AWS 和 Hugging Face 发布基础模型工程工具包

图：改编自 NVIDIA 博客《AI 的三大扩展定律详解》。

综合来看，这些扩展范式推动基础模型生命周期——预训练、后训练和推理——逐渐形成趋同的基础设施需求：紧密耦合的加速器计算、高带宽低延迟网络，以及分布式存储后端。同时，它们也提升了编排对于资源管理的重要性，以及应用层和硬件层可观测性的重要性，后者用于维护集群健康状态，并在大规模环境中诊断性能异常。

另一个关键趋势是，基础模型生命周期日益依赖覆盖模型开发框架、集群资源管理和运维工具的开源软件（OSS）生态系统。在集群层，资源管理通常由 Slurm 和 Kubernetes 等系统提供。模型开发和分布式训练通常使用 PyTorch、JAX 等框架实现。监控和可视化——即可观测性——通常使用 Prometheus 采集指标，并使用 Grafana 进行可视化和告警；它们作为运维层，位于基础设施和资源管理层之上。图 1 展示了这种分层架构：硬件基础设施为资源编排提供支撑，资源编排进而支持 ML 框架，而可观测性则贯穿所有层级。

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

本文面向参与基础模型训练和推理的机器学习工程师与研究人员，尤其关注构建在 OSS 框架之上的工作流。本文分析了 AWS 基础设施——包括多节点加速器计算、高带宽低延迟网络、分布式共享存储及相关托管服务——如何在基础模型生命周期的各个阶段与常见 OSS 技术栈协同工作。其主要目标是提供必要的技术基础，帮助读者理解贯穿预训练、后训练和推理阶段的系统瓶颈与扩展特性。作为本系列的导论，本文将呈现整体系统架构，重点介绍 AWS 基础设施组件与支撑大规模分布式训练和推理的 OSS 工具之间的集成点。

AWS 构建模块

本系列接下来的内容将分析这种分层架构如何在 AWS 上落地，并依次介绍基础设施、资源编排、ML 软件栈和可观测性。以下各节将预览每一层。

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

如图 1 所示，基础设施以三个相互耦合的构建模块为核心：拥有大容量设备内存的加速计算、用于集合通信的宽带宽互连，以及用于数据和检查点的可扩展分布式存储。

加速计算构成了大规模基础模型预训练、后训练和推理的基础。AWS 通过 Amazon EC2 加速计算实例（包括 Amazon EC2 P 实例系列）提供多代 NVIDIA GPU。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 容量与带宽，以及节点内部和节点之间的互连带宽。

作为一阶近似，以每秒浮点运算次数（FLOPS）衡量的 Tensor Core 峰值吞吐量，有助于在统一维度上比较这些加速器。下表汇总了单块 GPU 在稠密 BF16/FP16 和 FP8 Tensor 运算下的峰值吞吐量，以及 HBM 容量和 HBM 带宽；其中使用的 SXM/HGX 级规格与基于 NVSwitch/NVLink 的多 GPU 节点相匹配。

注：NVIDIA 产品表通常会报告“启用稀疏性”时的 Tensor 吞吐量；本表报告的是稠密吞吐量。在适用情况下，依据 NVIDIA 针对 HGX 级平台的说明，稠密吞吐量按稀疏吞吐量的一半计算（NVIDIA）。DGX 数据为系统级数据；B200 的 HBM 容量和带宽值通过将 DGX 总量除以八，换算为单块 GPU 的数值（NVIDIA）。

随着模型规模扩大，单步耗时往往主要受集合通信和内存数据移动影响，而非原始计算吞吐量，因此需要明确核算纵向扩展和横向扩展的带宽。对于多 GPU 实例，GPU 通信分为两种模式。内部纵向扩展（NVLink/NVSwitch）在节点内提供高带宽、低延迟的 GPU 间连接，使 all-reduce、all-gather 等集合通信操作无须经过主机网络栈即可执行。外部横向扩展（EFA）则在节点之间提供绕过操作系统的网络通信能力；AWS 将其作为 Amazon EC2 UltraClusters 的构建模块，使通信密集型集合操作能够跨越数千个实例。下表汇总了这些实例类型的关键规格：

注：为与其他带宽指标保持一致，EFA 带宽通过除以 8 从 Gbps 转换为 GB/s；可参阅 EC2 加速计算网络规格。NVLink 和 EFA 带宽数据均以每个实例的聚合值表示，而非每条链路的数值；有关相应的节点内互连和网络特性，请参阅 P5 实例系列页面和 P6 实例系列页面。

