T
traeai
登录
返回首页
AI HOT 精选

在AWS上进行基础模型训练与推理的核心构建模块

9.2Score
在AWS上进行基础模型训练与推理的核心构建模块

TL;DR · AI 摘要

AWS 提供了支持大规模基础模型训练与推理的完整技术栈,涵盖高性能计算、网络、存储及开源软件集成,适配 NVIDIA 提出的“三大缩放定律”。

核心要点

  • NVIDIA 的三大缩放定律包括预训练、后训练(如 SFT 和 RL)和推理时计算,需统一基础设施支持。
  • AWS 使用 Trn1/Inf2 实例、EFA 网络和 FSx for Lustre 存储构建低延迟高带宽分布式系统。
  • 主流开源工具链(PyTorch、Kubernetes、Prometheus/Grafana)在 AWS 上实现高效协同与可观测性。

结构提纲

按章节快速跳转。

  1. 基础模型的扩展已从单纯增加算力演变为涵盖预训练、后训练与推理时计算的三维体系。

  2. NVIDIA 提出的三类缩放路径——模型规模、后训练优化与测试时计算显著影响最终性能。

  3. AWS 通过加速计算实例、EFA 高速网络和 FSx for Lustre 分布式存储支撑大规模训练负载。

  4. PyTorch、JAX、Kubernetes 和 Slurm 等开源框架在 AWS 上实现资源调度与模型开发协同。

  5. Prometheus 采集指标、Grafana 可视化告警,形成贯穿全栈的监控与诊断能力。

  6. AWS 将硬件、网络、存储与开源生态整合,为大模型全生命周期提供一致性架构支持。

思维导图

用一张图看清主题之间的关系。

查看大纲文本(无障碍 / 无 JS 友好)
  • AWS 基础模型构建模块
    • 三大缩放定律
      • 预训练缩放
      • 后训练优化
      • 推理时计算
    • 核心基础设施
      • Trn1/Inf2 实例
      • EFA 网络
      • FSx for Lustre
    • 软件栈集成
      • PyTorch / JAX
      • Kubernetes / Slurm
      • Prometheus / Grafana

金句 / Highlights

值得收藏与分享的关键句。

  • 扩展不再是一条曲线——性能不仅依赖预训练,也通过后训练和推理时计算提升。

    第 1 段

    ⬇︎ 下载 PNG𝕏 分享到 X
  • 三个关键组件是具备大显存的加速计算、高带宽互联和分布式共享存储。

    基础设施章节

    ⬇︎ 下载 PNG𝕏 分享到 X
  • 可观测性贯穿所有层级,使用 Prometheus 收集指标,Grafana 进行可视化与告警。

    可观测性章节

    ⬇︎ 下载 PNG𝕏 分享到 X
  • PyTorch 和 JAX 用于模型开发,Kubernetes 和 Slurm 负责集群资源管理。

    开源生态章节

    ⬇︎ 下载 PNG𝕏 分享到 X
  • 该架构实现了在整个模型生命周期中对大规模加速器资源的高效利用。

    结论部分

    ⬇︎ 下载 PNG𝕏 分享到 X
  • 图1展示了硬件如何支撑编排层,进而支持ML框架,可观测性覆盖各层。

    图1说明

    ⬇︎ 下载 PNG𝕏 分享到 X
#AWS#基础模型#分布式训练#PyTorch#可扩展性
打开原文

标题:在 AWS 上进行基础模型训练与推理的构建模块

来源链接:https://huggingface.co/blog/amazon/foundation-model-building-blocks

发布时间:2026-05-11T23:18:26.540Z

Markdown 内容:

