混合精度训练 - 亚马逊 SageMaker AI

本文属于机器翻译版本。若本译文内容与英语原文存在差异,则一律以英文原文为准。

混合精度训练

SageMaker 模型并行度 (SMP) 库 v2 通过与 Transformer Engine 等 PyTorch FSDP开源框架集成,开箱即用地支持混合精度训练。要了解更多信息,请参阅以下主题。

使用变形引擎FP8在 P5 实例上进行混合精度训练

从 SageMaker 模型 parallelism (SMP) 库 v2.2.0 开始,该SMP库与 Transformer E ngin e 集成,并支持开箱即FP8用的混合精度训练,同时保持与的兼容性。PyTorch FSDPMixedPrecision这意味着您可以同时使用混合精度训练和 Tr PyTorch FSDP ansformer Engine 进行FP8训练。对于 Transformer Engine 的FP8训练功能不支持的模型层,这些层会回退到 PyTorch FSDP混合精度。

注意

SMPv2 FP8 支持以下 Hugging Face Transformer 型号:

  • GPT-neoX(在 SMP v2.2.0 及更高版本中可用)

  • Llama 2(在 SMP v2.2.0 及更高版本中可用)

  • Mixtral 8x7b 和 Mixtral 8x22b(在 v2.5.0 及更高版本中可用)SMP

注意

本关于 P5 功能的FP8培训有以下库 SageMaker 和库组合: PyTorch

  • SageMaker Python SDK v2.212.0 及更高版本

  • PyTorch v2.2.0 及更高版本

FP8(8 位浮点精度)是一种数据类型,已成为加速LLM模型深度学习训练的另一种范式。随着 NVIDIA H100 GPUs 支持FP8数据类型的发布,您可以从配备 H100 的 P5 实例的性能改进中受益GPUs,同时通过FP8混合精度训练加速分布式训练。

FP8数据类型进一步分支到 E4M3 和 E5M2 格式。E4M3 具有更高的精度,但动态范围有限,是模型训练中前向传递的理想选择。E5M2 具有更宽的动态范围,但精度有所降低,更适用于精度要求不高、动态范围更宽的后向传递。因此,我们建议您使用混合FP8策略配方来有效地利用这些特征。

对于半精度数据类型(FP16和BF16),全局损失缩放技术(例如静态损失缩放或动态损失缩放)可以处理因半精度舍入梯度而导致的信息丢失所产生的收敛问题。但是,的动态范围甚至FP8更窄,全局损失缩放技术还不够。此时,我们需要一种更精细的按张量缩放技术。延迟缩放是一种根据之前迭代中观察到的一些张量的最大绝对值来选择缩放因子的策略。这种策略需要权衡取舍;它利用了FP8计算的全部性能优势,但需要内存来保存张量的最大值历史记录。要了解有关延迟扩展策略的更多信息,请参阅论文《深度学习FP8格式》。

实际上,在 P5 实例的所有训练场景中使用FP8都很有帮助。我们强烈建议FP8尽可能启用以提高训练绩效。

SMPv2 开箱即用支持 Transformer Engine。因此,在 SageMaker AI (ml.p5.48xlarge) 的 P5 实例上使用 SMP v2 运行FP8训练时,你唯一需要做的就是导入torch.sagemaker训练脚本并继续使用原生 Transformer Engine Python 包。要了解有关使用 Transformer Engine 进行一般FP8训练的更多信息,请参阅 Transf FP8 ormer Engine 文档中的与NVIDIA变形引擎一起使用 以下代码片段显示了用于导入SMP库和在训练脚本FP8中进行设置的代码行应该是什么样子。

import torch.sagemaker as tsm import transformer_engine.pytorch as te from transformer_engine.common.recipe import DelayedScaling, Format # Initialize the SMP torch.sagemaker API. tsm.init() # Define a transformer model and wrap it with the torch.sagemaker.transform API. from transformers import AutoModelForCausalLM model = AutoModelForCausalLM.from_config(ModelConfig) model = tsm.transform(model) # Enable E4M3 during forward pass, E5M2 during backward pass. fp8_format = Format.HYBRID # Create an FP8 recipe. fp8_recipe = DelayedScaling(fp8_format=fp8_format, amax_history_len=32, amax_compute_algo="max") # Enable FP8 autocasting. with te.fp8_autocast(enabled=True, fp8_recipe=fp8_recipe, fp8_group=tsm.state.world_process_group): out = model(inp) loss = out.sum() loss.backward()

要查找在 P5 实例上使用 SMP v2 进行FP8训练的实际示例,请参阅在 P5 实例上使用 llama-v2(或 GPT-neoX)加速 SageMaker PyTorch FSDP训练中的示例笔记本。FP8

使用半精度数据类型进行混合精度训练 PyTorch FSDP

SMPv2 支持PyTorch FSDPMixedPrecision在 P4 和 P5 实例上训练作业。 PyTorch FSDP为混合精度提供了各种配置,以提高性能和减少内存。

注意

这种带有该 PyTorch FSDP功能的混合精度训练可在以下库 SageMaker 和库组合中使用。 PyTorch

  • SMPv2.0.0 及更高版本

  • SageMaker Python SDK v2.200.0 及更高版本

  • PyTorch v2.0.1 及更高版本

为混合精度配置模型的标准方法是在中创建模型float32,然后允许通过传递MixedPrecision策略将参数FSDP转换为float16bfloat16即时强制转换参数,如以下代码片段所示。有关在dtype中更改参数、缩减或缓冲区以实现混合精度的选项的更多信息 PyTorch,请参阅PyTorch文档PyTorch FSDPMixedPrecisionAPI中的。

# Native PyTorch API from torch.distributed.fsdp import MixedPrecision dtype = torch.bfloat16 mixed_precision_policy = MixedPrecision( param_dtype=dtype, reduce_dtype=dtype, buffer_dtype=dtype ) model = FSDP( model, ..., mixed_precision=mixed_precision_policy )

请注意,某些模型(例如 Hugging Face Transformers Llama 模型)期望缓冲区为 float32。要使用 float32,请在定义 dtype 对象的一行中将 torch.bfloat16 替换为 torch.float32