使用 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 實例上開始優化的實用步驟。