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

混合精度训练

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

使用 Transformer Engine 在 P5 实例上进行 FP8 混合精度训练

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

FP8（8 位浮点精度）是一种数据类型，是加速 LLM 模型深度学习训练的另一种范式。随着支持 FP8 数据类型的英伟达™（NVIDIA®）H100 GPU 的发布，您可以从配备 H100 GPU 的 P5 实例的性能提升中获益，同时利用 FP8 混合精度训练加速分布式训练。

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

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

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

SMP v2 支持开箱即用的 Transformer Engine。因此，在 SageMaker 人工智能 (ml.p5.48xlarge) 的 P5 实例上使用 SMP v2 运行 FP8 训练时，你唯一需要做的就是导入torch.sagemaker训练脚本并继续使用原生 Transformer Engine Python 包。要了解有关使用 Transformer Engine 进行 FP8 训练的更多信息，请参阅 NVIDIA Transformer Engine 文档中的使用 Transformer Engine 的 FP8。下面的代码片段显示了在训练脚本中导入 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 实例上使用 FP8 加速 SageMaker PyTorch FSDP 训练 Llama-v2 （或 GPT-NeoX）中的示例笔记本。

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

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

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

# 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。