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

AWS ML Blog 2026-06-26T09:28:37.062031

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

在 Amazon SageMaker AI 上结合 NVIDIA Blackwell GPUs（图形处理器）优化模型训练，改变了大型 AI 模型的实际可行性。如果你今天在训练大型模型，很可能遇到了熟悉的限制：受 GPU 内存限制的批次大小、为避免内存溢出错误而缩短的序列长度，以及随着扩展而增加通信开销的模型分片。Blackwell 的扩展内存和新精度格式直接缓解了这些限制。配备 8 个 Blackwell GPU 的 P6-B200 实例已在 Amazon SageMaker AI Training jobs 中可用，你可以使用 Flexible Training Plan（灵活训练计划）预订容量，享受可预测的访问、成本管理和自动化资源管理。Amazon SageMaker AI training jobs 通过自动配置和管理底层计算基础设施及资源，让你能够大规模训练 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 参数的 Transformer。训练配置采用 PyTorch Fully Sharded Data Parallel (FSDP)，这是一种分布式训练技术，将模型参数、梯度和优化器状态分片到多个 GPU 上，以训练超出单个 GPU 内存容量的模型。结果涵盖了多种配置，包括不同批次大小、序列长度和精度格式，展示了不同方法何时能获得最佳结果。

内存管理

Blackwell 的扩展内存（B200 上 180 GB，B300 上 268 GB）为你提供了三个优化空间：更大的批次大小、简化的模型分片和更长的序列长度。

更大的批次大小 减少了 GPU 间的梯度同步步骤数量，从而提高了整体吞吐量。
简化的模型分片 成为可能，因为每个 GPU 拥有更多内存，这意味着对于某些模型，你可以降低模型并行度，甚至完全消除它。分片越少，GPU 间通信开销就越小。
更长的序列长度 允许模型在单次处理中处理更多上下文，这对于长距离依赖任务至关重要。

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

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

例如，在图1中，我们比较了使用MXFP8精度、序列长度8K时，对一个10亿参数LLM的三种训练配置。在不使用激活检查点（BS=1）的情况下，吞吐量约为6K tokens/sec，但峰值内存高达15.5 GB。在相同批次大小下启用激活检查点，内存大幅降至2.3 GB（因为中间激活被重新计算而非存储），但由于重新计算开销，吞吐量也略有下降。关键收益体现在第三个柱状图中：启用激活检查点并将批次大小提升到16后，释放的内存允许更大的批次，将吞吐量推高至约51K 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提供最高的吞吐量，但其性能优势主要随着大型模型和推理工作负载（无需主权重）而扩展。

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

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

原文：
* 小模型（最多约 140 亿参数）： 在此模型规模下，低精度格式（FP8、MXFP8、NVFP4）相比 FP16 都能带来相似且适度的吞吐量提升，因为量化开销会抵消部分速度优势。调整批次大小通常比选择精度格式能带来更有意义的收益。建议从 FP8 开始以获得更高吞吐量。与 MXFP8 或 NVFP4 相比，其开销更低，通常是大多数小模型工作负载的良好默认选择。请注意，使用默认的 TransformerEngine 设置时，低精度格式比 FP16 占用更多内存，因为 TransformerEngine 会以更高精度保留权重并在运行时进行转换。如果内存受限且优化器支持，可以使用 quantized_model_init 直接将权重以 FP8 存储，从而将内存降低到 FP16 水平以下。
* 大模型（约 140 亿参数以上）： 这是低精度格式发挥最大作用的场景。FP8 通常在吞吐量和内存效率之间提供了良好的平衡。虽然 MXFP8 理论上内存效率更高，但其转置开销在实践中会部分抵消这一优势。然而，如果工作负载的收敛稳定性或数值精度是优先考虑因素，那么 MXFP8 可能是更好的选择，因为其更细粒度的量化方案往往能比 FP8 更可靠地保持模型精度。对于内存是主要瓶颈的大模型，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)）的集成。

前提条件

在开始之前，请确认您已具备：

一个 AWS 账户，具有创建 Amazon SageMaker AI 训练作业、访问 Amazon ECR 以及为 Amazon SageMaker AI 执行创建 IAM 角色的权限。
通过弹性训练计划或托管 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 wheel 提供 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 命令）。有关详细的构建说明，请参阅 Adapting your own Docker container to work with Amazon SageMaker AI。

步骤 4：确保容量

通过 Flexible Training Plan（灵活训练计划）预订容量，以标准费率获得可预测的访问权限；或者使用 Managed Spot Training（托管 Spot 训练）进行成本优化的工作负载。对于需要容量预留以确保连续可用性的生产训练任务，请使用 Flexible Training Plans；对于实验性和容错性工作负载，如果成本降低的收益大于中断风险，请使用 Spot。Spot 实例可能会中断，因此请确保您的训练脚本定期将检查点保存到 Amazon S3。如果您在 estimator 配置中提供了 checkpoint_s3_uri，Amazon SageMaker AI 会自动恢复中断的 Spot 任务。在 SageMaker AI 控制台上创建 Flexible Training Plan，以获得可预测的容量，并选择 'training-job' 作为目标资源。如果您的任务可以容忍重启，并且成本降低是优先事项，您可以选择 Spot，但需要首先提高 ml.p6-b200.48xlarge Spot 训练任务使用的配额。注意：Flexible Training Plans 会预留容量，并在计划有效期内产生费用，无论训练任务是否正在运行。在创建计划之前，请查看定价详情。

步骤 5：提交训练任务

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

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

training_plan_arn = "" # 替换为您的训练计划 ARN
ecr_image = "" # 替换为您的 ECR 镜像 URI

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

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

estimator = Estimator(

查看原文