长期以来,基础模型中的“扩展”主要意味着一件事:在预训练上投入更多计算资源,能力随之提升。这种直觉得到了实证研究的支持,例如 Kaplan 等人 (2020),该研究指出,当扩大模型参数量数据集规模训练计算量时,损失会呈现出可预测的幂律趋势。在实践中,这些趋势证明了持续投资大规模加速器硬件以及维持其高效利用所需的分布式基础设施是合理的。但前沿技术已经发展——如今的扩展不再是一条单一曲线。NVIDIA 提出的“从一条到三条扩展定律”框架明确指出,在预训练之外,性能越来越多地通过后训练(例如监督微调(SFT)和基于强化学习(RL)的方法)以及推理时计算(“深度思考”、搜索/验证、多采样策略)来实现扩展。

图片 2: 三重扩展定律图表-1280x720

图示: 改编自 "AI 的三条扩展定律详解"(NVIDIA 博客)。

总体来看,这些不同的扩展方式推动基础模型生命周期——包括预训练、后训练和推理——趋向于统一的基础设施需求:紧密耦合的加速器计算资源、高带宽低延迟网络,以及分布式存储后端。它们也提升了资源管理中编排的重要性,以及应用层和硬件层可观测性的关键作用,以维护集群健康并在大规模下诊断性能问题。

另一个重要趋势是,基础模型生命周期越来越依赖一个开源软件(OSS)生态系统,涵盖模型开发框架、集群资源管理和运维工具。在集群层面,资源管理通常由 SlurmKubernetes 等系统提供。模型开发和分布式训练通常使用 PyTorchJAX 等框架实现。监控与可视化(即可观测性)则常借助 Prometheus 进行指标采集,并结合 Grafana 实现可视化和告警,作为基础设施与资源管理之上的运维层。图 1 展示了这一分层架构,说明硬件基础设施如何支撑资源编排,进而支持机器学习框架,而可观测性则贯穿所有层级。

图片 3: 构建模块简介

_图 1:用于基础模型训练与推理的开源软件栈分层架构_

本文面向参与基础模型训练与推理的机器学习工程师和研究人员,特别关注基于开源软件(OSS)框架构建的工作流。文章分析了 AWS 基础设施——包括多节点加速计算、高带宽低延迟网络、分布式共享存储及相关托管服务——如何在整个基础模型生命周期中与常见的 OSS 软件栈协同工作。主要目标是为理解跨越预训练、后训练和推理各阶段的系统瓶颈及扩展特性提供技术基础。本篇引言性文章概述整体系统架构,重点强调 AWS 基础设施组件与 OSS 工具之间的集成点,这些集成点支撑着大规模分布式训练与推理。

AWS 构建模块

本系列后续内容将探讨上述分层架构如何在 AWS 上实现,依次涵盖基础设施、资源编排、机器学习软件栈和可观测性。以下部分将对每一层进行预览。

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

如图 1 所示,基础设施由三个紧密关联的构建模块组成——具备大容量设备内存的加速计算、用于集合通信的高带宽互连,以及用于数据和检查点的可扩展分布式存储。

加速计算是大规模基础模型预训练、后训练和推理的基础。AWS 作为其 Amazon EC2 加速计算实例 的一部分,提供了多代 NVIDIA GPU,包括 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。在这些不同代际中,主要的扩展维度是峰值张量吞吐量、HBM 容量与带宽,以及节点内和跨节点的互联带宽。

作为一阶近似,以每秒浮点运算次数(FLOPS)衡量的峰值张量核心吞吐量有助于将这些加速器置于一个统一的比较基准上。下表总结了 SXM/HGX 类规格下每个 GPU 在 BF16/FP16 和 FP8 密集型张量操作中的峰值吞吐量,以及 HBM 容量和 HBM 带宽,该规格与基于 NVSwitch/NVLink 的多 GPU 节点一致。

| GPU(代表性变体) | BF16/FP16 张量峰值(密集型) | FP8 张量峰值(密集型) | FP4 张量峰值(密集型) | 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 产品规格表通常报告“启用稀疏性时”的张量吞吐量;本表列出的是密集型吞吐量。如适用,密集型吞吐量按稀疏型吞吐量的一半计算,遵循 NVIDIA 对 HGX 平台的指导说明 (NVIDIA)。DGX 数据为系统级总量;B200 的 HBM 容量和带宽值通过将 DGX 总量除以 8 得出,表示为单个 GPU 的数值 (NVIDIA)。_

