Optimize model training on Amazon SageMaker AI with NVIDIA Blackwell

在 Amazon SageMaker AI 上使用 NVIDIA Blackwell 优化模型训练 | 人工智能

在 Amazon SageMaker AI 上使用 NVIDIA Blackwell 优化模型训练

使用 NVIDIA Blackwell GPU 在 Amazon SageMaker AI 上优化模型训练，改变了大型 AI 模型的可行性。如果你今天正在训练大型模型，你可能正在应对一组熟悉的限制：由于 GPU 内存限制，批量大小受限；为了防止内存溢出错误，序列长度被缩短；随着规模扩大，模型分片增加了通信开销。Blackwell 扩展的内存和新的精度格式直接减少了这些限制。配备 8 个 Blackwell GPU 的 P6-B200 实例现已在 Amazon SageMaker AI 训练任务中提供，你可以通过灵活训练计划预订容量，实现可预测的访问、成本管理和自动资源管理。Amazon SageMaker AI 训练任务通过自动提供和管理底层计算基础设施和资源，使你能够以大规模训练 ML 模型，从而专注于数据和算法，而不是基础设施操作。

本文将向你展示如何在 Amazon SageMaker AI 上配置训练任务，以充分利用 Blackwell 在 AWS 上的架构。你将学习如何选择能够利用 Blackwell 扩展内存的批量大小和序列长度，为不同模型规模（1B 到 64B 参数）选择合适的精度格式，并战略性地应用激活检查点。到文章结尾时，你将拥有一个实用的框架，用于调整训练配置并在 P6-B200 实例上启动分布式训练任务。

正确配置的 Blackwell 训练任务可以在不进行激进分片的情况下处理更大的批量大小，从而减少通信开销并提高吞吐量。更长的序列长度对于处理长距离依赖任务变得可行。使用合适的精度格式，以前需要多节点设置的模型可以在单个 8-GPU 节点上运行，这意味着更快的迭代周期、更少的网络开销和更低的基础设施成本。

了解 NVIDIA Blackwell

在配置训练任务之前，了解 Blackwell 与之前 GPU 世代的不同之处是有帮助的。Blackwell 的双芯片架构和第五代 Tensor Cores 提供了开箱即用的多 GPU 训练性能提升。NVLink 5 互连提供高达 1.8 TB/s 的双向 GPU 到 GPU 带宽，而 B200 的更大 HBM 容量和更高的内存带宽有助于减少大型批量、长序列和分布式训练任务的内存压力。

本文中的示例使用单节点 8-GPU 训练，适用于参数范围从 1B 到 64B 的变压器模型。训练配置使用 PyTorch 全部分片数据并行（FSDP），这是一种分布式训练技术，它将模型参数、梯度和优化器状态在 GPU 之间进行分片，以训练超出单 GPU 内存的模型。结果涵盖了多种配置，包括不同批量大小、序列长度和精度格式，以展示不同方法何时能带来最佳效果。

内存管理

Blackwell 扩展的内存（B200 上的 180 GB，B300 上的 268 GB）使你在以下三个方面可以进行优化：更大的批量大小、简化模型分片和更长的序列长度。

• 更大的批量大小可以减少 GPU 之间的梯度同步步骤数量，从而提高整体吞吐量。

• 每块 GPU 拥有更多的内存，使得简化模型分片成为可能，因为您可能能够减少模型并行度，甚至对某些模型完全消除并行度。更少的分片意味着 GPU 之间的通信开销更少。

• 更长的序列长度允许模型在单次处理中处理更多的上下文，这对于处理长距离依赖任务至关重要。

如果吞吐量是您的主要目标，请从调整批量大小开始。如果通信开销是瓶颈，请首先简化分片。如果您的任务需要长距离上下文，请优先考虑序列长度。批量大小和序列长度都会增加内存消耗，找到有效的平衡点非常重要。

激活检查点有助于在内存使用和计算之间取得平衡。它通过在反向传播过程中重新计算中间激活值（而不是存储它们）来减少 GPU 内存使用，从而以增加计算时间（根据模型架构通常增加 10-30% 的开销）为代价。释放的内存可以重新用于更大的批量大小或更长的序列。由于计算开销因工作负载而异，请对特定配置进行基准测试，以了解检查点的权衡后再决定是否启用。

