AWS 上基础模型训练与推理的构建模块：深度解读

背景

随着基础模型（Foundation Models）的规模不断扩张，其生命周期——包括预训练、后训练（Post-training）和推理（Inference）——正面临着日益复杂的工程挑战。正如 NVIDIA 博客中提到的“AI 的三大缩放定律”所示，这些不同的缩放阶段正在推动基础设施需求趋于收敛：即需要紧密耦合的加速器计算能力、高带宽低延迟的网络连接，以及分布式存储后端。

与此同时，为了在大规模集群中维持系统健康并诊断性能瓶颈，资源编排（Orchestration）的重要性日益凸显，应用层和硬件层的可观测性（Observability）也变得不可或缺。另一个关键趋势是，基础模型的生命周期越来越依赖于一个涵盖模型开发框架、集群资源管理和操作工具的开源软件（OSS）生态系统。

本文旨在为参与基础模型训练与推理的机器学习工程师和研究人员提供技术基础，深入分析 AWS 基础设施（包括多节点加速器计算、高带宽低延迟网络、分布式共享存储及相关托管服务）如何与常见的 OSS 堆栈相互作用。文章重点介绍了系统架构的整体概览，强调了 AWS 基础设施组件与支撑大规模分布式训练和推理的 OSS 工具之间的集成点。

核心内容

开源软件栈的分层架构

在集群层面，资源管理通常由 Slurm 和 Kubernetes 等系统提供；模型开发和分布式训练则普遍在 PyTorch 和 JAX 等框架中实现。在可观测性方面，通常使用 Prometheus 进行指标收集，并使用 Grafana 进行可视化和告警，它们作为操作层位于基础设施和资源管理之上。这种分层架构展示了硬件基础设施如何支持资源编排，进而赋能机器学习框架，而可观测性则贯穿所有层级。

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

AWS 的基础设施由三个紧密耦合的构建模块组成：具有大显存的加速计算、用于集体通信的高带宽互连，以及用于数据和检查点的可扩展分布式存储。

1. 加速计算（Accelerated Compute）

加速计算是大规模基础模型预训练、后训练和推理的基础。AWS 提供了多代 NVIDIA GPU 作为其 Amazon EC2 加速计算实例的一部分：

• P5 实例系列：
  • p5.48xlarge：配备 8 块 NVIDIA H100 GPU。
  • p5.4xlarge：配备单块 H100 GPU，适用于较小规模的工作负载。
  • p5e.48xlarge / p5en.48xlarge：变体版本，配备 NVIDIA H200 GPU。
• P6 实例系列：
  • 引入了 NVIDIA Blackwell B200 架构（p6-b200.48xlarge）以及 Blackwell Ultra B300（p6-b300.48xlarge）。

在这些代际中，主要的缩放轴包括峰值 Tensor 吞吐量、HBM（高带宽内存）容量与带宽，以及互连带宽（节点内和节点间）。

性能指标说明：

• 峰值吞吐量：以每秒浮点运算次数（FLOPS）衡量。需要注意的是，NVIDIA 产品表通常报告“带稀疏性”的 Tensor 吞吐量，而本文表格报告的是密集（Dense）吞吐量。根据 NVIDIA 对 HGX 类平台的指导，密集吞吐量通常被视为稀疏吞吐量的一半。
• DGX 数据：DGX 数据为系统级数据；B200 的 HBM 容量和带宽值是通过将 DGX 总值除以 8 得出的单 GPU 值。

随着模型规模扩大，步长（Step time）往往由集体通信和内存移动而非原始计算吞吐量主导，因此需要明确计算 Scale-up（节点内）和 Scale-out（节点间）的带宽。

2. 网络连接（Network）

多 GPU 实例的 GPU 通信跨越两个阶段：

• 内部 Scale-up（NVLink/NVSwitch）：提供节点内高带宽、低延迟的 GPU 到 GPU 连接，使 all-reduce 和 all-gather 等集体操作无需经过主机网络栈即可执行。
• 外部 Scale-out（EFA）：提供跨节点的 OS 旁路（OS-bypass）网络。AWS 将其作为构建模块用于 Amazon EC2 UltraClusters，其中通信密集的集体操作跨越数千个实例。

Elastic Fabric Adapter (EFA) 演进： EFA 是 Amazon EC2 的网络接口，使用可扩展可靠数据报（SRD）协议提供 OS 旁路远程直接内存访问（RDMA）能力。通过 Libfabric API 允许应用直接与网络设备通信，绕过操作系统内核，从而降低延迟并提高分布式训练中的集体操作吞吐量。

• EFAv2：配备于 Amazon EC2 P5 和 P5e 实例。
• EFAv3：配备于 P5en 实例，相比 EFAv2 将数据包延迟降低了约 35%。
• EFAv4：配备于 P6 实例，相比 EFAv3 在集体通信性能上额外提升了 18%。

注：EFA 带宽已从 Gbps 转换为 GB/s 以与其他带宽指标保持一致。NVLink 和 EFA 带宽数据显示为每实例的聚合值，而非每条链路的值。

3. 分布式存储（Storage）

在大规模场景下，分布式训练（流式传输语料库和写入多 TB 检查点）以及大规模推理（暂存权重和管理 KV 缓存增长）都推动了一种分层存储体系结构：

• 本地 NVMe SSD：用于热数据。在主要的多 GPU 实例中，本地 NVMe 作为实例存储（临时存储）提供，原始容量为 30.72 TB（8 × 3.84 TB NVMe SSD）。
• Lustre：用于共享的高吞吐量访问。Amazon FSx for Lustre 提供 Lustre 作为完全托管服务，作为一个并行文件系统，具备每秒 TB 级的吞吐量、数百万 IOPS 和亚毫秒级延迟。通过数据仓库关联（Data Repository Associations），它可以与 Amazon S3 集成，支持训练数据集的懒加载和检查点的自动导出以确保持久性。
• Amazon S3：用于持久化存储。

在集群规模上，这些实例部署在 Amazon EC2 UltraClusters 中。第二代 UltraClusters（例如 P5 UltraCluster）在一个可用区（Availability Zone）内将数千个加速实例配置为单个紧密放置的集群，并使用 petabit 级无阻塞网络进行互连。

关键要点

• 基础设施收敛：基础模型的预训练、后训练和推理阶段正趋向于相同的基础设施需求：紧密耦合的加速器、高带宽低延迟网络和分布式存储。
• 开源生态依赖：集群管理（Slurm/K8s）、模型开发（PyTorch/JAX）和可观测性（Prometheus/Grafana）构成了支撑大规模 AI 工作负载的核心开源软件栈。
• 计算能力迭代：AWS 提供了从 H100/H200 到 Blackwell B200/B300 的多代 GPU 实例，主要关注点在于峰值 Tensor 吞吐量、HBM 容量/带宽以及互连带宽。
• 通信瓶颈关键性：随着模型变大，集体通信和内存移动成为性能瓶颈，而非单纯的计算能力。因此，节点内（NVLink）和节点间（EFA）的带宽至关重要。
• EFA 性能持续优化：AWS 的 EFA 网络接口经历了从 v2 到 v4 的迭代，显著降低了延迟（v3 降低 35%）并提升了集体通信性能（v4 提升 18%）。
• 分层存储策略：采用“本地 NVMe（热数据） + Lustre（共享高吞吐） + S3（持久化）”的分层存储架构，以平衡速度、容量和成本。
• UltraClusters 架构：通过 Amazon EC2 UltraClusters，AWS 能够在单个可用区内提供数千个实例的紧密耦合集群，并利用