随着模型规模扩大,步长时间往往由集体通信和内存移动主导,而非原始计算吞吐量,因此需要明确考虑纵向扩展(scale-up)和横向扩展(scale-out)的带宽。对于多 GPU 实例,GPU 通信涉及两种模式。内部纵向扩展(NVLink/NVSwitch) 提供节点内高带宽、低延迟的 GPU 间连接,使 all-reduce 和 all-gather 等集合操作无需经过主机网络栈即可执行。外部横向扩展(EFA) 提供跨节点的绕过操作系统的网络通信能力,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 的一种网络接口,使用 可扩展可靠数据报(SRD) 协议提供绕过操作系统的远程直接内存访问(RDMA)功能。通过允许应用程序经由 Libfabric API 直接与网络设备通信(绕过操作系统内核),EFA 降低了延迟,并提升了分布式训练中集合操作的吞吐量。

不同代的 EFA 可用于不同的实例系列。Amazon EC2 P5 和 P5e 实例配备了 EFA 第 2 版(EFAv2)。在 P5en 实例上提供的 EFA 第 3 版(EFAv3)相比 EFAv2 将数据包延迟降低了约 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 中,后者在一个可用区内部署数千个加速实例,形成单一且紧密布局的集群,并通过拍比特级无阻塞网络互连。

图 4: ec2-ultraclusters-gen2

图示: 第二代 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 |

_Note: UltraServer EFA 带宽由 AWS 报告的太比特每秒(Tbps)换算为 GB/s(÷8);详见 P6e-GB200 UltraServers 公告P6 实例系列页面。_

资源编排:Slurm 与 Kubernetes

当训练任务跨越数百或数千个加速器时,手动资源管理将变得不可行。例如,一个需要 512 个 GPU 的训练任务必须同时调度 64 个八卡节点(P 实例),并在完成或失败时原子性地释放资源。SlurmKubernetes 都通过控制平面架构来应对这一挑战:一个集中式调度器维护集群状态并做出资源分配决策,而工作节点则执行被分配的工作负载。

图 5: slurm-k8s-highlevel-arch

_图 2:在 AWS 上基于 Slurm 和 Kubernetes 的资源编排高层架构_

Slurm(资源管理的简单 Linux 工具)是高性能计算领域主流的工作负载管理器,其基于模块化插件架构构建,允许独立配置调度算法、拓扑模型、资源类型和计费后端。它的调度模型将资源组织成分区(节点的逻辑分组),通过 sbatch 接收作业提交,并使用 srun 在已分配节点上启动并行任务,实现同步启动。对于分布式训练尤为关键的是,Slurm 以作业级别进行调度——在任何任务启动前原子性地分配整个多节点作业。回填调度器 可在空闲槽位中启动低优先级作业,而不会延迟高优先级作业;多因素优先级系统 则综合考量公平份额使用率、作业等待时间和 QOS 等级,以跨租户排序队列。Slurm 还通过插件支持拓扑感知放置,该插件可建模网络交换机层级结构——在 AWS 上编码 EFA 网络拓扑,使作业尽可能部署在交换跳数最少的节点上——并通过其通用资源(GRES)接口提供原生GPU 调度,用于跟踪 GPU 类型并强制设备亲和性。

AWS 提供了多种基于 Slurm 编排的部署选项。AWS ParallelCluster 是一款开源集群管理工具,可自动在 EC2 上部署 Slurm 集群,处理主节点供应、计算资源池扩展以及与共享存储的集成。AWS 并行计算服务 (PCS) 提供另一种选择,提供托管的控制平面。针对分布式训练工作负载,Amazon SageMaker HyperPod 支持 Slurm 模式,并具备专为大规模训练优化的功能,例如持续节点健康监控作业自动恢复功能

