AWS基础模型训练与推理的构建模块
本文分析了AWS在基础模型预训练、后训练和推理中的基础设施组件,包括GPU实例、弹性网络适配器(EFA)、Lustre文件系统及UltraCluster/UltraServer架构,并强调开源软件在资源管理和监控中的作用。
文章情报
要点
- 基础模型扩展已从单一预训练扩展到后训练和测试时计算三个尺度。
- AWS提供从H100到B300的多代GPU实例,并配有NVLink和EFA网络。
- Lustre共享存储和S3持久化构成分层存储体系。
- UltraServer通过扩展NVLink域减少跨节点通信开销。
为什么重要
这条新闻值得关注,因为基础模型扩展已从单一预训练扩展到后训练和测试时计算三个尺度。
技术影响
可能影响模型选型、推理成本、产品能力和评测基准。
本文旨在为机器学习工程师和研究人员提供关于在AWS上进行基础模型训练和推理的技术基础。文章主要基于开源软件栈,分析AWS基础设施如何与常见工具链集成,以应对大规模分布式训练和推理中的系统瓶颈和扩展特性。
基础模型的扩展定律已经发生变化。过去主要依赖预训练阶段的算力投入,但现在后训练(如监督微调、基于强化学习的方法)和测试时计算(如长思考、搜索验证、多样本策略)也显著影响性能。AWS的基础设施需要支持这三个阶段的协同需求,包括紧密耦合的加速计算、高带宽低延迟网络以及分布式存储后端。
在基础设施层,AWS提供多种GPU实例。P5实例基于H100 GPU,P5e和P5en基于H200,P6则采用Blackwell B200和B300。这些实例的峰值张量吞吐量、HBM容量和带宽以及互连带宽是主要扩展维度。例如,H100的BF16/FP16张量峰值达到0.9895 PFLOPS,而B300则达到2.25 PFLOPS,并支持FP4精度(13.5 PFLOPS)。设备内存从H100的80 GB HBM3到B300的288 GB HBM3e,带宽也从3.35 TB/s提升至8 TB/s。
多GPU实例利用NVLink进行机内扩展,提供高带宽GPU直连,避免主机网络栈开销。对于跨节点扩展,AWS的弹性网络适配器(EFA)提供内核旁路的RDMA通信,支持可扩展可靠数据报(SRD)协议。EFA已有多个版本:EFA v2(P5实例)、EFA v3(P5en实例,延迟降低35%)和EFA v4(P6实例,集合同性能再提升18%)。
存储方面,AWS采用分层架构:本地NVMe SSD用于热数据,Lustre并行文件系统(通过Amazon FSx for Lustre托管)提供共享高吞吐访问,Amazon S3用于持久化存储。Lustre与S3通过数据仓库关联(Data Repository Associations)实现懒加载数据集和自动检查点导出。
在大规模集群层面,Amazon EC2 UltraClusters将数千个加速实例部署在同一可用区,通过PB级非阻塞网络互连。针对通信密集型工作负载(如MoE模型的专家并行),NVLink域大小成为关键约束。为此,AWS推出UltraServers,通过专用加速互连将多个实例的NVLink域扩展。例如,P6e-GB200 UltraServer基于NVIDIA GB200 NVL72平台,在单个NVLink域内包含最多72个Blackwell GPU和13.4 TB HBM3e。其组件实例p6e-gb200.36xlarge提供4个GPU和200 GB/s EFA带宽,组合后UltraServer总EFA带宽达1,800 GB/s。
文章还强调开源软件栈在集群资源管理(Slurm、Kubernetes)、模型开发框架(PyTorch、JAX)和可观测性(Prometheus、Grafana)中的核心地位。AWS基础设施与这些工具深度集成,为整个基础模型生命周期提供端到端支持。后续系列文章将更详细地探讨各层的具体实现。