在 AWS 上基礎模型訓練與推論的構建模組

基礎模型的規模化已從單純關注預訓練計算，演變為涵蓋預訓練、後訓練（SFT 與 RL）以及測試時計算（搜尋與驗證）的三管齊下策略。此變化需要一套以緊密耦合的加速器計算、高頻寬低延遲網路以及分散式儲存後端為核心的收斂基礎設施策略。

基礎設施：計算、網路與儲存

AWS 提供分層式基礎設施，旨在最小化集體通信與記憶體移動的瓶頸，這些瓶頸往往比純粹的計算吞吐量更主導步驟時間。

加速計算

AWS 透過 Amazon EC2 P 系列實例提供多代 NVIDIA GPU：

  • P5 實例： 配備 NVIDIA H100 GPU（p5.48xlarge）與 H200 GPU（p5e.48xlarge/p5en.48xlarge）。
  • P6 實例： 引入 NVIDIA Blackwell B200（p6-b200.48xlarge）與 B300（p6-b300.48xlarge）架構。

網路與互連

GPU 通信分為兩種模式以最佳化資料移動：

  • 內部擴展（NVLink/NVSwitch）： 單節點內的高頻寬、低延遲連接。
  • 外部擴展（EFA）： Elastic Fabric Adapter（EFA）使用 Scalable Reliable Datagram（SRD）協議提供 OS‑bypass RDMA 以進行跨節點通信。EFAv3 相較於 EFAv2 可降低約 35% 的封包延遲，EFAv4 在集體通信效能上較 EFAv3 再提升 18%。
  • UltraServers： Amazon EC2 UltraServers（例如 P6e-GB200）將 NVLink 範圍延伸至單一實例之外，最多可在同一 NVLink 域內容納 72 顆 Blackwell GPU，減少通信離開 NVLink 纖維的頻率。

分層儲存

為處理多 TB 級別的檢查點與龐大語料庫，AWS 採用三層階層結構：

  1. 本機 NVMe SSD： 暫存實例儲存，用於熱資料（最高 30.72 TB 原始容量）。
  2. Amazon FSx for Lustre： 受管的平行檔案系統，提供高聚合吞吐量與次毫秒延遲。
  3. Amazon S3： 用於持久保存與透過資料庫關聯（Data Repository Associations）延遲載入資料集。

資源編排：Slurm 與 Kubernetes

管理成千上萬的加速器需要集中式控制平面，以確保作業排程的原子性與資源效率。

Slurm（以 HPC 為中心）

Slurm 是 HPC 工作負載的主要管理器，於作業層級排程，確保在任務啟動前已分配所有必要節點。AWS 透過以下方式支援 Slurm：

  • AWS ParallelCluster： 開源工具，用於自動化 Slurm 叢集部署。
  • AWS Parallel Computing Service（PCS）： 受管的 Slurm 控制平面。
  • Amazon SageMaker HyperPod： 為 Slurm 模式加入持續節點健康監控與作業自動恢復功能。

Kubernetes（雲原生）

Kubernetes 雖擅長部署，但缺乏原生的作業層級原子性與拓撲感知。以下元件彌補此缺口：

  • Kueue： 管理作業層級的 gang admission 與多租戶配額。
  • Volcano 與 NVIDIA KAI Scheduler： 提供拓撲感知的 Pod 放置，以最佳化 NVLink 與 EFA 使用。
  • SageMaker HyperPod（EKS 模式）： 整合受管的 Kueue 與 Karpenter 以實現即時供應。它亦引入 checkpointless training，透過 EFA 的點對點狀態複製，在不需從儲存讀取多 TB 檢查點的情況下從失敗中復原。

機器學習軟體堆疊

效能由五層堆疊決定，從硬體驅動程式到高階框架皆有所影響。

低階啟用與執行環境

  • 核心驅動程式： NVIDIA GPU 驅動支援 GPUDirect RDMA，EFA 驅動則透過 libfabric 提供 OS‑bypass 網路功能。
  • CUDA 與 Kernels： CUDA Toolkit 13.x 支援 Blackwell 架構。效能進一步受 FlashAttention 等融合 kernel 與 Triton、NVIDIA CuTe 等可程式化工具鏈驅動。

通訊基礎層

  • NCCL： NVIDIA Collective Communications Library 實作拓撲感知的 all‑reduce 與 all‑gather 演算法。對於 Mixture‑of‑Experts（MoE）模型，all‑to‑all 集體通信對於在專家之間路由 token 至關重要。
  • aws-ofi-nccl： 將 NCCL 傳輸 API 映射至 libfabric 的插件，使 NCCL 能使用 EFA 的 SRD 協議。
  • NIXL： NVIDIA Inference Xfer Library 為分散式推論架構（將 prefill 與 decode 階段分離）提供點對點傳輸。

框架

  • PyTorch： 分散式工作負載的主要框架，使用 torch.distributed 與 FSDP2 進行參數與 optimizer 狀態的分片。
  • 分散式框架：
    • Hugging Face Transformers/Accelerate： 強調易用性與相容性。
    • NVIDIA Megatron Core/NeMo： 透過 3D 平行（tensor、pipeline、expert）追求最高效能。
    • veRL： 一個強化學習框架，使用 HybridFlow 在單一作業中混合訓練後端（FSDP2、Megatron）與推論引擎（vLLM、SGLang）。
    • vLLM 與 SGLang： 推論引擎，利用 PagedAttention 與 RadixAttention 最佳化 KV 快取管理與請求排程。

可觀測性與故障偵測

系統化的遙測是診斷成千上萬 GPU 瓶頸的必要條件。

遙測堆疊

  • Prometheus 與 Grafana： 度量收集與可視化的標準。AWS 提供 Amazon Managed Service for Prometheus（AMP） 與 Amazon Managed Grafana（AMG），以減少運維負擔。
  • DCGM‑Exporter： 暴露 NVIDIA GPU 度量。SM 活動（DCGM_FI_PROF_SM_ACTIVE）被視為比基本利用率更精確的計算效率指標。

健康監控

主動偵測硬體故障可防止訓練中斷。關鍵指標包括：

  • ECC 錯誤： 單位元錯誤（SBE）率上升往往預示雙位元錯誤（DBE）。
  • XID 事件： 如 XID 63（行重映射失敗）、XID 64（GPU 從匯流排掉線）以及 XID 94/95（受限/未受限錯誤）會立即觸發節點更換。

Sources