Kubernetes 采用声明式、API 驱动的方法:用户通过资源清单定义期望状态,控制器则负责将实际状态调整至与之匹配。尽管 Kubernetes 在模型部署方面表现出色,但其原生调度模型在紧密耦合的分布式训练场景中存在若干不足。Kubernetes 以 Pod 为单位进行调度;缺乏作业级别的原子性,可能导致多节点训练作业部分启动——某些进程已运行而其他仍处于 Pending 状态——从而浪费 GPU 资源或引发死锁。标准 Kubernetes 也缺少基于优先级回填的批处理队列语义,以及对网络架构拓扑(如 NVLink 域、EFA 互联)的内置感知能力,而这对于通信密集型集合操作的放置至关重要。

多个 Kubernetes 原生项目在不同层面弥补了这些缺陷。Kueue 作为默认调度器之上的准入控制器运行,管理作业级别的组作业准入、支持分层公平共享的多租户配额,以及基于优先级的抢占机制——同时将 Pod 放置交由底层调度器处理。VolcanoNVIDIA KAI Scheduler 采取了不同的方法,通过替换或增强默认调度器,将组作业调度直接与拓扑感知的 Pod 放置相结合——Volcano 作为通用批处理调度器,KAI Scheduler 则深度感知 NVLink/NVSwitch,实现面向 GPU 优化的放置策略。这些层次可以互补:Kueue 可负责准入和配额策略管理,而将已准入的作业交给具备拓扑感知能力的调度器进行具体放置。

对于在 AWS 上基于 Kubernetes 的编排,Amazon Elastic Kubernetes Service (EKS) 通过 NVIDIA 设备插件 提供支持 GPU 调度的托管式 Kubernetes。Amazon SageMaker HyperPod 还支持 EKS 模式,将 Kubernetes 编排与 HyperPod 针对训练任务的专用功能相结合。HyperPod EKS 在 EKS 基础上扩展了专为大规模基础模型训练设计的功能。任务治理 提供跨团队的计算资源分配和策略执行能力,并集成托管的 Kueue 实现准入控制,以及使用 Karpenter 实现按需节点供给。无检查点训练(Checkpointless training) 解决了传统基于检查点的容错机制所固有的恢复延迟问题。不同于周期性地将模型状态序列化到共享存储中,无检查点训练通过 GPU 之间持续的点对点状态复制来维护一致性。当发生故障时,存活节点通过基于 EFA 的通信重建丢失的状态,而不是从 FSx for Lustre 或 S3 中读取数 TB 大小的检查点文件。弹性训练(Elastic training) 使训练任务能够根据可用资源自动伸缩。当有额外加速器资源可用时(例如来自已完成的任务或新配置的容量),弹性任务可以扩展以利用这些资源;而当更高优先级的工作负载需要资源时,任务可在保留训练进度的同时收缩规模。

机器学习软件栈

分布式训练和推理涉及多个必须正确配置和协调的软件层。一个有用的模型将运行时栈分为五个层次,从靠近硬件的组件(这些组件必须正常工作才能运行任何内容)到框架级抽象(决定开发效率和模型吞吐量):硬件支持、加速器运行时与数学库、通信底层、机器学习框架,以及分布式训练/推理框架。

图 6:揭秘机器学习软件栈

_图 3:在 EC2 实例上进行分布式训练和推理的机器学习软件栈_

#### [](https://huggingface.co/blog/amazon/foundation-model-building-blocks#hardware-enablement-kernel-drivers) 硬件支持:内核驱动

最底层是 Linux 内核驱动程序,提供对硬件的直接访问。NVIDIA GPU 驱动暴露计算能力,并支持 GPUDirect RDMA,实现 GPU 与网络适配器之间的直接数据传输。GDRCopy 驱动(gdrdrv)支持 CPU 发起的低延迟 GPU 内存读写操作,被 NCCL 用于小消息传输。EFA 驱动 通过 libfabric API 提供操作系统旁路网络通信,而 Lustre 客户端 驱动则启用对 FSx for Lustre 并行文件系统的 POSIX 访问。