例如，在图 1 中，我们比较了使用 MXFP8 精度在 8K 序列长度下训练一个 10 亿参数的 LLM 的三种训练配置。不启用激活检查点（BS=1）时，吞吐量约为 6000 tokens/sec，但峰值内存高达 15.5 GB。在相同批量大小下启用激活检查点，内存显著下降至 2.3 GB（因为中间激活值被重新计算而不是存储），但由于重新计算的开销，吞吐量略有下降。关键的回报体现在第三个条形图中：启用激活检查点并将批量大小增加到 16，释放的内存允许使用更大的批量，将吞吐量提升至约 51000 tokens/sec（约为基准值的 8 倍），同时峰值内存上升至 22.8 GB，仍然在 GPU 限制范围内。

图 1. 启用和不启用激活检查点时的吞吐量差异

为了判断激活检查点是否适用于您的工作负载，请考虑您的模型大小和内存使用情况：

• 小型模型（最多约 140 亿参数）：通常不需要激活检查点。凭借 Blackwell 扩展的内存，大多数小型模型在不启用检查点的情况下也能舒适地运行。如果您处于该范围的上限并遇到内存压力，激活检查点会增加计算开销，但可以带来显著的内存节省，您可以将这些节省的内存重新投资于更大的批量大小。

• 大型模型（约 140 亿参数及以上）：在该模型大小下，内存消耗根据批量大小和序列长度在 87 GB 到 171 GB 之间变化。如果不启用激活检查点，大多数配置会因 CUDA 内存不足（OOM）错误而失败。当您添加检查点后，释放的内存允许您增加批量大小，从而在增加计算开销的情况下提高吞吐量。对于大型模型，检查点不是可选的，而是稳定训练的前提条件。

精度格式

Blackwell 的第五代 Tensor Cores 提供了对低精度格式（FP8、MXFP8 和 NVFP4）的硬件加速，使其主要作为吞吐量优化手段，而不是节省内存的技术。使用较低精度可以减少内存带宽需求，同时增加 GPU 每个周期可以执行的操作数量。然而，默认情况下，低精度训练在内存使用上大致是中性的，因为 Transformer Engine 会同时维护高精度的主权重（用于优化器更新）和量化副本，因此较低精度格式并不直接转化为更低的内存使用。量化本身会引入开销（在精度格式之间转换以及在内存中维护权重的多个副本），这意味着净收益取决于模型大小以及训练是计算受限还是内存受限。虽然 NVFP4 提供了最高的吞吐量，但其性能优势主要随着大型模型和推理工作负载而提升，在这些情况下不需要主权重。

对于计算受限的工作负载（通常是较小的模型），计算速度是限制因素，量化开销部分抵消了低精度带来的吞吐量提升。对于内存受限的工作负载（通常是较大的模型），数据移动是瓶颈，低精度格式的较小内存占用直接解决了这一限制，带来更显著的提升：

• 小型模型（最多约 140 亿参数）：在该模型规模下，低精度格式（FP8、MXFP8、NVFP4）在 FP16 上提供的吞吐量改进大致相似，因为量化开销抵消了速度优势。批量大小调整通常比精度格式选择带来更显著的提升。建议从 FP8 开始，以获得更高的吞吐量。与 MXFP8 或 NVFP4 相比，FP8 的开销更低，通常对于大多数小型模型工作负载来说是一个不错的默认选择。请注意，默认的 TransformerEngine 设置下，低精度格式使用的内存比 FP16 更多，因为 TransformerEngine 会将权重保存在更高精度，并在运行时动态转换。如果内存是限制因素且你的优化器支持，可以使用 quantized_model_init 将权重直接存储在 FP8 中，从而将内存使用量降低到 FP16 以下。

• 大型模型（约 140 亿参数及以上）：这就是低精度格式产生最大影响的地方。FP8 通常在吞吐量和内存效率之间提供了良好的平衡。虽然 MXFP8 在理论上更节省内存，但其转置开销在实践中部分抵消了这一优势。然而，如果你的工作负载对收敛稳定性或数值精度有较高要求，MXFP8 可能是更好的选择，因为其更细粒度的量化方案更可靠地保留了模型的准确性。对于内存是主要瓶颈的大型模型，NVFP4 可以带来额外的吞吐量提升，因为其矩阵乘法速度优势随着模型规模的增大而提升。实现这些提升需要实质性的工程投入。建议使用 Megatron Core 提供的框架级配方，这些配方提供了经过验证的 NVFP4 配置，而不是从头开始实现。

NVIDIA 的 TransformerEngine 处理了实现的复杂性：自动混合精度切换、融合内核和动态损失缩放。在投入生产之前，通过跟踪不同格式下的损失曲线来验证收敛性，以确认所选精度满足准确性要求。

