AWS 上基础模型训练与推理的构建模块

基础设施：计算、网络与存储

加速计算

  • 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）使用可扩展可靠数据报（SRD）协议提供 OS 绕过 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 绕过网络功能。
  • CUDA 与内核： CUDA Toolkit 13.x 支持 Blackwell 架构。性能还受融合内核（如 FlashAttention）和可编程工具链（如 Triton、NVIDIA CuTe）的推动。

通信底层

  • NCCL： NVIDIA 集体通信库实现了面向拓扑的 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 为解耦的推理架构（将预填充和解码阶段分离）提供点对点传输。

框架

  • PyTorch： 分布式工作负载的主要框架，使用 torch.distributed 和 FSDP2 对参数和优化器状态进行分片。
  • 分布式框架：
    • Hugging Face Transformers/Accelerate： 注重易用性和兼容性。
    • NVIDIA Megatron Core/NeMo： 通过 3D 并行（张量、流水线和专家）实现最高效率。
    • 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