AWS 上基础模型训练与推理的构建模块
Building Blocks for Foundation Model Training and Inference on AWS
NVIDIA提出基础模型扩展已从单一预训练曲线转变为三条定律:预训练、后训练(SFT、RL)和测试时计算(长思考、搜索)。AWS分析了其基础设施(EC2 P5/P6实例、H100/B200 GPU、EFA网络、FSx for Lustre存储)如何与开源软件栈(Slurm、Kubernetes、PyTorch、NCCL、vLLM)交互,支持大规模分布式训练和推理。文章涵盖四层架构:计算/网络/存储基础设施、资源编排、ML软件栈和可观测性(Prometheus、Grafana、DCGM),由Aman Shanbhag、Pavel Belevich和Keita Watanabe撰写。
](https://huggingface.co/KeitaWatanabe)
长期以来,基础模型中的"扩展"主要意味着一件事:在预训练上投入更多算力,能力就会提升。这一直觉得到了经验研究的支持,例如 Kaplan et al. (2020),该研究报告了在扩展模型参数、数据集大小和训练算力时,损失函数存在可预测的幂律趋势。在实践中,这些趋势证明了持续投资大规模加速器容量以及维持其高效利用所需的分布式基础设施是合理的。但前沿已经演变——扩展不再是单一曲线。NVIDIA 的"从一条到三条扩展定律"框架有效地强调了,除了预训练之外,性能还越来越多地通过后训练(例如,监督微调(SFT)和基于强化学习(RL)的方法)以及测试时计算("长思考"、搜索/验证、多样本策略)来扩展。
图: 改编自 "AI's Three Scaling Laws, Explained"(NVIDIA 博客)。
综合来看,这些扩展机制将基础模型的生命周期——预训练、后训练和推理——推向趋同的基础设施需求:紧密耦合的加速器计算、高带宽低延迟网络,以及分布式存储后端。它们也提升了编排在资源管理中的重要性,以及应用和硬件层面的可观测性,以维护集群健康并在大规模下诊断性能问题。
另一个关键趋势是基础模型生命周期越来越依赖一个开源软件(OSS)生态系统,该生态系统涵盖模型开发框架、集群资源管理和运维工具。在集群层面,资源管理通常由 Slurm 和 Kubernetes 等系统提供。模型开发和分布式训练通常使用 PyTorch 和 JAX 等框架实现。监控和可视化——即可观测性——通常使用 Prometheus 进行指标收集,使用 Grafana 进行可视化和告警,定位为基础设施和资源管理之上的运维层。图 1 展示了这种分层架构,显示了硬件基础设施如何支持资源编排,资源编排又如何支持 ML 框架,而可观测性则贯穿所有层。
图 1:用于基础模型训练和推理的开源软件栈的分层架构
本文面向参与基础模型训练和推理的机器学习工程师和研究人员,特别关注基于 OSS 框架的工作流。它分析了 AWS 基础设施——包括多节点加速器计算、高带宽低延迟网络、分布式共享存储以及相关的托管服务——如何与基础模型生命周期中的常见 OSS 栈进行交互。主要目标是为理解跨越预训练、后训练和推理的系统瓶颈和扩展特性提供技术基础。这篇介绍性文章揭示了整体系统架构,强调了 AWS 基础设施组件与支撑大规模分布式训练和推理的 OSS 工具之间的集成点。
AWS 构建模块
本系列的其余部分将探讨这种分层架构如何在 AWS 上实现,依次涵盖基础设施、资源编排、ML 软件栈和可观测性。以下各节预览每一层。
基础设施:计算、网络和存储
如图 1 所示,基础设施由三个耦合的构建模块锚定——具有大设备内存的加速计算、用于集合通信的宽带宽互连,以及用于数据和检查点的可扩展分布式存储。
加速计算构成了大规模基础模型预训练、后训练和推理的基础。AWS 提供多代 NVIDIA GPU,作为其 Amazon EC2 加速计算实例 的一部分,包括 Amazon EC2 P 实例系列。P5 实例系列 包括配备八块 NVIDIA H100 GPU 的 p5.48xlarge、配备单块 H100 GPU 用于较小规模工作负载的 p5.4xlarge,以及配备 NVIDIA H200 GPU 的 p5e.48xlarge/p5en.48xlarge 变体。P6 实例系列 引入了 NVIDIA Blackwell B200 架构,包括 p6-b200.48xlarge 和配备 Blackwell Ultra B300 的 p6-b300.48xlarge。在这些代际中,主要的扩展轴是峰值 Tensor 吞吐量、HBM 容量和带宽,以及互连带宽(节点内和节点间)。
作为一阶近似,峰值 Tensor Core 吞吐量——以每秒浮点运算次数(FLOPS)衡量——有助于将这些加速器放在一个共同的轴上进行比较。下表总结了每个 GPU 在密集 BF16/FP16 和 FP8 Tensor 操作上的峰值吞吐量,以及 HBM 容量和 HBM 带宽,使用的是与基于 NVSwitch/NVLink 的多 GPU 节点一致的 SXM/HGX 级规格。
| GPU(代表性变体) | BF16/FP16 Tensor 峰值(密集) | FP8 Tensor 峰值(密集) | FP4 Tensor 峰值(密集) | HBM 容量 | HBM 带宽 |
|---|---|---|---|---|---|
| H100 (SXM) | 0.9895 PFLOPS | 1.979 PFLOPS | — | 80 GB HBM3 | 3.35 TB/s |
| H200 (SXM) | 0.9895 PFLOPS | 1.979 PFLOPS | — | 141 GB HBM3e | 4.8 TB/s |
| B200 (HGX, 每 GPU) | 2.25 PFLOPS | 4.5 PFLOPS | 9 PFLOPS | 180 GB HBM3e | 8 TB/s |
| B300 (HGX, 每 GPU) | 2.25 PFLOPS | 4.5 PFLOPS | 13.5 PFLOPS | 288 GB HBM3e | 8 TB/s |
注:NVIDIA 产品表通常报告"带稀疏性"的 Tensor 吞吐量;此表报告密集吞吐量。在适用的情况下,密集吞吐量取稀疏吞吐量的一半,遵循 NVIDIA 对 HGX 级平台的指导(NVIDIA)。DGX 数据为系统级;B200 HBM 容量和带宽值通过将 DGX 总数除以八表示为每 GPU 值(NVIDIA)。
随着模型规模的扩大,步时间通常由集合通信和内存移动主导,而非原始计算吞吐量,这促使需要明确的纵向扩展和横向扩展带宽核算。对于多 GPU 实例,GPU 通信跨越两个领域。内部纵向扩展(NVLink/NVSwitch) 提供节点内高带宽、低延迟的 GPU 到 GPU 连接,使得 all-reduce 和 all-gather 等集合操作无需经过主机网络栈即可执行。外部横向扩展(EFA) 提供跨节点的 OS 旁路网络,AWS 将其用作 Amazon EC2 UltraClusters 的构建模块,其中通信密集型的集合操作跨越数千个实例。下表总结了这些实例类型的关键规格:
| 实例类型 | GPU | GPU 数量 | GPU 内存 | NVLink | NVLink 带宽(聚合) | EFA | EFA 带宽(聚合) |
|---|---|---|---|---|---|---|---|
| p5.4xlarge | H100 | 1 | 80 GB HBM3 | — | — | v2 | 12.5 GB/s |
| p5.48xlarge | H100 | 8 | 640 GB HBM3 | 第 4 代 | 7.2 TB/s | v2 | 400 GB/s |
| p5e.48xlarge | H200 | 8 | 1,128 GB HBM3e | 第 4 代 | 7.2 TB/s | v2 | 400 GB/s |
| p5en.48xlarge | H200 | 8 | 1,128 GB HBM3e | 第 4 代 | 7.2 TB/s | v3 | 400 GB/s |
| p6-b200.48xlarge | B200 | 8 | 1,440 GB HBM3e | 第 5 代 | 14.4 TB/s | v4 | 400 GB/s |
| p6-b300.48xlarge | B300 | 8 | 2,100 GB HBM3e | 第 5 代 | 14.4 TB/s | v4 | 800 GB/s |
注:EFA 带宽从 Gbps 转换为 GB/s(÷8),以与其他带宽指标保持一致;请参阅 EC2 加速计算网络规格。NVLink 和 EFA 带宽数据以每实例聚合值而非每链路值显示;请参阅 P5 实例系列页面 和 P6 实例系列页面 了解相应的节点内互连和网络特性。
Elastic Fabric Adapter (EFA) 是 Amazon EC2 的一种网络接口,它使用 Scalable Reliable Datagram (SRD) 协议提供 OS 旁路远程直接内存访问(RDMA)能力。通过使应用程序能够通过 Libfabric API 直接与网络设备通信——绕过操作系统内核——EFA 减少了延迟,并提高了分布式训练中集合操作的吞吐量。
不同实例系列提供多代 EFA。Amazon EC2 P5 和 P5e 实例配备 EFA 版本 2 (EFAv2)。与 EFAv2 相比,P5en 实例上提供的 EFA 版本 3 (EFAv3) 将数据包延迟降低了约 35%。P6 实例上提供的 EFA 版本 4 (EFAv4) 相对于 EFAv3 进一步提升了 18% 的集合通信性能。
在大规模下,分布式训练(流式传输语料库和写入数 TB 的检查点)和大规模推理(暂存权重和管理 KV 缓存增长)都促使采用分层存储架构——本地 NVMe SSD 用于热数据,Lustre 用于共享高吞吐量访问,以及 Amazon S3 用于持久化存储。
在本系列的主要多 GPU 实例中,本地 NVMe 作为实例存储(临时) 提供,具有 30.72 TB 原始容量(8 × 3.84 TB NVMe SSD);请参阅 EC2 加速计算实例存储规格。
Lustre 是一个开源的、符合 POSIX 标准的分布式文件系统,广泛用于高性能计算(HPC),以提供跨多个客户端具有高聚合吞吐量的共享命名空间。Amazon FSx for Lustre 将 Lustre 作为完全托管服务提供,并将其暴露为一个并行文件系统,能够提供每秒 TB 级的吞吐量、数百万的 IOPS 和亚毫秒级的延迟。数据存储库关联(Data Repository Associations)支持与 Amazon S3 集成,支持训练数据集的延迟加载和检查点的自动导出以实现持久性。
在集群规模下,这些实例部署在 Amazon EC2 UltraClusters 中,该集群将数千个加速实例作为一个紧密放置的集群在一个可用区内进行配置,并使用 PB 级非阻塞网络将它们互连。
图: 第二代 Amazon EC2 UltraClusters(示例 P5 UltraCluster)。
对于每步通信强度高的工作负载(例如,MoE 模型中的专家并行,其中 all-to-all token 分发跨越许多 GPU),NVLink 域的大小可能成为一个一阶约束。作为内部纵向扩展轴的延伸,增加 NVLink 域可以减少性能关键的通信必须离开 NVLink 结构的频率。
Amazon EC2 UltraServers 通过专用的加速器互连连接多个组件实例,将 NVLink 域扩展到单个 EC2 实例之外。AWS 报告称,P6e-GB200 UltraServers 基于 NVIDIA GB200 NVL72 平台构建,在一个 NVLink 域内暴露多达 72 块 Blackwell GPU 和 13.4 TB 的聚合 HBM3e。在更大规模下,EFA 仍然是多 UltraServer 作业的跨节点结构,但增加域内 GPU 数量可以减少性能关键的通信必须离开 NVLink 结构的频率。
这些系统由 NVIDIA Grace–Blackwell 超级芯片构建而成,该芯片通过缓存一致的 NVLink-C2C 将 Grace CPU 内存和 Blackwell GPU HBM 耦合在一起,无需显式的主机-设备拷贝即可实现对 CPU 和 GPU 附加内存的直接访问。在实践中,这可以扩展 GPU 工作负载可用的有效内存(例如,将较冷的模型状态或 KV 缓存放置在 CPU 附加内存中),同时避免 PCIe 级别的拷贝开销,尽管延迟更高且带宽低于本地 HBM。
P6e-GB200 UltraServers 的组件实例类型是 p6e-gb200.36xlarge,它提供四块 GPU 和 Elastic Fabric Adapter (EFA) v4 网络。下表总结了每实例和组合的 UltraServer 配置。
| 实例类型 | GPU | GPU 数量 | GPU 内存 | 内存带宽 | NVLink | NVLink 带宽 | EFA | EFA 带宽 |
|---|---|---|---|---|---|---|---|---|
| p6e-gb200.36xlarge | GB200 NVL72 | 4 | 740 GB HBM3e | — | — | — | v4 | 200 GB/s |
注:p6e-gb200.36xlarge EFA 带宽从已发布的聚合 EFA 网络(4 × 400 Gbps)转换为 GB/s(÷8);请参阅 EC2 加速计算网络规格。
| UltraServer | 组件实例类型 | GPU 数量(NVLink 域) | HBM3e(聚合) | EFA | EFA 带宽 |
|---|---|---|---|---|---|
| u-p6e-gb200x36 | p6e-gb200.36xlarge | 36 | 6.7 TB | v4 | 1,800 GB/s |
| u-p6e-gb200x72 | p6e-gb200.36xlarge | 72 | 13.4 TB | v4 | 3,600 GB/s |
注:UltraServer EFA 带宽从 AWS 报告的每秒兆兆位(Tbps)转换为 GB/s(÷8);请参阅 P6e-GB200 UltraServers 公告 和 P6 实例系列页面。
资源编排:Slurm 和 Kubernetes
当训练跨越数百或数千个加速器时,手动资源管理变得难以处理。例如,一个需要 512 个 GPU 的训练作业必须同时协同调度 64 个八 GPU 节点(P 实例),并在完成或失败时原子地释放资源。Slurm 和 Kubernetes 都通过控制平面架构解决了这一挑战:一个集中式调度器维护集群状态并做出分配决策,而工作节点执行分配的工作负载。
图 2:AWS 上基于 Slurm 和基于 Kubernetes 的资源编排的高层架构
Slurm(Simple Linux Utility for Resource Management)是高性能计算中占主导地位的工作负载管理器,构建于模块化插件架构之上,允许独立配置调度算法、拓扑模型、资源类型和计费后端。其调度模型将资源组织成分区(节点的逻辑分组),通过 sbatch 接受作业提交,并通过 srun 在分配的节点上同步启动并行任务。对于分布式训练至关重要的是,Slurm 在作业级别进行调度——在任何任务启动之前,原子地分配整个多节点作业。回填调度器 在空闲槽位中启动低优先级作业,而不会延迟高优先级作业,而多因素优先级系统权衡公平份额使用、作业年龄和 QOS 层级,以跨租户对队列进行排序。Slurm 还通过模拟网络交换机层次结构的插件支持拓扑感知放置——在 AWS 上,编码 EFA 结构拓扑以将作业共置于交换机跳数最少的节点上——并通过其通用资源(GRES)接口支持原生 GPU 调度,该接口跟踪 GPU 类型并强制设备亲和性。
AWS 为基于 Slurm 的编排提供了多种部署选项。AWS ParallelCluster 是一个开源集群管理工具,可自动在 EC2 上部署 Slurm 集群,处理头节点配置、计算集群扩展以及与共享存储的集成。AWS Parallel Computing Service (PCS) 提供了一种替代方案,提供托管控制平面。专门针对分布式训练工作负载,Amazon SageMaker HyperPod 支持 Slurm 模式,并具有针对大规模训练定制的附加功能,例如连续节点健康监控和作业自动恢复功能。
Kubernetes 采用声明式、API 驱动的方法:用户通过资源清单指定期望状态,控制器将实际状态与期望状态协调一致。虽然 Kubernetes 擅长模型部署,但其原生调度模型在紧密耦合的分布式训练方面暴露了几个差距。Kubernetes 在 Pod 级别进行调度;没有作业级别的原子性,多节点训练作业可能部分启动——某些 rank 正在运行,而其他 rank 仍处于 Pending 状态——浪费 GPU 或导致死锁。Vanilla Kubernetes 还缺乏具有基于优先级的回填的批处理队列语义,以及用于放置通信密集型集合操作的内置网络结构拓扑(NVLink 域、EFA 互连)感知。
几个 Kubernetes 原生项目在不同层面解决了这些差距。Kueue 作为默认调度器之上的准入控制器运行,管理作业级别的 gang 准入、具有层次化公平共享的多租户配额以及基于优先级的抢占——同时将 Pod 放置委托给底层调度器。Volcano 和 NVIDIA KAI Scheduler 采用不同的方法,替换或增强默认调度器,以将 gang 调度直接与拓扑感知的 Pod 放置集成——Volcano 作为通用批处理调度器,KAI Scheduler 具有深入的 NVLink/NVSwitch 感知,用于 GPU 优化的放置。这些层是互补的:Kueue 可以管理准入和配额策略,同时将已准入的作业传递给拓扑感知调度器进行放置。
对于 AWS 上基于 Kubernetes 的编排,Amazon Elastic Kubernetes Service (EKS) 提供托管的 Kubernetes,并通过 NVIDIA device plugin 支持 GPU 调度。Amazon SageMaker HyperPod 也支持 EKS 模式,将 Kubernetes 编排与 HyperPod 的训练特定功能相结合。HyperPod EKS 使用专为大规模基础模型训练设计的功能扩展了 EKS。任务治理提供跨团队的计算分配和策略执行,集成了托管的 Kueue 用于准入控制和 Karpenter 用于即时节点配置。无检查点训练解决了传统基于检查点的容错中固有的恢复延迟问题。无检查点训练不是定期将模型状态序列化到共享存储,而是跨 GPU 维护持续的端到端状态复制。当发生故障时,幸存的节点通过基于 EFA 的通信重建丢失的状态,而不是从 FSx for Lustre 或 S3 读取数 TB 的检查点。弹性训练使作业能够根据资源可用性自动扩展。当额外的加速器可用时(例如,来自已完成的作业或新配置的容量),弹性作业可以扩展以利用它们;当更高优先级的工作负载需要资源时,作业可以收缩,同时保持训练进度。
ML 软件栈
分布式训练和推理涉及多个软件层,这些层必须正确配置和协调。一个有用的模型将运行时栈视为五层,从硬件相邻组件(必须正确运行才能运行任何东西)到框架级抽象(决定程序员生产力和模型吞吐量)排序:硬件启用、加速器运行时和数学库、通信底层、ML 框架,以及分布式训练/推理框架。
图 3:EC2 实例上用于分布式训练和推理的 ML 软件栈
硬件启用:内核驱动
在基础层,Linux 内核驱动提供直接的硬件访问。NVIDIA GPU 驱动暴露计算能力,并支持 GPUDirect RDMA 用于 GPU 和网络适配器之间的直接数据传输。GDRCopy 驱动 (gdrdrv) 支持低延迟的 CPU 发起的 GPU 内存拷贝,NCCL 将其用于小消息传输。EFA 驱动 通过 libfabric API 提供 OS 旁路网络,而 Lustre 客户端 驱动支持对 FSx for Lustre 并行文件系统的 POSIX 访问。
加速器运行时、编译器和内核库
CUDA 平台为 GPU 计算提供编程模型和运行时。针对 CUDA 编译的应用程序可以在 NVIDIA GPU 上启动内核、管理设备内存并协调跨多个设备的执行。当前版本是 CUDA Toolkit 13.x,支持 Blackwell 架构(计算能力 10.x)。
现代训练和推理性能越来越由专门的优化库和自定义内核驱动,而不仅仅是通用的供应商原语。像 FlashAttention 这样的内核将注意力融合为一次内存高效的传递,减少了 HBM 流量并提高了吞吐量。许多团队还编写针对其确切模型调整的形状和精度专用的融合内核(例如,layernorm/residual/activation、量化 GEMM、MoE 分发、KV 缓存操作)。这是通过可编程工具链实现的,例如 Triton(Python GPU 内核编译器)和 NVIDIA 的 CuTe(张量布局和 warp 级 DSL),以及像 CUTLASS 这样的库提供高度优化的 GEMM 和融合构建块。在实践中,这个内核和编译器层通常与 ML 框架一样决定端到端性能。
通信底层:NCCL 和传输插件
多 GPU 训练依赖于高效的集合通信。NVIDIA Collective Communications Library (NCCL) 实现了集合操作——all-reduce、all-gather、reduce-scatter、all-to-all、broadcast 和点对点 send/receive——使用拓扑感知算法,利用 NVLink 进行节点内通信,利用网络传输进行节点间流量。NCCL 动态检测通信拓扑,并根据消息大小和可用带宽选择 ring 或 tree 算法。虽然数据并行和张量并行策略主要依赖 all-reduce 和 all-gather,但具有专家并行的 Mixture-of-Experts (MoE) 模型依赖 all-to-all 集合在 GPU 之间路由 token:一个 dispatch all-to-all 将每个 token 发送到托管其分配专家的 GPU,一个 combine all-to-all 将专家输出返回到原始 GPU(NVIDIA Developer Blog)。由于每个 GPU 都与专家并行组中的每个其他 GPU 交换数据,all-to-all 通信量随专家数量扩展,并且在高专家并行度下可能成为主要瓶颈。
在 AWS 上,NCCL 的节点间通信通过 aws-ofi-nccl 插件启用,该插件将 NCCL 的传输 API 映射到 libfabric 接口。这允许 NCCL 利用 EFA 的 OS 旁路和 Scalable Reliable Datagram (SRD) 协议,而无需更改应用程序。
对于推理工作负载,集合操作并不能捕获所有通信模式。分离式推理架构——将 prefill 和解码阶段分离到不同的 GPU 池——需要高效的点对点数据移动,特别是在实例之间传输 KV 缓存状态。NVIDIA Inference Xfer Library (NIXL) 通过提供统一的 API 来满足这一需求,用于跨内存层级(HBM、DRAM、NVMe、分布式存储)和互连(NVLink、InfiniBand、Ethernet)的点对点传输。NIXL 与 NVIDIA Dynamo 等推理框架集成,并支持包括 UCX 和 GPUDirect Storage 在内的后端。
ML 框架:PyTorch
基础模型开发的两个主要框架是 PyTorch 和 JAX。JAX 通过 XLA 采用 SPMD(单程序多数据)方法,其中相同的程序跨设备执行,并自动进行数据分布和集合降低。本系列专注于 PyTorch,它在开源生态系统中得到更广泛的采用,并构成了下面讨论的分布式训练和推理框架的基础。
PyTorch 提供带有 GPU 加速的张量计算、自动微分和灵活的 eager 执行模型。对于分布式工作负载,PyTorch 的 torch.distributed 模块提供了核心原语:用于集合通信的进程组,以及包括 Distributed Data Parallel (DDP) 和 Fully Sharded Data Parallel (FSDP2) 在内的分布式数据并行抽象。DDP 跨 GPU 复制模型并通过 all-reduce 同步梯度,而 FSDP2 使用 ZeRO 算法中的技术跨工作节点分片参数、梯度和优化器状态,从而能够训练超过单个 GPU 内存容量的模型。
分布式训练和推理框架
顶层包括构建在 PyTorch 之上的框架,为大规模分布式训练和推理提供更高级别的抽象。对于训练,三类框架在复杂性-性能权衡中解决了不同的点。以下是一些示例。
Hugging Face Transformers 提供 Trainer 类,通过 Accelerate 内置支持分布式训练,后者抽象了 DDP、FSDP 和 DeepSpeed。此路径优先考虑易用性和广泛的模型兼容性,使其适用于微调和中等规模的训练,其中配置简单性比最大吞吐量更重要。
NVIDIA Megatron Core 针对大规模下的最大效率,实现了 3D 并行(张量、流水线和专家并行),并通过 Transformer Engine 进行了 FP8 混合精度等优化。NeMo Framework 构建在 Megatron Core 之上,为预训练和微调提供端到端工作流。
对于基于人类反馈的强化学习(RLHF)和相关后训练方法,veRL(Volcano Engine Reinforcement Learning)提供了一个灵活的框架,实现了包括 PPO、GRPO 和 REINFORCE++ 在内的算法。veRL 的 HybridFlow 架构允许在同一作业中混合训练后端(FSDP2、Megatron)和推理引擎(vLLM、SGLang),通过在 actor 和 rollout 组件之间在内存中共享模型权重来避免权重同步开销。
对于推理服务,vLLM 实现了 PagedAttention,将 KV 缓存作为分页虚拟内存进行管理,以减少碎片并实现更高的批处理大小。SGLang 通过 RadixAttention 扩展了这一点,实现了跨请求的自动前缀重用、一个将 CPU 调度与 GPU 计算重叠的零开销批处理调度器,以及一个基于预测缓存命中率路由请求的缓存感知负载均衡器。两个框架都支持张量并行,用于服务超过单 GPU 内存的模型,并且都与 NVIDIA Dynamo 集成,用于分离 prefill 和解码阶段的分离式服务架构。
可观测性
可观测性是调试和操作大规模分布式训练系统的先决条件。当训练作业停滞或吞吐量下降时,从业者需要了解原因是硬件故障、网络拥塞、存储瓶颈还是应用级低效。在本系列讨论的基础设施规模下——数千个 GPU、PB 级互连带宽和 TB 级检查点数据——挑战从简单的监控转变为系统化的遥测数据收集、存储和分析。可观测性涵盖三类遥测数据:基础设施指标(GPU、网络、存储)、工作负载指标(训练吞吐量、队列延迟)以及用于主动故障检测的告警。
核心栈:Prometheus 和 Grafana
Kubernetes 和 HPC 环境中可观测性的事实标准结合了用于指标收集的 Prometheus 和用于可视化和告警的 Grafana。Prometheus 基于拉取模型运行,定期抓取指标导出器暴露的 HTTP 端点。收集的指标存储在时间序列数据库(TSDB)中,并通过 PromQL(一种用于聚合、过滤和告警规则评估的灵活查询语言)进行查询。Grafana 将 Prometheus 作为数据源使用,渲染仪表板并根据 PromQL 表达式触发告警。
对于生产部署,Amazon Managed Service for Prometheus (AMP) 提供一个完全托管的、与 Prometheus 兼容的时间序列数据库,能够每秒摄取数百万个样本,而无需操作员管理存储、复制或高可用性。Amazon Managed Grafana (AMG) 提供一个托管的 Grafana 工作区,与 AMP 原生集成,并通过 IAM Identity Center 提供 AWS 身份验证。这些服务共同消除了运维开销,同时保持与现有 Prometheus 导出器和 Grafana 仪表板的兼容性。
GPU、网络和应用遥测
DCGM-Exporter 以 Prometheus 格式暴露 NVIDIA GPU 指标,包括利用率、内存使用、功耗、温度和硬件健康指标,如 ECC 错误和 XID 事件。对于训练工作负载,SM 活动 (DCGM_FI_PROF_SM_ACTIVE) 通常比基本利用率指标更能准确衡量计算效率。
EFA 暴露驱动级统计信息(字节、数据包、重传、超时),有助于诊断分布式训练中的集合操作瓶颈。aws-ofi-nccl 插件将 NCCL 桥接到 libfabric 接口,操作员可以结合 EFA 计数器和 NCCL 诊断信息 (NCCL_DEBUG=INFO) 来隔离网络层问题。
Amazon FSx for Lustre 暴露客户端侧指标,包括吞吐量和元数据延迟,而应用级指标(训练步时间、每秒 token 数、损失值;推理的 TTFT、token 间延迟)可以通过 Prometheus 客户端库导出。
GPU 健康监控和告警
主动故障检测可防止硬件问题蔓延为长时间的训练中断。典型的工作流监控 DCGM 健康指标,并在错误计数超过阈值时触发告警。ECC 单比特错误(SBE)在数量较少时可能可以容忍,但加速的 SBE 速率通常先于双比特错误(DBE)或其他故障。XID 63(行重映射失败)、XID 64(GPU 脱离总线)和 XID 94/95(受控/非受控错误)通常需要立即更换节点。
GPU Health - Cluster 仪表板(Grafana 仪表板 ID 21645)提供了常见 GPU 错误模式的参考可视化。该仪表板聚合了所有集群节点的 ECC 错误、XID 事件、热违规和行重映射状态,使操作员能够在硬件影响训练作业之前识别故障硬件。
图 4:GPU Health - Cluster 仪表板,显示 GPU 错误模式和实例报告
本系列的第 5 部分将全面介绍可观测性架构,包括指标收集策略、仪表板配置以及用于大规模维护集群健康的告警模式。
结论
从单一的预训练扩展定律到三个互补机制——预训练、后训练和测试时计算——的转变并没有分散基础设施需求;反而强化了它们。所有三种机制都需要紧密耦合的加速器计算、高带宽低延迟网络和可扩展的分布式存储,主要区别在于工作负载配置文件和资源调度模式。
本文揭示了在 AWS 上满足这些需求的四层架构:基础设施构建模块(EC2 P 实例、EFA 网络和分层存储)、资源编排(Slurm 和 Kubernetes 与 SageMaker HyperPod)、ML 软件栈(从内核驱动和 CUDA 到 NCCL 再到 PyTorch)以及可观测性(Prometheus、Grafana 和 GPU 健康监控)。每一层都约束并启用其上层——一个配置错误的驱动或饱和的网络链路可以像次优的并行策略一样有效地瓶颈一个原本调优良好的训练运行。
理解这些集成点是诊断性能瓶颈并在基础模型生命周期中做出明智扩展决策的基础。
作者
Aman Shanbhag 是 NVIDIA MARS MLOps 团队的 AI 性能和基础设施工程师,他帮助研究团队构建可扩展、高性能的 ML 训练和推理系统。他之前曾在 AWS 担任专家解决方案架构师,支持全球客户在 AWS 上进行 ML 训练和推理优化。Aman 拥有莱斯大学的计算机科学、数学和创业学位,专注于 AI 基础设施、性能优化以及分布式训练和推理。
Pavel Belevich 是 Amazon Web Services GenAI ML Frameworks 团队的高级应用科学家。他将自己在分布式训练和大模型推理方面的研究应用于生产规模的实际客户工作负载。在加入 AWS 之前,Pavel 曾在 PyTorch Distributed 团队工作,为 FSDP 和 Pipeline Parallelism 等核心分布式训练技术做出了贡献。在 AWS,他致力于 MoE 通信模式和大规模服务/训练工作流。他还通过关于专家并行和大模型系统的技术深度分享定期分享最佳实践。
Keita Watanabe 是 Amazon Web Services GenAI ML Frameworks 团队的首席解决方案架构师,专门从事 ML 系统性能工程,并支持全球客户在 AWS 上进行 ML 训练和推理优化。他的背景是机器学习研究和开发。在加入 AWS 之前,Keita 在乐天担任研究科学家,开发了一个基于图像的产品搜索系统。Keita 拥有东京大学的科学博士学位。