#### [](https://huggingface.co/blog/amazon/foundation-model-building-blocks#accelerator-runtime-compilers-and-kernel-libraries) 加速器运行时、编译器与内核库

CUDA 平台为 GPU 计算提供编程模型和运行时环境。基于 CUDA 编译的应用程序可以在 NVIDIA GPU 上启动内核、管理设备内存,并协调跨多个设备的执行。当前版本为 CUDA Toolkit 13.x,支持 Blackwell 架构(计算能力 10.x)。

现代训练和推理性能越来越依赖于专门优化的库和自定义内核,而不仅仅是通用的厂商原语。例如 FlashAttention 这类内核将注意力操作融合为一次内存高效的处理过程,减少 HBM 流量并提升吞吐量。许多团队还会编写针对特定形状和精度优化的融合内核(如 layernorm/residual/activation、量化 GEMM、MoE 分发、KV-cache 操作),以精确匹配其模型需求。这得益于可编程工具链的支持,例如 Triton(Python GPU 内核编译器)和 NVIDIA 的 CuTe(张量布局与线程束级 DSL),以及像 CUTLASS 这样的库提供了高度优化的 GEMM 和融合构建模块。实际上,这一内核与编译器层对端到端性能的影响往往不亚于机器学习框架本身。

#### [](https://huggingface.co/blog/amazon/foundation-model-building-blocks#communication-substrate-nccl-and-transport-plugins) 通信底层:NCCL 与传输插件

多 GPU 训练依赖于高效的集合通信。NVIDIA 集合通信库(NCCL) 实现了多种集合操作——包括全归约(all-reduce)、全收集(all-gather)、规约散列(reduce-scatter)、全对全(all-to-all)、广播(broadcast)以及点对点发送/接收——采用拓扑感知的算法,利用 NVLink 实现节点内通信,并通过网络传输处理节点间流量。NCCL 动态检测通信拓扑结构,并根据消息大小和可用带宽选择环形或树形算法。虽然数据并行和张量并行策略主要依赖 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 开发者博客)。由于在专家并行组中每个 GPU 都需要与其他所有 GPU 交换数据,all-to-all 的通信量随专家数量增长,在高程度专家并行下可能成为主要瓶颈。

在 AWS 上,NCCL 的跨节点通信通过 aws-ofi-nccl 插件实现,该插件将 NCCL 的传输 API 映射到底层 libfabric 接口。这使得 NCCL 能够利用 EFA 的操作系统旁路和可扩展可靠数据报(SRD)协议,而无需修改应用程序代码。

对于推理工作负载,集合操作无法涵盖所有的通信模式。解耦式推理架构——将预填充(prefill)和解码(decode)阶段分离到不同的 GPU 资源池中——需要高效的点对点数据传输,特别是用于在实例之间传递 KV 缓存状态。NVIDIA 推理传输库(NIXL) 提供了一个统一的 API,支持跨内存层级(HBM、DRAM、NVMe、分布式存储)和互连技术(NVLink、InfiniBand、Ethernet)进行点对点传输,满足了这一需求。NIXL 可与 NVIDIA Dynamo 等推理框架集成,并支持 UCX 和 GPUDirect Storage 等后端。

#### [](https://huggingface.co/blog/amazon/foundation-model-building-blocks#ml-frameworks-pytorch) 机器学习框架:PyTorch

目前基础模型开发的两大主流框架是 PyTorchJAX。JAX 借助 XLA 采用 SPMD(单程序多数据)方式,使同一程序在多个设备上执行,并自动完成数据分发和集合操作的降低。本系列聚焦于 PyTorch,因其在开源生态中应用更广泛,并构成了下文讨论的分布式训练与推理框架的基础。