Elastic Fabric Adapter（EFA）是 Amazon EC2 的一种网络接口，它使用 Scalable Reliable Datagram（SRD）协议提供绕过操作系统的远程直接内存访问（RDMA）能力。EFA 允许应用通过 Libfabric API 直接与网络设备通信并绕过操作系统内核，从而降低延迟，并提升分布式训练中集合操作的吞吐量。

不同实例系列提供了多代 EFA。Amazon EC2 P5 和 P5e 实例配备 EFA version 2（EFAv2）。P5en 实例提供的 EFA version 3（EFAv3）相比 EFAv2 可将数据包延迟降低约 35%。P6 实例提供的 EFA version 4（EFAv4）在集合通信性能方面，相较 EFAv3 可进一步提升 18%。

在大规模环境中，无论是分布式训练（流式读取语料库并写入数 TB 的检查点），还是大规模推理（暂存权重并管理不断增长的 KV cache），都需要采用分层存储体系：使用本地 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 中。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，带宽也低于本地 HBM。

P6e-GB200 UltraServers 的组件实例类型为 p6e-gb200.36xlarge，可提供四个 GPU 和 Elastic Fabric Adapter（EFA）v4 网络。下表汇总了单实例配置和组合后的 UltraServer 配置。

注意：p6e-gb200.36xlarge 的 EFA 带宽由官方公布的 EFA 聚合网络带宽（4 × 400 Gbps）换算为 GB/s（÷8）；请参阅 EC2 加速计算网络规范。

注意：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 已经运行，而另一些仍处于 Pending 状态——从而浪费 GPU 或导致死锁。原生 Kubernetes 还缺少支持基于优先级进行回填的批处理队列语义，以及对网络结构拓扑（NVLink 域、EFA 互连）的内置感知能力，因而难以针对通信密集型集合操作进行合理放置。

多个 Kubernetes 原生项目从不同层面解决了这些不足。Kueue 作为准入控制器运行在默认调度器之上，负责管理作业级组调度准入、采用分层公平共享机制的多租户配额，以及基于优先级的抢占，同时将 pod 放置委托给底层调度器。Volcano 和 NVIDIA KAI Scheduler 采用了不同的方法，通过替代或增强默认调度器，将组调度直接与拓扑感知的 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 的检查点。弹性训练使作业能够根据资源可用性自动扩缩容。当有更多加速器可用时（例如其他作业完成后释放了资源，或新预置的容量已经就绪），弹性作业可以扩展并使用这些资源；当优先级更高的工作负载需要资源时，作业可以收缩，同时保持训练进度。

分布式训练和推理涉及多个必须正确配置并协调运行的软件层。一个实用的模型是将运行时栈划分为五层，按照从硬件邻近组件（它们必须正常工作，其他一切才能运行）到框架级抽象（它们决定程序员生产力和模型吞吐量）的顺序排列：硬件支持、加速器运行时与数学库、通信底层、机器学习框架，以及分布式训练/推理框架。

图 3：EC2 实例上用于分布式训练和推理的机器学习软件栈

在最底层，Linux 内核驱动程序提供对硬件的直接访问。NVIDIA GPU 驱动程序公开计算能力，并支持 GPUDirect RDMA，使 GPU 与网络适配器之间能够直接传输数据。GDRCopy 驱动程序（gdrdrv）支持由 CPU 发起、往返 GPU 内存的低延迟复制，NCCL 使用它来传输小型消息。EFA 驱动程序通过 libfabric API 提供绕过操作系统的网络通信，Lustre 客户端驱动程序则支持通过 POSIX 访问 FSx for Lustre 并行文件系统。

CUDA 平台提供 GPU 计算所需的编程模型和运行时。针对 CUDA 编译的应用程序可以在 NVIDIA GPU 上启动内核、管理设备内存，并协调多个设备之间的执行。当前版本为 CUDA Toolkit 13.x，支持 Blackwell 架构（计算能力 10.x）。

现代训练和推理的性能越来越多地由专用优化库和自定义内核驱动，而不仅仅依赖通用的厂商原语。FlashAttention 等内核将注意力计算融合为一次内存高效的处理过程，从而减少 HBM 流量并提高吞吐量。许多团队还会编写针对特定形状和精度进行专门优化的融合内核（例如 layernorm/residual/activation、量化 GEMM、MoE dispatch、KV-cache ops），以适配其具体模型。这得益于 Triton（Python GPU 内核编译器）和 NVIDIA CuTe（张量布局及 warp 级 DSL）等可编程工具链，而 CUTLASS 等库则提供经过高度优化的 GEMM 和融合构建块。在实践中，内核和编译器层对端到端性能的影响往往与机器学习框架本身同样显著。