并非所有工作负载都能从激进的优化中受益。如果您的模型在内存限制内可以顺利训练，并且使用 FP16 即可满足吞吐量要求，那么使用降低精度格式所带来的额外复杂性可能不值得投入工程努力。从基准测量开始，然后仅优化那些可以测量的瓶颈。

在 Amazon SageMaker AI 上使用 Blackwell 进行训练的入门

前面的章节涵盖了关键决策：您可用的内存大小、激活检查点是否适合您的模型规模，以及哪种精度格式适合您的工作负载。接下来的章节将使用 Amazon SageMaker AI 训练任务来实践这些决策。

Amazon SageMaker AI 提供了一个完全托管的环境，用于在 Blackwell 实例上进行分布式训练，处理实例配置、容器编排，并与 AWS 服务（如 Amazon Simple Storage Service（Amazon S3）、Amazon CloudWatch、Amazon Elastic Container Registry（Amazon ECR）和 AWS Identity and Access Management（AWS IAM））进行集成。

先决条件

在开始之前，请确认您具备以下条件：

• 一个具有创建 Amazon SageMaker AI 训练任务、访问 Amazon ECR 和为 Amazon SageMaker AI 执行创建 IAM 角色权限的 AWS 账户。

• 通过灵活训练计划或托管 Spot 训练访问 ml.p6-b200.48xlarge 实例。在开始之前，请检查您的服务配额。

• 本地安装了 Docker、Python 3.9 或更高版本，以及 SageMaker Python SDK。

• 熟悉 PyTorch 和 FSDP。如果您是 FSDP 的新手，请参阅《分布式训练入门》。

启动训练任务

要启动训练任务，请完成以下步骤。

#### 步骤 1：创建脚本

从 NVIDIA TransformerEngine 仓库中的 FSDP 示例下载 fsdp.py 文件。此脚本实现了 FSDP 训练，并接受命令行参数作为超参数。

#### 步骤 2：创建入口脚本

准备一个 train.sh 文件，用于配置 torchrun 并启动训练脚本：

#!/bin/bash

# SageMaker 将超参数作为环境变量传递（SM_HP_<NAME>）
PRECISION=${SM_HP_PRECISION:-"mxfp8"}
NUM_LAYERS=${SM_HP_NUM_LAYERS:-10}
BATCH_SIZE=${SM_HP_BATCH_SIZE:-8}
SEQ_LENGTH=${SM_HP_SEQ_LENGTH:-2048}

NUM_GPUS=$(nvidia-smi --list-gpus | wc -l)
torchrun --standalone --nnodes=1 --nproc-per-node="$NUM_GPUS" \
    fsdp.py --no-defer-init --precision "$PRECISION" \
    --num-layers "$NUM_LAYERS" --checkpoint-layer "transformerlayer" \
    --batch-size "$BATCH_SIZE" --seq-length "$SEQ_LENGTH"

#### 步骤 3：构建并推送容器

构建一个自定义的 Docker 容器，该容器扩展了 AWS Deep Learning Containers（DLC），包含 fsdp.py 和 train.sh，并安装了 TransformerEngine 2.11。DLC 提供了一个经过验证的基础镜像，其中包含 PyTorch 和 Blackwell 兼容所需的 CUDA 库。您可以使用以下 Dockerfile：

FROM 763104351884.dkr.ecr.us-west-2.amazonaws.com/pytorch-training:2.9.0-gpu-py312-cu130-ubuntu22.04-sagemaker

# 安装 Transformer Engine
RUN pip install --upgrade --no-build-isolation transformer_engine[pytorch]==2.11.0
# 为预构建的 flash-attn 轮子提供 libcudart.so.12
RUN pip install nvidia-cuda-runtime-cu12

# 使链接器能够找到它
ENV LD_LIBRARY_PATH=/usr/local/lib/python3.12/site-packages/nvidia/cuda_runtime/lib:$LD_LIBRARY_PATH

COPY fsdp.py /opt/ml/code/fsdp.py
COPY train.sh /opt/ml/code/train.sh

ENV SAGEMAKER_SUBMIT_DIRECTORY /opt/ml/code ENV SAGEMAKER_PROGRAM train.sh

构建完成后，如果尚未拥有 Amazon ECR 私有仓库，请创建一个，并将镜像推送到 Amazon ECR（注意从输出中获取 repositoryUri，用于 docker tag 和 push 命令）。有关详细的构建说明，请参阅将您自己的 Docker 容器适配为与 Amazon SageMaker AI 兼容。