PyTorch 提供了支持 GPU 加速的张量计算、自动微分以及灵活的即时执行(eager-execution)模型。对于分布式任务,PyTorch 的 torch.distributed 模块提供了核心原语:用于集合通信的进程组,以及分布式数据并行抽象,包括 分布式数据并行(DDP)完全分片数据并行(FSDP2)。DDP 在多个 GPU 上复制模型并通过 all-reduce 同步梯度;而 FSDP2 使用 ZeRO 算法中的技术,在工作节点间对参数、梯度和优化器状态进行分片,从而支持训练超出单个 GPU 显存容量的大模型。

#### [](https://huggingface.co/blog/amazon/foundation-model-building-blocks#distributed-training-and-inference-frameworks) 分布式训练与推理框架

顶层是由基于 PyTorch 构建的高级抽象框架组成,用于大规模分布式训练与推理。在训练方面,三类框架分别针对复杂性与性能之间的不同权衡点。以下是几个示例:

Hugging Face Transformers 提供了 Trainer 类,通过 Accelerate 内置支持分布式训练,后者对 DDP、FSDP 和 DeepSpeed 进行了抽象。该路径强调易用性和广泛的模型兼容性,适用于微调和中等规模训练场景,尤其适合配置简洁性比最大吞吐量更重要的情况。

NVIDIA Megatron Core 致力于在大规模下实现最高效率,实现了三维并行(张量、流水线和专家并行),并包含通过 Transformer Engine 支持 FP8 混合精度等优化功能。NeMo 框架 在 Megatron Core 基础上构建,提供预训练和微调的端到端工作流。

对于基于人类反馈的强化学习(RLHF)及相关后训练方法,veRL(火山引擎强化学习)提供了一个灵活的框架,实现了 PPO、GRPO 和 REINFORCE++ 等算法。veRL 的 HybridFlow 架构允许在同一任务中混合使用不同的训练后端(如 FSDP2、Megatron)和推理引擎(如 vLLM、SGLang),并通过在 actor 和 rollout 组件之间共享内存中的模型权重,避免了权重同步开销。

对于推理服务,vLLM 实现了 PagedAttention,将 KV 缓存作为分页虚拟内存进行管理,以减少内存碎片并支持更大的批处理规模。SGLang 在此基础上进一步扩展,引入了 RadixAttention 以实现跨请求的自动前缀复用、零开销批调度器(可重叠 CPU 调度与 GPU 计算)以及基于预测缓存命中率路由请求的缓存感知负载均衡器。两个框架均支持张量并行,用于服务超出单个 GPU 内存容量的模型,并且都集成了 NVIDIA Dynamo,以支持将预填充和解码阶段分离的 disaggregated 服务架构。

可观测性

可观测性是大规模调试和运行分布式训练系统的前提。当训练任务停滞或吞吐量下降时,实践者需要能够判断问题根源是硬件故障、网络拥塞、存储瓶颈还是应用层效率低下。在本系列讨论的基础架构规模下——数千个 GPU、PB 级互联带宽和 TB 级检查点数据——挑战已从简单的监控转变为系统化的遥测数据采集、存储与分析。可观测性涵盖三类遥测数据:基础设施指标(GPU、网络、存储)、工作负载指标(训练吞吐量、队列延迟)以及用于主动故障检测的告警机制。

#### [](https://huggingface.co/blog/amazon/foundation-model-building-blocks#core-stack-prometheus-and-grafana) 核心栈: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 仪表板的兼容性。

#### [](https://huggingface.co/blog/amazon/foundation-model-building-blocks#gpu-network-and-application-telemetry) 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 数、损失值;推理阶段的首 token 时间 TTFT、token 间延迟)可通过 Prometheus 客户端库导出。

#### [](https://huggingface.co/blog/amazon/foundation-model-building-blocks#gpu-health-monitoring-and-alerting) GPU 健康监控与告警