多 GPU 训练依赖高效的集合通信。NVIDIA Collective Communications Library（NCCL）通过可感知拓扑的算法实现了集合操作，包括全归约（all-reduce）、全收集（all-gather）、归约散射（reduce-scatter）、全交换（all-to-all）、广播（broadcast）以及点对点发送/接收；这些算法利用 NVLink 进行节点内通信，并使用网络传输机制处理节点间流量。NCCL 会动态检测通信拓扑，并根据消息大小和可用带宽选择环形或树形算法。数据并行和张量并行策略主要依赖全归约与全收集，而采用专家并行的混合专家（Mixture-of-Experts，MoE）模型则依赖全交换集合操作在 GPU 之间路由 token：分发全交换（dispatch all-to-all）将每个 token 发送到托管其分配专家的 GPU，合并全交换（combine all-to-all）再将专家输出返回给原始 GPU（NVIDIA Developer Blog）。由于专家并行组中的每个 GPU 都会与其他所有 GPU 交换数据，全交换通信量会随专家数量增加而扩展，并且在专家并行度较高时可能成为主要瓶颈。

在 AWS 上，NCCL 的节点间通信通过 aws-ofi-nccl 插件实现，该插件将 NCCL 的传输 API 映射到 libfabric 接口。这样，NCCL 无需更改应用程序即可利用 EFA 的操作系统旁路能力和可扩展可靠数据报（Scalable Reliable Datagram，SRD）协议。

对于推理工作负载，集合操作并不能涵盖所有通信模式。解耦式推理架构会将预填充（prefill）和解码（decode）阶段分配到不同的 GPU 池，因此需要高效的点对点数据传输，尤其是在实例之间传输 KV 缓存状态。NVIDIA Inference Xfer Library（NIXL）通过提供统一 API 来满足这一需求，可在不同内存层级（HBM、DRAM、NVMe、分布式存储）和互连方式（NVLink、InfiniBand、Ethernet）之间执行点对点传输。NIXL 可与 NVIDIA Dynamo 等推理框架集成，并支持包括 UCX 和 GPUDirect Storage 在内的后端。

基础模型开发领域的两大主流框架是 PyTorch 和 JAX。JAX 通过 XLA 采用 SPMD（Single Program Multiple Data，单程序多数据）方式，在多个设备上执行同一个程序，并自动完成数据分发和集合操作降级。本博客重点介绍 PyTorch，因为它在开源生态系统中的应用更为广泛，并且是下文讨论的分布式训练和推理框架的基础。

PyTorch 提供支持 GPU 加速的张量计算、自动微分以及灵活的即时执行模型。对于分布式工作负载，PyTorch 的 torch.distributed 模块提供核心原语：用于集合通信的进程组，以及包括 Distributed Data Parallel（DDP）和 Fully Sharded Data Parallel（FSDP2）在内的分布式数据并行抽象。DDP 在多个 GPU 上复制模型，并通过全归约同步梯度；FSDP2 则使用源自 ZeRO 算法的技术，在多个工作进程之间对参数、梯度和优化器状态进行分片，从而能够训练超出单个 GPU 显存容量的模型。

最上层由构建在 PyTorch 之上的框架组成，这些框架为大规模分布式训练和推理提供更高级别的抽象。在训练方面，有三类框架分别面向复杂度与性能权衡中的不同位置。以下是一些示例。

Hugging Face Transformers 提供 Trainer 类，通过 Accelerate 内置支持分布式训练；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 集成，以支持将预填充阶段与解码阶段分离的解耦式服务架构。

可观测性是大规模分布式训练系统调试和运维的先决条件。当训练作业停滞或吞吐量下降时，实践者需要判断原因究竟是硬件故障、网络拥塞、存储瓶颈，还是应用层效率问题。在本系列所讨论的基础设施规模下——数千个 GPU、数拍比特的互连带宽以及数 TB 的检查点数据——挑战会从简单监控转变为系统化的遥测数据收集、存储和分析。可观测性涵盖三类遥测数据：基础设施指标（GPU、网络、存储）、工作负载指标（训练吞吐量、队列延迟），以及用于主动检测故障的告警。

在 Kubernetes 和 HPC 环境中，可观测性的事实标准是结合使用 Prometheus 进行指标收集，并使用 Grafana 进行可视化和告警。Prometheus 采用拉取模型，定期抓取指标导出器暴露的 HTTP 端点。收集到的指标存储在时间序列数据库（TSDB）中，并通过 PromQL 进行查询；PromQL 是一种灵活的查询语言，可用于聚合、筛选以及评估告警规则。Grafana 将 Prometheus 作为数据源，根据 PromQL 表达式呈现仪表板并触发告警。

对于生产部署，Amazon Managed Service for Promet