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

本文基于 Hugging Face Blog 发布的文章《Building Blocks for Foundation Model Training and Inference on AWS》，深入剖析了在 AWS 基础设施上运行大型基础模型（Foundation Models）全生命周期（预训练、后训练、推理）所需的底层技术架构。文章重点阐述了开源软件栈（OSS）与 AWS 硬件基础设施之间的交互关系，特别是计算、网络和存储这三个核心构建模块如何协同工作以应对大规模分布式 AI 工作负载的挑战。

背景

随着基础模型规模的指数级增长，其生命周期中的预训练、后训练和推理阶段正趋向于收敛的基础设施需求。这种趋势主要由三大扩展定律驱动，它们共同推动了对以下基础设施要素的极高要求：

1. 紧密耦合的加速器计算：提供极高的峰值吞吐量。
2. 高带宽、低延迟的网络：确保节点间通信效率。
3. 分布式存储后端：支持海量数据读写和检查点持久化。

与此同时，资源管理的编排能力以及应用层和硬件层的可观测性（Observability）变得至关重要，以维持集群健康并在大规模环境下诊断性能瓶颈。

另一个关键趋势是基础模型生命周期对开源软件（OSS）生态系统的日益依赖。这一生态系统涵盖了模型开发框架、集群资源管理和操作工具。在集群层面，资源管理通常由 Slurm 和 Kubernetes 等系统提供；模型开发和分布式训练常基于 PyTorch 和 JAX 等框架实现；而监控和可视化（即可观测性）则常使用 Prometheus 进行指标收集，Grafana 进行可视化和警报，作为位于基础设施和资源管理层之上的操作层。

本文旨在为参与基础模型训练和推理的机器学习工程师和研究人员提供技术基础，重点分析 AWS 基础设施（包括多节点加速器计算、高带宽低延迟网络、分布式共享存储及相关托管服务）如何与常见的 OSS 栈交互，从而理解跨越预训练、后训练和推理的系统瓶颈和扩展特性。

核心内容

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

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

1. 加速计算：AWS EC2 实例家族

加速计算是大规模基础模型预训练、后训练和推理的基础。AWS 提供多代 NVIDIA GPU 作为 Amazon EC2 加速计算实例的一部分，主要包括 P 系列和最新的 P6 系列：

• 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）吞吐量。对于 HGX 类平台，密集吞吐量通常被视为稀疏吞吐量的一半。
• 扩展轴：随着模型规模扩大，步长（step time）往往由集体通信和内存移动主导，而非原始计算吞吐量。因此，显式的 Scale-up（节点内）和 Scale-out（节点间）带宽核算变得至关重要。

2. 网络互连：从 NVLink 到 EFA

多 GPU 实例的通信跨越两个主要机制：

• 内部 Scale-up（NVLink/NVSwitch）：
  • 提供节点内高带宽、低延迟的 GPU 间连接。
  • 使 All-Reduce 和 All-Gather 等集体操作无需穿越主机网络栈即可执行。
• 外部 Scale-out（EFA, Elastic Fabric Adapter）：
  • 提供跨节点的 OS-Bypass（绕过操作系统内核）网络。
  • AWS 将其作为 Amazon EC2 UltraClusters 的构建模块，支持数千个实例间的高通信强度集体操作。

EFA 版本演进：

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

EFA 使用可扩展可靠数据报（SRD）协议提供远程直接内存访问（RDMA）能力。通过 Libfabric API，应用程序可以直接与网络设备通信，绕过操作系统内核，从而降低延迟并提高分布式训练中的吞吐量。

3. 存储层级：NVMe、Lustre 与 S3

在大规模场景下，分布式训练（流式处理语料库、写入多 TB 检查点）和大规模推理（暂存权重、管理 KV Cache 增长）需要一个分层存储体系：

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

Amazon EC2 UltraClusters： 在集群规模上，这些实例部署在 Amazon EC2 UltraClusters 中。UltraClusters 在一个可用区（Availability Zone）内配置数千个紧密放置的加速实例，并使用 PB 级无阻塞网络进行互连，以支持极高通信强度的工作负载。

关键要点

• 基础设施收敛：基础模型的生命周期（预训练、后训练、推理）正趋向于对紧密耦合的加速器计算、高带宽低延迟网络和分布式存储的共同需求。
• 开源软件栈的核心地位：Kubernetes/Slurm 用于资源管理，PyTorch/JAX 用于模型开发，Prometheus/Grafana 用于可观测性，构成了现代 AI 基础设施的操作层。
• 计算性能的关键指标：除了峰值 FLOPS，HBM 容量/带宽以及互连带宽（NVLink/EFA）是决定扩展性能的关键因素。
• 通信瓶颈的缓解：
  • 节点内通信依赖 NVLink/NVSwitch 实现低延迟集体操作。
  • 节点间通信依赖 EFA（Elastic Fabric Adapter）实现 OS-Bypass RDMA，EFAv3 和 EFAv4 分别显著降低了延迟并提升了集体通信性能。
• 分层存储架构：
  • 本地 NVMe 用于热数据（如临时缓存）。
  • FSx for Lustre 提供高吞吐的共享命名空间，并与 S3 集成以实现数据持久化和懒加载。
  • S3 作为最终的持久化存储层。
• UltraClusters 的价值：