#### 第 4 步：确保容量

通过灵活训练计划保留容量，以标准费率获得可预测的访问，或使用托管 Spot 训练以优化成本的工作负载。对于需要预留容量以确保持续可用性的生产训练运行，请使用灵活训练计划；对于可以容忍重启且成本降低优先于中断风险的实验和容错工作负载，请使用 Spot。Spot 实例可能会被中断，因此请确保训练脚本定期将检查点保存到 Amazon S3。如果在估算器配置中提供 checkpoint_s3_uri，Amazon SageMaker AI 会自动恢复被中断的 Spot 作业。在 SageMaker AI 控制台中创建灵活训练计划以获得可预测的容量，并选择“training-job”作为目标资源。如果您的任务可以容忍重启且成本降低是首要考虑因素，可以选择 Spot，但需要首先提高 ml.p6-b200.48xlarge Spot 训练任务的使用配额。注意：灵活训练计划会预留容量，并在计划期间无论训练任务是否正在运行都会产生费用。在创建计划之前，请查看定价详情。

#### 第 5 步：提交训练任务

将占位符值替换为实际的训练计划 ARN 和 ECR 镜像 URI，然后从本地开发环境或 SageMaker AI 笔记本实例运行以下代码：

from sagemaker.estimator import Estimator from sagemaker import get_execution_role from sagemaker.debugger import ProfilerConfig

training_plan_arn = "<your-training-plan-arn>" # 替换为您的训练计划 ARN ecr_image = "<your-ecr-image-uri>" # 替换为您的 ECR 镜像 URI

根据您的工作负载调整这些值

precision = "mxfp8" num_layers = 10 batch_size = 8 seq_length = 2048

estimator = Estimator( image_uri=ecr_image, role=get_execution_role(), base_job_name='blackwell-training', instance_count=1, instance_type='ml.p6-b200.48xlarge', hyperparameters={ "precision": precision, "num-layers": num_layers, "batch-size": batch_size, "seq-length": seq_length, }, profiler_config=ProfilerConfig(disable_profiler=True), training_plan=training_plan_arn)

estimator.fit()

对于托管 Spot 训练，请将 training_plan 替换为 use_spot_instances=True，设置 max_run 和 max_wait，并添加 checkpoint_s3_uri 以实现自动恢复。

#### 第 6 步：监控训练任务

Amazon SageMaker AI 会自动将日志流式传输到 Amazon CloudWatch。在 SageMaker AI 控制台中，导航到训练任务并选择您的任务以找到 CloudWatch 日志组。打开 /aws/sagemaker/TrainingJobs 并查找 [rank 0] 行以查看损失值和吞吐量。要确认已加载的精度格式，请查找类似“Using FP8 recipe”或“MXFP8 enabled”的消息。

如果任务因 CUDA 内存不足（OOM）错误而停止，日志中会显示分配大小。如果没有已经这样做，请减少批量大小或序列长度，或添加激活检查点。

为了避免测试结束后产生持续费用，请在 Amazon SageMaker AI 控制台中停止任何正在运行的训练任务（请注意，这不会取消灵活训练计划）。警告：以下清理步骤将永久删除资源且无法恢复。在继续操作之前，请确认已备份所有需要的数据。删除您的 Amazon ECR 仓库（这将永久删除容器镜像），删除存储在 Amazon S3 中的任何训练工件（这将永久删除训练数据和检查点），并删除 CloudWatch 日志组（这将永久删除训练日志）。删除为 Amazon SageMaker AI 创建的 IAM 执行角色。

## 结论

在本文中，您学习了如何使用 Amazon SageMaker AI 训练任务在 NVIDIA Blackwell GPU 上优化 AI 模型训练。您配置了批量大小和序列长度，以利用 Blackwell 扩展的内存，根据模型大小应用激活检查点，并选择了适合您工作负载的精度格式。您还使用 TransformerEngine 设置了自定义容器，通过灵活训练计划确保了计算资源，并在 ml.p6-b200.48xlarge 实例上启动了分布式训练任务。

从 1B 到 64B 参数的 Transformer 模型在结合这些优化时表现出一致的提升。关键是了解您的工作负载是计算受限还是内存受限，然后逐步应用更改，以便衡量每个更改的影响。

如果您准备开始，请查阅 Amazon SageMaker AI 文档以查看实例选项和配置详情，或购买灵活训练计划以保留 Blackwell 资源用于下一次训练运行。如果您对特定工作负载有疑问，请联系 AWS 支持或您的 AWS 账户团队。

## 作者简介

'"`