主动故障检测可防止硬件问题演变为长时间的训练中断。典型的工作流程是监控 DCGM 健康指标,并在错误计数超过阈值时触发告警。少量 ECC 单比特错误(SBE)可能是可容忍的,但 SBE 速率快速上升通常预示着双比特错误(DBE)或其他故障的发生。XID 63(行重映射失败)、XID 64(GPU 脱离总线)和 XID 94/95(受控/非受控错误)通常意味着需要立即更换节点。

Grafana 仪表板 ID 21645:GPU Health - Cluster dashboard 提供了常见 GPU 错误模式的参考可视化视图。该仪表板聚合了集群所有节点上的 ECC 错误、XID 事件、热违规和行重映射状态,使运维人员能够在硬件故障影响训练任务之前识别出问题设备。

![图片 7](https://huggingface.co/blog/amazon/figs/gpu-health.png)_图 4:GPU Health - Cluster 仪表板展示 GPU 错误模式及实例上报情况_

本系列第五部分将全面介绍可观测性架构,包括指标采集策略、仪表板配置以及用于维护大规模集群健康的告警模式。

结论

从单一的预训练扩展法则转向三个互补的计算阶段——预训练、后训练和测试时计算——并未分散基础设施需求,反而强化了这些需求。这三个阶段均需要紧密耦合的加速器计算能力、高带宽低延迟的网络以及可扩展的分布式存储,主要区别在于工作负载特征和资源调度模式。

本文揭示了 AWS 上满足这些需求的四层架构:基础设施构建模块(EC2 P系列实例、EFA 网络和分层存储)、资源编排(Slurm 和 Kubernetes 配合 SageMaker HyperPod)、机器学习软件栈(从内核驱动和 CUDA 到 NCCL 再到 PyTorch),以及可观测性(Prometheus、Grafana 和 GPU 健康监控)。每一层都会对上层形成约束并提供支持——一个配置错误的驱动程序或饱和的网络链路,其造成瓶颈的效果与次优的并行策略一样显著。

理解这些集成点是诊断性能瓶颈并在基础模型生命周期中做出明智扩展决策的基础。

作者

[Aman Shanbhag](https://www.linkedin.com/in/aman-shanbhag/) 是 NVIDIA MARS MLOps 团队的一名 AI 性能与基础设施工程师,致力于帮助研究团队构建可扩展、高性能的机器学习训练和推理系统。此前,他曾在 AWS 担任专业解决方案架构师,为全球客户在 AWS 上优化机器学习训练和推理提供支持。Aman 毕业于莱斯大学,拥有计算机科学、数学和创业学学位,专注于 AI 基础设施、性能优化以及分布式训练与推理。

[Pavel Belevich](https://www.linkedin.com/in/pbelevich/) 是亚马逊云服务(Amazon Web Services)GenAI 机器学习框架团队的高级应用科学家。他将自己在分布式训练和大模型推理方面的研究成果应用于实际生产规模的客户工作负载。加入 AWS 之前,Pavel 曾在 PyTorch 分布式团队工作,为 FSDP 和流水线并行等核心分布式训练技术做出了贡献。在 AWS,他专注于专家混合模型(MoE)的通信模式以及大规模服务/训练工作流。他还通过深入的技术分享,定期传播关于专家并行和大模型系统的最佳实践。

[Keita Watanabe](https://www.linkedin.com/in/keitawatanabe/) 是亚马逊云服务(Amazon Web Services)GenAI 机器学习框架团队的首席解决方案架构师,专注于机器学习系统性能工程,并为全球客户在 AWS 上的机器学习训练和推理优化提供支持。他的背景是机器学习的研究与开发。加入 AWS 之前,Keita 在乐天(Rakuten)担任研究科学家,开发基于图像的商品搜索系统。Keita 拥有东京大学理学博士学位。

AI 可能会生成不准确的信息,请核实重要内容