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