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

使用 NVIDIA Blackwell GPU 在 Amazon SageMaker AI 上优化模型训练，改变了大型 AI 模型的可行性。如果您今天训练大型模型，很可能在应对一系列熟悉的约束：受 GPU 内存限制的批大小、为避免内存不足错误而缩短的序列长度，以及随着规模扩大而增加通信开销的模型分片。Blackwell 的扩展内存和新的精度格式直接减少了这些约束。配备 8 个 Blackwell GPU 的 P6-B200 实例现已在 Amazon SageMaker AI 训练作业中可用，您可以通过 Flexible Training Plan 预订容量，获得可预测的访问、成本管理和自动化资源管理。Amazon SageMaker AI 训练作业通过自动配置和管理底层计算基础设施和资源，让您能够大规模训练 ML 模型，从而专注于数据和算法，而非基础设施运维。

本文向您展示如何在 Amazon SageMaker AI 上配置训练作业，以充分利用 Blackwell 架构的优势。您将学习如何选择批大小和序列长度以利用 Blackwell 的扩展内存，为模型大小（1B 到 64B 参数）选择合适的精度格式，并策略性地应用激活检查点。最终，您将获得一个实用的框架来调整训练配置并在 P6-B200 实例上启动分布式训练作业。

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

了解 NVIDIA Blackwell

在配置训练作业之前，了解 Blackwell 与以往 GPU 代次的不同之处会有所帮助。Blackwell 的双芯片架构和第五代 Tensor Core 为多 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 序列长度的 1B 参数 LLM 的三种训练配置。没有激活检查点时（批大小=1），吞吐量约为 6K tokens/sec，但峰值内存高达 15.5 GB。启用相同批大小的激活检查点后，内存急剧下降到 2.3 GB（因为中间激活被重新计算而非存储），但吞吐量也因重新计算开销而略有下降。关键回报在第三个柱中：启用激活检查点并将批大小提升到 16 后，释放的内存允许更大的批，将吞吐量推高至约 51K tokens/sec（约为基线的 8 倍），而峰值内存上升到 22.8 GB，仍在 GPU 限制内。

要决定激活检查点是否适合您的工作负载，请考虑您的模型大小和内存使用情况：

• 小模型（最多约 14B 参数）：通常不需要激活检查点。借助 Blackwell 的扩展内存，大多数小模型无需它即可轻松适应。如果您在此范围的上限运行并遇到内存压力，激活检查点会增加计算开销以换取有意义的内存节省，您可以将这些内存重新投入到更大的批大小中。
• 大模型（约 14B+ 参数）：在此模型大小下，内存消耗范围从 87 到 171 GB，具体取决于批大小和序列长度。没有激活检查点时，大多数配置会因为 CUDA 内存不足 (OOM) 错误而失败。添加检查点后，释放的内存让您可以增加批大小，从而尽管有额外的计算开销，吞吐量仍能提高。对于大模型，检查点不是可选的；它是稳定训练的先决条件。

精度格式

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

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

• 小模型（最多约 14B 参数）：在此模型大小下，降低精度格式（FP8、MXFP8、NVFP4）在 FP16 基础上都提供类似、适度的吞吐量改进，因为量化开销抵消了速度优势。调整批大小通常比选择精度格式带来更有意义的增益。从 FP8 开始以提高吞吐量。它的开销低于 MXFP8 或 NVFP4，通常是大多数小模型工作负载的良好默认选择。请注意，使用默认的 TransformerEngine 设置时，降低精度格式比 FP16 使用更多内存，因为 TransformerEngine 保持权重为高精度并动态转换。如果内存是约束且优化器支持，请使用 quantized_model_init 将权重直接存储为 FP8，从而将内存降低到 FP16 水平以下。
• 大模型（约 14B+ 参数）：这是降低精度发挥最大作用的地方。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)）的集成。

先决条件

在开始之前，请确认您拥有：

• 具有创建 Amazon SageMaker AI 训练作业、访问 Amazon ECR 以及为 Amazon SageMaker AI 执行创建 IAM 角色权限的 AWS 账户。
• 通过 Flexible Training Plan 或受管 Spot Training 访问 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 passes hyperparameters as environment variables (SM_HP_)
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=8

# Launch training with torchrun
torchrun --nproc_per_node=$NUM_GPUS fsdp.py \
  --precision $PRECISION \
  --num-layers $NUM_LAYERS \
  --batch-size $BATCH_SIZE \
  --seq-length $SEQ_LENGTH

步骤 3：构建 Docker 映像并推送到 Amazon ECR

创建一个 Dockerfile，基于 PyTorch 映像并包含 TransformerEngine，然后将映像推送到您的 Amazon ECR 存储库。

步骤 4：使用 SageMaker Python SDK 启动训练作业

使用 SageMaker Python SDK 中的 Estimator 类指定映像、实例类型和超参数。

结论

通过采用上述策略——利用更大的批大小、更长的序列长度、适当的精度格式和激活检查点——您可以显著提高训练吞吐量并降低成本。Blackwell 架构在 SageMaker AI 上提供了强大的能力，通过精心的配置，您可以为 AI 模型训练树立新的性能标杆。本文提供的框架为您提供了在 P6-B200 实例上开始优化的实用步骤。