哈佛-CS249r-机器学习系统第二卷-二-

哈佛 CS249r:机器学习系统第二卷(二)

原文:Machine-Learning-Systems-Vol2

译者:飞龙

协议:CC BY-NC-SA 4.0

ML 中的 Incast 问题

拥塞控制和自适应路由管理穿越结构内部的流量。另一种故障模式发生在边缘,那里的流量模式本身会淹没单个端口,无论结构内部提供多少容量。

Incast 是一种多对一 ML 集群流量模式,大量发送方同时向单个接收端口传输数据,将来自多个源的线速流量集中到单个交换机队列,导致缓冲区溢出,即使结构其余部分未拥塞。

  1. 重要性:在 AllReduce 的 Reduce 阶段,每个参与的 GPU 同时向同一聚合点发送梯度。以 256 个发送方每个 50 GB/s 瞄准一个交换机端口为例,瞬时 Incast 达到 12.8 TB/s,比 400 Gb/s 端口的吸收能力高出数个数量级。32 MB 交换机缓冲区可容纳约 640 μs 的单条 400 Gb/s 出站流,但在 256 倍 Incast 下,扣除单个 50 GB/s 出站端口后,净填充率约为 12.75 TB/s,因此约 2.5 μs 后发生溢出。这种突发溢出会触发 PFC 级联或丢包,从而使全集群的 L[lat] 升高。

  2. 区别:与通用拥塞(发生在聚合流量超出链路容量的共享内部链路上)不同,Incast 是端点瓶颈:即使每条内部脊叶链路完全未拥塞,它也会发生,因为瓶颈是单个目标端口,而非结构内部。

  3. 常见误区:一个常见的误解是 Incast 罕见或不可预测。因为 ML 训练在每次反向传播时同步所有 GPU,Incast 是每层梯度计算结束时的确定性事件——每个训练步骤发生数百次,必须通过分层交错 AllReduce 或基于树的集合算法在架构层面缓解,而不应视为边缘情况。

由于同步通信模式,ML 训练在结构上易受 Incast 影响。当一层完成反向计算时,成千上万个节点同时发起 AllReduce,瞄准相同的交换机端口。在我们跨 1,024 个 GPU 的 175B 模型训练中,每个 AllReduce 涉及每个节点同时注入数据,产生的突发可能瞬间超出特定交换机端口的结构容量。生产集群通过三种互补技术缓解此问题:

  • 分层交错:一旦反向传播产生梯度桶,立即触发其 AllReduce。最后一层的梯度在归约时,早期层仍在计算,从而使通信交错进行,而不是作为单个突发释放。

  • 算法选择:选择限制任意单端口扇入的集合操作。分层和 Rail-local 归约分阶段聚合梯度,使没有交换机端口必须一次性吸收所有发送方,这与每个节点瞄准单一聚合点的扁平归约形成对比。带宽最优的 Ring AllReduce 本身是点对点的,不会产生单一聚合热点;Incast 压力来自许多此类流在超预订端口汇聚,拓扑感知放置和 Rail 优化布线可缓解此问题。

  • 服务质量(QoS)分类:将梯度流量标记为最高服务类别,以确保后台存储或管理流量不会延迟同步路径。

这些缓解措施共同在时间、空间和优先级上分散了突发流量,但它们并不能消除拥塞或尾延迟。问题变为如何构建能够在现实条件下仍能正常运行的端到端集群。

第五级:集群设计与案例研究

堆栈的最终级别是集群设计,其中电线、传输、拓扑和拥塞控制结合成一个连贯的系统。集群设计的目标是提供一个端到端的梯度总线(Gradient Bus),使成千上万的分布式 GPU 感受到就像一台单机器。在此级别,抽象层会坍塌为具体的工程决策:应购买哪种电缆,如何布线机架,应部署哪种协议,以及如何验证所构建的织构是否真正能够提供训练作业所期望的带宽。两个代表性的大规模架构 — 一个基于 InfiniBand,一个基于以太网 — 通过说明定义机器学习基础设施的权衡,阐明了这种权衡(NVIDIA 2023; Meta Engineering 2024; Gangidi et al. 2024)。

GPU 到 GPU 的“梯度总线”

在设计良好的集群中,网络织构充当系统总线的感知调度器的扩展,而非被动管道。节点内(NVLink)带宽大约是节点间(InfiniBand/RoCE)带宽的 9 倍,而隐藏这一鸿沟的主要机制是通信-计算重叠。在反向传播过程中,最后一层的梯度最先被计算。系统不会等待整个反向传播完成,而是在 GPU 继续计算前一层梯度的同时,立即触发这些梯度的异步 AllReduce。如果反向传播需要 500 毫秒的计算时间,而 AllReduce 需要 300 毫秒,则完美重叠可以将整个通信成本隐藏在计算之后,从而使实际开销降低至 max (0, 300 − 500) = 0。在实际中,依赖链和资源竞争会将此效率限制在 60–80 百分比之间,留下一个最后一英里问题:最先计算(即在反向传播最后计算)的第一层梯度没有后续计算来掩盖其传输时间,因而其原始传输时间会暴露在关键路径上。

图 3.2 中所示的带宽层次结构决定了用于缓解此鸿沟的并行策略。张量并行性需要巨大带宽以进行频繁的激活交换,因此被限制在单个节点的 NVLink 域内。流水线并行性涉及管道阶段之间激活的点到点传输,跨越节点间的 InfiniBand 链路。数据并行性通过梯度累积和重叠可以容忍较低带宽,因此可伸展到整个织构上。

案例研究:NVIDIA DGX SuperPOD

NVIDIA DGX SuperPOD 架构采用 NDR InfiniBand 网络将 DGX H100 节点互连,作为梯度总线概念的具体实现(NVIDIA 2023)。每个节点充当密集计算岛,通过 NVSwitch 将八颗 H100 GPU 互连,以提供 900 GB/s 的内部带宽。在外部,每个 GPU 与一块 ConnectX-7 NIC 配对,提供 400 Gb/s 的注入带宽。在标准可扩展单元(32 节点,256 GPU)中,这可产生 12.8 TB/s 的总注入带宽(256 × 50 GB/s),确保织构能够以加速器生成梯度的速度接收梯度。

此架构明确实现了本章构建的五层模型。在第 1 层,它通过在机架内使用无源铜线(DAC)并在仅对脊椎连接使用有源光纤来最小化延迟。在第 2 层,它依赖 InfiniBand 原生的基于信用的流量控制来保证无丢失介质,而不受以太网 PFC 脆弱性的影响。第 3 层实现了基于轨优化的胖树:GPU 轨道在可扩展单元内对齐,使得同轨流量保持局部,而交叉轨流量穿过脊椎层。在第 4 层,基于硬件的自适应路由可以在脊椎路径之间进行选择,以改善跨秩通信的双割行为。在第 5 层,设计是模块化的:多个 SuperPOD 可扩展单元通过额外的交换机连接以形成更大的集群(NVIDIA 2023)。就我们的 1750 亿参数模型而言,物理基础设施将由大约四个 SuperPOD 可扩展单元通过布线连接而成,从而使大约 1,024 颗 GPU 能够作为单一同步仪器协同工作。

案例研究:Meta Grand Teton

Meta 披露了两个基于 Grand Teton 的 24,576 块 H100 集群:一个在 Arista/OCP 以太网织构上使用 RoCE,另一个使用 NVIDIA Quantum-2 InfiniBand,两者均具有 400 Gb/s 端点。Meta 将两者均用于 Llama 3 训练,并报告称在 RoCE 集群上正在进行 Llama 3 训练(Meta Engineering 2024)。采用以太网的主要动机是运营规模和供应链韧性:通过使用以太网,Meta 可以从多个供应商处采购交换机,并使用与其前端服务机群共享的相同光纤基础设施和管理工具,从而避免专用 InfiniBand 岛屿带来的运营孤岛。

要使以太网在该规模下像专用 HPC 织构一样运行,需要在第 4 层进行显著的工程工作。Meta 的 RoCE 研究描述了用于长时长停顿事件的 PFC 看门狗、超越普通 ECMP 的迭代路由设计、面向未来的流让(flowlet)切换实验,以及一次 400G 部署经验:在该经验中,DCQCN 调整证明相当困难,以至于 Meta 决定不使用 DCQCN,而依赖于 PFC 加上更高层次的集体控制(Gangidi et al. 2024)。尽管此架构能够为大规模梯度传输提供高线路速率带宽,但它接受了一项权衡:小消息延迟可能仍然高于 InfiniBand,因为 RoCE 部署依赖于以太网缓冲、QoS、PFC/ECN/DCQCN 调整、路由策略以及交换机/NIC 实现选择,而不是 InfiniBand 原生基于信用的织构语义。对于带宽占主导的巨型模型,这种权衡是可以接受的;对于对延迟敏感的 MoE 路由,则需要仔细的算法补偿来应对该惩罚。

InfiniBand(NVIDIA DGX SuperPOD)和 RoCE(Meta Grand Teton,Google)在生产环境中的共存反映了一种真实的权衡,而不是明显的胜者。InfiniBand 能提供更低的尾延迟和更简单的无损配置。RoCE 能提供更低的交换机成本和多供应商灵活性。对于训练运行(其迭代时间以秒计),时延差异通常会被噪声吸收。对于推理服务(具有严格 SLO),时延差异可能会变得重要。

设计空间正在趋同。NVIDIA 的 Spectrum-4 以太网交换机融入了 InfiniBand 启发的自适应路由和拥塞控制。博通的 Memory DCS 芯片增加了对 RDMA 优化交换的硬件支持。两种生态系统之间的区别正在缩小,尽管它尚未完全消失。

这些案例研究通过展示物理线路、传输协议、拓扑和拥塞控制如何成为一个生产织构,完成了五层堆栈的闭环。剩余的运营问题不再是如何构建单个快速集群,而是如何在不牺牲训练所需可预测性能的前提下,在团队和工作负载之间共享这一昂贵的基础设施。

网络虚拟化

生产环境的 ML 集群极少专用于单个训练任务,这就给高效的多租户架构带来了巨大的经济动力。一台价值 3 亿美元的超级计算机如果因为无法安全隔离并发工作负载而闲置了 30%,就意味着 9 亿美元的资本浪费。要回收这部分效用,网络必须支持三个正交维度的虚拟化:带宽分区(保证最小吞吐量)、时延确定性(防止队头阻塞)和安全隔离(防止租户间的内存窥探)。诸如 SR-IOV(单根 I/O 虚拟化)和虚拟通道之类的技术,将训练任务与物理线缆解耦,正如虚拟机管理程序将操作系统与 CPU 解耦一样。对于我们的 175B 模型,这意味着训练任务可以可靠地消耗集群 BW[bisect](分割带宽)的 80%,而一个高优先级的推理服务和一个后台数据预处理任务共享剩余的 20%,网络结构强制执行硬边界,防止预处理器的突发流量阻塞梯度更新。

SR-IOV:硬件 NIC 虚拟化

云服务商可通过单根 I/O 虚拟化,向虚拟化的 GPU 实例提供接近裸金属的 RDMA 性能。该标准允许物理网卡呈现为多个独立的虚拟功能(VF)。每个 VF 拥有自己的硬件队列、门铃寄存器和 DMA 映射。通过为每个虚拟机或容器分配专用 VF,硬件创建了一条 DMA 操作的直接路径,完全绕过了宿主机内核和虚拟机管理程序。这种直通架构对 ML 训练至关重要,因为当网卡、虚拟机管理程序和放置策略配置正确时,它能将虚拟化开销降低到很低水平,通常在裸金属性能的百分之几以内。对集群管理的影响是直接的:VF 成为可分配的网络资源,因此放置逻辑必须在 GPU 之外,预留网卡容量和 QoS 策略,而不应将网络视为无界的共享资源池。

然而,硬件级隔离仍需显式的带宽策略。SR-IOV 暴露的多个 VF 共享网卡的物理资源;管理员可添加逐 VF 或逐组的 QoS 策略来限制或保证带宽,例如在 400 Gb/s 网卡上为 8 个 VF 各分配 50 Gb/s。对于在多租户集群上训练的 175B 参数模型,这种配置分区可防止邻近租户窃取带宽,但也对峰值吞吐量设定了硬性上限。训练任务的架构必须适应在此切片内运行,因为任何软件优化都无法突破配置的 VF 限制。

流量隔离与服务质量

设想共享集群上最坏的竞争场景:我们的 175B 模型正处于对延迟敏感的 AllReduce 操作中途,而邻近任务发起了大规模检查点保存。若无严格隔离,一次突发的 100 GB 写入(以满 400 Gb/s 线速需 2 秒)可能会饱和共享的脊柱链路,引入排队延迟,使 AllReduce 时间增加 50% 甚至更多。为防止这种“吵闹邻居”效应,高性能网络结构依赖服务质量(QoS)机制在数据包层面强制执行公平性。

主要工具是 InfiniBand 中的虚拟通道(VL)(或 RoCE 中的流量类别),它在单条物理链路上提供多达 16 条独立的逻辑通道。通过将不同流量类型映射到不同的 VL,网络确保某通道的饱和事件不会阻塞另一通道的进展。每个 VL 维护独立的基于信用的流控:若检查点的存储流量填满其缓冲区,交换机仅暂停该特定通道。梯度更新被标记为高优先级服务等级,继续通过其预留通道畅通无阻地流动。在以太网侧,增强传输选择(ETS)为每个流量类别提供类似的带宽保障,而先进的交换机 ASIC 可将其转发表和缓冲池划分为隔离的网络切片,确保某租户切片内的拥塞不会触发另一租户切片的 PFC 暂停。

虚拟化解决了共享问题,但增加了性能诊断的难度。当多租户集群上的训练任务变慢时,原因可能是物理链路退化、吵闹邻居超出其带宽配额,或 QoS 策略配置错误。系统性监控对于区分这些情况至关重要。

监控与调试

ML 集群中的网络性能问题隐蔽而棘手,因为它们表现为静默浪费而非显性故障。一块退化的光模块导致有效带宽降低 10%,可能仅使每次训练迭代慢 2%–3%,这种漂移极易被检查点或数据加载的自然方差掩盖。在 1,024 张 GPU 上进行为期 30 天的训练运行中,这种不可见的拖累累计约达 15,000–22,000 个浪费的 GPU 小时,按 3 美元/GPU-小时 计算约燃烧 4.5 万–6.6 万美元,却不会触发任何告警。传统的 IT 监控工具如 SNMPICMP ping 测量的是连通性,而非 RDMA 所需的持续吞吐量。有效的可观测性要求采取三层方法:物理监控(FEC 错误、信号衰减)、传输监控(PFC 暂停帧、重传率)和应用监控(NCCL 算法带宽)(Jeaugey 2017;Gangidi et al. 2024)。只有关联这三层的信号,运维人员才能发现“训练变慢”实际上是由某机架中一根退化的线缆导致的。

链路级遥测

在物理层,硬件计数器回答的是线缆是否仍按假设模型工作。每个 RDMA 网卡和交换机都维护这些计数器,可通过 perfquery 实用程序访问 (Linux man-pages project 2026)。PortXmitData(端口发送数据)和 PortRcvData(端口接收数据)提供原始字节计数,用于计算实时链路利用率,而 PortXmitDiscards(端口发送丢弃)记录交换机丢弃的数据包——在正确配置的无损网络中,丢弃应恰好为零。物理信号完整性通过 SymbolErrorCounter(符号错误计数器)跟踪,其因松动线缆或脏污光模块导致的比特级错误而递增;LinkDownedCounter(链路中断计数器)记录链路翻转,常由间歇性硬件故障或光模块过热触发。

这些指标通常由中心化监控系统(例如带 Grafana 仪表板的 Prometheus)以 10–30 秒间隔轮询,以捕捉静默退化。常见故障模式包括:QSFP56 线缆协商速率为 HDR(200 Gb/s)而非 NDR(400 Gb/s),或因单个连接器针脚损坏而静默从 4 通道降为 2 通道。链路仍保持“Up”且功能正常,但带宽减半。在拥有 1,024 张 GPU、3,000 多条链路的集群中,0.1% 的瞬时组件退化率意味着预期约有 3 条退化链路,且至少有一条链路退化的概率极高。由于 Ring AllReduce 等分布式训练算法是同步的,单条退化链路会将整个任务限制在该落后者的速度上,将微小的硬件故障转化为海量的空闲算力周期浪费。

带宽与延迟验证

排除可见链路故障后,验证工作必须测试训练作业实际可用的带宽和延迟。健康的 NDR InfiniBand 链路在扣除编码与协议开销后,可逼近原始 50 GB/s 的速率,但唯一有用的数字是特定主机对上交付的有效载荷带宽。运维人员依赖定期健康检查,常在选定节点对间运行 ib_write_bw 等 perftest 工具,并将结果汇总成全配对带宽矩阵(linux-rdma project 2026)。该热力图能直观显示网络中特定脊叶交换机或线缆束表现不佳的“冷点”,从而在作业调度前实施针对性维护。

针对对延迟敏感的同步场景,ib_write_lat 用于测量小 RDMA 写操作的往返时间。直连节点的基线 NDR 延迟应保持在 2 μs 以下。延迟超过 5 μs 提示交换机缓冲区拥塞或路由不均衡,而飙升至 100 μs 以上通常意味着 RDMA 路径已回落到 TCP/IP 仿真模式——这是一种灾难性的配置错误。在启动大规模训练作业前,稳健的验证包括应用层测试,如 nccl-tests,以验证网络能否在工作负载使用的特定集合拓扑(环形或树形)上维持预期的 AllReduce 带宽。这确保了在分配昂贵的算力资源前,物理网络现状与理论设计相符。

系统化调试工作流

当训练作业报告吞吐量低于预期时,诊断顺序应遵循分层模型,而非从检查线缆开始。以下序列从应用症状出发,逐步向物理原因推进:

  1. 检查 GPU 利用率:使用 dcgminvidia-smi 排除计算瓶颈。若 SM 利用率达 100%,网络不是瓶颈。然而,利用率低并不自动意味着网络有罪:输入管道饥饿(数据加载或存储)会产生相同症状,因此在调查网络前,先确认数据路径能持续喂饱加速器。

  2. 检查 NCCL 日志:设置 NCCL_DEBUG=INFO 以揭示所选网络传输、节点间检测带宽及任何向较慢协议的回落情况。

  3. 运行点对点测试:在作业的特定节点间使用 ib_write_bw。在 Ring AllReduce 中,单个退化链路即可成为整个环的瓶颈。

  4. 检查 PFC/ECN 计数器:检查路径上交换机的计数器。持续的 PFC 活动表明存在持久拥塞,应在调度器或路由层面调查。

  5. 验证物理层:检查符号错误和 CRC 计数,以识别故障的收发器或线缆。

该诊断序列遵循本章的分层模型:应用症状定位故障路径,传输计数器显示是涉及拥塞还是回落,物理遥测确认线缆本身是否在退化。

您的 175B 模型训练作业已在 512 块 GPU 上运行了 3 天。您注意到迭代时间已从 4.2 秒逐渐增加到 4.8 秒(慢了 14%)。nvidia-smi 报告的 GPU 利用率已从 92% 下降到 85%。

突破上限的网络技术

监控能让已部署网络“诚实”,但当铜缆在高信令速率下逼近物理上限时,可观测性无法突破这一上限。监控所揭示的物理上限,正是新型网络值得关注的原因:只有当它们能推动同步主干无法跨越的边界时,才具有意义。随着集群规模向 10 万节点演进,持久的设计变量是每比特功耗、传输距离,以及算力、内存与网络容量的比率,而非特定路线图上挂着的产品名称。

协议融合与链路密度

首个压力点是可靠性。超以太网联盟(UEC)旨在解决这一约束,致力于让以太网表现得更像 HPC 网络,同时不继承 RoCEv2 的 PFC 脆弱性。其超以太网传输设计融合了多路径包喷洒、灵活排序、多传输交付服务、拥塞控制变更,以及面向 AI/HPC 流量的遥测(Ultra Ethernet Consortium 2025)。目标是在保留以太网普及性与多厂商经济性的同时,将其故障行为推向 InfiniBand 无损网络模型。包修剪、链路层重传等具体机制是这一大方向的实现细节,应依据同一持久标准评估:网络能否在拥塞成为同步训练拖后腿的“落后者”前,在本地完成修复。

更高的端口速率同时攻克跳数与原始带宽。XDR InfiniBand 将单端口推至 800 Gb/s,以太网路线图包含 800GbE 与 1.6TbE 链路速率(NVIDIA 2026c; Ethernet Alliance 2025)。搭载 1.6 Tb/s 端口的以太网交换机芯片,在当前 Tomahawk 6 级设计中将聚合容量推至 102.4 Tb/s(Broadcom Inc. 2025)。单个 1.6 Tb/s 端口提供的带宽相当于 4 个 400G 端口。密度提升让架构师得以扁平化拓扑:在基数、端口速率与部署约束允许时,原本需要三级交换机的集群可能仅需两级,将收发器数量减半,并通过移除整整一级交换与 FEC 开销来降低尾延迟。

光互联

![垂直传输距离阶梯,标注四个层级:封装级(毫米级)、DAC(1 至 3 米)、AOC(3 至 30 米)、光纤(100 米)。]

每增加一米,都将网络从铜缆推向光学。

光互联针对铜缆在数据率提升时愈发恶化的每比特功耗与传输距离问题。在每通道 112 Gb/s(NVLink 4.0 与 InfiniBand NDR 采用的 PAM-4 信令速率)下,铜缆 SerDes 收发器功耗约 7–10 pJ/比特,传输距离限制在 2–3 米,超过此距离信号完整性退化将超出均衡器的恢复能力。

224 Gb/s 级别将铜缆推向极限。在 224 Gb/s 下,SerDes 功耗约翻倍至 15–20 pJ/比特,无源铜缆传输距离收缩至 1–1.5 米。有源电缆(含中继芯片以在线缆中段重塑信号)虽能延长距离,但会增加延迟和功耗。

共封装光学(CPO)通过将光收发器直接集成在交换机或加速器封装内,而非置于可揔拔线缆模块末端,解决了这些限制。在 CPO 中,硅光子芯片与处理器共享同一封装基板,在芯片边界完成电光转换。光信号随光速在光纤中传输,在数据中心相关的数十米距离内衰减可忽略,并在接收端封装转回电信号。

CPO 通过三个独立的物理效应改变链路预算:

  • 更短的电路径:将光收发器从交换机面板移至封装基板,将 SerDes 到光学的路径从厘米级缩短至毫米级,降低了电信号调理功耗。

  • 距离稳定的带宽:光纤在数据中心相关的最高 100 米距离内无距离相关的带宽衰减,省去了中继器,实现了更灵活的物理布局。

光缆轻量化

光纤比铜缆更轻、更细,简化了高密度机柜中的线缆管理。

这些效应共同将架构从受铜缆限制的布局问题,转变为受光学限制的封装和功耗问题。

快速估算使功耗红利具体化。

问题:计算将一个 51.2 Tb/s 交换机从插拔式收发器迁移到共封装光学(CPO)的功耗节省。

  1. 插拔式架构:128 端口 × 20 W = 2.56 kW,仅用于光学模块。

  2. CPO 架构:128 个引擎 × 10 W = 1.28 kW。

  3. 节省:每台交换机节省 1.28 kW 功耗。

系统洞察:在拥有 1,000 台交换机的集群中,插拔式光模块仅用于传输光信号就消耗 2.56 MW。CPO 将这笔“网络税”减半,节省 1.28 MW,重定向的电力足以供约 1,800 块 H100 GPU 使用。在 BW[bisect] 瓶颈处,可持续性不是一种选择;它是由面板热极限驱动的架构要求。

对于 10,000-GPU 集群,消除插拔式模块可节省超过 1 MW 的功耗,重定向给计算。CPO 还移除了收发器作为独立的现场可更换单元,消除了一个常见的机械故障点。

对于 ML 基础设施,CPO 可通过缩小节内和节间带宽的差距来扁平化带宽层级。如果节间链路实现与 NVLink 级本地结构相当的带宽(每 GPU 数百 GB/s),则限制张量并行性仅限于单节点内的约束将放宽。这将启用新的并行策略,使张量并行跨越两到四个节点,而不再局限于单个 8-GPU 节点,从而可能提高超大模型的扩展效率。

产品路线图不如其说明的约束持久:随着链路速率上升,光学器件向硅靠拢以控制功耗和传输距离。将光学集成到交换机封装或加速器封装中的产品,是为相同物理距离支付更低功耗和延迟的具体尝试。特定产品代际是否成功不如系统方向重要:电气路径缩短,光学路径起始于更靠近芯片的位置,封装级热学成为网络设计的一部分。

然而,CPO 引入了新挑战。光学器件对温度敏感(激光波长随温度漂移,需要主动热管理),将光子学与 1,000 W GPU 集成在同一封装内,为光学器件创造了恶劣的热环境。硅光子学和 CMOS 晶体管的制造工艺相似但不完全相同,需要单独的制造步骤,增加了封装成本和复杂性。这些挑战正通过混合集成方法解决,即将光子学芯片放置在封装基板较冷的区域,与 GPU 芯片热隔离。

解耦与可组合架构

传统节点捆绑了固定比例的计算、内存和网络,第 2.7.1 节 说明了当一个工作负载需要每 GPU 更多内存容量,而另一个需要更多网络带宽时,这种比例为何是浪费的:CXL 内存池化让一个机箱中的处理器以加载/存储语义访问另一个机箱中的内存,解耦设计将计算、内存和网络解耦为独立可组合的资源池。本章关注的是面料层面的后果。一旦内存池化和资源组合跨越节点边界,节内 NVLink 域和节间 InfiniBand 域之间的刚性区别就软化为单一内存面料,同步主干从固定的节点边界转移到可组合的边界。

这种转变并未废除本章阐述的物理规律。一个组合的虚拟节点在每次跨池访问时仍支付相同的 α-β 成本,受相同的链路功耗每比特限制,并具有相同的可观测性要求,以防止过订阅路径伪装成缓慢的加速器。系统层面的教训是开销的转移:解耦重新定位了同步主干,而非消除它。可组合面料能否工作,取决于决定固定节点面料能否工作的相同面料属性。

谬误与陷阱

为 ML 设计和运营高性能面料,需要摒弃传统数据中心网络的假设。以下谬误和陷阱捕捉了导致训练停滞和集群生产力下降的最常见错误。

谬误带宽越大,训练越快。

工程师假设从 HDR(200 Gb/s)升级到 NDR(400 Gb/s)将带来成比例的收益,但 α-β 模型(第 3.4.4 节)揭示这仅在带宽主导模式下成立。对于小模型或流水线并行的小消息阶段,由交换机跳数和 FEC 主导的延迟项 α(约 1 μs)决定性能。如果工作负载受延迟限制,带宽增加 2 倍可能仅带来线速增益的一小部分,却增加大量功耗和收发器成本。

以 10 KB 消息为例(控制同步的典型大小)。在 200 Gb/s InfiniBand 上耗时 1.90 μs。升级到 400 Gb/s 将其降至 1.70 μs,尽管链路速率翻倍,性能仅提升 10.5%。

陷阱将 InfiniBand 视为只是更快的以太网。

采购团队仅凭链路速率对比 InfiniBand 和高速以太网,得出在相同 Gb/s 下两者等价的结论。InfiniBand 和以太网在架构上不同,不仅仅是速度。InfiniBand 提供内核旁路(RDMA)、硬件管理的流控和基于信用的拥塞避免,所有这些都是标准化且可预测的。以太网依赖软件管理的 TCP/IP 栈,延迟高出数个数量级,且需要依靠 PFC/ECN 近似来在 RDMA 下实现无损行为。在 InfiniBand 和以太网之间的选择是系统架构决策,而非带宽选择:它决定了故障模式是由硬件信用界定,还是由软件配置纪律界定。

谬误无损以太网与 InfiniBand 一样可靠。

RoCE over Ethernet 在大传输中能达到与 InfiniBand 相当的吞吐量,但其“无损”特性是由 PFC(第 3.6.1 节)维持的一种近似。配置错误的交换机或固件 Bug 可触发 PFC 风暴,其中 PAUSE 帧在环路中传播,冻结整个面料。InfiniBand 基于信用的流控天生免疫此类级联,因为它基于逐跳缓冲区可用性运行,而非被动信号。部署 RoCE 的团队必须在面料测试和监控上投入大量工程精力,以避免由“无损”死锁引起的多日停机。

陷阱因为大部分流量是本地的,所以接受过订阅。

在通用云中,脊叶层 4:1 或 8:1 的过订阅很常见,因为流量是随机的。然而,ML 训练是批量同步的。当 AllReduce 开始时,屏障的协调强制每个节点同时注入满带宽。正如分析所示(第 3.5.2 节),即使 4:1 的过订阅也会使整个同步慢 4 倍。对于同步占 30% 时间的 3 亿美元集群,这种“成本优化”浪费了超过 1.421 亿美元的空闲 GPU 周期。

谬误TCP 带宽测试足以验证 ML 面料。

标准基准测试如 iperf 测量的是基于内核的 TCP/IP 性能。由于 TCP 需要由 CPU 管理缓冲区拷贝和上下文切换,它通常会在 20–40 Gb/s 处达到上限,与线速无关。ML 训练使用 RDMA,它完全绕过内核。在 iperf 中表现为 30 Gb/s“损坏”的链路,在 ib_write_bw 中可能完全健康并能提供 390 Gb/s 的带宽。验证协议必须使用 perftest 套件中的 RDMA 专用工具,以匹配工作负载的数据路径。

陷阱:依赖自适应路由而非拓扑感知放置

自适应路由将流量分发到可用路径上,但它无法创造不存在的带宽。如果调度器将一个 1,024 GPU 的作业放置在两个超额订阅的脊叶组之间,自适应路由会平衡流量,但组间链路仍会成为瓶颈。拓扑感知放置(详见第 8.3 节)与自适应路由是互补的:前者确保带宽存在,后者确保其被高效利用。

谬误:网络结构问题总表现为硬性故障

ML 集群中的网络性能问题通常表现为微妙的训练减速,而非错误。由于作业持续运行而梯度到达延迟,由间歇性拥塞导致的 10% 吞吐量下降,可能在被发现前浪费数千 GPU 小时。

陷阱:将 PFC 和 ECN 计数器视为可选的生产细节

运维人员必须对 PortXmitDiscards 和 PFC 暂停帧速率配置告警。这些计数器的逐渐增加,往往是光模块故障或路由不平衡的先兆,最终会导致作业失败。

总结

图 3.1 所示的五层模型构建了我们对高带宽结构的分析框架,从信号传输的物理层面向上延伸至仓库级集群的架构。该框架揭示:网络性能是物理传输距离、传输协议、拓扑结构和拥塞控制之间相互作用的产物,而非孤立的链路速度。第 1 层(线路)确立了 PAM4 编码和 FEC 强制规定了不可消除的延迟下限,限制了集群直径。在第 2 层(传输),InfiniBand 原生的基于信用的流控与 RoCE 对 PFC 的依赖提供了不同的可靠性保证,α-β 模型量化了分布式集合操作中固有的带宽-延迟权衡。

第 3 层(拓扑)展示了非阻塞 Fat-tree 和链路优化设计如何提供全局 AllReduce 模式所需的结构性 BW[bisect]。然而,第 4 层(行为)表明仅有结构是不够的:在 BSP 训练的同步世界中,尾延迟是主要约束,这需要 DCQCN 和 HPCC 等主动拥塞控制机制来防止因 Incast 导致的停顿。第 5 层(集群设计)将这些层面整合进生产架构,如 NVIDIA SuperPOD 和 Meta Grand Teton,展示了虚拟化和多租户如何让这些巨型设备得以安全共享。

在物理极限处,铜缆的约束和以太网的运维复杂度驱动了新的互联设计。超以太网联盟(UEC)标准、共封装光学(CPO)和 CXL 内存池化,都旨在扁平化拓扑或降低移动数据的功耗代价。然而,正如我们的监控工作流讨论所强调的,没有任何技术能消除对严格可观测性的需求。无论是调试单个退化的收发器,还是优化多租户调度器,将物理计数器与应用层吞吐量相关联的能力,仍然是防范机器学习集群中无声浪费的终极保障。

这种分层理解的实用价值在于诊断精度。当分布式训练作业表现不佳时,抱怨总是“网络慢”,但慢因很多:单个链路退化的故障收发器、跨子网级联的 PFC 风暴、一种拓扑瓶颈饿死某种通信模式却能很好服务另一种。理解五层模型的工程师能隔离故障层,关联物理计数器与传输行为,并将拓扑限制与拥塞控制误配置区分开来。这种跨越抽象层推理的能力——从 SerDes 信号完整性到集群级 BW[bisect]——将常规故障排查与真正的系统工程区分开来。

同样重要的是,α-β 成本模型为在采购硬件和布线机架之前做架构决策提供了定量词汇。在 InfiniBand 与 RoCE 之间、Fat-tree 与链路优化拓扑之间、400G 与 800G 链路速度之间的选择,都是取决于消息大小分布、集合算法和物理链路特性之间相互作用的百万美元级决策。本章开发的框架使从业者能够以分析严谨性评估这些权衡,而不只是依赖供应商基准测试。

  • 链路速度不等于结构速度:全局 AllReduce 受限于最窄的二分切割,而非最快的标称端口。拓扑决定了购买的加速器吞吐量有多少变成了有用的训练吞吐量,有多少变成了空闲的硅片。

  • 延迟和带宽的绑定方式不同α-β 模型将启动开销与每字节传输开销分离,n* = αβ 标志着机制转换。消息大小分布而非供应商峰值数字,决定了是优化软件延迟还是硬件带宽。

  • 无损转移了风险:RDMA 需要无损结构以避免昂贵的重传停顿。InfiniBand 原生提供此特性,而 RoCE 依赖 PFC、ECN、DCQCN 和 HPCC,用尾延迟和运维复杂度换取硬件灵活性。

  • 拓扑必须匹配流量:Fat-tree 购买灵活的二分带宽,链路优化设计加速同秩 AllReduce,Dragonfly 或 Torus 设计在布线和局部性间做不同权衡。正确的结构取决于工作负载是否强调 AllReduce、AllToAll 或多租户共享。

  • 遥测节省 GPU 小时:PFC 计数器、链路错误率、带宽基线和应用吞吐量必须跨层关联。没有这种可观测性,退化的收发器或拥塞风暴会悄无声息地将高端集群变成等待加速器的队列。

链路速度是人人引用的数字,也是不应盲信的错误指标。决定一个集群是否能作为一台机器计算的是二分带宽——即集群一半与另一半通信的速率,因为全局 AllReduce 只能快到结构最窄切割的速度。一个饿死该切割的拓扑,会将昂贵的硅片变成空闲的硅片,在等待导线无法及时传输的梯度。α-β 模型给等待定了价,将消息为启动支付的延迟与为完成支付的带宽分离开来。这就是本章为何将互联视为计算机的一部分,而非计算机之间的管道:在舰队规模下,结构决定了有多少硅片值得购买。

网络结构现在将计算节点绑定成一个舰队,每一个字节的梯度数据和激活张量都流经它。然而,舰队还需要移动海量数据集和巨大的检查点。第 4 章 探讨并行存储系统和数据加载架构,它们让舰队持续获得数据供给。

  • 为什么二分带宽而非标称链路速度,决定了有多少加速器算力变得有用?

数据存储

车队存储燃料管线层级,数据从对象存储经 NVMe、DRAM 和 HBM 层级流向加速器工作节点。

目的

为什么存储会成为阻碍加速器发挥潜力的隐形瓶颈?

加速器的计算速度快于存储的供数据速度。高端加速器内部处理数据的速度达到每秒 TB 级,但单个本地驱动器仅能提供每秒 GB 到几十 GB 级的吞吐量,而分布式存储系统引入的延迟会累积,导致加速器空闲等待数据到达。这种不匹配在单独测量加速器性能的基准测试中是不可见的,但在真实工作负载中却占据主导地位:训练数据必须持续流式传输,检查点必须可靠保存,推理时必须加载模型权重。加速器能消费的数据量与存储能提供的数据量之间的差距,在各个层面塑造了系统架构:它迫使人们仔细关注数据格式、缓存策略和管道设计——如果存储能跟上计算,这些原本是不必要的。那些在不解决存储问题的情况下优化加速器利用率的组织会发现,他们昂贵的硬件只能以极小部分的产能运行,因为没人规划过数据路径。用 术语来说,数据存储是一个计算-通信协同设计问题:当存储层级无法以加速器消费的速率供应数据时,车队中最快的计算资源就会空闲等待。

  • 解释 ML 存储工作负载如何通过流式处理、检查点突发和元数据压力颠覆数据库假设

  • 利用管道方程和目标加速器利用率计算所需训练带宽

  • 对比内存、闪存、并行文件、对象和归档存储层级的带宽、延迟、容量和成本

  • 设计预取、分片、缓存和数据局部性策略,防止车队级加速器“饿死”

  • 评估加速器直连和 CPU 旁路存储路径在延迟、扩充和局部性瓶颈方面的表现

  • 选择检查点暂存和复制策略,在暂停时间、恢复风险和存储成本间取得平衡

  • 将检索索引和合成数据管道作为具有独特延迟、一致性和治理约束的存储工作负载进行评估

燃料管线

密集加速器节点可集成八张 GPU,提供每秒 PFLOP 级的聚合算力,InfiniBand 网络可将数千个此类节点以每秒数百 Gb 级互联。如图 1.13 所示的车队栈中,数据存储通过提供燃料供应——训练数据、模型权重、优化器状态和置于距加速器适当位置的中间检查点——完成了这套物理基础。没有燃料的引擎只是昂贵的雕塑。工程问题在于如何足够快地输送燃料,使 1,000 个加速器从不“饿死”。

三个并排的胶囊标签 D、A 和 I。I 胶囊填充蓝色,D 和 A 胶囊为灰色,标记基础设施轴为当前聚焦活动轴。

存储是车队栈的基础设施轴。

考虑存储分析的运行示例。一个 1750 亿参数的语言模型在 1.5 万亿 token 的文本上训练:压缩源文件约 3 TB,表示为 4 字节 token ID 后约 6 TB。每个训练周期按随机种子决定的打乱顺序读取每个 token 一次。不存在传统意义上的“热数据”子集主导访问;每字节每轮恰好被消费一次。同时,每个加速器在约 200 ms 内处理完其本地批次,然后等待下一个。如果存储无法在 200 ms 窗口内交付数据,加速器就会空闲,组织为产生热量而非梯度的硅片买单。

问题具有欺骗性,因为存储技术已大幅进步。NVMe 驱动器实现 7 GB/s 顺序吞吐,这在十年前似乎遥不可及。然而,加速器进步得更快。H100 GPU 从其 HBM 消费数据的速率为 3.35 TB/s,约为单个 NVMe 驱动器供数据速度的 478.6 倍。存储交付与加速器消费之间的差距是核心存储张力,无法靠任何单一技术解决。相反,它需要一个存储层级体系,每层精心匹配 ML 生命周期的特定阶段,通过预取和流水线隐藏延迟来连接。

存储问题从根本上讲是物理与经济的博弈。物理决定了:离加速器更近(物理距离和互联跳数)的数据交付更快,但容量更小。经济决定了:更便宜的存储能容纳更多数据,但距离更远。工程艺术在于构建一条管道来弥合这些约束,通过从低层级快速抽取数据保持昂贵的顶层满载,使加速器从未察觉到延迟。由此产生的设计问题是定量的:每层必须多快、管道必须多深、哪些字节值得保留在靠近加速器的位置。

对数刻度上的两级**:高耸的蓝色级代表检查点写入 7.56 PB,远超微小的每周期训练数据读取 6 TB 级,显示检查点移动的字节数比数据集多约一千倍。

检查点写入量比训练数据读取量大约 1,000 倍。

该示例的标准训练数据在层级中的总占用约 9 TB,结合了上文引入的压缩语料库和分词分片。额外的打乱或打包变体可能增加暂存占用,但基于 4 字节 token ID 的每周期训练读取量为 6 TB。模型每 10 分钟生成约 1.75 TB 检查点(共 1,750 GB:350 GB 权重加 1.4 TB Adam 优化器状态)。在 256 节点上进行 30 天训练期间,存储系统必须每周期交付 6 TB 分词训练数据,吸收 7.6 PB 检查点写入,并为评估运行暂存模型权重。这些数字为后续章节定下基调,将抽象原则落地为具体工程约束。

这些数字强制确定了存储路径。ML 访问模式首先颠覆了传统存储系统背后的假设;这种颠覆强制建立从 HBM 到冷归档的层级;层级随即要求管道方程、直连数据路径和经济模型,以决定哪些字节属于各层。检查点、检索索引和合成数据溯源是同一燃料管线问题的变体:存储系统必须在加速器请求之前,将正确的数据表示放置在正确的距离上。

ML 工作负载如何颠覆存储假设

一位数据库管理员转岗到 ML 基础设施团队,会发现 ML 存储工作负载颠覆了他们赖以生存的几乎每一条存储设计原则。175B 运行示例使这种颠覆具体化。每个训练周期按打乱顺序读取每个 token 恰好一次。下一周期再次打乱并再次全部读取。不存在传统意义上的“热数据”,也没有 80/20 法则——即一小部分数据占绝大多数访问。标准存储优化之所以失效,正是因为它们假设情况相反。

传统存储与 ML 工作负载

传统存储系统是为了服务事务型数据库而演变而来的,这类工作负载的特点是小块随机访问、强一致性和中等带宽。数据库服务器可能会每秒发出数千次 4 KB 读取请求来响应用户查询。数十年来,行业一直针对这种模式进行优化,开发了复杂的缓存算法、预写日志,以及针对小块随机访问进行调优的 RAID 配置。每一项优化都假设最近访问的数据极有可能在不久后再次被访问。

ML 工作负载系统性地颠覆了这些假设。训练数据访问主要是顺序的,流式读取跨越数百 TB 的数据集。单次访问的粒度很大(兆字节而非千字节),因为模型以批次消费图像或文本序列。一致性要求有所放宽,因为略微陈旧的特征很少影响模型质量。然而,带宽需求极高:数百 GB/s,持续数天或数周。存储系统的优化目标与 ML 实际需求之间的错位,造就了我们所说的I/O 墙(原则):当存储吞吐无法以加速器消费训练数据的速度提供数据时,GPU 就会空转,无论其算力多强。一个足以支撑 8 张 GPU 的存储系统,在 64 张 GPU 面前就会成为瓶颈,使得限制因素变成了数据管线,而非模型。第 B.5.2 节将这堵墙归类为一个约束,它位于车队级诊断框架中通信与计算的交汇处,因此存储工程师可以确认,治本之策是增加带宽,而非采用更快的加速器。

分片分配中的生日悖论

一个简单的分片分配计算展示了即便聚合存储容量看似充足,这种瓶颈也会涌现。

问题:某数据集在共享文件系统上被拆分为 1000 个分片。若 32 个工作进程各自随机选取一个分片以开启下一个 epoch,那么至少有两个 GPU 在同一存储服务器上“碰撞”并导致性能瓶颈的概率是多少?

数学原理:这是概率论中“生日问题”的一个变体。

  1. 无碰撞概率:≈ e^(-n[workers]²/(2*K[shards])) = e^(-32²/2000) ≈ 0.60。

  2. 资源竞争概率:1 − 0.60 = 40%

系统洞察:即使分片数量很大,生日悖论近似公式 1 − e^(-n²/(2*k)) 表明,当 n = 32 个工作进程、k = 1,000 个分片时,存储分片争用的概率高达 40.1%:即同一个 epoch 内,多个工作进程落在同一存储服务器上形成“热点”。在分布式集群中,这些碰撞会产生尾延迟:整个集群必须等待共享同一磁盘的那两张 GPU 完成任务。为解决此问题,生产环境的数据加载器采用全局打乱和确定性分片分配,确保工作进程在存储结构中完美分布,从而消除“惊群效应”。

硬件鸿沟:算力与存储带宽

分片碰撞是 I/O 墙的一个运行时症状;这堵墙背后的硬件趋势,是算力吞吐与存储带宽之间日益扩大的鸿沟。

2016 年至 2024 年间,加速器 Tensor Core 峰值吞吐大幅增长,但具体倍数取决于对比是否固定精度,或是遵循各代产品支持的最低训练/推理精度。同期,单块 NVMe 顺序带宽增长缓慢,从约 3.5 GB/s 增至 14 GB/s。因此,算力与存储带宽之比显著恶化。若此趋势持续,存储层级必须引入新层级(持久内存、CXL 连接存储),或从根本上改变数据管线架构(近存计算、存内计算),以防止 I/O 墙成为训练吞吐的决定性约束。持久的规划启示是:算力提升不会自动带动存储路径同步进化。

图 4.1 通过在同一时间尺度上并排追踪 GPU 吞吐与存储带宽,直观精确地展示了这一不断扩大的鸿沟。

图 4.1:存储-算力鸿沟:2016 年至 2024 年 GPU 峰值 FP16 吞吐(蓝色,左轴)与 NVMe 顺序读带宽(橙色,右轴),均采用对数刻度。GPU 吞吐增长了 236 倍,而存储带宽仅增长了 4 倍。阴影区域突出了数据管线工程必须通过预取、缓存和格式优化来弥合的不断扩大的差距。

图 4.1 说明了为何 I/O 墙是本章的决定性约束。GPU 吞吐增长(236 倍)与存储带宽增长(4 倍)近 60 倍的比率意味着,每一代新 GPU 都在加大对数据管线的压力。这种增长差异——即差距历年来扩大的速度——不同于后续章节量化的静态层级间带宽悬崖,后者衡量的是某一时刻差距有多陡峭。若缺乏多层存储层级、预取和格式优化,栈顶昂贵的加速器将花费更多时间等待数据,而非进行计算。本章剩余部分将探讨存储层级的每一层如何应对这一鸿沟的不同侧面。

访问模式反转:随机与顺序

首个反转在于访问模式。数据库工作负载呈现随机访问模式,受益于寻道时间优化;而 ML 训练执行大规模顺序扫描。一个训练 epoch 按打乱算法生成的顺序读取每个样本一次。这种模式更类似视频流,而非数据库查询。针对随机 IOPS 优化的存储系统在 ML 工作负载上浪费了能力,而针对顺序吞吐优化的系统则表现出色。这种差异在量化上极为显著:Gen4/Gen5 NVMe 驱动器顺序读可达约 7–14 GB/s,而 4 KB 小块随机读可能降至 0.5 GB/s 左右,错误的访问模式导致数量级到 30 倍的性能损失。在硬盘上惩罚更为严重,机械寻道时间导致随机访问吞吐比顺序访问降低 100 倍。

ML 训练所需的打乱操作,因随机梯度下降的要求而增加了复杂性。随机梯度下降要求每个小批次近似均匀地从训练分布中抽取;按相关顺序呈现样本(如同一文档的所有样本、同一类别的所有图像)会偏移梯度估计,拖慢或破坏收敛。真正的全局打乱虽能完美满足此要求,但需要跨整个数据集进行随机访问,从而破坏存储硬件所需的顺序访问模式。对于 PB 级语料库,逐样本的随机寻道在 I/O 上是不可接受的:假设 1 万亿个 4 字节令牌作为独立元素存储在 NVMe 上,按随机方式访问,以驱动器的随机读 IOPS 速率读取将耗费数千小时,而顺序流式读取仅需数小时。务实的解决方案因此是一种折衷:打乱大数据分片的顺序,然后在每个分片的本地缓冲区内打乱样本。这在保证训练收敛所需的充分随机性的同时,保留了存储硬件所要求的顺序 I/O 模式。分片大小决定了权衡:更大的分片提供更丰富的分片内打乱多样性,但需要更多内存用于打乱缓冲区。

ML 存储工作负载反转

工作集大小

传统应用程序表现出时间局部性:Web 服务器反复访问热门页面,而持有前 10% 内容的缓存可服务于 90% 的请求。ML 训练数据集是被均匀访问的。我们 6 TB 的序列化语料库没有“热门” token;每个 token 每个 epoch 仅被消费一次。任何切实可行的缓存大小仅能容纳数据集的一小部分,并且每个样本在被访问时实际上都是“冷”的。这种缺乏时间局部性的特性使得传统的缓存策略失效。即使是数据库系统的主力 —— LRU(最近最少使用)缓存,在均匀访问的数据上也能达到 0% 的命中率,因为当下一个 epoch 再次访问某个样本时,它早已被驱逐以便为其他样本腾出空间。

这一均匀性的例外是多任务或课程学习,在特定训练阶段,数据集的某些子集会被更频繁地访问。在课程学习中,训练器从“简单”样本开始,逐步引入更难的样本。这创造了一个确实表现出局部性的临时工作集,而 NVMe 层的本地缓存可以利用这种结构。然而,对于绝大多数大规模预训练工作负载而言,访问模式实际上是均匀的,因此存储系统必须为全数据集流式读取而设计,而非热点子集缓存。

写入模式

事务系统生成持续的小写入流,每次写入都立即持久化。ML 系统在保存检查点时会生成偶发的大规模写入。一个包含优化器状态的 175B 参数模型检查点大约占用 1,750 GB。每 10 分钟保存一次会产生持续数秒的集中突发写入,从而饱和带宽,随后是漫长的空闲期。这些突发写入共同构成了检查点风暴:一个同步的检查点写入事件,并行文件系统必须在不干扰正在进行的训练读取的情况下吸收它。这种突发写入模式极具挑战性,因为集群中的所有节点同时写入其检查点分片。如果 128 个节点每个写入 14 GB,并行文件系统将在单次突发中接收 1.8 TB 的写入,这必须在训练管线恢复之前完成。表 4.1 将这些反转整合为采购原则:针对流式读取和突发写入进行优化,而非数据库式的随机 IOPS。

表 4.1:ML 工作负载反转传统存储假设

| 工作负载模式 | 传统假设 | ML 现实 |

| --- | --- | --- |

| 访问模式 | 随机访问 | 顺序流式读取 |

| 工作集 | 适合缓存 | 超过所有缓存层级 |

| 写入模式 | 持续小写入 | 突发大写入 |

| 读/写比例 | 均衡 | 阶段依赖(100:1 到 1:0) |

| 局部性 | 强时间局部性 | 无局部性(均匀采样) |

数据库针对可缓存工作集优化随机 IOPS,而 ML 训练则顺序流式读取超出所有缓存层级的数据集。

正如表 4.1 所示,这些反转还有一个第四个更微妙的维度:读/写比例随生命周期阶段发生剧烈变化。在训练期间,读主导写,比例可达 100:1 以上,因为系统持续流式读取数据并偶尔保存检查点。在易故障集群的检查点密集阶段,写入可能短暂地占主导地位。在数据预处理期间,读写都很繁重,且访问模式更类似于 MapReduce 作业而非训练循环:分词、去重和洗牌会扫描一次原始语料库(大量顺序读取),写入中间产物(大量顺序写入),然后再次扫描这些产物以供下一阶段使用。单个 1 TB 文本语料库通常在首个训练 epoch 开始前就会扩展为 6 TB 的分词分片,与随后的只读训练阶段相比,这以不同方式加载每个存储层级。没有单一存储配置能针对所有阶段进行优化,这也是 ML 系统需要多层级层次结构而非单一存储技术的原因。

训练与推理

对比训练和推理工作负载时,会出现第五个反转。训练顺序读取数据集并以突发方式写入检查点。推理则相反,它在启动时一次性读取模型权重(可能数百 GB 的大规模顺序读取),随后在正常运行期间不再进行任何存储 I/O,因为模型完全驻留在 HBM 中。推理的存储挑战在于冷启动延迟:扩容或从故障恢复时,将模型从存储加载到 HBM 所需的时间。因为模型只读一次便常驻内存,约束条件从持续流式吞吐量转变为单次大容量读取,其耗时完全取决于哪个层级提供权重,并行文件系统比本地 NVMe 慢数倍。第 4.6.3 节 推导了 175B 模型的具体加载时间。对于有严格可用性要求的服务工作负载,这种冷启动时间推动设计走向在主机 DRAM 中保持热副本,或使用模型分片跨多个存储设备并行化加载。

当从单用户扩展到数千并发推理请求时,存储挑战从单流吞吐量问题转变为大规模扇出分发问题。拥有 100 个 175B 参数模型副本的服务集群需要在集群中分发 35 TB 的模型权重。部署新模型版本(模型发布)时,必须更新所有 100 个副本,触发一场必须在几分钟内完成的 35 TB 数据分发事件,以最大程度减少服务中断。这类似于训练中的检查点风暴,但方向相反:不是许多节点同时向中心位置写入,而是许多节点同时读取相同数据。存储系统必须在继续无降级地为现有模型版本服务推理请求的同时,维持这种用于模型分发的突发读取带宽。

采购启示

这些反转对系统采购和架构有直接影响。基于数据库时代启发式方法为 ML 分配存储的组织,将在随机 IOPS 上过度投资(ML 不需要),在顺序带宽上投资不足(ML 迫切需要),且无法解释检查点产生的突发写入模式。

设计练习

你正在为一个拥有 512 个 GPU 的新 ML 训练集群设计存储子系统。主要工作负载将在 10 TB 文本数据集上训练大语言模型。

了解哪个存储层级能持续供给加速器对于系统设计至关重要。图 4.2 绘制了维持 GPU 充分利用所需的 I/O 吞吐量与模型大小的关系,并叠加了各存储层级的水平上限。在 NVMe 上限以下,本地 SSD 可维持训练供给;在 NVMe 与并行文件系统上限之间,仅 Lustre 级存储足以胜任;超出并行文件系统上限后,工作负载进入存储瓶颈区,没有任何单一层级能独立维持加速器运行。

图 4.2:存储瓶颈区

图 4.2:存储瓶颈区:维持 GPU 充分利用所需的 I/O 吞吐量 (GB/s) 与模型大小 (参数量, 100M 至 175B) 的关系。红色曲线显示存储需求随模型大小增长;水平虚线标记常见存储层级的上限(NVMe 7 GB/s,对象存储 50 GB/s,Lustre 100 GB/s)。计算受限区与存储瓶颈区之间的过渡地带,是存储而非计算限制训练吞吐量的区域。

ML 存储层级

系统架构师必须对存储进行组织,以服务于那些同时要求计算具有每秒 TB 级带宽、数据集具有 PB 级容量以及检查点具有极高耐久性的工作负载。没有任何单一技术能同时满足这三项要求。HBM 提供带宽但不提供容量。对象存储提供容量和耐久性但不提供带宽。解决方案是采用多层层级架构:将少量快速、昂贵的存储放置在靠近加速器的位置,将大量缓慢、廉价的存储放置在边缘。每一层之所以存在,是因为它解决了物理学(带宽和延迟受距离加速器的远近支配)与经济学(随着容量增加,每比特成本下降)之间的特定矛盾。该层级扩展了学生在计算机体系结构课程中遇到的经典处理器内存层级(寄存器、L1/L2 缓存、DRAM),增加了大规模数据系统特有的、位于 DRAM 之下的层级。表 4.2 揭示了 ML 系统必须应对的极端带宽差异。

Three stacked horizontal bars on a log scale, longest at top: HBM at 3.35 TB/s, host DRAM at 200 GB/s, and NVMe at 7 GB/s, showing bandwidth dropping sharply across the top three storage tiers.

带宽在顶层三个存储层级(从 HBM 到本地 NVMe)中下降了约 479 倍。

| 存储层级 | 典型容量 | 带宽 | 延迟 | 成本 ($/GB) |

| --- | --- | --- | --- | --- |

| GPU HBM | 80 GB | 3.35 TB/s | ~100 ns | ~15.00 |

| 主机 DRAM | 512 GB–2 TB | 200 GB/s | ~100 ns | ~3.00 |

| 本地 NVMe SSD | 4–30 TB | 7–25 GB/s | ~100 μs | ~0.10 |

| 并行文件系统 | 100+ PB | 1+ TB/s 聚合 | ~1 ms | ~0.03 |

| 对象存储 | 极大的存储池 | 100 GB/s 聚合 | ~50 ms | ~0.02 |

| 归档/冷存储 | 极大的存储池 | 1 GB/s | 分钟到小时 | ~0.004 |

表 4.2:ML 系统的扩展内存层级:HBM 与对象存储之间大约 30 倍的聚合带宽差距(一旦单个推理实例从共享对象端点拉取数据,单客户端的差距会大得多),驱动了跨多个层级进行复杂预取和缓存的需求。

表 4.2 中有一点值得注意:HBM 相对于主机 DRAM 的优势在于带宽,而非延迟。这两种技术的随机访问延迟均处于 ~80–150 ns 范围内;HBM 通过硅中介层上数千条并行走线和 3D 堆叠实现其带宽领先优势(10–20 倍),而非更快的存储单元。相邻层级之间的延迟断层仅在 NVMe 及以下层级出现,此时距离加速器的物理距离从毫米级增长到米级。

存储层级原则支配着本章中的每一个设计决策。随着数据离加速器越远,存储性能下降而容量增加:表 4.2 中的每一层级都会使带宽下降 10–100 倍,同时容量增加 10–100 倍。数据格式选择、缓存策略、预取缓冲区大小调整和分层策略的存在,都是为了管理数据在层级中的上行移动,确保加速器永不“饥饿”。

图 4.3 将这六个层级映射为空间层级,展示了带宽如何在远离加速器的每一步降低,而容量如何增加。

图 4.3:ML 存储层级:存储层级分为热、温、冷三个频段,从 GPU HBM(每秒 TB 级带宽、GB 级容量)一直到磁带归档(亚 GB/秒级带宽、超大容量),每 GB 成本呈反比变化。该图在文中发展的五层模型之外又增加了两个层级——传统 SATA SSD 层级和磁带归档层级——以展示权衡的完整跨度;共享对象存储和并行文件系统层级标记为红色,因为它是使大型集群“饥饿”的瓶颈,也是本章其余部分致力于隐藏其带宽断层的层级。该图使用示例数值使带宽与容量的权衡一目了然;本章全文使用的标准每层数值来自表 4.2,当图表与表格不一致时,该表为真理来源。

图 4.3 中的金字塔体现了一个基本权衡:层级中的每一步向下都在用带宽换取容量和成本。这种权衡并非随意设定;它反映了数据邻近性的物理规律。HBM 与加速器位于同一硅中介层上,通过数千条以毫米计的并行走线连接。主机 DRAM 通过跨越厘米级的 PCIe 通道通信。NVMe 通过 PCIe 连接器跨越电路板。并行文件系统穿越数米的线缆和网络交换机。对象存储可能跨越数据中心间数公里的光纤。在每一层级,不断增加的物理距离直接转化为增加的延迟、降低的单连接带宽和降低的每字节成本(因为相同介质可在更低密度下存储更多数据)。

工程挑战在于确保数据以足够快的速度向上流经金字塔,使顶层(HBM)在加速器需要时永不为空。回到我们的持续示例:持久语料库以 3 TB 压缩源副本加 6 TB 分词训练分片的形式驻留在对象存储(第 4 层)中,但加速器需要在 200 ms 内将每个批次加载到 HBM(第 0 层)。数据必须通过中间层级提升,在逐步更快的存储中暂存,以便当加速器请求批次时,数据已在主机 DRAM 中就绪,仅需一次 PCIe 传输即可到达 HBM。

层级间的带宽断层

相邻层级之间的带宽比率揭示了层级结构中每次转换的严重程度。HBM 与主机 DRAM 之间,比率约为 16.8×(3.35 TB/s 对比约 200 GB/s)。主机 DRAM 与 NVMe 之间,比率约为 7.1×(200 GB/s 对比 4 驱动器 RAID-0 提供的约 28 GB/s)。NVMe 与并行文件系统之间,比率取决于每节点分配:若 1 TB/s 聚合 PFS 服务于 256 个节点,每个节点获得约 4 GB/s,比本地 NVMe 降低 7×。并行文件系统与对象存储之间,比率通常为 10× 或更高,取决于并发客户端数量和网络带宽。

这些“带宽断崖”有一个关键含义:流水线不能简单地实时“流式穿越”层级结构。如果加速器以 3.35 TB/s 从 HBM 消费数据,而下一层级仅能提供 200 GB/s,那么 HBM 可在 50 ms 内耗尽,但从主机 DRAM 回填却需 400 ms。加速器避免停顿的唯一途径,是它所需的下一批次数据在完成当前批次前已预先置于 HBM 中。这正是为什么层级结构中的每一层都作为其上层的预取缓冲区:主机 DRAM 为 HBM 缓冲数据,NVMe 为主机 DRAM 缓冲数据,并行文件系统为 NVMe 缓冲数据,而对象存储则是终极事实来源。每个缓冲区都必须足够深,以吸收其下层的延迟和带宽抖动。

带宽算术也解释了为什么增加集群规模会产生存储压力。配备 8 张 GPU 的单节点大约需要 4 到 40 GB/s 的存储带宽(取决于负载)。256 个此类节点的集群需要 1,000 到 10,000 GB/s。10,000 个节点的集群需要 40 到 400 TB/s。在上限端,即使拥有 1,000 台对象存储服务器(即下文并行文件系统层级中引入的数据服务节点)且聚合吞吐达 1 TB/s 的世界级并行文件系统,也无法满足需求,架构必须依赖本地 NVMe 缓存来分担共享存储的负载。这种集群级压力的严重程度急剧依赖于数据模态。

文本训练带宽与图像训练带宽对比。

存储带宽需求随数据模态剧烈波动。

问题

对于文本训练与图像训练,一个 2,048-GPU 集群需要多少聚合存储带宽?

设置

带宽需求完全取决于数据模态。

文本训练

对于文本训练,需求出奇地低。采用典型的每 GPU 4,096 token 批次大小和 200 ms 步长,聚合带宽为:

2,048 GPUs × 4,096 tokens/GPU × 4 bytes/token ÷ 0.2s ≈ **167.8 MB/s**

单个网络附加存储节点即可轻松满足。

图像训练

对于图像训练,情况发生剧变。采用常见的每 GPU 256 批次大小(ImageNet 224 × 224,约 150 KB/张),聚合带宽激增:

2,048 GPUs × 256 images/GPU × 150 KB/image ÷ 0.2s ≈ **393.2 GB/s**

这比文本负载高出约 2,300 倍,需要高性能并行文件系统。这一根本差异驱动层级设计:文本训练是数据量大带宽轻的,受限于总数据集大小和检查点写入;图像训练是带宽重的,受限于存储系统向加速器供数据的能力。

层级间的带宽断崖也影响各层级的数据格式。在 HBM 层,数据必须为加速器可直接计算的格式:float16 张量、打包的 token ID 或预处理后的特征向量。在 NVMe 层,数据可采用更紧凑的格式(压缩 JPEG、字典编码的分词文本),因为 CPU 有时间在加速器处理上一批次时解码。在对象存储层,最大压缩有利于最小化存储成本和传输时间,即使解压会增加 CPU 开销。从压缩存储到计算就绪张量的格式转换,是流水线“增值”工作的一部分,将原始字节转换为加速器所需的表示。这种转换发生在主机 DRAM 中,这正是为什么主机 DRAM 作为流水线的关键暂存区。

多模态训练加剧了格式挑战,它在单一模型中融合文本、图像、音频和视频。每种模态的数据特征截然不同:一个文本 token 占 4 字节,一张高分辨率图像占 150 KB,一个短视频片段占 10 MB。它们的压缩特性也不同,需要不同的增强流水线。多模态训练作业必须管理多条并行数据流,每条流有自己的带宽特征和预取需求。存储层级必须为所有模态的带宽需求之和进行规划,而非仅针对主导模态。对于结合 3 TB 文本、50 TB 图像和 200 TB 视频的训练作业,视频模态在存储容量和 I/O 带宽需求上压倒性地占据主导,即使文本模态对模型质量贡献更大。这种存储成本与训练价值之间的不对称,是多模态系统设计中反复出现的挑战。

下行始于实际消耗字节的计算之处:HBM,这一层级足够快以供给加速器算术运算,却因稀缺而无法容纳完整训练语料。

第 0 层:GPU HBM

整个系统中最受限的资源,同时也是最稀缺的。高带宽内存(HBM)是权重和激活值在主动计算期间唯一能驻留的存储层级。正如第 2.3.1 节所述,HBM 是一种 3D 堆叠内存技术,将 DRAM 芯片垂直堆叠于加速器之上,通过数千个硅通孔连接,在 H100 上提供 3.35 TB/s 聚合带宽。该带宽约为本章 200 GB/s 主机 DRAM 数值的 17 倍,约为单张 7 GB/s NVMe 驱动器的 480 倍。正是这种带宽使大规模深度学习成为可能:涉及数十亿参数的矩阵乘法,要求每次前向和反向传播都从内存读取这些参数。

此层级的约束在于容量,而非带宽。H100 提供 80 GB HBM,足以容纳 FP16 下的 400 亿参数模型(每参数 2 字节),却远不及我们贯穿示例的 1750 亿参数模型。要理解此约束的严重性,请看训练该 1750 亿模型的内存预算。FP16 模型权重消耗 350 GB。自适应矩估计优化器维护两个额外状态(动量和方差)为 FP32,消耗 175 × 10⁹ × 4 × 2 = 1,400 GB。单批次的激活值视序列长度和批次大小而定,还可能消耗 100 到 400 GB。总内存占用达约 1.85–2.15 TB,视激活大小而接近或超过单张 H100 HBM 容量的 25 倍。这个 25 倍是未分区的占用,这正是表 4.3 中分区策略强制执行的原因。从存储层级角度看,HBM 是所有下层服务的目的地。数据流水线的目的是确保 80 GB HBM 始终包含加速器下一步所需的数据,而非它一秒前所需的数据。

由于 HBM 容量相对于模型和数据集尺寸极其有限,加速器以批次方式处理数据。每个批次占用一部分 HBM,持续一个前向-反向传播的时长,随后被丢弃以腾出空间给下一个。批次必须供给的速率,设定了所有下层的带宽需求。

批处理生命周期与存储层级溢出动态

批处理生命周期在存储层级中产生了向下的“溢出”动态。当仅权重无法装入单个 HBM 池时,系统只有两种基本选择:将活跃权重拆分到多个加速器,或将部分状态保留在 HBM 之外并在需要时取回。以 1750 亿参数模型为例,FP16 权重占用 350 GB,仅权重存储就需要至少五块 80 GB H100 加速器,且不留激活值或优化器状态的空间。Adam 优化器的 FP32 动量和方差还需额外 1.4 TB。将这一占用跨集群分片使得训练成为可能,但每个分片都会产生更多低层流量和更多协调开销。关键在于:HBM 稀缺性迫使层级的其余部分存在。

HBM 内存预算的极薄边际

极薄的边际是大模型训练的显著特征。表 4.3 展示了在单块拥有 80 GB HBM 的 H100 GPU 上训练 1750 亿参数模型的 HBM 内存预算。表格将分区视为存储布局:模型权重在节点内八向拆分,而优化器和梯度状态跨 256 个节点分片。目的不在于推导分区算法,而在于展示 HBM 面临的存储压力。

| 组件 | 每 GPU 大小 |

| --- | --- |

| 模型权重 (FP16, 8-way shard) | 43.75 GB |

| 优化器状态 (节点分片) | 0.7 GB |

| 激活值 (随序列长度变化) | 10–20 GB |

| 梯度缓冲区 (FP16, 节点分片) | 0.2 GB |

| 通信缓冲区 (NCCL) | 2–4 GB |

| 总计占用 | 56.6–68.6 GB |

表 4.3:1750 亿参数训练的 HBM 内存预算:将 1750 亿模型在节点内八向拆分、优化器和梯度状态跨 256 节点分片时的每 GPU 分配。优化器和梯度行是在两级分区后的每 GPU 分片。最终预算占用了 80 GB HBM 设备的大部分,留给传入数据管道的缓冲有限。从主机内存取下一个批次的任何延迟都可能使加速器“饿死”。

即使采用激进分区,总内存占用在十进制吉字节基础上也达 56.6–68.6 GB —— 按二进制容量计算约占 80 GB HBM 池的 65.9%–79.9%。扣除碎片和框架工作空间后,留给传入数据管道的缓冲极其有限。从主机内存取下一批次的任何延迟都会导致加速器空转,使系统中最昂贵的资源产生热量而非梯度。

HBM 内的批处理生命周期展示了该层级存储的瞬时性。当新训练批次经 PCIe 从主机 DRAM 到达时,它被置于 HBM 中预分配的输入缓冲区。前向传播读取输入数据、读取模型权重(跨批次持久化)并将激活值写入 HBM。反向传播读取激活值、计算梯度并写入梯度更新。优化器步骤读取梯度和模型权重、计算更新后的权重并写回。优化器步骤后,输入批次和激活值不再需要,其 HBM 区域被释放供下一批次使用。输入批次在 HBM 中的完整生命周期(从到达到释放)跨度为单个训练步骤:通常 100 到 500 毫秒。相比之下,模型权重和优化器状态在整个训练运行期间持久驻留在 HBM 中,占用固定分配且无法回收给批次数据。

从数据管道视角看,第 0 层不是待管理的存储层,而是待满足的约束。管道的目的是确保加速器下一个所需的批次在当前批次计算完成前已驻留在 HBM 中。到达晚了,加速器停顿;到达早了,会消耗本可存放激活值的 HBM。“刚好及时”与“刚好太晚”之间的张力定义了管道的缓冲管理策略,我们将在第 4.4 节(存储管道方程)中量化分析。

第 1 层:主机 DRAM

在 HBM 之下一级,主机 DRAM 充当数据管道的暂存区。到达加速器的每字节训练数据均先经过主机 DRAM(除非使用 GPUDirect Storage 绕过,见第 4.5 节)。典型训练节点包含 512 GB 至 2 TB 系统内存,由主机 CPU 及其外设共享。虽然主机 DRAM 与加速器间的带宽受限于 PCIe Gen 5(单向约 64 GB/s,双向 128 GB/s)或 NVLink(900 GB/s),主机 DRAM 在 ML 存储层级中扮演三大关键角色。

运行于主机 DRAM 的数据加载器管道遵循四个有序阶段:

  1. 读取:I/O 线程从 NVMe 或网络存储读取压缩数据到读取缓冲区。

  2. 解码:解码线程解压数据,如图像的 JPEG 解码或文本的解压。

  3. 增强:增强线程应用变换,包括图像的随机裁剪、翻转、归一化,或文本的分词与序列打包。

  4. 整理:整理阶段将单个样本组装成批次,并置于锁页内存以便高效 DMA 传输至加速器。

最终的放置至关重要,因为锁页内存(在 PyTorch 中通过 DataLoaderpin_memory=True 暴露)防止操作系统在 DMA 传输完成前换出缓冲区,允许 NVMe 控制器直接写入物理地址,PCIe DMA 引擎无需额外拷贝即可到达。各阶段并发运行,形成重叠 I/O、CPU 计算和数据传输的管道。该管道的效率决定主机 DRAM 能否跟上加速器的需求。

主机 DRAM 最关键的功能是作为预取缓冲区。CPU 数据加载器从低层(NVMe 或网络存储)读取数据,解码压缩格式(JPEG、gzip),应用增强(随机裁剪、翻转、色彩抖动),并在主机 DRAM 中组装张量。当加速器完成批次 i 处理时,批次 i + 1 应已在主机内存组装就绪,可传输至 HBM。该预取缓冲区的深度决定了管道能吸收多少 I/O 方差而不停顿。

推荐负载对主机 DRAM 提出不同需求:托管可能超过 100 GB 的嵌入表,远大于 HBM 容量。这些表驻留在主机 DRAM 中,通过查找仅获取当前批次所需行。主机 DRAM 与 HBM 间的带成为此类负载的关键瓶颈,这也是为什么某些系统使用带 RDMA 的 CPU 端 DRAM 跨网络服务嵌入查找。

主机 DRAM 还提供 CPU 管道所需的增强工作空间。数据增强操作(调整图像大小、文本分词、加噪)在 CPU 上执行,需要临时内存存放中间结果。对 256 张图像的批次应用五种增强(每图 150 KB)需每增强阶段数十 MB 工作空间。虽单批次占用不大,但多个数据加载工作进程并行运行时会累积。

基于 GPU 的增强方法权衡

一些增强流水线已从 CPU 迁移至 GPU 执行,使用如 NVIDIA 数据加载库(DALI)之类的库在加速器自身上执行图像解码和增强。此方法消除了 CPU 增强瓶颈,并降低了主机 DRAM 带宽需求,因为压缩数据(较小)被传输到 GPU 而不是解码数据(较大)。权衡在于,GPU 上的增强会消耗本可用于训练的 HBM 容量和计算周期。对于 compute-bound 工作负载(GPU 已被矩阵乘法饱和),GPU 基础的增强是适得其反的。对于 I/O-bound 工作负载(GPU 正在等待数据),通过将工作从瓶颈(CPU)转移到具有空闲容量的资源(GPU),它可以提高整体吞吐量。

内存竞争和容量规划

三者以微妙的方式相互作用。预取缓冲区和增强工作区竞争相同的物理 DRAM,而嵌入表会占用本可用作更深预取队列的容量。一个拥有 512 GB DRAM 并托管 200 GB 嵌入表的节点,仅剩 312 GB 用于预取和增强。如果数据加载器使用 8 个工作进程,每个维护 1 GB 的解码缓冲区,则有效预取容量会进一步下降。系统架构师必须通过分析每个流水线阶段的实际内存消耗并相应配置 DRAM 来平衡这些竞争需求。

NUMA 拓扑与性能

DRAM 的物理布局具有性能影响,在单插座基准测试中是不可见的,但在生产规模中至关重要。多插座服务器表现出非一致内存访问(NUMA)拓扑,其中每个 CPU 插座拥有可以以全带宽访问的“本地” DRAM,以及连接到另一个插座、带宽大约只有其一半的“远程” DRAM。在一个双插座 DGX 节点中,每个插座分配四个 GPU,运行在插座 0 上的数据加载器线程若将其预取缓冲区分配在插座 1 的 DRAM 上,则每次缓冲区访问将面临大约 2× 的带宽惩罚。解决方案是 NUMA-aware 分配:将每个数据加载器工作进程固定到与其服务的 GPU 处于相同的 CPU 插座(相同的 NUMA 域),并使用 numactllibnuma 确保内存分配保持本地。正确的 NUMA 绑定在双插座系统上可提高数据加载吞吐量 30–50%,这一提升在开发环境中是不可见的,但在每个利用率百分点的提升都能转化为每日数千美元时则至关重要。

带宽断崖与加速器饥饿

主机 DRAM 带宽与 HBM 带宽之间的差距是层次结构中的第一个主要断崖:使用本章节的 200 GB/s 主机-DRAM 参考值和 3.35 TB/s HBM 数值,大约为 16.8×。主机到 GPU 的互连则增加另一个断崖,实际传输带宽取决于节点使用的是 PCIe 还是 NVLink。未能从较低层级维持主机 DRAM 数据供给会立即导致加速器饥饿,因为加速器无法直接从 NVMe 或网络存储获取数据。在我们的运行示例中,256 节点集群要求每个节点的主机 DRAM 持续输出已解码、已增强的批次,以便通过 PCIe 传输。如果 NVMe-to-DRAM 读取管道落后甚至只有几百毫秒,预取缓冲区将耗尽,加速器将空闲直至下一个批次到达。

二级:本地 NVMe

当工作集超过主机 DRAM 时,系统将降至二级:直接连接到计算节点的本地 NVMe 驱动器。NVMe⁴⁵ 提供了专为固态驱动器设计的高性能协议,每个驱动器可实现 7 GB/s 的顺序吞吐量。四个驱动器以 RAID-0 配置时,单个节点可维持大约 28 GB/s 的顺序读取(未考虑开销),足以在理想顺序条件下,约 3.6 分钟内从本地磁盘流式传输 6 TB 序列化数据集。

在机器学习训练中,本地 NVMe 作为一个 温暖缓存,存储从分布式存储获取的数据分片。此设计使工作进程能够在多个 epoch 中重新读取样本,而无需重新穿越网络。对于 PB 级数据集的多 epoch 训练,每个 epoch 从对象存储重新获取数据的网络出口成本将是禁止性的(正如我们在第 4.6 节中所量化的那样)。在作业开始时从共享存储填充本地 NVMe,随后仅从本地读取,既消除了成本也消除了延迟。

温暖缓存模式需要仔细的容量规划。具备四个 7.68 TB NVMe 驱动器的训练节点可提供约 30.7 TB 的本地存储。对于我们的运行示例,6 TB 序列化数据集可舒适地容纳在单个节点的本地存储中,仍有空间用于检查点暂存和临时增强缓冲区。一个多模态训练作业,结合 20 TB 图像、10 TB 文本和 5 TB 音频,总计 35 TB,超出本地容量约 4.3 TB,迫使管道至少部分数据流从并行文件系统读取。设计权衡在于:是预留更多本地 NVMe(增加节点成本),还是接受依赖网络的读取(可能导致延迟尖峰)。

NVMe 的内部并行性是其相比传统存储吞吐量优势的关键。NVMe 规范支持多达 65,535 个 I/O 队列,每个队列最多可有 65,536 条未完成命令。一个拥有 32 个工作进程、每个发出异步读取的数据加载器,可保持 NVMe 控制器的内部管道饱和。相比之下,NVMe 所取代的遗留 AHCI 协议仅支持一个包含 32 条命令的队列,无论底层介质的能力如何,都会在协议层面限制并行性。这种架构差异解释了为什么 NVMe 即使在使用相同 NAND 闪存的情况下,也能提供比 SATA SSD 高 10–50 倍的吞吐量。

为顺序 I/O 设计的数据格式

本地 NVMe 上数据格式的选择对有效吞吐量具有显著影响。考虑存储 128 万张图像数据集的三种方法。

第一种方法将每张图像作为单独的 JPEG 文件存储在目录层次结构中。这种格式很自然地适用于数据收集(下载一张图像,保存一个文件),但在训练 I/O 方面是敌对的。每次 open() 系统调用在内核的虚拟文件系统(VFS)层中都有大约 10–50 μs 的固定开销。以每秒 8,000 张图像的速度,仅开销就占用了每秒 80–400 ms 的 CPU 时间。更糟的是,目录结构迫使文件系统为每个文件维护一个 inode,消耗了本可用于数据读取的元数据资源。

第二种方法将所有图像打包到少量的大型二进制文件中(如 HDF5、LMDB 或带索引文件的原始连接张量)。每个文件包含数千张连续存储的图像,并有一个单独的索引将样本 ID 映射到文件内的字节偏移。数据加载器寻找所需偏移并直接读取样本。这消除了每个文件的元数据开销,并使每个二进制文件内部实现顺序访问成为可能。其缺点是数据集不再具有人类可读性,且修改单个样本需要重写整个文件。

第三种方法使用由 WebDataset 推广的基于 tar 的归档格式。每个样本存储为一组相关文件(图像、标签、元数据)内的标准 POSIX tar 归档。tar 格式支持无需单独索引的顺序迭代,因为每个文件的头部包含其大小,使读者能够跳到下一个样本。此格式兼具单独文件的简洁性(每个样本自描述)和大型二进制文件的顺序 I/O 效率。tar 归档同样是有效的超文本传输协议(HTTP)字节范围目标,使其可直接从对象存储流式传输,无需本地暂存。

对于我们的运行示例,1.5 万亿 token 的数据集通常存储为 256 MB 到 4 GB 的二进制分片集合,每个分片包含一段连续的 token 化文本序列。数据加载器打开一个分片,将其顺序读入缓冲区,并在缓冲区内迭代 token。当缓冲区耗尽时,加载器打开下一个分片。如果 6 TB 的序列化语料库被分割成 3 GB 的分片,则每个 epoch 的 open() 总调用次数即为分片数量(约 2,000 个),而非 token 数量(数万亿)。元数据操作约减少了 750,000,000 倍,这使得从本地 NVMe 和远程存储进行流式传输在训练规模下成为可能。

在 NVMe 层面,压缩代表了 I/O 带宽与 CPU 周期之间的关键权衡。对于 I/O 受限的流水线(即 NVMe 驱动器无法满足加速器的需求),激进压缩是有益的:zstd 9 级压缩可实现约 4:1 的压缩比,但每 CPU 核心解压速度仅为 0.5 GB/s。对于 CPU 受限的流水线(解码和增强已饱和主机处理器),则倾向于轻量压缩:zstd 1 级压缩可提供约 3:1 的压缩比,但每核心解压速度高达 1.5 GB/s。在一个 28 GB/s 的四驱动器 NVMe RAID-0 阵列上,zstd-1 可提供 84 GB/s 的非压缩数据有效吞吐量,而 zstd-9 可提供 112 GB/s,但需要 3 倍的 CPU 核心专门用于解压。因此,最优压缩级别并非数据的属性,而是流水线瓶颈的属性,且当集群配置变更时(增加 GPU 会将瓶颈推向 I/O,从而倾向于更重的压缩),最优级别也会随之改变。

标准 ImageNet 训练流水线

标准的 ImageNet 训练流水线直观地展示了错误格式的代价。

问题:在 ImageNet(128 万张,平均约 150 KB)上进行 ResNet-50 训练作业,目标吞吐为 1000 张图片/秒。问题是应在 HDD 还是 NVMe 上使用单独的 JPEG 文件。

计算

  1. 原始带宽:1000 张/秒 × 150 KB = 150 MB/s。

  2. HDD 现实:一块 7200 RPM 的硬盘驱动器(HDD)在随机读取下可提供 100 IOPS。维持 1000 张/秒所需的 IOPS 是磁盘提供能力的 10 倍。

  3. 结果:在 HDD 上随机读取单个文件会使 GPU“挨饿”,在解码、寻道和文件系统开销之前,利用率被限制在约 10%。

系统洞见:流水线必须使用 NVMe(随机访问约 100 μs,比对象存储快三个数量级),或将数据集转换为顺序格式(TFRecord/WebDataset)以实现顺序吞吐量。

I/O 墙与本地检查点暂存

该层级的核心挑战是 I/O 墙:单块 NVMe 驱动器比 HBM 慢约 500 倍,即使是四驱动器 RAID-0 条带化,也仍慢约 120 倍。弥合这一差距需要流水线技术(I/O 与计算重叠,详见 4.4 节),以及日益重要的 GPUDirect Storage(详见 4.5 节)以完全绕过 CPU 开销。该层级的 I/O 墙具有隐蔽性,因为 NVMe 性能按历史标准来看非常优秀。习惯了 HDD 时代 100 MB/s 吞吐量的存储工程师,可能会认为本地 RAID-0 条带提供的 28 GB/s 非常充裕。然而,相对于加速器的“食量”,28 GB/s 只是涓涓细流。弥合差距的唯一途径是将存储读取与计算彻底重叠,使加速器完全感知不到存储延迟的存在。

本地 NVMe 也是本地检查点暂存的主要层级。保存模型检查点时,最快的策略是以全带宽写入本地 NVMe(最小化训练流水线暂停时间),然后异步复制到共享存储以保证持久性,因为每节点分片直接写入竞争激烈的并行文件系统所需时间,比写入本地驱动器要长几倍。最优检查点频率取决于集群故障率和检查点写入时间,但存储侧的设计目标已很明确:通过分级暂存以全带宽写入本地 NVMe,再在后台复制到共享存储,从而最小化 T_write。4.7 节将详细推算这些时序。

利用本地缓存恢复利用率

同一本地 NVMe 层级还通过恢复因远程读取延迟而损失的利用率,证明了其价值。

问题:某视觉模型训练流水线每步耗时 800 ms。由于网络拥塞,从共享并行文件系统获取数据会增加 150 ms 的 I/O 等待。在每个节点添加本地 NVMe SSD 能在多大程度上提升 GPU 利用率?

计算:GPU 利用率(η_hw)是步耗时中用于计算的时间占比。

  1. 仅远程:800 ms / (800 ms + 150 ms) ≈ 84.2%。

  2. 本地缓存:利用预取至本地 NVMe 将暴露的 I/O 等待降至接近零。

    • 新利用率:800 ms / (800 ms + 10 ms) ≈ 98.8%。

系统洞见:本地存储是 GPU 利用率乘数器。在此场景中,向一台价值 30,000 美元的 GPU 节点添加一块 500 美元的 NVMe 驱动器,恢复了先前因 I/O 等待而浪费的 14.6% GPU 算力。在一个拥有 1,024 张 GPU 的集群中,这种利用率提升相当于在不购买更多加速器的情况下,增加了 149 张 GPU 价值的有效工作量。在 ML 基础设施中,本地 NVMe 并非辅助组件;它是将昂贵算力与不可预测的共享存储解耦的物理缓冲区。

SSD 耐久性

该层级的一个实际顾虑是 SSD 耐久性。NAND 闪存在单元退化前只能承受有限的写擦周期。企业级 NVMe 驱动器的额定值为在 5 年寿命期内每天 1 到 3 次盘写入(DWPD)。对于一块 7.68 TB、1 DWPD 的驱动器,这意味着每天可吸收 7.68 TB 的写入量,终生约 14 PB。ML 训练工作负载以读为主(数据集写一次、读多次),这对 SSD 耐久性有利。然而,检查点写入可能很密集:如果每个节点每 10 分钟保存一个 4 GB 的检查点分片,那就是每天 576 GB 的检查点写入量,完全在 1 DWPD 的预算之内。风险出现在当本地 NVMe 同时用作检查点写入和数据集缓存的暂存缓冲区时:初始数据集暂存加上重复的检查点保存的总写入量,必须保持在驱动器的耐久额定值范围内。

本地 NVMe 在单节点内提供高带宽和低延迟,但分布式训练要求每个节点访问相同的数据集并看到相同的检查点。这种共享命名空间的需求无法仅靠节点本地存储满足,这引出了层级体系中的下一层。

第三层:并行文件系统

超越单节点,工作负载需要一个共享命名空间,使所有工作节点能访问相同的数据集,并使持久化检查点全局可见。这就是并行文件系统(PFS)的角色。⁴⁶

并行文件系统(PFS)

并行文件系统(PFS) 是面向 ML 训练集群的分布式存储架构,它将数据跨许多存储服务器条带化,以提供超越任何单一设备容量的聚合吞吐量。

  1. 意义:PFS 随存储服务器(对象存储服务器,OSS)数量线性聚合 BW_io。一个拥有 20 个 OSS 节点、每个节点提供 10 GB/s 的 Lustre 集群,可提供 200 GB/s 的聚合带宽,而单个 NAS 服务器上限为 10 GB/s。这使得训练作业加载一个 10 GB 的条带化分片仅需约 50 ms,而非 1 秒。这种聚合带宽直接减少了铁律中的 D_vol/BW 项。

区别

与网络附加存储(NAS)不同,后者的每个 I/O 请求都通过单台服务器路由,并行文件系统(PFS)客户端从专用的元数据服务器(MDS)接收条带位置元数据,然后直接从多个 OSS 节点并行读取数据:MDS 和 OSS 路径在架构上是分离的,因此数据带宽随 OSS 数量扩展,而元数据操作随 MDS 数量扩展。

常见误区

一个常见的误解是,只要添加足够多的 OSS 节点,PFS 就拥有无限吞吐量。实际上,Lustre MDS 每秒大约处理 100,000–300,000 次元数据操作;在 10,000 个工作进程各自打开一个小文件的情况下,MDS 不到 1 秒就饱和了,成为序列化点,导致整个集群空转,无论存在多少 OSS 节点。

Lustre 风格并行文件系统的架构将传统文件系统共同处理的两个关注点分离开来(Schwan 2003)。元数据服务器(MDS)管理命名空间:文件创建、目录列表、权限检查和锁管理。对象存储服务器(OSS)⁴⁷管理实际的数据块,每个 OSS 服务于每个大文件的一条条带。当客户端打开一个 10 GB 的训练分片时,MDS 告诉客户端哪些 OSS 节点持有哪些条带,客户端随即从所有这些节点并行读取。一个拥有 100 个 OSS 节点、每个节点提供 10 GB/s 的 Lustre⁴⁸ 部署,可提供高达 1 TB/s 的聚合带宽。

这种抽象架构对性能调优有具体的影响。当训练作业在 Lustre 上创建新的数据集目录时,管理员会为该目录配置 stripe count(文件分布所跨越的 OSS 节点数量)和 stripe size(写入每个 OSS 的块大小)(Lustre Wiki 2017)。对于大型顺序读取的训练分片,管理员可能会选择更宽的条带化和兆字节级的条带大小,以便单个文件能从多个 OSS 节点汲取带宽;Lustre 的条带化指南将宽条带化视为对超大文件或由许多客户端访问的文件有用,而较小的文件通常使用较少的条带以减少开销(Lustre Wiki 2019)。一个 4 GB 的数据分片以 4 MB 的条带大小跨 100 个 OSS 节点条带化,每个 OSS 大约放置 10 个条带。当数据加载器读取大的连续范围时,Lustre 客户端可向持有这些条带的 OSS 节点发出并行读取请求,聚合其带宽。客户端还维护自己的预读缓冲区,在应用程序处理当前条带时预取接下来的几个条带。这种文件系统级的预读不同于数据加载器的应用级预取缓冲区;这两层预取叠加,提供了深度的延迟隐藏,使得到 OSS 节点的物理距离对训练过程几乎透明。

元数据路径与数据路径的分离,正是实现 ML 工作负载所需聚合带宽的关键。在 GPFS⁴⁹(IBM Spectrum Scale)中,架构采取了不同方法:它将数据和元数据条带化到共享磁盘上,并通过基于令牌的分布式锁协调并发访问(Schmuck and Haskin 2002)。其结果并非简单的 MDS/OSS 分离;而是一种共享磁盘设计,其元数据和数据布局可分布,而锁协议保证跨客户端的一致性。代价是锁管理的复杂性更高,必须扩展到数千个节点。

这种分离造成了一个关键瓶颈:小文件问题。如果数据集由数百万个 10 KB 的图像组成,且作为单独的文件存储,元数据负载会在数据链路饱和之前很久就淹没 MDS。每个 open() 系统调用都需要一次元数据查找、锁获取和属性获取。当 10,000 个工作进程同时对不同文件调用 open() 时,MDS 成为序列化点。存储系统的吞吐量降至 MDS 处理元数据操作的速率,通常为每秒几十万次,远低于数据路径所能提供的速率。

小文件问题

小文件问题 是一种 ML 数据加载病理现象,即数百万个单独的小文件淹没存储系统的元数据服务器。

意义

它将有效 I/O 带宽(BW[io])降至理论额定值的一小部分,因为每个文件都需要自己的元数据操作(openstatclose)。当 10,000 个工作进程同时访问小文件时,元数据服务器成为序列化点,导致整个集群空转。

区别

与批量数据吞吐量(衡量比特率)不同,小文件问题是一个元数据延迟(L[lat])问题:瓶颈在于请求的频率,而非数据的大小。

常见误区

一个常见的误解是,购买更快的 SSD 就能“修复”此问题。实际上,这是一个格式问题:解决方法是将样本打包进大型顺序容器(例如 TFRecord、Parquet),以在数千个样本间摊销元数据操作。

在生产规模下,元数据瓶颈可能拖垮一个原本高带宽的文件系统。

场景:一个大规模训练集群从预处理后的顺序数据集迁移到作为单独文件存储的原始图像。

设置:并行文件系统预置了 500 GB/s 的聚合数据带宽,但实际交付的吞吐量不足额定值的 1%,因为元数据服务器无法维持 ML 数据加载器产生的每秒数百万次 stat()open() 调用。将 2 亿张图像的数据集转换为 50,000 个大型 tar 文件,将元数据操作减少了四个数量级。

系统教训:在大规模下,即使存储结构拥有充足的数据带宽,元数据操作也可能成为第一个瓶颈。

针对小文件问题的设计应对是条带化聚合。条带化将单个大文件分布到多个 OSS 节点,使客户端能并行地从所有节点读取。条带大小(通常每个 OSS 1 到 4 MB)决定了粒度:一个 4 GB 文件以 1 MB 条带大小跨 100 个 OSS 节点条带化,每个节点放置 40 MB,顺序读取可同时饱和所有 100 条数据路径。条带数(参与的 OSS 节点数量)可按文件或按目录配置,允许管理员针对不同访问模式调优带宽。训练数据分片受益于最大条带化;小型配置文件受益于最小条带化,以避免跨多个节点协调的开销。

条带大小与工作负载访问模式的交互决定了实际吞吐量。如果训练数据加载器顺序读取一个分片,PFS 客户端从 OSS-A 读取条带 1,从 OSS-B 读取条带 2,从 OSS-C 读取条带 3,依此类推,自然地将负载分布到持有该文件条带的所有 OSS 节点。如果读取大小小于条带大小,每次读取由单个 OSS 节点服务,客户端无法从并行读取中获益。如果读取大小跨越多个条带,客户端同时向多个 OSS 节点发出并行读取。对于读取兆字节级块(整批图像,或 4 MB 的标记化文本块)的 ML 工作负载,读取大小通常超过条带大小,从而实现全带宽聚合。

聚合通过减少 MDS 必须跟踪的文件数量来补充条带化。小样本被捆绑进大型顺序文件(如 WebDataset⁵⁰ tar 归档或 TFRecord 序列),以在数千个样本间摊销元数据成本。对一个 4 GB tar 文件的单次 open() 即可访问 40,000 个样本,与单独文件相比,元数据负载降低了 40,000 倍。

并行文件系统和存储层次结构

一致性和协调

并行文件系统提供了可靠检查点所需的协调共享命名空间和锁管理(Schmuck and Haskin 2002)。当检查点保存完成时,随后读取该检查点的每个节点都必须看到相同的已完成状态。这一保证对于故障恢复至关重要:如果一个节点发生故障并重新启动,它必须读取幸存节点写入的检查点,而不是一个部分刷新的版本。在规模上实现这种一致性需要 MDS 锁管理器和分布式 OSS 写入路径之间的仔细协调,这也是检查点写入比训练读取更昂贵的原因之一。

一致性模型对机器学习工作负载有细微影响。对于训练数据读取,强一致性不是必需的,因为数据是不可变的:一旦数据集被预处理并上传到并行文件系统,它就永不会被修改。多个工作进程同时读取同一个分片时可以不加锁地进行,因为没有并发写入者来制造冲突。这种只读访问模式使 PFS 能够以接近理论带宽的速度提供训练数据。相比之下,检查点写入需要独占锁以防止在写入过程中出现部分读取,而锁的获取和释放会为每个检查点操作增加延迟。一些系统通过将检查点写入新文件而不是覆盖之前的文件来缓解这一问题,以存储空间换取减少的锁竞争。

负载均衡和吞吐量

并行文件系统的性能不仅取决于 OSS 节点的数量,还取决于它们之间负载的平衡。如果一个训练作业读取一个跨越 10 个 OSS 节点条带化的大型分片,而剩余的 90 个节点处于空闲状态,则该作业只能达到系统总带宽的 10%。相反,如果 100 个训练作业各自读取跨越全部 100 个 OSS 节点条带化的不同分片,则总带宽将接近理论最大值。因此,在集群中调度数据访问模式是一个系统级优化机会,能够显著影响实际吞吐量。

尾延迟和预防性请求

在数千个节点的规模下,尾延迟(Dean and Barroso 2013)主导系统性能。其数学原理令人深思。如果一个训练步骤需要来自 100 个存储服务器的数据,而每个服务器因后台维护、垃圾回收或网络抖动而有 1% 的概率变慢,那么所有 100 个服务器快速响应的概率为 0.99¹⁰⁰ = 0.366,这意味着超过 63% 的训练步骤将经历至少一次慢响应。最慢的响应决定了步骤的完成时间,因为数据并行训练要求所有工作进程在集体通信阶段开始之前完成他们的批处理。

系统通过预防性请求来缓解尾延迟:在等待一个可配置的超时时间(通常是中间响应时间)后,客户端会向同一数据条带的不同副本发送冗余读取请求。首先到达的响应被使用;第二个被丢弃。如果存储系统将每个条带复制到两个 OSS 节点,那么两个副本都变慢的概率为 0.01² = 0.0001,这样慢速步骤的比例从 63% 降低到不到 1%。冗余请求的成本——一次额外的网络读取——相比于等待期间消耗数百瓦功率的空闲加速器的成本,可以忽略不计。

预防性请求的有效性取决于复制的数据布局。如果两个副本位于共享同一网络交换机的 OSS 节点上,则交换机故障会导致两个副本同时不可用。有效的预防需要故障域感知放置:副本应位于不同的机架,通过不同的顶部机架交换机连接,以确保任何单个组件的故障最多只影响一个副本。

共享基础设施挑战

噪 y 邻居问题

在生产集群中,并行文件系统是一个共享实用程序,同时为具有不同 I/O 模式的数十个并发训练作业提供服务。在任何给定时刻,一个大型语言模型作业可能正在顺序流式传输文本分片,一个计算机视觉作业可能正在读取随机图像分片,而来自第三个作业的检查点风暴可能正在饱和写入带宽。这种并发导致了噪 y 邻居问题:一个作业的 I/O 模式会严重降低所有其他作业的性能。一个消耗 PFS 内部网络带宽不成比例份额的检查点保存会导致其他作业的读取延迟激增,可能使它们的训练管道停滞。PFS 管理员通过 I/O 调度政策(服务质量等级、每个作业的带宽配额)来缓解这一问题,但这些政策增加了管理复杂性,并可能降低单个作业的峰值吞吐量。

命名空间隔离

一个相关的挑战是命名空间隔离。不同的团队和工作负载需要不同的存储配置。一个在数百万张 100 KB 图像上进行训练的团队需要一个具有高条带数量和小条带大小的目录,以最大化元数据性能。一个处理大型视频文件的团队需要较少且较大的条带,以优化顺序带宽。一个团队的工作负载条带化配置错误可能会创建热点,从而降低整个共享文件系统的性能。共享的影响很容易量化:一个具有 1 TB/s 总带宽的 PFS,如果被 10 个并发训练作业共享,则平均每个作业仅能获得 100 GB/s。如果一个作业的检查点风暴在 10 秒内消耗了 400 GB/s,那么剩余的九个作业将共享 600 GB/s,相比平均值减少了 33%,这可能触发它们加速器管道中的数据停滞。

容量规划示例

考虑一个训练集群,包含 512 个节点,每个节点运行 8 个 GPU。训练作业需要 400 GB/s 的总读取带宽。

第 4 层:对象存储

并行文件系统解决了活跃训练集群的共享命名空间问题,但其每字节成本对于组织的完整数据 holdings 来说过高。ML 组织多年积累的 PB 级原始训练数据、历史检查点和预处理数据集需要一个针对容量和耐久性而非带宽优化的层级。

在 PB 级规模下,对象存储提供了层次结构中耐用且成本效益高的基础。诸如 Amazon S3 和 Google Cloud Storage 之类的服务通过对象命名空间组织数据,其中对象通过键而非 POSIX 目录路径检索(Amazon Web Services 2026c)。与按层次目录组织数据并具备丰富元数据操作(重命名、链接、权限继承)的文件系统不同,对象存储将每个对象视为由唯一键标识的不透明 blob。这种简化消除了在规模上困扰并行文件系统的元数据复杂性,并启用了水平扩展到非常大的容量池。

对象存储通过删除编码[⁵²]和地理复制实现“十一个九”[⁵¹]的耐久性(99.999999999%),这使得对象存储成为训练数据和检查点归档的自然事实真相层,同时仍然需要版本控制、生命周期策略、访问控制和恢复测试来处理存储服务耐久性模型之外的运营故障。

删除编码

删除编码 是一种用于 ML 数据集和检查点归档的存储保护方案,它将一个对象分割为 k 个数据块和 m 个奇偶校验块,确保原始对象可以从任意 k 个剩余片段中恢复。

1. 重要性

它以远低于复制的存储开销实现了极端的数据持久性(例如,“十一个 9”)。例如,4+2 码存储四个数据片段和两个校验片段,因此它能在 1.5× 的存储开销下容忍两个不可用片段,而三副本复制的开销则是

2. 区别

与复制(存储完整副本)不同,擦除编码使用数学编码(例如,Reed-Solomon)将冗余信息分布在不同的故障域(磁盘、机架或站点)中。

3. 常见误区

一个常见的误解是擦除编码是“免费的持久性”。实际上,它是一种延迟-计算权衡:故障后重建数据需要额外的 CPU 周期,并增加了存储读取的尾部延迟(L_lat)。

擦除编码背后的工程原理阐释了存储系统中的一个反复出现的主题:通过对应用程序不可见的冗余来实现极端持久性。当 ML 训练作业从对象存储读取一个分片时,存储服务会透明地从可用节点读取 k 个片段,重建原始分片,并将其返回给客户端。如果某个片段位于故障磁盘上,服务会读取额外的编码片段并重建丢失的数据。客户端从未观察到故障,但存储系统付出了额外的 I/O、网络和解码工作,这导致对象存储的尾部延迟高于本地 NVMe。

对象存储的成本优势显著:按代表性的 $0.02/GB/月 计算,其每字节成本比 HBM 便宜 750×,比本地 NVMe 便宜 。对于 100 TB 的训练数据集,按该价格假设,对象存储每年成本约 $24,000,而本地 NVMe 约为 $120,000。这一成本优势使对象存储成为组织训练数据湖的常见归宿,原始数据、预处理后的数据集和归档的模型工件均可持久驻留于此。

延时劣势同样显著。对象存储的 GET 请求可能产生数十毫秒的延迟;本章设计示例使用 50 到 100 ms 作为代表性规划范围,而非通用的服务保证。该延迟是提供持久性的多层间接寻址的结果。当客户端请求一个对象时,存储服务必须在分布式元数据索引中查找对象位置,确定要读取哪些擦除编码片段,从潜在的不同存储节点检索片段,并重建原始对象。对于大对象(传输时间占主导),此开销不可见;但对于小对象(单次请求延迟占主导),则是灾难性的。

解决方案是并行文件系统中使用的相同聚合模式:样本使用 WebDataset (Aizman et al. 2019) 或 Mosaic Streaming (Databricks 2026) 等格式打包成多 GB 的分片,将高延迟随机访问转化为高带宽顺序流式传输。数据加载器不再发起数十亿次小的 GET 请求,而是并行发起数百次大读取,每次获取一个包含数千个样本的完整分片。在足够的并行度和同地网络容量下,对象存储可维持高聚合吞吐量,足以供养某些工作负载的大规模训练集群。流式模式通过为每个数据加载器工作进程分配一组分片并让其顺序读取来工作。在每个分片内,样本连续存储,因此工作进程按顺序解码它们并将其在本地缓冲区中打乱。这种方法在保持存储层顺序 I/O 的同时,实现了跨数据集的伪随机访问。

面向对象存储的流式数据加载器架构在重要方面区别于本地存储加载器。本地存储加载器可使用内存映射 I/O 将数据集视为虚拟内存区域,依赖操作系统的缺页机制按需加载数据。这种方法虽优雅,但与对象存储不兼容,因为对象存储不支持内存映射所需的 POSIX 文件系统接口。取而代之,对象存储加载器必须显式管理 HTTP 连接,为分片内的特定字节范围发起范围 GET 请求,处理瞬态故障重试,并管理用于已解码样本的本地缓冲池。WebDataset 和 Mosaic Streaming 等库封装了这种复杂性,向训练循环提供简单的迭代器接口,同时在内部管理 HTTP 传输和缓冲。

将这种流式架构应用于我们的运行示例,使设计更加具体。分词后的训练集存储在类 S3 的对象存储中,约为 2,0003 GB 的分片。流式架构将这 256 个计算节点中的每一个分配一组唯一的约 8 个分片。在每个节点内,8 个数据加载器工作进程发起并发 HTTP Range 请求,从其分配的分片中获取 256 MB 的数据块。跨集群来看,这产生了 2,048 个并发请求,当对象存储和网络路径按此配置时,可实现 50–100 GB/s 范围内的聚合吞吐量,轻松满足文本模型训练所需的约 167.8 MB/s。在此模式下,瓶颈往往从对象存储的带宽转移到计算 VPC 与对象存储端点之间的网络带宽。将计算和存储同地部署在同一云可用区的组织可以降低跨网络延迟,直接减少所需的预取缓冲区深度,并提高整体流水线效率。

对象存储暴露了大规模请求并行性。每个 GET 请求都是独立的,可由云提供商基础设施内的不同存储节点服务。一个拥有 1,024 个节点的训练集群,每个节点运行 8 个数据加载器工作进程发起并发请求,会产生 8,192 个同时进行的 GET 请求。如果每个请求获取一个 256 MB 的分片,聚合吞吐量取决于云提供商的网络容量、内部带宽和请求速率配额。实际限制通常是计算集群与对象存储服务之间的网络带宽,而非存储服务的内部介质带宽。

许多对象存储服务提供强读后写一致性,这意味着一旦检查点写入对象存储,后续读取即可看到完整数据。情况并非一直如此:直到 2020 年 12 月,Amazon S3 对覆盖 PUT 和 DELETE 操作仅提供最终一致性,这意味着在更新或删除现有对象后立即读取可能返回陈旧数据。新对象的 PUT 始终是强一致的,但更新“最新”指针或覆盖现有工件的检查点工作流易受攻击:从故障中恢复的节点可能读取陈旧的检查点并破坏训练状态。Amazon S3 于 2020 年 12 月对所有操作过渡到强读后写一致性,填补了 S3 风格检查点保留的这一空白 (Amazon Web Services 2020),而更广泛的设计教训是验证流水线所用对象存储的一致性语义。这一一致性属性使对象存储既适合作为训练数据源(原始数据集),也适合作为长期检查点保留(本地 NVMe 暂存后的归档副本),前提是流水线验证其所依赖的服务语义。

海量 ML 数据集使其极易遭受静默数据损坏:存储介质或网络传输过程中发生的微妙位翻转,可能引入几乎无法诊断的训练伪影。对象存储服务提供按对象校验和,可在上传、复制或下载时验证完整性;例如 Amazon S3 记录了包括 CRC 系列和加密哈希在内的校验和选项(Amazon Web Services 2026a)。并行文件系统通常依赖底层 RAID 或擦除编码来保证物理完整性,但可能不提供应用可见的端到端校验和。稳健的大规模 ML 流水线会在预处理期间为每个数据分片计算并存储校验和,在训练工作进程首次加载分片时进行验证。与在损坏数据上训练的成本相比,验证开销微不足道。对于模型检查点,完整性验证更为关键:权重或优化器状态中的单个位翻转就可能破坏模型,导致恢复后训练静默发散。框架和检查点库可包含完整性验证,但流水线不应在未测试检查点格式的情况下假设其存在。

存储安全是生产 ML 基础设施的头等大事。训练数据集通常包含敏感信息(个人身份数据、专有文本、授权图像),需要访问控制。对象存储通过身份与访问管理(IAM)角色、按存储桶策略以及静态和传输中数据的集成加密,提供细粒度安全策略。并行文件系统提供 POSIX 权限和访问控制列表(ACL),但通常缺乏合规所需的审计日志。对于受《通用数据保护条例》(GDPR)和《加州消费者隐私法案》(CCPA)等数据保护法规约束的组织,存储架构必须确保记录数据访问、满足删除请求(因不可变预处理分片而变得复杂),并确保模型检查点不会无意中记忆受保护信息。因此存储设计必须同时保留吞吐量和治理元数据。

从 Tier 4 到 Tier 5 的过渡主要由经济因素驱动:当检索延迟和治理要求允许时,很少访问的数据应移至归档存储,因为存储成本可降低一个数量级。

Tier 5:归档与冷存储

最后一层为合规和可审计性提供长期保存。归档服务(如 S3 Glacier)专为极少访问的数据设计:往年训练日志、被取代的模型检查点、法规合规审计追踪。按每 GB 每月 $0.004 的代表性价格计算,100 TB 归档存储每年约 $4,800,比上文假设的对象存储 $0.02/GB/月便宜 5 倍。代价是检索延迟以分钟到小时计,使该层不适合运营访问。

归档存储的主要价值在于组织记忆。当两年前训练的模型在生产环境中表现出意外行为时,能否恢复确切的训练数据、超参数和中间检查点以进行取证分析,取决于归档存储。为节省成本而激进清理旧数据的组织,往往在最糟糕的时刻发现,自己无法复现或解释已部署模型的行为。

自动化生命周期策略根据访问近期性管理层间数据迁移。举例说明此类自动化的表达方式,基于时间的 ML 基础设施规则可能按如下方式迁移检查点:最近三个检查点保留在并行文件系统以便快速恢复;超过 72 小时的检查点迁移至对象存储;超过 90 天的检查点迁移至归档存储;超过两年的检查点删除,除非标记为法规保留。S3 生命周期规则为对象生命周期提供相同的通用迁移/归档/删除机制(Amazon Web Services 2026b)。后文的运行示例采用该思路的基于计数的变体(第 4.7 节),迁移规则依据检查点数量而非时间。没有自动化,团队要么在昂贵层囤积数据(浪费预算),要么过度删除数据(丧失可复现性)。生命周期策略以声明式规则编码组织的成本-恢复权衡,无需人工干预。

以 175B 模型训练运行为例,归档层保存完整训练谱系:原始 3 TB 压缩源语料库、6 TB 序列化分片集、所有预处理脚本、最终检查点,以及抽样的中间检查点子集。若组织为保留谱系预算保守 100 TB,两年归档成本约 $9,600(100 TB × $0.004/GB/月 × 24 个月),远低于训练集群一天的 GPU 时间。归档成本与计算成本的巨大不对称,使保留成本相对于其记录的运行极为低廉;归档失败通常反映生命周期所有权缺失、保留策略不清或检索流程疏忽,而非原始容量成本。

归档存储服务提供多种检索速度层,各有不同定价。以 S3 Glacier 风格服务为例,标准检索可能在数小时内以中等成本完成,加急检索可能在数分钟内以较高成本完成,批量检索则以更低检索价格换取更高延迟(Amazon Web Services 2026d)。检索层选择取决于用例紧迫性。生产事件取证分析需要更快检索;年度合规审计可使用较慢的批量检索。设计生命周期策略需要了解存储成本、各类数据的预期检索频率和紧迫性。

合规要求可能强制组织为已部署模型保留训练数据溯源记录。例如《欧盟 AI 法案》要求高风险 AI 系统提供文档,包括训练数据构成和预处理信息。在要求的保留期内,归档存储通常是保留模型从原始数据到预处理再到最终检查点的完整谱系的经济可行层。合规存储成本通常远低于文档不足导致的运营和法律成本。

数据格式全景

数据格式选择决定存储层级是否能发挥理论带宽,还是在元数据开销下崩溃。同一多 TB 数据集用不同格式存储,在相同硬件上可产生 100 倍以上的 I/O 吞吐量差异。大规模机器学习要求数据结构化以实现高吞吐顺序访问,而非为了人类可读性或事务便利性。

数据格式与存储访问模式

核心的设计问题是工作负载必须承担哪种访问成本:逐文件元数据、不必要的特征读取,还是顺序分片遍历。

行式格式

行式格式(Row-oriented formats),如 CSV、JavaScript Object Notation (JSON) Lines,或单独图像文件目录(JPEG、PNG),最直观且人类可读。它们易于检查和调试,适合初始数据探索或完全放入内存的小型数据集(通常 10 GB 以下)。但在大规模训练中,它们对 I/O 性能是灾难性的。每个样本是一个单独的文件,需要独立的文件系统操作(openstatreadclose),产生显著的逐样本开销,往往主导实际读取数据的时间。

列式格式

第二种选择解决不同的存储问题:列式格式(columnar formats)如 Apache Parquet 和 Apache Arrow,当表格特征可在传输前选择时,避免不必要的读取。它们按列而非行组织数据,因为列内数据通常类型相同、熵较低,从而实现极佳压缩。这种结构还支持谓词下推(predicate pushdown),存储引擎仅读取查询所需的特定列(特征),避免不必要的 I/O。这些格式适合表格 ML 任务和特征工程,但对图像、音频等非结构化数据不那么自然,后者被视为原子二进制大对象。

顺序流式格式

第三种选择直接针对训练访问模式:顺序流式格式(sequential streaming formats)将逐样本元数据降低到打开分片的成本。TFRecord、Mosaic 的 StreamingDataset 格式(Databricks 2026)和 WebDataset(基于标准 tar 归档)将许多样本分组为大型连续二进制文件,称为分片。这将打开文件的成本分摊到数千个样本上。数据加载器打开单个分片并顺序流式传输其内容,最大化带宽。基于 tar 的格式还支持 HTTP 字节范围请求,使其成为直接从对象存储流式传输数据而无需先下载整个文件的理想选择。

流式格式内部的一个关键权衡是索引的有无。TFRecord 和 HDF5 常用单独的索引文件记录分片内每个样本的字节偏移。这支持分片内的高效随机访问,但引入元数据依赖:索引必须读取并驻留内存。WebDataset 无索引:tar 归档中每条记录包含自身的头部指定大小,允许纯顺序流式传输,无需外部元数据。代价是无索引格式无法高效定位到分片内任意样本,这对课程学习或主动学习等需要非顺序数据访问的工作负载很重要。

压缩集成在不同格式间也有差异。Parquet 在列级集成压缩,利用类型同质性实现高压缩比。WebDataset 和 TFRecord 通常将压缩留给单个样本:图像的 WebDataset 存储已压缩的 JPEG 文件,文本数据集可能对每个分片应用 zstd 压缩。这种选择影响 NVMe 读取背景下讨论的 CPU-I/O 权衡,因为逐样本解压必须由主机 CPU 处理。

对于我们的运行示例,1.5 万亿 token 数据集存储为约 2,000 个分片,每个约 3 GB,采用自定义二进子格式。每个分片包含一个小头部(token 计数、词表大小、字节序),随后是连续的 4 字节 token ID 数组。当暂存到 POSIX 存储(如本地 NVMe 或并行文件系统)时,数据加载器内存映射分片并按位置直接索引 token 数组,以零逐样本元数据开销实现近峰值读取带宽。当从对象存储直接读取相同分片时,加载器使用 HTTP 范围 GET 流式传输和显式本地缓冲区,而非内存映射。表 4.4 按存储路径施加的访问成本比较格式选择。

表 4.4:ML 训练数据格式比较

| 格式 | 逐样本开销 | 随机访问 | 流式传输 | 压缩 |

| --- | --- | --- | --- | --- |

| 单独文件 | ~50 μs (open/stat/close) | 是 | 差 | 逐文件 |

| Parquet | ~1 μs (行组查找) | 是 (行组) | 良好 | 列级 |

| TFRecord | ~0.1 μs (索引查找) | 有索引时 | 优秀 | 逐记录 |

| WebDataset (tar) | ~0 (顺序) | 否 | 优秀 | 逐样本 |

| 原始二进制 (tokens) | 0 | 是 (字节偏移) | 优秀 | 无需 |

表 4.4:ML 训练数据格式比较:逐样本开销决定存储硬件带宽是被实现还是浪费在元数据操作上。

表 4.4 的比较将管道方程中的数据量项量化,将格式选择转化为带宽规模问题。

数据管道方程

每个训练作业最终都会提出具体的规模问题:管道必须提供多少存储带宽,加速器才不会停顿?所需速率取决于训练配置,任一方向出错都代价高昂。存储带宽不足会饿死加速器;过度配置则浪费预算在部分闲置的存储容量上。数据管道吞吐量方程(公式 4.1)量化管道必须维持的带宽,使加速器永不等待数据:

\[\\text{BW}\_{\\text{required}} = N\_{\\text{GPU}} \\times \\eta\_{\\text{target}} \\times \\frac{D\_{\\text{vol,batch}}}{T\_{\\text{iteration}}} \\qquad(4.1) \]

其中 \(N\_{\\text{GPU}}\) 为加速器数量,\(\\eta\_{\\text{target}}\) 为目标利用率(通常 0.8 到 0.95),\(D\_{\\text{vol,batch}}\) 为单个本地批次大小(字节),\(T\_{\\text{iteration}}\) 为一次前向-反向传播耗时。

该方程揭示控制存储需求的四个杠杆:

  • 加速器数量:减少 \(N\_{\\text{GPU}}\) 直接降低带宽需求,但更少加速器也降低训练吞吐量。

  • 目标利用率:降低 \(\\eta\_{\\text{target}}\) 减少所需带宽,但接受更低利用率浪费昂贵硬件。

  • 批次数据量:减少 \(D\_{\\text{vol,batch}}\) 降低每次迭代数据量,但更小批次可能损害模型收敛。

  • 迭代时间:增加 \(T\_{\\text{iteration}}\) 给存储更多时间交付每个批次,但更慢迭代延长总训练时间。

实践中,这些杠杆无一免费;管道必须按训练配置要求的带宽进行工程化。

作为单独的图像训练示例,考虑 256 GPU 使用 ImageNet 规模图像、200 ms 迭代、80% 目标利用率的训练。该配置需要 39.3 GB/s 聚合存储吞吐量。带宽必须在整个训练期间持续维持,可能持续数天或数周。若存储系统任何时刻交付低于 39.3 GB/s,加速器开始空闲。对于本章 175B 语言模型运行示例,token 化文本的每批次带宽需求低得多,但运行期间的总数据量和检查点流量大得多,将瓶颈从峰值训练数据带宽转移到数周内的持续吞吐量和检查点移动。

不足的后果由数据停顿比(data stall ratio)捕获,定义于公式 4.2 为每个训练步骤中加速器等待数据的时间比例:

\[\\text{Data Stall \\%} = \\frac{T\_{\\text{step}} - T\_{\\text{compute}}}{T\_{\\text{step}}} \\times 100 \\qquad(4.2) \]

数据停顿分析与流水线优化

流水线方程

其中当 I/O 与计算不重叠时,T[step] = T[compute] + T[I/O],或采用流水线时 T[step] = max(T[compute], T[I/O]) = T[compute] + max(0, T[I/O]T[compute])。

问题

某 GPU 处理一个批次需 200 ms,但存储系统交付下一个批次需 250 ms。每个训练步骤中加速器空闲的比例是多少?将 I/O 与计算流水线化后,该停顿能减少多少?

无流水线(顺序执行 I/O 再计算)

T[step] = T[I/O] + T[compute] = 250 + 200 = 450 ms

\[\\text{停顿 \\%} = \\frac{250}{450} = 55.6\\% \]

有流水线(I/O 与计算重叠)

T[step] = max(200, 250) = 250 ms

\[\\text{停顿 \\%} = \\frac{250 - 200}{250} = \\mathbf{20\\%} \]

系统洞察

即便采用流水线,加速器仍有 20% 的时间处于空闲状态。要彻底消除停顿,下一个批次必须在 200 ms 的计算窗口内到达,方法包括:提高有效存储带宽、利用更深的预取缓冲区隐藏尾部延迟,或使用更多并行 I/O 流。

数据停顿率作为诊断指标

数据停顿率提供了一项诊断指标,存储工程师可借此判断训练作业是计算受限还是 I/O 受限。

  • 停顿率低于 2% 表明存储流水线健康:加速器几乎将所有时间花在计算上。

  • 停顿率在 2% 到 10% 之间表明流水线勉强够用,但若训练配置变更(更多 GPU、更小批次、更快模型)将会退化。

  • 停顿率超过 10% 表明存在明显的存储瓶颈,浪费了大量计算预算。

B.5.1 节将相同的绿、黄、红三个区间映射到更广泛的车队诊断中,因此存储工程师可将高停顿率解读为红灯信号——计算处于空闲等待,并将存储与其他导致加速器“挨饿”的因素并列考量。

PyTorch 的 DataLoader 分析器和 NVIDIA Nsight Systems 等分析工具可直接测量数据停顿率,方法是对比加速器等待数据的时间与计算时间。

持续监控要求

存储系统不能是“部署即不管”的组件;它需要持续监控,以防无声瓶颈侵蚀训练效率。需实时跟踪的关键指标包括:

  • 分级带宽利用率(NVMe、网络、对象存储)

  • 预取缓冲区深度(已排队并就绪的批次数量)

  • 每 GPU 的数据停顿率

  • 并行文件系统的元数据操作速率

  • NVMe 驱动器健康指标(磨损等级、温度、错误计数)

告警阈值应在停顿反映到训练损失曲线之前触发。在一个 1000 GPU 的集群中,若 5% 的数据停顿率增加持续一周未被察觉,将浪费超过 16,800 美元的空闲算力(1000 GPU × 2 美元/GPU-小时 × 168 小时 × 5%),这使得全面的存储监控成为一级经济考量,而非可有可无的运维事项。

多级分析方法

这种高层监控必须配合细粒度分析,以精确定位瓶颈源头。分析分为三个层级:

  1. 应用层级:NVIDIA Nsight Systems 提供时间线视图,以微秒精度显示加速器因等待数据而空闲的确切时刻与时长。

  2. 流水线层级:PyTorch 内置的 DataLoader 分析器报告每个工作线程的吞吐量、批次处理时间和数据队列深度,识别慢速工作线程或预取不足。

  3. 设备层级iostatnvme-cli 提供原始硬件带宽和延迟指标。针对并行文件系统,Lustre 的 lctl get_param 和 GPFS 的 mmpmon 给出存储服务器利用率、元数据操作速率及客户端缓存命中率的实时统计。

最高效的调试方法是结合三个层级:将应用层级的加速器停顿与特定的流水线工作线程及饱和的底层存储设备相关联,从而识别瓶颈层级,并指导优化工作集中在影响最大的地方。

流水线压力测试

一种常见且代价高昂的失败模式是:在小型本地数据集(1 GB)上开发和测试数据流水线,直到全规模生产(1 TB+)时才发现流水线无法维持所需吞吐量。根原因往往是小规模下不可见、但在负载下致命的瓶颈:

  • 数据增强函数中的 Python 全局解释器锁(GIL)争用

  • 随数百万样本累积的文件描述符泄漏

  • 连续运行数小时后导致预取缓冲区变慢的内存碎片问题

成熟的团队会隔离地对数据流水线进行基准测试,在目标数据量下测量至少 10 到 20 分钟的持续吞吐量,然后再将其接入训练循环。这种流水线压力测试能捕获系统级瓶颈,否则这些瓶颈会在生产运行数天后表现为神秘的训练减速。

面向工作负载的带宽概况

即使使用相同集群,不同 ML 工作负载也呈现截然不同的带宽概况。

  • 图像分类(ImageNet 上的 ResNet-50):产生高带宽需求,因为每个批次包含数百张大图(每张 150 KB),且迭代时间短(100 至 200 ms)。

  • 语言模型预训练:单批次带宽需求较低,因为标记化文本紧凑(每个 token 为 4 字节整数,4096 token 的批次每序列仅占 16 KB),但整个训练过程的总数据量巨大(数万亿 token)。

  • 推荐模型训练:产生混合带宽需求:嵌入查找需要对嵌入表进行高 IOPS 随机访问,而稠密层消费标准顺序训练数据。

每个工作负载的带宽概况决定了哪个存储层级是瓶颈,以及优化工作应集中在何处。

数据停顿前沿

图 4.4 形象展示了存储带宽与 GPU 利用率之间的关系。S 型曲线揭示了一个急剧转折:在停顿阈值以下,即使微小的带宽不足也会导致利用率大幅下降;在阈值以上,增加带宽收益递减。图中技术的位置说明了为何单一层级不足:硬盘和网络文件系统深陷停顿区,单块 NVMe 驱动器恰处于利用率仍陡增的“膝点”位置,唯有 NVMe RAID 阵列跨越阈值进入计算饱和区。这种分布正是分级存储架构必不可少的原因,因为没有任何单一层级能同时提供大规模集群饱和所需的容量与带宽。

图 4.4:数据停顿前沿:GPU 利用率随存储带宽呈陡峭 S 型曲线变化。在停顿阈值附近(典型 8 GPU 节点约 10 GB/s),利用率从等待主导跨越到计算主导。常见存储技术在曲线上的位置解释了为何分级架构至关重要。

扩展挑战

流水线方程还揭示了一个随集群规模扩大而恶化的扩展挑战。随着 N[GPU] 增加,所需带宽 BW[required] 线性增长,但存储系统的聚合带宽不会自动匹配扩展。向共享并行文件系统的集群添加更多计算节点,会增加对固定存储资源的需求。最终,存储系统饱和,饱和点之后每增加一个 GPU 都不会带来额外训练吞吐量,却增加成本。该饱和点是给定存储配置下集群扩展的实际上限,定量识别它需要将公式 4.1 应用于特定存储系统的实测带宽。

场景

某大型 GPU 训练集群在理论上拥有足够的原始存储带宽,但实测的模型 FLOPS 利用率(MFU)仍远低于目标。

故障模式

性能剖析显示,尽管网络和存储带宽指标看起来正常,GPU 却频繁处于空闲状态。一个常见的根本原因是元数据放大:数据加载器在打开每个文件前都会发起一次 stat() 系统调用,导致数千个工作进程的元数据服务器负载成倍增加。

后果

在数千个工作进程中,逐文件的元数据检查每分钟可产生数百万次元数据请求,淹没并行文件系统的元数据服务器。GPU 随后在文件打开操作上停滞,而这种延迟在标准的 I/O 带宽计数器中是不可见的。

系统洞见

在作业启动时预计算文件清单,并让加载器仅读取一次,从而在热路径上消除逐文件的元数据检查。存储瓶颈往往是元数据瓶颈,而非字节吞吐瓶颈。

流水线与预取

对抗数据停滞的主要武器是 流水线:CPU 准备批次 i + 1 的同时,GPU 处理批次 i。图 4.5 展示了这种重叠。当 I/O 时间小于计算时间时,流水线可完全隐藏存储延迟,加速器永不停滞。

![流水线数据加载](https://github.com/OpenDocCN/dsai-notes-pt3-zh/tree/master/docs/hav-cs249r-mlsys-vol2/img/file69.svg)

图 4.5:流水线数据加载:CPU 准备批次 i + 1 的同时,GPU 处理批次 i。流水线可隐藏存储延迟,但前提是预取缓冲区足够深,能够吸收方差。

从 图 4.5 中得出的关键洞见是:流水线将串行瓶颈转化为并行重叠,但仅当预取缓冲区足够深以吸收 I/O 延迟的方差时才成立。当 I/O 时间稳定时,流水线效果完美。但在实践中,它们并不稳定。存储系统表现出 I/O 抖动:由多种因素导致的读取延迟变化。NVMe 驱动器在内部垃圾回收(控制器重组 NAND 闪存块)或磨损均衡操作期间会经历偶发的延迟尖峰。当多个训练作业争抢同一 OSS 节点、MDS 处理元数据操作突发、或后台擦洗检测并修复比特错误时,并行文件系统会出现延迟尖峰。对象存储的延迟可能在跨区域复制、版本对象垃圾回收或负载均衡器重新平衡期间飙升。

一次缓慢的 I/O 可能在流水线中产生一个“气泡”,并向前传播,导致加速器停滞。如果 CPU 正在准备批次 i + 1 而读取耗时是预期的两倍,当 GPU 完成批次 i 时,批次 i + 1 尚未就绪。GPU 空闲等待读取完成,流水线落后一个批次。如果后续读取也很慢,流水线将永远无法恢复。防御手段是预取缓冲区:CPU 不再仅准备下一个批次,而是维护一个包含 Q[prefetch] 个预加载批次的队列。

最小缓冲区深度 Q[prefetch] 必须能在队列不耗尽的情况下吸收最坏情况下的 I/O 延迟:

\[Q\_{\text{prefetch,min}} = \left\lceil \frac{T\_{\text{I/O,p99}}}{T\_{\text{compute}}} \right\rceil \qquad(4.3) \]

公式 4.3 中的预取深度公式告诉我们,必须同时有多少个批次在处理中(在飞行中),才能将 I/O 延迟隐藏在计算之后。该公式假设长期数据供给率至少与加速器的消费率持平;预取吸收的是延迟方差,而非持续的带宽赤字。当平均吞吐量足够但 P99 延迟超过计算时间时,更深的预取可防止停滞。当持续 I/O 吞吐量低于需求时,解决之道是更快的存储、更多并行流、更小的批次或本地暂存。公式使用 P99 I/O 延迟而非平均延迟,因为一次缓慢读取即可耗尽缓冲区并导致加速器停滞;按平均值规模化保证了大规模下的频繁停滞。

若 99 分位数 I/O 耗时 500 ms 而计算耗时 200 ms,则 Q[prefetch,min] = 3 个批次,建议安全边际为 5。实践中,PyTorch 的 DataLoader 等数据加载器使用 prefetch_factornum_workers 参数控制此深度。设置 prefetch_factor=2 配合 4 个工作进程可创建 8 个批次的缓冲区,这通常足以应对 NVMe 支持的流水线,但对于 P99 延迟可能超过 500 ms 的对象存储支持的流水线可能显得不足。

![序列图显示 500 ms 的 P99 I/O 延迟跨越三个 200 ms 计算窗口,意味着需要三个批次的预取深度。](https://github.com/OpenDocCN/dsai-notes-pt3-zh/tree/master/docs/hav-cs249r-mlsys-vol2/img/file70.svg)

P99 I/O 延迟决定了所需的预取深度。

为了说明预取缓冲区深度的内存成本,考虑一个大批次文本流水线:每个 GPU 处理一个打包的 Token 批次约需 200 ms,整理后的批次每 GPU 占用约 40 MB。从本地 NVMe 读取时,该批次的 P99 I/O 延迟约为 50 ms。最小预取深度为 ⌈50/200⌉ = 1 个批次,安全边际 2 即足够。从并行文件系统读取时,因网络抖动和争用,P99 I/O 延迟上升至约 200 ms,要求最小深度为 ⌈200/200⌉ = 1,需安全边际 3 以应对偶发的数百毫秒离群值。从对象存储读取时,P99 延迟可能超 500 ms,要求深度至少为 3,安全边际 5 或更多。这些数字直接转化为主机 DRAM 消耗:按每批次 40 MB 计,深度为 5 的预取缓冲区每 GPU 消耗 200 MB,每节点 8 GPU 共消耗 1.6 GB。若深度为 1,同一节点仅需 320 MB。存储层级直接决定了预取缓冲区的内存成本。

深度预取的代价是内存:每个缓冲批次占用主机 DRAM。一个包含 256 张 224 × 224 × 3 字节图像的批次(解码后)约占 37 MB。8 个此类批次的预取缓冲区消耗 300 MB,对于拥有 512 GB DRAM 的节点微不足道。然而,对于每个批次包含数百万 Token 的大批次语言模型训练,预取缓冲区可能增长至数 GB,与嵌入表及其他 DRAM 消费者竞争资源。

预取与存储层级的交互构建了对抗停滞的分层防御。第一层是主机 DRAM 中的预取缓冲区,吸收 NVMe 读取的 I/O 方差。第二层是本地 NVMe 热缓存,吸收来自并行文件系统的网络方差。第三层是并行文件系统的预读缓存,吸收 OSS 集群中磁盘寻道的方差。每一层以内存或存储容量为代价增加延迟容忍度。系统设计者的任务是确保所有层的总深度超过常规使用中最底层最坏情况的延迟尖峰。以我们的持续示例为例,若并行文件系统偶发 500 ms 延迟尖峰而计算每批次需 200 ms,本地 NVMe 热缓存通过本地服务读取完全消除了该风险,因此预取缓冲区只需吸收微秒级的 NVMe 延迟方差,而非毫秒级。

某训练集群运行 1,024 个 GPU,批次大小为 128 张图像(压缩后每张 150 KB,每次迭代 150 ms)。

多工作进程数据加载

单个 CPU 核心即使在存储硬件足够快时也无法为高吞吐量加速器提供数据,因为存储读取和 GPU 传输之间需要 CPU 工作。考虑图像训练所需的吞吐量:将 150 KB JPEG 图像解码为 224 × 224 × 3 原始张量需要 Huffman 编码系数的熵解码、离散余弦逆变换 (DCT)、色度上采样和像素格式转换。此过程在服务器 CPU 核心上每张图像大约需要 1 到 5 ms。解码后,增强添加了进一步的 CPU 工作:随机裁剪需要计算裁剪坐标并复制子区域;水平翻转需要按反转列顺序复制;颜色抖动需要逐像素乘法和加法。每项增强每张图像增加 0.5 到 2 ms。

以每秒每 GPU 1,000 张图像和每节点 8 个 GPU 计算,单个核心每秒需要解码并增强 8,000 张图像,吞吐量是单个核心可维持能力的 10× 到 40×。解决方案是 multi-worker data loading(多工作进程数据加载),其中 W 个工作进程各自从存储读取、解码、增强并独立入队批次。有效 I/O 吞吐量随 W 大约线性扩展,直到达到三个瓶颈之一:存储设备带宽饱和、主机与设备间 PCIe 总线饱和,或工作进程耗尽主机 CPU 周期。

在 PyTorch 的 DataLoader 中,每个工作进程是一个拥有自己文件描述符和内存空间的独立进程。num_workers 参数控制 W。设置 W 过低会使 GPU 饥饿;设置 W 过高会在上下文切换和争用上浪费 CPU 资源。一个好的启发式方法是从 W = 4 × G 开始,其中 G 是每节点 GPU 数量,并分析数据停滞比率,调整直到停滞低于 2%。对于拥有 8 个 GPU 的节点,这产生 32 个工作进程,在拥有 64 个 CPU 核心的系统上为 PyTorch 运行时、GPU 间通信线程和操作系统开销留下 32 个核心。数据加载与训练编排之间的 CPU 资源划分本身就是一项容量规划工作。

多工作进程加载与存储层级的交互很重要。从本地 NVMe 读取时,工作进程可以向同一 RAID 阵列发起并发读取而无竞争,因为 NVMe 的深命令队列(高达 64K 未完成命令)在硬件层面处理并行性。从并行文件系统读取时,工作进程将读取分布在不同 OSS 节点上,自然聚合带宽。从对象存储读取时,工作进程发起并发 GET 请求,每个请求指向不同分片,在 HTTP 层面实现并行性。

多工作进程加载的一个微妙陷阱是 shuffle quality vs. I/O efficiency(打乱质量与 I/O 效率)。完美打乱要求每个批次包含从整个数据集均匀抽取的样本,这意味着跨所有分片的随机访问。然而,随机访问破坏了存储系统优化的顺序 I/O 模式。实际折衷方案是 shard-level shuffling(分片级打乱)结合 within-shard shuffling(分片内打乱):加载器在 epoch 开始时打乱分片列表,将连续分片组分配给每个工作进程,然后打乱每个分片本地缓冲区内的样本。此方法在保持存储层面顺序 I/O 的同时为训练收敛提供足够随机性。实证研究证实,只要分片足够大(通常 256 MB 或更大)以提供足够的分片内多样性,分片级打乱对大多数工作负载产生的训练损失曲线与完全随机打乱无法区分。

打乱缓冲区大小在内存与随机性之间产生权衡。较大缓冲区提供更好随机性(因为从更大池中抽取样本)但消耗更多主机 DRAM。以我们的文本数据为例,打乱缓冲区 10,000 序列,每序列 2,048 个 token,每个 token 表示为 4 字节整数,大约消耗 80 MB。相对于 512 GB 主机 DRAM 可忽略不计。对于图像数据,每个解码图像占 150 KB,10,000 张图像的打乱缓冲区消耗 1.5 GB,仍可管理但占预取预算的非平凡比例。数据加载器设计者必须平衡打乱缓冲区大小与预取深度,因为两者争夺相同主机 DRAM 容量。

从存储读取经过解码、增强到传输到加速器的数据加载流水线,代表了训练引擎的完整燃料输送系统。当任何阶段成为瓶颈时,加速器就会饥饿。

Data locality and placement(数据局部性与放置)

在分布式训练集群中,数据分片在存储层级的放置决定流水线性能。每个训练工作进程需要高效访问其分配的数据集部分。核心权衡在于共享存储的灵活性与本地存储的原始速度之间。从节点本地 NVMe 驱动器读取通常比必须遍历网络到达并行文件系统或对象存储的读取快 10–100× 且延迟更低。

最简单策略是 static placement(静态放置)。在训练作业开始时,编排器将固定数据分片子集分配给每个节点。节点通过从共享 PFS 复制分片到本地 NVMe 驱动器来暂存数据。在训练运行的剩余时间里,所有数据读取均为本地读取,最大化 I/O 带宽。此方法对单数据集训练运行高度有效,但引入不灵活性:如果数据集变更或作业需要不同数据子集,必须重新暂存分片。

替代方案是 dynamic placement(动态放置),按需从共享存储获取分片。这提供最大灵活性,任何节点可随时访问任何分片。当总数据集大小超过集群本地 NVMe 存储总量,或作业涉及随时间改变活跃子集的数据采样策略时,动态放置至关重要。代价是性能:每次读取产生网络延迟并消耗共享存储带宽。

更高级方法是 locality-aware scheduling(局部性感知调度)。调度器跟踪哪些数据分片缓存在哪些节点的本地 NVMe 驱动器上。启动新作业或替换故障节点时,它倾向于已缓存所需数据的节点,而非视所有节点为可互换。这利用了 data gravity⁵³(数据引力)原则,将计算与数据共置以最小化传输时间。存储要点很简单:仅当放置系统知道缓存位置时,缓存数据才能节省时间。对于在同一数据集上运行重复实验的团队,局部性感知调度将暂存时间从分钟减少到零。

对于分布在 256 个节点上的 6 TB 序列化数据集,每个节点负责约 23.4 GB 数据。如果节点故障,其替代节点必须从 PFS 暂存该分片子集。以无竞争 PFS 读取速度 4 GB/s 计算,暂存约需 5.9 秒;在恢复竞争下,每节点 400 MB/s 有效速率将相同复制延长至约 58.6 秒。如果编排器能将替代工作负载调度到已从先前运行缓存数据的空闲节点上,暂存时间为零且训练即时恢复。规模化时此优化累积:在每天发生 10 次节点故障的 10,000 节点集群中,局部性感知调度每次故障可节省约 1 分钟暂存时间,竞争情况下每天节省约 10 分钟聚合集群空闲时间。

Data versioning and pipeline orchestration at scale(数据版本控制与大规模流水线编排)

数据版本控制与 GPUDirect Storage

了解数据存放位置只是故事的一半;在大规模场景下,我们还必须确定每次运行使用的数据版本以及数据在流水线中的流向。版本控制将存储层级转变为一个可复现系统:机群不仅需要知道读取了哪些字节,还需要知道在运行开始前,是哪些预处理代码、洗牌种子、模式定义和完整性检查产生了这些字节。

虽然原型 A(GPT-4)处理的是包含万亿 Token 的静态版本化数据集,但原型 B(DLRM 大规模部署),即 DLRM 工作负载,处理的是动态、高速的特征流。对于推荐系统而言,“真实标签”随着用户的点击和交互每秒都在变化。这迫使架构从简单的基于文件的存储转向特征存储架构,后者必须解决时间点正确性问题:确保用于训练的检索特征与模型在推理时刻所见的完全一致,且不泄露未来信息。

当输入是实时特征而非不可变文本分片时,这一可复现性要求变得更难满足。特征存储架构通过维护一个不可变的特征账本或“时间旅行”查询引擎来解决时间点正确性问题,但这显著增加了存储复杂性。

在万卡 GPU 规模下,数据集不再是静态文件集合,而是动态流,必须以绝对精度捕获。基础模型训练运行可能持续数周或数月;在我们的 30 天运行示例中,接近运行结束时看到的数据在组成和顺序上必须与第 1 天看到的数据完全相同。否则,我们无法区分模型退化和输入分布的静默变化。数据集版本控制通过对每个数据分片的加密哈希值、用于分词的确切预处理代码以及控制全局洗牌的随机种子进行快照,提供了必要的“数据版 Git”能力。这种严谨性对于审计追踪和调试大规模训练中特有的微妙“Loss 峰值”至关重要,我们必须能够重放导致发散的确切批次,以判断原因是损坏的数据样本还是优化器中的数值不稳定。

在 PB 级规模下,数据流水线必须使每个转换都可复现。原始文档经过清洗、分词、打包成训练样本,并分片存储到各存储层级。每个阶段写入一个不可变制品,每个制品记录其满足的输入版本、代码版本、模式定义和统计检查。一旦这些依赖关系显式化,流水线就形成了一个有向无环图:后续制品依赖前序制品,且任何阶段都不允许重写自己的历史。当训练作业启动时,集群管理器会验证每个工作节点的本地 NVMe 分片是否与为该运行声明的数据快照匹配。通过将数据流水线视为版本化制品图而非简单的文件传输,系统消除了基础模型训练中主要故障来源——静默数据漂移。

除了流水线编排,一项硬件优化消除了最常见的瓶颈之一:CPU 作为存储与加速器之间中介的角色。

GPUDirect Storage 与 CPU 绕过

存储层级可以提供足够的聚合带宽,但当 CPU 充当数千次小传输的中介时,GPU 仍可能处于等待状态。在传统路径中,每次分片读取都会在到达 GPU HBM 之前,经过主机 DRAM 和软件层;当本地或 RDMA 连接的 NVMe 成为瓶颈时,GPUDirect Storage (GDS) 会移除这条 CPU 中介路径。先理解传统路径,再看 GDS 优化就清晰了。

传统加载训练数据的数据路径遵循三跳:存储 → 主机 DRAM → GPU HBM。数据首先从 NVMe 读取到主机 DRAM 的内核缓冲区(由 NVMe 控制器发起 DMA 传输),然后复制到用户空间缓冲区(数据加载器的张量,通过 CPU 执行的 memcpy),最后通过 PCIe 总线传输到 GPU 内存(使用编程 GPU DMA 引擎的 cudaMemcpycudaMemcpyAsync 调用)。每一跳都增加延迟并消耗 CPU 资源。CPU 必须编排每次传输、管理缓冲区分配,并处理来自存储设备和 GPU 的中断。对于单次传输,此开销可忽略。但当 8 个 GPU 每秒各需数千次小传输时,仅数据移动的聚合 CPU 负载就可能饱和多个核心,留给同样对流水线至关重要的数据增强的 CPU 算力不足。

GPUDirect Storage (GDS) 是一项在 NVMe 存储设备与 GPU 内存之间建立直接 DMA 路径的技术,绕过主机 CPU 和系统 DRAM。

  1. 重要意义:消除了通过系统内存的“反弹缓冲区”。在下述小传输模型中,这将单次传输延迟降低 75%,使本地传输路径快 4 倍。它允许 GPU 将 NVMe 链路速度(例如 7 GB/s)跑满,同时降低 I/O 的 CPU 利用率。

  2. 区别:与传统 I/O 中每个字节都必须由 CPU 处理并存入内核缓冲区不同,GDS 在存储控制器与加速器之间提供直接内存访问。

  3. 常见陷阱:常见的误解是 GDS 让所有存储“更快”。实际上,它仅加速本地或 RDMA 连接的 NVMe;无法消除网络附属文件系统或对象存储的物理延迟。

传统 I/O 路径中的软件层导致了延迟开销。当数据加载器对 NVMe 支持的文件调用 read() 时,调用会遍历应用运行时、Python/C++ 边界、操作系统的虚拟文件系统 (VFS) 层、文件系统驱动(ext4、XFS)、块层,最后到 NVMe 驱动。每层为参数验证、锁获取和缓冲区管理增加几微秒开销。在返回路径上,NVMe 控制器触发中断,中断处理程序唤醒阻塞线程,数据从内核缓冲区复制到用户空间。小读取的总往返开销为 10–50 μs,主要由上下文切换和缓冲区复制主导,而非实际 NVMe 访问时间。

GDS 为数据路径消除了这些软件层。cuFile API 向 NVMe 控制器注册一块 GPU 内存区域,建立直接 DMA 映射。后续读取将数据从 NVMe 传输到 GPU 内存,无需经过内核的文件系统或块层。CPU 的角色简化为发出 DMA 命令并检查完成状态,仅需几微秒工作,而传统路径需要 50+ 微秒。

当训练流水线已优化且剩余瓶颈是 CPU 传输中介能力时,GDS 的延迟降低最为关键。在一个 8 GPU 节点中,每个 GPU 运行 4 个工作进程的数据加载器,CPU 必须管理 32 个并发 I/O 流。按 120 μs/次传输计算,CPU 在中断处理和缓冲区管理上花费大量时间。GDS 将这部分工作卸载给硬件 DMA 引擎,释放 CPU 核心用于数据增强和其他预处理任务。图 4.6 对比了两种数据路径,展示了 GDS 如何消除主导单次传输开销的 CPU 中介复制。

![](https://github.com/OpenDocCN/dsai-notes-pt3-zh/tree/master/docs/hav-cs249r-mlsys-vol2/img/file71.svg)

图 4.6:传统路径 vs. GPUDirect Storage 路径:传统路径(上)需要两次复制和 CPU 参与。GDS 路径(下)使用 DMA 直接从 NVMe 传输数据到 GPU 内存,消除了作为瓶颈的 CPU。

正如图 4.6 所示,GDS 带来的吞吐量提升在读取许多小对象(如解码后的图像块)且单次传输开销占主导地位的工作负载中最为显著。对于大型顺序读取(如流式传输数 GB 级的训练分片),吞吐量提升相对温和,因为传输时间占主导地位,而非建立连接的开销。通用原则是 GDS 消除了每次传输固定的开销,因此每秒传输次数越多的工作负载获益越大。

问题:一个配备 8 个 GPU 的训练节点以每 GPU 8,000 张/秒的速度加载 150 KB 图像(总计 64,000 张/秒)。对比传统 I/O 与 GDS 下的 CPU 负载。

传统路径:每张图像需要从 NVMe 到 DRAM 的一次 DMA,一次从内核空间到用户空间的主机拷贝,以及一次通过 PCIe 传输到 GPU 的操作。在 64,000 张/秒、每张图像耗费 120 μs CPU 时间的条件下,CPU 每墙钟秒消耗 7.68 秒 CPU 时间,仅数据移动就占用了约 8 个核心的处理能力。

GDS 路径:每张图像直接从 NVMe 通过 DMA 传输到 GPU。在每张图像耗费 30 μs CPU 时间(用于发起 DMA)的条件下,CPU 每墙钟秒消耗 1.92 秒 CPU 时间,从而释放约 6 个核心用于数据增强。

系统洞察:GDS 不仅降低了延迟,更重要的是将 CPU 利用率从无增值的数据拷贝转移到了能提升模型质量的数据增强上。

GDS 存在实际限制。并非所有 NVMe 控制器都支持点对点 PCIe 传输。数据必须采用 GPU 可直接消费的格式,这意味着数据必须已解码或可由 GPU 解码(原始像素值、Token ID 或 GPU 解码器支持的格式),而非任意压缩的 JPEG 或 gzip 流。实践中,许多管线采用混合方法:对压缩格式进行 CPU 侧解码,随后通过普通的页锁定内存传输到 GPU;而 GDS 则保留用于从受支持的本地或远程存储直接读取 GPU 可消费数据。

从 CPU 中介 I/O 到直接存储访问的架构演进,反映了 ML 系统的一个更广泛趋势:尽可能将 CPU 从关键数据路径中移除。3.4.1 节介绍的 RDMA 将 CPU 从网络数据传输中移除,GDS 对存储传输做到了同样的事情。RDMA 与 GDS 的联合效应是:数据可以从远程存储流经网络结构进入 GPU 内存,关键路径上无需一条 CPU 指令。CPU 因此得以执行其具有独特价值的任务:数据增强、管线编排和错误处理。

GDS 设计原则同样适用于检查点写入。在传统路径中,检查点数据从 GPU HBM 经由主机 DRAM(通过 PCIe)再到 NVMe(通过 CPU 中介写入)。使用 GDS 时,检查点可直接从 GPU HBM 写入 NVMe,消除了主机 DRAM 拷贝。对于 1,750 GB 的检查点,这省去了一整次完整数据拷贝,并减少了 CPU 在写入路径上的介入,进一步降低了 T[write]。

完整数据路径

结合 GDS 与多级层级,我们现在可以追踪运行示例的完整数据路径。训练作业开始时,源自 3 TB 压缩源语料库的 6 TB 序列化分片集从对象存储调度至集群的本地 NVMe 缓存,设置阶段每个节点接收其分配的分片子集。数据加载工作进程从 NVMe 读取已分词的分片,在页锁定的主机 DRAM 中组装打包序列,并将这些批次传输给加速器。启用 GDS 时,已处于 GPU 可消费布局的 Token ID 可直接加载进 GPU 内存,无需 CPU 介入;压缩格式仍需在使用前经由 CPU 或 GPU 解压。

训练期间,管线进入稳态。CPU 数据加载工作进程持续从本地 NVMe 读取,填充主机 DRAM 中的预取队列。GPU 通过 PCIe DMA 从该队列拉取批次。计算阶段(前向传播、反向传播、优化器步骤)消费批次并在 HBM 中更新模型权重。每 10 分钟,训练框架发起一次检查点:模型权重与优化器状态从 GPU HBM 写入本地 NVMe(无论是通过主机 DRAM 还是经由 GDS),随后一个后台线程异步将本地检查点复制到并行文件系统。

训练作业结束时,最终模型检查点从并行文件系统提升至对象存储以供长期保留。作业清理阶段删除本地 NVMe 副本。若模型部署用于推理,则从对象存储加载至服务集群的 HBM,完成数据在层级间的全生命周期移动。

此生命周期中的数据拷贝总数颇具启示性。每个训练样本历经:对象存储 → NVMe(调度),NVMe → 主机 DRAM(读取),主机 DRAM → GPU HBM(传输)。每个检查点历经:GPU HBM → NVMe(本地保存),NVMe → 并行文件系统(异步复制),并行文件系统 → 对象存储(长期保留)。每一次拷贝都消耗带宽并增加延迟。工程目标是在关键路径(训练循环)上最小化拷贝,而在非关键路径(调度与归档)上容忍额外拷贝。

30 天训练运行期间移动的数据量揭示了一个反直觉的现实。将 6 TB 序列化数据集从对象存储单次调度至本地 NVMe 仅产生适度的传输量。然而,检查点过程生成的数据流大得多。在 256 个节点上每 10 分钟创建约 1,750 GB 检查点,持续 30 天,系统将产生约 4,320 个检查点,总计约 7.6 PB 状态必须穿越存储层级。对于大模型,检查点数据移动远超训练数据移动:6 TB 训练数据集可能仅被读取数次(每轮一次),而 7.6 PB 检查点数据是全新生成的,使检查点 I/O 成为主要存储工作负载。这一洞察解释了为何检查点调度策略(本地写入、异步复制)对整体存储设计的影响,可能大于训练数据管线优化。物理数据路径决定了字节在何处移动;剩下的问题是每次移动的代价是什么。

存储经济学

100 TB 训练数据集可驻留于层级的任何一层,但根据所选层级及访问方式不同,成本相差数个数量级。ML 存储经济学不仅关乎选择最便宜的层级,更关乎最小化数据交付总成本,其中包含存储成本、传输成本以及闲置加速器的机会成本。

考虑存储 100 TB 训练数据集。不同层级的年存储成本差异巨大:

  • 对象存储 (S3 Standard):$24,000/年

  • 本地 NVMe (预置):$120,000/年

  • 归档存储 (Glacier):$4,800/年

  • HBM 等效 (假设以 $15/GB 将 100 TB 存储于 GPU 显存):$1,500,000 硬件摊销成本

本地 NVMe 与对象存储间 5 倍的成本差距(而 HBM 到归档的完整跨度超过 3,000 倍)解释了为何层级必须存在:数据应尽可能驻留在最廉价的层级,仅在需要时上移,用毕即下移。这一成本梯度反映了底层物理特性。更快的存储需要更昂贵的材料(HBM 采用硅通孔的 3D 堆叠硅),每比特访问能耗更高,且必须物理靠近加速器(这限制了单节点可预置的容量)。更廉价的存储使用通用组件(归档用标准硬盘),闲置时能耗极低,且可部署在任何有网络连接的位置。

存储成本条形图,重复出口流量占主导。

重复出口流量,而非存储,主导了云成本。

然而,存储成本仅是整个方程的一部分。数据传输成本 可能会占主导地位,尤其是在云环境中。在此所使用的代表性跨边界出口假设下,从对象存储读取 100 TB 数据集到计算实例会产生 9,000 美元的传输费用。对于多轮训练(读取数据集 10 次),仅出口费用就会超过 90,000 美元,超过了年度存储成本。这种倒置——读取 数据的成本高于存储 数据的成本——驱动了将数据缓存在本地 NVMe 而不是在每个 epoch 从对象存储流式传输的架构决策。

当考虑到加速器空闲时间的机会成本时,经济分析进一步变化。一套每 GPU 每小时 2 美元、共 1,000 个 GPU 的集群每天成本为 48,000 美元。如果存储系统容量不足导致加速器利用率从 90% 降至 70%,组织每天将损失 9,600 美元的浪费算力。在为期 30 天的训练过程中,这相当于 288,000 美元的算力损失,这一金额足以资助并行文件系统升级或额外的 NVMe 容量。从经济上看合理的做法是将存储投资视为提升算力利用率的推动者,而非一项支出:每花一美元改进存储吞吐量,将带来数美元的算力生产率提升。

存储的可见成本($/GB/月)只是冰山一角。表面之下隐藏的成本往往超过存储成本本身:

  • 出口费用:代表性的跨边界出口费用为 0.09 美元/GB,读取 100 TB 数据集(每 epoch 一次,共 10 次 epoch)仅产生的出口费用即为 90,000 美元。

  • IOPS 费用:对象存储按请求收费。S3 标准 GET 请求每 1,000 次请求约收费 0.0004 美元,因此读取 1 亿个单独文件的数据集(一次读取)在数据传输或检索费用之前,仅 GET 请求费用约为 40 美元,尽管存储成本为每月 2,300 美元。

  • 加速器空闲成本:如果存储停顿导致每 GPU 每小时 2 美元、共 1,000 个 GPU 的集群的 GPU 利用率从 90% 降至 70%,则损失的算力成本为每天 9,600 美元(每日 48,000 美元支出的 20%),远高于任何存储节省。

  • 检索费用:归档存储除了出口费用外还会收取检索费用(Glacier 为 0.02 美元/GB)。将 100 TB 数据集从 Glacier 恢复将产生 2,000 美元的检索费用加上 9,000 美元的出口费用。

从经济上看合理的策略通常会在存储上投入更多(本地 NVMe 缓存、并行文件系统容量)以降低传输成本并防止加速器停顿。

成本冰山揭示了:原始存储容量可能仅占总支出的一半以下;复制、元数据操作、API 调用和跨区域传输等因素在规模上可能共同主导账单。在这些隐藏的成本项中,出口往往是导致架构决策逆转的关键因素。与按存储数据量增长的容量费用不同,出口费用随着读取的数据量增长,而机器学习训练在每个 epoch 会读取完整数据集。

问题

一个团队在 S3 上存储的 50 TB 图像数据集上训练一个视觉模型。训练持续 20 个 epoch。需决定是每个 epoch 从 S3 流式传输,还是将数据暂存到本地 NVMe

选项 A:从 S3 流式传输

  • 存储:50 TB × 0.023 美元/GB/月 × 12 个月 = 13,800 美元/年

  • 出口:50 TB × 20 个 epoch × 0.09 美元/GB = 90,000 美元/次训练

  • 每年 4 次训练的总成本:13,800 美元 + 360,000 美元 = 373,800 美元/年

选项 B:暂存到本地 NVMe

  • 存储(S3 作为事实来源):13,800 美元/年

  • NVMe 容量(集群中的 50 TB):50 TB × 0.10 美元/GB/月 × 12 = 60,000 美元/年

  • 出口(每次训练暂存一次):50 TB × 4 次 × 0.09 美元/GB = 18,000 美元/年

  • 总计:13,800 美元 + 60,000 美元 + 18,000 美元 = 91,800 美元/年

系统洞察

在每次训练 20 个 epoch 的情况下,本地 NVMe 缓存每年可节省 282,000 美元,足以资助 9 额外的 H100 GPU(按列表价格计算)。考虑到每年 60,000 美元的预置 NVMe 成本和每次训练一次的暂存出口费用,盈亏平衡点仅略高于每次训练 4 个 epoch;到第五个 epoch 时,本地缓存每次训练的成本低于每个 epoch 从 S3 流式传输的成本。

该笔记本将冰山转化为决策规则:最便宜的存储层级并不总是最便宜的传输路径。对象存储可以仍作为事实来源,但重复读取使得出口、请求开销和加速器空闲成本成为存储设计的一部分。多轮视觉训练可能在第五个 epoch 证明在本地 NVMe 上暂存数据是合理的,而单轮语言模型预训练可能选择将语料库保留在对象存储中并优化流式传输路径。因此,架构的选择应基于访问模式、读取次数和延迟敏感度,而非仅仅基于 $/GB/月。

成本分析还延伸到并行文件系统层级,该层级在性能和价格上介于 NVMe 和对象存储之间。能够提供 1 TB/s 聚合带宽的并行文件系统大约需要 100 个 OSS 节点,每个节点包含多个 NVMe 驱动器和高带宽网络连接。如此系统的资本和运营成本可能超过每年 1,000 万美元。只有当替代方案——从对象存储流式传输并支付出口费用——更昂贵,或者当对象存储 alone 无法满足检查点写入的延迟需求时,这一费用才能得到证明。盈亏平衡分析取决于组织的工作负载组合:运行单个长时间训练作业的团队可能发现,将数据暂存在本地 NVMe 并完全绕过并行文件系统是最便宜的选择;而运行许多并发短任务的团队则可能受益于并行文件系统提供的共享命名空间。

是否自建还是购买存储基础设施的决策取决于在本地部署的并行文件系统的高资本支出与云对象存储的持续运营支出之间的权衡。1 PB 的 Lustre 部署硬件成本大约为 300–500 万美元,此外每年还需额外支出 500,000 美元用于电力、冷却和工程支持。在云对象存储中以 0.02 美元/GB/月存储 1 PB 数据,仅存储成本每年为 240,000 美元。然而,访问数据的出口费用往往占据总成本的主导:以每年 10 次完整读取该 PB 数据计算,将产生 900,000 美元的数据传输费用。在此假设下,年度云交付成本为 1,140,000 美元;扣除 500,000 美元的本地运营成本后,净节省为每年 640,000 美元。由此得出的硬件盈亏平衡点大约为 4.7–7.8 年,而非 18–24 个月。要实现 18–24 个月的回收期,需要更高强度的工作负载(每年更多次完整数据集读取)、较低的资本支出,或者非出口驱动因素,如检查点延迟和共享命名空间需求。就我们目前的例子而言,6 TB 的序列化语料库远小于 1 PB,因此本地存储方案更多受延迟、重复实验和检查点暂存的驱动,而非仅仅依赖一年的出口节省。

硬件退化带来了对存储刷新周期的需求。NVMe 驱动器的写入耐久度有一个额定值,通常以五年保修期内的每日驱动器写入次数(DWPD)来衡量。许多大规模机器学习训练工作负载在初始数据集摄入后主要是读取密集型,因此训练集群中的驱动器可能超过其额定写入寿命。因此,刷新周期的主要驱动因素可能是密度而非磨损:较新的驱动器通常在类似价格点下提供更大容量。将集群从 7.68 TB 升级到 15.36 TB 的 NVMe 驱动器将使本地缓存容量翻倍,使更大的数据集能够直接在计算节点上暂存,从而在稳态读取时减少对并行文件系统的依赖。

分层策略

分层即放置策略

前文所述的成本结构将分层转化为一种放置策略:仅当数据的访问频率和延迟要求能够证明合理时,才将数据保存在离加速器足够近的位置。分层策略 将该策略编码到层级结构中。

热层

速度最快、成本最高的层级是 热层(本地 NVMe + 主机 DRAM),它存放正在被运行中的训练作业主动消费的数据。数据集在作业启动时从共享存储暂存至本地 NVMe。该层按节点配置,不对集群共享,因此多个作业或团队同时需要的数据不能仅驻留于此。

温层

这种共享访问需求正是 温层(并行文件系统)存在的理由,它存放由多个作业或团队访问的数据集、用于故障恢复的共享检查点,以及正在活跃开发中的模型工件。并行文件系统提供了多租户访问所需的共享命名空间,并保证强一致性。

冷层

温层之下是 冷层(对象存储),它作为组织所有训练数据的规范存储库。数据集在对象存储中编写和版本控制,需要训练时再提升至更热的层级。对象存储还充当检查点的持久备份,这些检查点从本地 NVMe 暂存,经由并行文件系统传输而来。

归档层

出于合规或复现目的必须保留,但在正常运营期间预计不会被访问的数据,会降级至 归档层(Glacier 或同类服务)。生命周期策略会在可配置周期后(通常为上次访问后 90 至 365 天)自动将数据从冷层转移至归档层。

自动化分层迁移

层级间的数据迁移应通过生命周期策略实现自动化。设计良好的分层系统会在训练作业请求数据时自动将其从冷层提升至温层,并在配置周期内无作业访问时将数据从温层降级至冷层。人工分层会带来运维负担,并不可避免地导致过度配置(数据滞留在昂贵层级)或配置不足(需要时数据不可用)。

运行示例:分层实战

回到我们的运行示例,观察分层的实际运作。3 TB 的源语料库和 6 TB 的序列化训练分片永久驻留在对象存储(冷层)中。当训练作业被调度时,编排系统会触发“数据暂存”阶段,在第一个训练节点启动前,将分词后的分片复制到并行文件系统(温层)。随后,每个节点的数据加载器在首个 epoch 期间将其分配的分片复制到本地 NVMe(热层)。后续 epoch 完全从本地 NVMe 读取,不产生任何网络流量。作业完成后,清理进程会删除本地 NVMe 副本,并在可配置的宽限期后清除并行文件系统中的副本。从暂存、训练到清理的整个生命周期均由策略管理,而非由训练工程师手动操作。

经济效益

自动化分层的经济效益随时间复利增长。某组织并行运行 50 个训练作业,每个作业使用 10 TB 数据,若所有数据就地保留,将需要 500 TB 并行文件系统容量。作业完成后的自动降级可将稳态并行文件系统用量降至 100 TB,按并行文件系统的每 TB 单价计算,每年可节省 144,000 美元。这些节省在单次实验中不可见,但在整个机群累积起来却十分可观。

数据暂存模式

暂存决策决定了作业在何处支付数据移动的代价:训练开始前、首个 epoch 期间,或每个 epoch 期间。三种常见模式针对启动延迟与稳态吞吐量之间的不同权衡:

  • 预暂存:在第一次训练迭代开始前,将整个数据集从对象存储复制到本地 NVMe。这能带来最佳的稳态性能,因为所有读取均为本地读取,但会导致最差的作业启动时间;以每节点 500 MB/s 的速度在 256 个节点上暂存 10 TB 数据集,每节点耗时约 80 秒,共享并行文件系统带宽可能将该延迟延长至数分钟。

  • 按需暂存:仅当数据加载器首次访问分片时,才将其复制到本地 NVMe。这消除了启动延迟,但使首个 epoch 承担全部网络读取开销,适用于多 epoch 训练,因为后续的本地读取可分摊初始暂存成本。

  • 流式读取:数据不落本地存储,每个 epoch 均从并行文件系统或对象存储读取。这具有零启动延迟和零本地存储需求,但每个 epoch 都要支付全部网络读取代价,适用于单 epoch 预训练或数据集大于本地 NVMe 容量的场景。

分层是一种放置策略,因此下一个问题是分层边界可在何处移动。新兴设备正在移动 DRAM、NVMe 和归档存储之间的边界,而推理工作负载为不同的目标复用相同的层级结构:降低冷启动延迟和管理服务时的工作集,而非最大化 epoch 吞吐量。

移动分层边界的存储技术

虽然六层层级结构代表了常见的生产基线,但内存和存储技术可以填补现有层级间的带宽缺口,重塑大规模 ML 的有效存储架构。层级结构中最大的未填补缺口位于主机 DRAM 与本地 NVMe 之间。Compute Express Link (CXL) 是一种开放标准互连,允许 CPU、加速器和内存设备以缓存一致性语义共享内存,精准瞄准这一缺口。CXL 连接内存提供的带宽介于本地 DRAM 与 NVMe 之间(约 30–60 GB/s),延迟也处于中间水平(200–500 ns)。对于 ML 工作负载,CXL 内存可作为主机 DRAM 与本地 NVMe 之间的“1.5 层”,为嵌入表和预取缓冲区提供海量容量池,且无 NVMe 的延迟惩罚。在我们的运行示例中,CXL 连接内存可将有效主机内存层从 512 GB 扩展至每节点 4 TB 以上,足以容纳压缩后的源语料库或大型本地分片缓存,但仍不足以在单节点上容纳完整的 6 TB 序列化训练副本。

计算存储驱动器 (CSD)

当 CPU 端解码而非存储带宽成为管线瓶颈时,会出现另一种瓶颈。计算存储驱动器 (CSD) 通过在 SSD 控制器上直接嵌入处理元件(FPGA 或简化 CPU 核心)来填补这一缺口,使解压、过滤或格式转换能在存储设备上完成,而不消耗主机 CPU 周期。对于 CPU 端解码成为瓶颈的 ML 管线,计算存储可将 CPU 完全从数据路径中移除,通过卸载直连 DMA 无法处理的工作来补充 GDS

持久内存

检查点写入仍受 SSD 写入延迟限制,这是 Young-Daly 公式中惩罚的 T[write] 项。持久内存 尽管经历了英特尔 Optane 停产等市场变动,仍是直接攻克该项的架构理念。字节寻址的非易失性内存在延迟和容量上均介于 DRAM 和 NVMe 之间,可变革检查点存储:具备 SSD 级持久性且写入延迟逼近 DRAM,可将 T[write] 从毫秒级降至微秒级,使频繁、低开销的检查点更切合实际。

存储层级的演进

这些技术尚未在最大规模 ML 集群中广泛部署,但它们展示了存储层级演进的方向。基本原则保持不变:更快、更昂贵的存储更靠近加速器,更慢、更便宜的存储位于外围。变化的是层级的粒度和层级间缺口的大小

推理工作负载的存储

推理的存储需求

推理的存储需求与训练从根本上有所不同。训练持续读取数据集并以突发方式写入检查点。推理在启动时一次性读取模型权重,在正常运行期间不再进行存储 I/O 操作。对于推理服务,数据集被传入的用户请求流取代,主要的存储挑战从大规模数据集的持续吞吐量转移到了加载模型本身的延迟上。

模型加载延迟

推理存储最关键的指标是模型加载延迟,通常称为冷启动时间:将模型从存储加载到加速器 HBM 所需的时长。对于我们的 1750 亿参数语言模型,仅权重在 FP16 格式下就占用 350 GB。从单个高性能 NVMe 驱动器以 14 GB/s 的速度顺序加载需要 25 秒,这对于面向用户的应用来说是不可接受的延迟。从共享并行文件系统(PFS)以每节点 4 GB/s 的速率加载需要近 88 秒。从对象存储以 1 GB/s 的速度加载,延迟则接近六分钟。

为减少冷启动时间,模型会被分片并并行加载。如果 350 GB 的模型跨 8 个 NVMe 驱动器条带化,每个驱动器以 7 GB/s 的速度加载 43.75 GB 的分片,总加载时间降至约 6 秒。当跨 8 个 GPU 使用张量并行时,每个 GPU 仅加载其 43.75 GB 的分片,通过从主机 DRAM 读取可在 4 秒内完成。对于对延迟最敏感的应用,组织会维护热副本:在主机 DRAM 中预先加载一份或多份模型权重副本,可在两秒内将其传输至 HBM。这以 DRAM 容量换取近乎即时的冷启动时间。

KV 缓存管理

在自回归生成过程中,服务器会保留每个请求的注意力键和值记录,这些由先前的 token 生成,以便在生成下一个 token 时无需重新计算整个前缀。该记录称为 KV 缓存,其大小随序列长度增长,并可能消耗大量 HBM。对于长序列(12.8 万 token 或更多),将 KV 缓存卸载到主机 DRAM 或 NVMe 可以扩展有效上下文窗口,但会增加每个生成 token 的延迟,因为缓存块必须在注意力机制使用前移回 HBM。存储层面的启示是,推理有其自身的工作集,不仅仅是模型加载问题。

虽然模型服务有其自身的经济核算,但训练阶段引入了完全不同、以突发写入为主的 I/O 挑战:通过检查点存储来保持模型状态。

检查点存储

经过一千个 GPU 训练六小时后,电源故障会让检查点存储成为快速恢复与重复数小时工作之间的分水岭。关键的存储问题在于:最近一次检查点保存有多快、它驻留在何处、以及能否被读回。检查点是存储层级中最苛刻的写工作负载。它们也是后果最严重的:硬件故障后丢失检查点意味着重复数小时甚至数天的训练。因此,存储架构必须将 T[write](训练流水线暂停以保存检查点的时间)降至最低。

采用 Adam 优化器的 1750 亿参数模型生成的检查点约为 1,750 GB。检查点包含 350 GB FP16 模型权重、1,400 GB 用于动量和方差的 FP32 优化器状态、学习率调度器状态、随机数生成器状态以及当前数据加载器位置。集群中的每个 GPU 同时保存其检查点分片,形成存储系统必须吸收且不干扰正在进行的训练读取的检查点风暴。

对比朴素复制检查点写入时间与 ZeRO-3 分片检查点写入时间的边际阶梯图。

分片将复制的检查点风暴转化为分片写入。

检查点风暴

检查点风暴是指 ML 训练集群中所有节点同时保存模型状态时发生的同步网络和存储流量突发。

重要性

风暴的规模按 T[write] = R × per-replica state/BW[fabric] 缩放,其中 R 是同时写入的副本数量。对于 700 亿参数模型,每个副本持有 140 GB FP16 权重、140 GB 梯度和 840 GB 优化器状态。在 1,024 个副本上使用朴素数据并行时,该完整副本状态的朴素检查点会产生 1146.9 TB 的同步写入。在 100 GB/s 网络带宽下,该写入耗时约 11,468.8 秒,即每次检查点事件超过 191.1 分钟的训练停顿。ZeRO-3 是一种优化器状态分片方案,每个节点仅存储其 1/N 的状态切片而非完整副本,将写入状态总量压缩至约 1.12 TB(140 GB 权重 + 140 GB 梯度 + 840 GB 优化器状态,分摊到 1,024 个副本),使 T[write] 降至约 11.2 秒。若省略分片优化,此停顿项可能远超两次检查点间的计算时间。

区别

与一般的 I/O 争用(随机且不可预测)不同,检查点风暴是同步且周期性的:因为训练编排器在固定步数后触发检查点,每个节点在同一时刻写入。其可预测性使其破坏性更强(所有节点同时竞争)且更易处理(可通过错峰调度或异步序列化在设计上预防)。

常见误区

常见的误解是每 100 步检查点一次属于“低开销”。在 700 亿规模下,此假设灾难性地错误:若每训练步耗时 10 秒而检查点风暴耗时 11,468.8 秒,每 100 步检查点意味着大部分运行时间花费在检查点 I/O 而非有效训练上。

检查点风暴的严重程度取决于两个因素:单节点分片大小(由模型规模和分片策略决定),以及第一写入层级(通常为本地 NVMe)的吞吐量。一个具体的 1750 亿参数示例表明,为何本地写入阶段(而非到并行文件系统的持久化拷贝)决定了暴露的训练暂停时间。

具体 175B 示例

问题:一个 256 节点集群每 10 分钟保存一次 1750 亿参数检查点。每个检查点总计 1,750 GB。采用 ZeRO-3 时,每个节点保存约 6.8 GB。

  1. 单节点写入本地 NVMe(4 块驱动器,每块 7 GB/s,合计 28 GB/s):6.8 GB ÷ 28 GB/s ≈ 0.24 秒。

  2. 异步拷贝至 PFS:256 × 6.8 GB ≈ 1.8 TB 总量。若 PFS 提供 1 TB/s 总带宽,风暴约 1.8 秒完成。

  3. 单节点 PFS 带宽:若 256 节点同时写入,每节点得 1000 GB/s ÷ 256 = 3.9 GB/s。

  4. 训练暂停:仅约 0.24 秒(本地 NVMe 写入时间)。PFS 拷贝与下一训练迭代重叠。

  5. 开销:每 600 秒暂停 0.24 秒 ≈ 0.04% 的训练时间损耗于检查点。

系统洞察:在总计 1 TB/s PFS 带宽假设下,分级暂存将检查点开销从约 1.8 秒的 PFS 直写降至亚秒级本地写入。关键洞见在于,训练流水线的 T[write]本地写入时间,而非持久化写入时间。

分级暂存策略

分级暂存策略通过两阶段写入最小化 T[write]。第一阶段,每个节点以满带宽将其检查点分片写入本地 NVMe。4 块 NVMe 提供 28 GB/s 总带宽下,单节点约 6.8 GB 分片(256 节点 ZeRO-3 分割 1,750 GB 总检查点)约 0.24 秒完成。训练流水线在本地写入完成后可立即恢复。

检查点存储与恢复

在第二阶段,一个后台进程会异步地将本地检查点复制到并行文件系统以实现持久化。此复制过程可与下一次训练迭代重叠,因此不会阻塞流水线。风险在于,如果在异步复制完成前节点发生故障,本地检查点将会丢失。缓解措施是在宣布保存完成前,将检查点复制到至少两个对等节点的 NVMe 驱动器上,即使其中一个节点发生故障也能提供持久性保障。

训练运行全生命周期内检查点消耗的总存储量巨大。一个拥有 1750 亿参数的模型,在为期 30 天的训练运行中每 10 分钟生成一个检查点,将产生约 4,320 个检查点,每个 1,750 GB,检查点数据总量约为 7.6 PB。保留所有检查点既非必要也非经济。本次运行采用的保留策略是:在并行文件系统上保留最近的三个检查点以便快速恢复;将每第 100 个检查点复制到对象存储以供长期审计;删除其余检查点。该策略将并行文件系统的检查点占用空间从 7.6 PB 降低至约 5.25 TB(三个存活的完整检查点),同时在对象存储中保留 43 个历史快照以供训练后分析。

增量检查点为减少 T[write] 提供了进一步的优化。增量检查点不保存每次检查点时的完整模型状态,而是仅保存自上一个完整检查点以来发生变化的参数。对于大部分参数被冻结的模型(例如仅更新最后几层的微调场景),增量检查点的大小可比完整检查点小几个数量级。代价是恢复复杂性:从增量检查点恢复需要将一系列增量更新应用到基础检查点上,这会延长恢复时间。大多数生产系统采用混合方法,频繁保存增量检查点,并定期(每 10 到 100 个增量)保存完整检查点。

检查点存储与系统的故障模型直接交互。从存储角度来看,设计目标是在下一个故障窗口开启前,确保至少存在一个检查点的持久副本,同时最小化 T[write](训练流水线暂停的时间)。通过分级暂存来减少 T[write],使更频繁的检查点成为可能。然而,仅最小化 T[write] 只有在检查点实际可读回的情况下才有意义;代价更高的故障是检查点已写入但从未验证过可恢复性。

案例研究:GitLab.com 数据库事故(2017 年 1 月 31 日)

背景:2017 年 1 月 31 日,约 23:30 UTC,一名 GitLab.com 站点可靠性工程师在处理数据库复制事故时,擦除了他们认为是次级 PostgreSQL 数据目录的内容,但此时他们已切换终端至生产主库 (GitLab 2017)。随后恢复的快照将数据库恢复到了约 17:20 UTC 的状态,这也是为何大约六小时的写入数据无法恢复的原因。

故障模式:GitLab.com 依赖五层恢复路径:pg_dump 备份至 S3、LVM 快照、Azure 磁盘快照、PostgreSQL 复制和 WAL 归档。所有五条路径均失效或缺失。由于 PostgreSQL 9.2/9.6 版本不匹配,S3 转储文件静默地为空;数据库服务器未启用 Azure 快照;复制正是那条已损坏的程序;LVM 快照虽存在但已过时六小时;WAL 归档从未实施。

后果:约六小时的最近提交和数据库状态无法恢复。GitLab 在 YouTube 上向约五千名同时在线观众直播了恢复过程,并发布了一个公开的 Google 文档跟踪进度,最终在约十八小时后从一个陈旧的预发布环境快照恢复了服务。

系统教训:存储可靠性是一种恢复属性,而非备份属性。只有当团队定期证明检查点、副本或转储文件能在事故条件下被定位、验证和恢复时,它们才算数。

该事件将检查点存储从一个容量问题转变为一个恢复设计问题。

示例场景:128 节点训练集群

一个由 128 个节点组成的训练集群每 10 分钟保存一个 1750 亿参数模型的检查点。每个检查点总计 1,750 GB,分布在所有节点上。

  • (此处将根据上下文给出计算细节)

分布式检查点协调

在包含 128 个节点的分片训练设置中,没有单个节点必然持有完整的训练状态。一个节点可能持有优化器状态的一个分片,而其他模型分区策略则将权重本身拆分到各设备上。因此,一个完整的检查点要求每个节点都保存其分片,且仅当所有分片都已持久写入时,检查点才算完成。这造成了一个协调问题:系统必须在训练恢复前确认所有 128 个节点均已完成写入。

最简单的方法是同步屏障:所有节点暂停训练,将其分片写入本地 NVMe,然后运行集群范围的确认操作以确认完成。仅当每个节点都确认其写入完成后,训练才会恢复。这种方法最大程度降低了不一致检查点的风险(即某些分片来自步骤 i,而其他分片来自步骤 i + 1),但它最大化了训练暂停时间,因为最慢的节点决定了屏障时间。

异步检查点通过在后台写入来减少训练暂停。每个节点将其分片快照到一个固定内存缓冲区(DRAM 内的快速副本),立即恢复训练,并在后台线程中将缓冲区写入 NVMe。该快照捕获了模型状态在某一时间点的一致性视图。风险在于,如果节点在异步写入期间发生故障,本地 NVMe 可能包含不完整的分片。缓解措施要求要么在认为检查点持久化前等待异步写入完成,要么在允许下一个训练步骤覆盖快照缓冲区前,将内存快照复制到对等节点。

在同步和异步检查点之间的选择取决于 T[write]T[iteration] 的比率。如果 T[write] 相对于 T[iteration] 很小(例如,1 秒对比每次迭代 200 毫秒,意味着检查点暂停消耗 5 次迭代),同步检查点是可以接受的。如果 T[write] 很大(例如,针对 5000 亿参数以上的模型为 30 秒),训练暂停将过于昂贵,异步检查点变得必不可少。

除了同步或异步方法的选择外,检查点格式优化可减少单节点 I/O 量。在最简单的数据并行设置中,每个节点持有模型权重的相同副本,因此朴素的检查点会从每个节点冗余地保存相同的数据。分片检查点协议通过保存分片而非完整副本来避免这种冗余。DeepSpeed 的零冗余检查点利用其 ZeRO 优化器状态分区:每个 H 个节点仅保存其 1/H 的优化器状态分片,系统保存一份合并的权重副本。PyTorch 的分布式检查点 (DCP) 协议遵循相同的存储原则,将每个秩的分片数据写入不同的文件。对于跨 256 个节点使用 ZeRO-3 的 1750 亿参数模型,结果是具体的:每个节点仅写入其约 6.8 GB 的分片(总状态 1,750 GB),将本地写入时间从分钟级降低到单个 NVMe 驱动器上约 1 秒,或在本地条带上远低于 1 秒。第 4.7 节 中的分级暂存推导已将这些单节点分片大小和写入时间贯穿于运行示例中,相同的算术逻辑也适用于分布式协调:暴露的流水线暂停即为本地 NVMe 写入,而持久复制到并行文件系统则与后续训练重叠。

随着检查点核算现已完全推导,我们可以按同一层级整合运行示例的存储需求,从而使训练数据集、检查点和归档谱系并排可见。

表 4.5 汇总了我们运行示例的完整存储图景:在 256 个节点上对 1750 亿参数模型进行 30 天的训练运行。

| 类别 || 容量 || 主存储层 |

| --- | --- | --- | --- |

| 训练数据集(压缩源) || 3 TB || 对象存储 |

| 训练数据集(分词后,每轮) || 6 TB || 对象存储 → NVMe 缓存 → 主机 DRAM |

| 模型权重 (FP16) || 350 GB || GPU HBM (分布式) |

| 优化器状态 (FP32) || 1,400 GB || GPU HBM (ZeRO 分区) |

| 单个检查点 (完整) || 1,750 GB || NVMe → PFS → 对象存储 |

| 所有检查点 (30 天, 10 分钟间隔) || 7.6 PB || NVMe (瞬态) → PFS (近期) → 对象存储 (归档) |

| 归档 (保留的检查点 + 数据集) || ~84.2 TB || Glacier |

表 4.5:1750 亿参数训练存储占用:30 天、256 节点训练运行中各数据类别的容量与主存储层。

流经层级的总数据量约为 7.6 PB 的检查点数据,加上每轮训练 6 TB 的数据读取量。对于单轮语言模型训练运行,检查点 I/O 以超过 1,260 倍的因子主导数据加载 I/O。

特征存储和模型注册表(管理特征服务和模型版本控制的运营系统)同样构建在此存储层级之上,依赖此处引入的相同分层、持久性和延迟边界。训练顺序流式读取其“燃料”并以突发方式写入检查点;而检索增强推理反转了该访问模式,用随机图遍历取代了顺序流式读取,因此相同的层级现在必须服务于其未针对性调优的工作负载。向量数据库(用于嵌入空间中近似最近邻搜索的专用存储系统)通过使该反转后的访问模式具备可操作性,为检索增强生成管线提供服务。

检索基础设施:向量索引

至此,存储层级已针对训练“燃料线”进行了优化:大分片、顺序读取和检查点突发写入。检索增强推理增加了一种存储工作负载,它属于本章范畴,因为它再次反转了访问模式。检索系统首先将文档转换为表示语义含义的向量,然后通过查找附近的向量并返回其源文档作为上下文来服务推理请求。存储系统不再通过 TFRecord 或 Parquet 等格式流式传输顺序 Blob,而是现在遍历高维向量图。

这种架构模式将主要存储瓶颈从顺序读带宽转移到了随机访问 IOPS。向量数据库必须遍历高维图索引,以在数十亿个嵌入中查找近似最近邻。性能受一项严格的物理约束支配:召回率 vs. 延迟 vs. 容量 的权衡。最大化召回率需要遍历索引的更大部分,从而推高 IOPS 和延迟。索引之所以庞大,是因为每个嵌入文档块占用大量存储:一个 1,536 维的 FP32 嵌入(每个值 4 字节)每个向量约消耗 6.1 KB,因此十亿文档语料库在包含用于快速导航的图邻接表之前,仅原始向量存储就需要约 6.1 TB。如果这些邻接表每个节点增加 50 到 100 字节,它们将增加约 50 GB 到 100 GB,使原始向量加最小图元数据在高端时对于十亿条目语料库约为 6.2 TB。完整的生产索引一旦包含标识符、图层级、分配器开销、已删除条目松弛空间和复制,可能会更大,但那些是实现开销,而非原始向量加邻接表的占用。在这个规模下,索引不再适合放入 DRAM,这是驱动基于磁盘替代方案的设计压力。基于内存的层级可导航小世界 (HNSW) 图索引在索引适合 DRAM 时使相同的图导航思路变得快速 (Malkov and Yashunin 2020),而混合搜索将向量相似性与词法或元数据过滤器结合,因此同时考验随机向量遍历和文档获取。基于磁盘的图索引(如 DiskANN)展示了十亿规模搜索如何在以高召回率和低查询延迟为目标的前提下,用 SSD 支持的图遍历权衡 DRAM 占用 (Subramanya et al. 2019)。

这使得检索基础设施成为训练燃料线的一个有力反例。相同的层级仍然适用,但目标发生了变化:训练存储通过预取隐藏顺序带宽,而向量搜索将存储预算花费在保持随机图遍历在服务延迟 SLO 内的索引层级上。只有当额外的 SSD 遍历不会抵消搜索更大语料库带来的质量增益时,基于磁盘的索引才具备经济性。因此,检索工作负载在反转目标的同时复用了相同的物理层级,从而闭合了燃料线打开的访问模式对比。

合成燃料线

当车队耗尽人工撰写的训练数据供应时,存储架构将再次变革,因为合成数据必须自带溯源。随着高质量人工生成文本变得更难扩展,车队遭遇数据墙:即额外训练样本必须越来越多地来自合成、模拟、策展或昂贵的人工收集的临界点 (Villalobos et al. 2022)。这种从简单收集到托管数据生成的转变,让不仅仅是负载字节流经层级。一条“合成燃料线”还必须存储每个样本的溯源链:由哪个模型生成、哪些过滤器接受或拒绝了它、哪个裁判或验证器模型对其打分、以及哪个下游运行消费了它。没有该谱系,系统就面临模型崩溃的风险,即模型因在其自身未经验证的错误上训练而退化 (Shumailov et al. 2024)。

问题:计算一个需要加密谱系和多模型验证的 1 TB 合成数据集的存储放大。

  1. 原始负载:1 TB。

  2. 溯源开销:40% 额外开销,用于谱系哈希、生成日志和奖励模型分数。

  3. 验证因子:为避免“自我中毒”,每个样本由 3 个独立的“裁判”模型验证。

  4. 结果:总占用 = 1 TB × 1.4 × 3 = 4.2 TB。

系统洞察:合成数据是经验证数据,而验证有 4.2× 的存储税。在机器学习车队中,存储从一个简单的比特桶变成了溯源引擎。一个无法存储合成 Token 背后的为什么的系统,冒着用自己的过去错误毒害车队未来的风险。

合成燃料线通过使谱系成为负载的一部分,闭合了本章的存储弧线。大规模训练、检查点、检索和合成数据生成考验不同的层级,但每种失效模式都源于同一个错误:将存储视为被动容量,而非控制吞吐、可复现性和模型质量的主动系统。

谬误与陷阱

存储设计错误是 ML 基础设施中代价最高的错误之一,因为它们发现得太晚(训练开始时),而补救需要昂贵的硬件升级或耗时的数据重新格式化。以下谬误和陷阱捕捉了经验丰富的存储工程师初次面对 ML 工作负载时最常犯的错误。

误区:NVMe 足够快,可直接供给 GPU 而无需流水线

单个 NVMe 驱动器提供 7 GB/s 的带宽,而 H100 从 HBM 消费数据的速率高达 3.35 TB/s。两者差距约为 478.6 倍。即便四块驱动器组建 RAID-0,在扣除开销前也仅能将差距缩小至约 120 倍。若没有流水线和预取机制来隐藏从 NVMe 到主机 DRAM 再到 HBM 的加载延迟,每次批次传输都会引入等同于传输时间的停顿。NVMe 速度虽属必要条件,但远非充分条件;完整的流水线架构(多工作器加载、预取缓冲区、异步传输)之所以存在,正是因为没有任何单一存储设备能匹配加速器的带宽。之所以产生这种混淆,是因为 NVMe 带宽以绝对值(GB/s)标称,听起来很惊人,但相关指标是存储带宽与加速器消费速率的比率,而这一比率不利达两到三个数量级。

陷阱:依赖预取来隐藏对象存储延迟

预取掩盖的是平均延迟,而非尾部延迟。在拥塞或跨区域故障转移期间,对象存储的 P99 延迟可能飙升至数百毫秒。深度为 Q[prefetch] 的预取缓冲区可吸收高达 Q[prefetch] × T[compute] 毫秒的方差。若尾部延迟尖峰超出此窗口,缓冲区将耗尽,导致加速器停顿。对于运行数周的训练作业,即使罕见的 P99.9 尖峰在大规模下每天也会发生数千次,从而显著降低利用率。缓解措施不是假设预取能解决所有问题,而是根据实测尾部延迟配置缓冲区深度,并在可能的情况下将数据暂存至本地 NVMe。

误区:海量小训练文件在大规模下无害

将每个训练样本存为单独文件(每张图片一个 JPEG,每个文本样本一个 JSON)符合数据采集直觉,却会给大规模训练带来灾难。每个文件都需要元数据操作(open()stat()close()),这些操作在元数据服务器上串行化。在 1 万个并发工作器下,元数据操作会在数据带宽饱和前很久就成为瓶颈。解决方案是在预处理阶段将样本聚合为大型顺序分片(TFRecord、WebDataset tar、Parquet)。这能将元数据操作减少 10,000 至 100,000 倍,并将随机访问模式转变为顺序流式读取。

陷阱:设计存储架构时忽略出站流量成本

云出站流量费用在开发阶段不可见,但在大规模下可能占据主导地位。按具有代表性的 $0.09/GB 跨边界出站费计算,一个从对象存储读取 100 TB 数据集 10 次的训练作业将产生 $90,000 的出站费用,远超上述对象存储价格假设下的年存储成本 $27,600。当读取跨越区域、供应商或互联网边界时会产生这些费用;同区域读取可能免费或打折,这正是成本隐藏至跨区域或多云拓扑使其不可避免时的原因。团队若在对象存储上用小数据集原型验证再放量,会发现其存储架构在经济上不可持续。对策是显性预算出站成本,并设计数据流水线以最小化跨层传输——通常在作业启动时将数据暂存至本地 NVMe,而非每个周期都从对象存储流式读取。

误区:平均写带宽足以满足检查点设计

检查点由所有节点同时写入,产生的突发流量可超过稳态带宽 10 倍以上。若并行文件系统按平均写负载规模化,将在检查点风暴期间成为瓶颈,拉长 T[write] 并降低有效训练吞吐量。并行文件系统必须按峰值突发带宽而非平均带宽预配,即使该突发大部分时间处于闲置。

陷阱:假设更快的加速器自动提升训练吞吐量

组织从 A100 升级到 H100 GPU 时,预期训练时间将按计算力提升比例缩短。若存储流水线在 A100 上已处于临界状态——仅能刚好跟上以避免停顿——则更快的 H100 计算只会更早撞上 I/O 墙。加速器升级缩短了 T[compute] 却未减少 T[I/O],导致数据停顿比上升。每次加速器升级都须伴随存储流水线审计,确保更强算力不至于暴露原本隐蔽的存储瓶颈。

误区:针对数据库优化的 RAID 配置适用于 ML 训练

面向数据库的 RAID 配置(如 RAID 5 或 RAID 6)以冗余优先于原始带宽,因奇偶校验计算导致 15–25% 的写性能损耗。对于训练数据不可变且已在持久对象存储中备份的 ML 工作负载,此开销纯属多余。对于供加速器使用的本地 NVMe 缓存,正确选择是 RAID 0(条带化),它在无奇偶校验开销下结合所有驱动器带宽,最大化顺序读吞吐。本地 NVMe 上的数据是缓存,而非数据源;因驱动器故障丢失它仅需重新暂存操作,不会造成数据丢失。

陷阱:为每个数据集快照创建全量副本

朴素的数据集版本控制在大规模下在财务和运维上均不可持续。对 10 TB 数据集,每个版本创建全量副本会迅速耗尽存储预算并拖慢实验周期。版本间若有 5% 差异,则存储了 9.5 TB 冗余数据。正确做法采用内容寻址存储或增量编码,仅存储变更样本。这将新版本存储成本从 10 TB 降至约 500 GB,实现 20 倍缩减。DVC(数据版本控制)和 lakeFS 等工具实现了这一模式,在不复制未变内容的前提下追踪数据集谱系。

误区:压缩总有益,因为它减少了读取数据量

这仅在流水线为 I/O 瓶颈时成立。压缩虽减少 I/O 体量,却因解压增加 CPU 负载。对于 CPU 瓶颈的流水线,复杂的数据增强已使主机处理器饱和,叠加解压工作会让瓶颈恶化而非改善,导致输送给加速器的吞吐量净下降。正确做法是先剖析流水线:若主机 CPU 是瓶颈,采用非压缩或轻度压缩格式(LZ4);若存储或网络 I/O 是瓶颈,采用高压缩比格式(zstd、gzip)。最优压缩级别取决于流水线的瓶颈所在,而非放之四海而皆准的常数。

陷阱:在剖析流水线瓶颈前就选定压缩格式

团队常为运维简便而标准化单一压缩格式,随后不顾访问模式将其套用于所有数据集。这种捷径掩盖了真问题:作业到底是在为搬运的字节付费,还是为解码消耗的周期付费。存储流水线应在测量了代表性批次上的 CPU 利用率、存储带宽、网络带宽及加速器空闲时间再选择压缩格式;否则该格式沦为又一固定假设,只是搬移瓶颈而非消除瓶颈。这些误区最终都指向同一核算练习:每字节驻留何处、移动多频、由哪层承担成本。

总结

ML 系统中的存储

ML 系统中的存储并非被动的数据仓库;它是一个主动的、多层级的流水线,其唯一目的就是让加速器的 HBM 时刻充满数据。从 HBM、主机 DRAM、本地 NVMe、并行文件系统、对象存储到冷归档存储,这一层级体系之所以存在,是因为没有任何单一技术能同时提供大规模训练所要求的带宽、容量和成本特性。HBM 与冷归档存储的聚合带宽相差大约三个数量级,而每个中间层级都作为一个暂存缓冲区,用于吸收加速器消费数据的速率与持久化存储供给数据的速率之间的错配。数据流水线吞吐量公式 BW_required = N_GPU × η_target × D_vol,batch / T_iteration 为每一层级的规模规划提供了定量基础:任何一级未能满足所需带宽都会导致昂贵的加速器空闲;过度配置则会造成资本浪费在闲置的存储容量上。

颠覆性的存储假设

ML 存储之所以极具挑战性,是因为这些负载颠覆了几十年存储系统设计中烙印的几乎所有假设。数据库针对随机 IOPS、可缓存的工作集和持续的小块写入进行优化。而 ML 训练要求在远超所有缓存级别容量的数据集上实现顺序流式吞吐,并伴随着巨大的检查点突发写入,这些写入会在几秒钟内饱和带宽,随后又归于寂静。在传统负载中无害的小文件元数据开销,当每个 epoch 需要打开、stat、关闭数百万个独立样本时,便成为主导性瓶颈。将样本聚合为大分片、选择 WebDatasetTFRecord 等顺序流式格式、以及针对 P99 尾延迟而非平均延迟设计预取流水线,所有这些都是对这些颠覆性访问模式的直接工程响应。GPUDirect Storage 将这种优化推向更远,它彻底将 CPU 从数据路径中移除,为训练流水线同样需要的数据增强工作释放了主机核心。

层级的经济学

层级的经济学同样影响深远。HBM 与归档存储之间数个数量级的成本差异强制要求分层策略,但数据交付的真实成本远不止每吉字节的存储价格。从对象存储移出数据的出站费用、元数据密集型访问模式的 IOPS 费用,以及加速器因等待缓慢检查点而空闲的机会成本,所有这些都计入总拥有成本。检查点暂存——模型先写入快速本地 NVMe,再异步复制到共享存储——正是精心设计的流水线如何将训练暂停时间与底层文件系统的性能限制解耦的典范。

诊断框架

内化存储层级的工程师将获得一个系统的训练性能诊断框架。当集群报告加速器利用率低时,直觉往往指向计算配置或模型本身。但在实践中,根因往往埋藏在数据路径中:数据加载器读取单个文件而非分片、预取缓冲区针对平均延迟而非尾延迟规模化、检查点策略在写入缓慢的并行文件系统时阻塞训练、或 NUMA 非感知的内存分配使有效 DRAM 带宽减半。诊断这些故障需要理解哪个层级是瓶颈以及为何成为瓶颈,而流水线公式和存储层级框架使这个问题变得可解。

并行策略的存储特征

这种诊断视角直接关联到训练作业如何在集群中分片,因为分片策略决定了施加于存储的需求模式。数据并行布局在每个节点复制模型,产生统一的读取模式,但每个节点都保存相同参数,导致巨大的检查点冗余。流水线式布局在各阶段间创建顺序依赖,早期阶段的数据加载停滞会级联向下,导致下游工作饥饿。专家式布局创建非均匀访问模式,不同加速器组在任何给定时间需要不同的数据子集。每种策略都施加独特的存储特征,这正是存储无法独立于其所服务的计算而设计的原因。

核心原则

  • HBM 是终点:每一级较低层级的存在都是为了让 GPU HBM 保持充满。HBM 与对象存储之间约 30 倍的聚合带宽差距(一旦单个实例从共享对象端点拉取数据,单客户端的差距更大)驱动着 ML 存储层级的每一个设计决策。

  • ML 负载颠覆存储假设:传统的缓存、IOPS 优化和小写入持久化模式对 ML 训练都得出错误答案。吞吐量 (GB/s) 比 IOPS 更重要,小文件使元数据成为首个瓶颈;聚合分片可将元数据操作减少数个数量级。

  • 流水线公式支配设计:所需带宽随 GPU 数量线性扩展,随迭代时间反比扩展。使用 BW_required = N_GPU × η_target × D_vol,batch / T_iteration 为每一层级确定规模。

  • 流水线隐藏的是平均延迟而非尾延迟:预取缓冲区必须按 P99 I/O 延迟而非平均延迟规模化,以防止大规模部署时加速器停顿。缓冲区深度 ⌈T_I/O,p99 / T_compute⌉ 是最小值。

  • GPUDirect Storage 消除 CPU 瓶颈:GDS 绕过数据路径中的 CPU,将每次传输的延迟降低 4 倍,并为增强工作释放 CPU 核心。

  • 检查点暂存最小化 T_write:先写入本地 NVMe,再异步复制到共享存储。这将检查点暂停时间与并行文件系统性能解耦。

  • 经济学驱动分层:本地 NVMe 与对象存储的 5 倍成本差异(以及从 HBM 到归档存储超过 3,000 倍的跨度)强制要求分层策略。出站费用常超存储费用;应为数据交付总成本设计,而非仅考虑存储成本。

总结

本卷其他章节优化的都是加速器所做的事。本章优化的是到达加速器的东西。整个存储层级的存在,是为了让大楼里最昂贵的资源不再空转,这重新定义了性能优化的战场:不在计算上——计算极少是瓶颈——而在于供养它的供应链。在舰队规模下,约束因素是数据到达的速率,而非处理的速率,这也是为何集群可能在每个加速器都健康的情况下仅跑到一半利用率。系统只能跑得和供给它的那一层一样快。

基础设施已完备。我们打造了以 petaFLOP/s 计算的加速器(第 2 章),用以 TB/s 移动梯度的网络互联将它们连接起来(第 3 章),并建立了供给它们数据的存储层级。下一个挑战不再是硬件问题,而是算法问题:如何将单个训练作业划分到这数千个资源上。第 5 章将探讨在舰队中分割计算的并行策略(数据、张量、流水线和专家并行)。

此处用于确保测验在部分开始前正确插入。

  • 针对特定负载,应如何设计分层存储流水线以保持加速器 HBM 被填满?

  • ML 负载如何颠覆数据库时代的存储假设,以及随之而来的数据格式选择是什么?

  • 何时应将数据预暂存至本地 NVMe,而非从并行文件系统或对象存储流式传输?

  • 检查点存储应如何在暂停时间、持久性、恢复验证和保留成本之间取得平衡?

分布式 ML 原则

协调集群:分布式逻辑

协调集群需要解决单机无需面对的问题。如果第一部分构建了物理基础设施,那么第二部分确立了分布式逻辑:将一组独立加速器转化为一个统一、整体的机器所需的算法、协议和协调策略。分布式学习的物理规律支配着这一转变,提供了基于第一性原理的推理,解释了为何增加 GPU 并不总能带来线性加速,以及为何通信而非计算往往是隐藏的瓶颈。

在此规模下,工程挑战是在并行化与协调之间做权衡。我们可以对模型进行分区以减轻单设备内存压力,但这必然会增加保持权重同步所需的通信量。我们可以扩展到数千个节点以缩短训练时间,但这同时会降低平均故障间隔时间(MTBF),使得容错成为强制性的系统组件,而非可选的运维奢侈品。这些原则分解了第一部分引入的通信税,并增加了规模带来的可靠性税。

规模经济学

第一个压力源于规模经济学本身:前沿能力需要算力,但每一次改进的成本都在逐步升高。

不变量:对于前沿基础模型,损失函数(ℒ)随算力(C)、数据集规模(D)和参数量(P)按幂律函数改进,γ 代表经验缩放指数。用于说明边际效应递减的仅算力简化公式为:ℒ(C) ∝ C^(−γ)

启示:在模型能力的前沿,规模是约束性需求,而非优化选项。因为损失随算力亚线性改进,每一次同等幅度的改进都需要成倍更多的算力;这种边际收益递减驱动了从单服务器训练向仓库级集群的转变。

步耗权衡

对规模的需求将每个训练步骤变成了更快的本地工作与更慢的全局协调之间的博弈。

不变量:跨 N 个工作节点或加速器完成一个训练步骤的时间,是可并行化计算、通信、同步以及可隐藏的重叠部分之和,其中 0 ≤ ToverlapTcomm(N) + Tsync(N)。

\[T_{\text{step}}(N) = \frac{T_{\text{compute}}}{N} + T_{\text{comm}}(N) + T_{\text{sync}}(N) - T_{\text{overlap}} \]

启示:扩展是在可并行化计算(在理想分区下随 N 缩减)、通信开销和同步开销之间的竞赛。为了高效扩展,算法必须减少或分摊 Tcomm(N) 和 Tsync(N)(例如通过梯度累积,它改变了有效批大小和收敛机制,或通过压缩,但这可能引入数值误差),并通过通信隐藏和流水线执行最大化 Toverlap

通信物理:延迟与带宽

该博弈中的通信项由互联网络的延迟和带宽决定,因此消息大小决定了哪种优化有效。

不变量:通信时间是固定延迟(α)、消息大小(n)和消息相关带宽(β)的函数。

\[T(n) = \alpha + \frac{n}{\beta} \]

启示:小消息——例如混合专家模型(MoE)路由元数据——受延迟限制;大消息(例如梯度)受带宽限制。优化策略必须匹配对应模式:融合小消息以分摊 α,压缩大消息以提升 β

容错与检查点

即便通信调优完毕,更大规模的集群故障率也足够高,以至于必须保护训练进度本身。

不变量:最优检查点间隔(τopt) 平衡了写入检查点的成本(Twrite)与因系统级故障导致的进度丢失重做的预期成本(MTBFsystem)。

\[\tau_{\text{opt}} = \sqrt{2 \cdot T_{\text{write}} \cdot \text{MTBF}_{\text{system}}} \]

启示:检查点并非“免费”。随着集群规模增长,MTBF 下降,迫使更频繁的检查点。这要求高带宽存储(突发缓冲区)以防止 I/O 主导训练时间。

C³ 分类法:开销守恒

容错是一条更广泛规则的一个例子:分布式系统很少消除开销,只是转移它。

不变量:分布式机器学习系统中的开销无法被消除,只能在计算、通信和协调(C³ 分类法)之间转移。减少其中一项必然会增加至少另一项。

启示:异步训练消除了协调屏障(Tsync(N) → 0),但引入了梯度陈旧,表现为额外的训练迭代(Tcompute↑)。流水线并行减少了通信量,但增加了流水线气泡时间(Tsync(N)↑)。分布式系统中没有免费的午餐;C³ 分类法揭示了系统在何处支付代价。

确立这些不变量后,第二部分阐述了集群如何协调分布式训练工作:将工作负载跨数千个加速器分区的并行策略(数据、张量、流水线和专家并行),将这些加速器绑定成一个连贯计算机的集合操作(AllReduceAllGatherAllToAll),吸收规模带来的必然故障的容错机制,以及将集群分配给工作负载的编排层。这些章节共同发展了将第一部分的物理集群转化为工作的分布式训练系统所需的算法和运维纪律。

分布式训练

展示大规模工作负载沿数据、张量和流水线轴跨加速器分区及同步路径的等轴测分布式训练地图。

目的

为何“更多硬件 = 更快训练”的线性逻辑会在剖分带宽墙前崩塌?

分布式训练表面看似简单:将工作拆分到各机器并合并结果。然而,随着机器学习集群规模增长,一种新的物理规律浮现。通信成本随机器数量扩大,而单机计算量缩减,直到同步开销占主导地位,增加硬件反而主动降低性能。当作业在单个加速器上训练时,设计优化目标是算术强度;当在 10,000 个加速器上训练时,设计优化目标变为通信强度。扩展天花板并非待修复的 Bug,而是可靠性鸿沟和通信计算比的基本属性:协调独立机器要求在比片上内存慢几个数量级的网络上移动 TB 级状态。分布式训练的艺术在于管理这种张力——分区工作以最小化协调税,重叠通信与计算以隐藏延迟,并选择在一致性与吞吐之间取得平衡的同步策略。缺乏这种理解,组织将在闲置等待梯度到达的硬件上浪费数百万,或因陈旧更新破坏优化而产出永不收敛的模型。分布式训练是 C³ 分类法的动态体现:每个设计选择都在并行化计算、网络受限通信和同步协调之间进行权衡。

  • 应用集群定律诊断分布式训练中的计算受限、内存受限、通信受限或协调受限模式

  • 计算数据并行训练作业的扩展效率、临界批大小限制和同步开销

  • 设计以额外通信换取加速器容量的内存分片方案

  • 根据硬件带宽层级和模型结构映射张量、流水线、专家和数据并行

  • 构建在保持吞吐和收敛稳定性的同时最小化气泡的微批次流水线调度

  • 利用过时性、掉队者容忍度、数值稳定性和集群异构性来选择同步和低精度策略

  • 为前沿模型训练、推荐模型训练和对齐训练的约束条件综合混合并行配置

为什么分布式训练是必要的

加速器层级、网络结构和存储管道构成了机群的物理基础。剩下的挑战是算法层面的:如何将单个训练作业跨数千个资源进行分区,同时又不丢失单一连贯优化过程的语义。通用扩展定律(原理)解释了为什么这种压力持续增加:前沿模型质量的提升需要不成比例地更多算力、数据和参数,因此训练问题最终会超出任何单台机器的承载能力。

如果一台加速器拥有 100 TB 内存和 1 EFLOPS 算力,分布式训练就变得没有必要。现实系统则施加了有限的 HBM 容量、有限的互联带宽和有限的故障预算,因此训练必须分区到许多独立芯片上。在图 1.13 所示的机群栈框架中,分布式训练代表分发层:将数学工作负载分区到物理机群的逻辑。此处定义的策略,包括数据、张量、流水线和混合并行,创造了互联网络必须承载的流量模式。

机群的物理特性

在优化算法之前,我们必须了解机器学习机群的物理约束。任何分布式训练作业的性能都受机群定律(原理)支配,该定律在第 1.5.2 节中引入,它将每步时间分解为:

T_step(N) = T_compute / N + T_comm(N) + T_sync(N) - T_overlap

关键项是通信计算比ρ = Tcomm/(T[compute]/N)。这个比率决定了一个集群表现为超级计算机,还是一堆空转的加热器。

两个等宽区域标记为计算和通信,中间有一条虚线 rho 等于 1 的阈值,分隔理想扩展和等待区域。

超过通信计算阈值后,扩展不再理想。

由此比率衍生出两种状态。在计算受限集群中,T[compute]/NTcomm,大部分时间花在矩阵乘法上;这是理想状态,典型场景是在稠密模型(如 ResNet)上使用大批量大小。在通信受限集群中,Tcomm ≈ T[compute]/N,大量时间花在等待梯度或激活值到达上,这是大语言模型(LLM)和 DLRM 风格推荐模型的常见状态,参数同步会饱和网络。第 B.5.2 节将此比率纳入诊断框架,根据测量带宽和算术强度将工作负载分类为计算受限、内存受限或通信受限,将上述定性状态转化为可重复的测试。

多机训练需求

三个具体信号表明何时分布式训练变成了必要而非仅仅是有益。第一个信号是内存耗尽:模型参数、优化器状态和激活存储超过了单设备容量。对于使用 Adam 的全混合精度训练,仅参数、梯度和优化器状态的预算,在激活值计入前,大约在 50 亿参数时就会超过 80 GB 加速器的容量;更大的 100-200 亿参数模型需要分片、卸载或其他节内存技术(Rajbhandari et al. 2020)。第 G.1 节记录了本章预算中使用的 80 GB H100 和 A100 容量数据,以便读者将每一个内存上限追溯到单一有文档记载的来源。

第二个信号是不可接受的训练时长。即使模型能放下,单设备训练可能需要数周或数月才能收敛,使得墙钟时间本身成为一种系统约束。GPT-3 的 1750 亿参数训练运行使用了 V100 GPU 集群(Brown et al. 2020),这说明了为什么日历(而不仅仅是 HBM 容量)迫使这种规模进行分布式训练。

第三个信号是数据集规模。当训练数据达到数 TB 量级,如在大规模视觉或语言建模任务中,单台机器不再是存储或输入吞吐的自然单元。分布式训练此时既是喂养模型的方式,也是容纳模型的方式。

分布式训练复杂性权衡

分布式通过引入单机内不存在的成本,改变了优化问题。三个复杂度维度决定并行策略是否可行,因为其中任何一个都可能成为瓶颈上限。通信开销是同步梯度的成本:对于拥有 P 个参数、分布在 N 个设备上的模型,all-reduce 操作每步必须传输约 2P(N − 1)/N 个梯度值,再乘以每个梯度元素的字节数,在商用网络上这可能主导计算时间。容错随着集群增长变得更难,因为单位时间的预期故障数随集群规模线性上升,而整个集群在一个区间内零故障存活的概率呈指数级衰减;若一个 100 节点集群的单节点小时存活率为 99.9%,集群级故障概率约为每小时 9.5%,对应 MTBF 约 10 小时。第 E.1.2 节从单节点存活率推导了这种级联效应至集群 MTBF,并通过完整示例演示了相同计算,以便读者复现为什么 MTBF 随集群增长而崩溃。算法稳定性完善了这组因素,因为数据并行带来的大批量大小影响收敛行为,需要单机训练不需要的学习率缩放和预热策略(Goyal et al. 2017)。综上,这些成本解释了为什么分布式训练是一个约束满足问题,而非硬件乘法练习。

从单机到分布式的过渡

为单机训练建立的系统化优化方法论可扩展至分布式环境,但需进行重要调整。性能分析除捕获计算和内存指标外,还必须捕获设备间通信模式和同步开销。解空间扩展为包含数据并行、模型并行、流水线并行和混合方法。图 5.1 可视化了这个三维配置空间。

图 5.1:**3D 并行立方体**:三个正交扩展轴的概念可视化:数据并行(复制模型)、流水线并行(分割深度)和张量并行(分割层)。大模型训练在这个立方体中占据特定坐标点(*d*, *p*, *t*),以平衡内存使用、计算效率和通信开销。

图 5.1 的关键洞见是总加速器数量 N[total] = d × p × t:每个轴相互独立,训练系统根据目标模型和集群的内存、计算及带宽约束,在这个立方体中选择特定坐标。

工程权衡:选择并行策略

选择正确的并行策略

选择正确的并行策略并非偏好问题;它是一个受参数量(P)、批次大小(B)和互联带宽支配的约束满足问题。表 5.1 量化了每种策略的并行通信成本,揭示了哪些方法在给定硬件拓扑下是物理上可行的。

| 策略 | 通信模式 | 通信量 | 硬件约束 |

| --- | --- | --- | --- |

| 数据并行 (DP) | AllReduce 梯度 | ∝ M[grad](梯度字节数) | 需高分切带宽 |

| 张量并行 (TP) | AllReduce 激活值 | ∝ B × N[L](层数) | 关键:需 NVLink |

| 流水线并行 (PP) | 点对点 (P2P) | ∝ B × S × d[model](激活值) | 低带宽(以太网足矣) |

表 5.1:并行通信成本:张量并行具有最高的通信频率(逐层),将其限制在节点内部(NVLink)使用。流水线并行具有最低的通信量(仅边界激活值),使其适合跨节点(以太网/InfiniBand)扩展。数据并行处于中间地带,但随着通信的梯度字节数 M[grad] 增长,其扩展性变差,最终需要内存高效的分片。

带宽要求对硬件布局施加了硬性约束:每种并行策略都必须与需要通信的 GPU 之间可用的互联相匹配,否则设计在离开白板前就已失败。快速的“信封背面”带宽可行性检查可在编写任何代码前捕获这种不匹配。

跨越由标准以太网连接的服务器机架进行张量并行将会停滞。通信量随批次大小 B 和层数 N[L] 缩放。它需要 NVLink 600 GB/s–900 GB/s 的吞吐量。对于跨机架扩展,设计必须切换到流水线或数据并行,以尊重网络的物理特性。

图 5.2 将此约束满足过程形式化为决策树,展示了模型大小和硬件拓扑如何决定可行的并行策略。

图 5.2:并行策略决策树:从模型的内存需求出发,该树指导实践者完成确定哪些并行策略在物理上可行的约束满足过程。关键分支点在于模型内存与单 GPU 容量的对比,以及通信带宽与并行需求的对比。叶节点标注了主导的硬件约束。

决策树揭示,并行策略的选择并非偏好,而是物理约束的结果。这是一个约束层面的预览:上述通信量公式指出了每种策略移动的内容,但其完整含义取决于本章其余部分将展开的张量、流水线和混合并行的机制。5.9 节在机制就绪后完成最终筛选。下一个问题是这些约束如何塑造真实集群上分布式训练步骤的机制。

分布式训练步骤

上一节的选择约束假设作业必须被拆分;现在的问题是分布式步骤必须保证什么。核心挑战在于确保 1,024 块完全独立运行的 GPU,在每次训练迭代结束时,就一组单一、数学上严谨的更新权重达成一致。

分布式训练 是一种训练方法论,它将优化循环划分到多个计算节点上(分发数据、模型层或单个张量操作),并通过同步通信原语协调其输出,以产生单一连贯的模型。

  1. 意义:当模型的内存需求超过单个加速器容量时,分布式训练变得必要。GPT-3(1750 亿参数)在 BF16 下需要约 350 GB——超过单张 H100 80 GB 容量的 4 倍。仅模型分片就需要至少 5 张 H100 来训练它,而大规模运行可能使用数百或数千个加速器以达到可接受的墙钟时间。

  2. 区别:与用于独立请求的分布式系统(Web 服务、数据库读取)不同,后者节点间不共享可变状态,分布式训练要求每个节点维护模型参数的一致视图——使得梯度同步成为强制性的协调步骤,而非可选的优化。

  3. 常见陷阱:一个常见的误解是分布式训练随节点数线性扩展。实际上,通信开销随集群规模和每步串行部分增长(阿姆达尔定律):若 30% 的步骤时间花在同步上,理论扩展上限为 1/0.30 ≈ 3×,无论增加多少加速器。

一个有用的心智模型将这些分布式策略视为循环变换,这是编译器用于优化顺序代码的相同概念工具箱。若将训练过程视为对数据和层的大规模循环,分布式策略就是由集群级编译器应用的循环变换。逻辑训练循环嵌套三个迭代器(epoch、批次、层),每种并行策略将其中一个跨设备展开:

数据并行是 parallel for 循环,将外层循环(批次维度)跨设备展开,使每个设备在不同数据索引上运行相同的代码体。张量并行是向量化,即单指令多数据流(SIMD),将内层循环(矩阵乘法)跨设备拆分,形成集群级 SIMD,其中 NVLink 充当向量寄存器文件。流水线并行是指令流水线,将顺序操作(层)跨设备拆分;正如 CPU 流水线分级取指、译码、执行,集群将第 1 层、第 2 层、第 3 层分级,以保持所有 ALU 忙碌。

无论策略展开哪个循环,分布式训练[⁵⁷]都将工作负载分散到必须协调以训练单一模型的机器上。协调意味着保持每个模型分片或副本在兼容的训练步骤上。基本屏障可使小规模研究运行保持有序,但长期运行的训练作业还需要超时、检查点和恢复机制,以免单个工作节点故障浪费数天算力。第 7 章将深入探讨这些可靠性工程挑战。

扩展路径的每个台阶都继承了前一阶段的挑战。单 GPU 训练仅需本地内存管理和前向/后向传播。扩展至节点内多 GPU 增加了高带宽通信,通过 NVLink[⁵⁸] 或带 NCCL[⁵⁹] 优化的 PCIe 处理,同时保持单机容错和调度。

跨越到多节点训练增加了网络通信开销、容错需求和集群编排。因为每个阶段都会叠加前一阶段的瓶颈,必须在横向扩展前优化单 GPU 性能,以免低效在整个集群中放大。尽管框架通过分片数据并行和通信库抽象了大量细节,高效实现分布式训练仍需精心的网络配置(InfiniBand 调优、拓扑感知路由)、通过集群调度器进行基础设施管理,以及调试非本地问题,如同步挂起和通信瓶颈。

尽管存在这种复杂性,但核心工作流在机制上是直观直接的;工程挑战在于如何使其在大规模下实现快速且可靠。反复出现的成本是跨设备聚合结果的梯度同步,随着系统规模扩大,这种开销会不断累积,这一点在 5.4 节 中得到了量化。

有四种方法可针对不同的约束条件。数据并行将训练数据划分到各台机器上,同时每台机器维护一份完整的模型副本,这使其成为适合单设备内存的模型的最简单方法。模型并行在参数超出单设备内存时,将模型本身拆分到各设备上。流水线并行将模型划分为按顺序处理微批次的阶段并发执行,相比于朴素的模型并行提高了利用率。混合方法集成多种策略,使得在任何单一方法都会失效的规模下进行训练成为可能。只有当前一种策略触及物理上限后,下一种策略才变得必要。

数据并行

最简单的方法是给每个 GPU 一份完整、相同的模型副本,并分配给它一份独特的数据切片。数据并行是分布式训练的自然起点,因为它只需要对单设备训练循环进行最小改动。

数据并行 是一种分布式训练策略,其中每个工作进程持有一份完整的模型副本,并处理一个独立的小批次分片,随后通过 AllReduce 同步梯度更新,使所有副本在每一步都应用相同的参数变更。

  1. 意义N 个工作进程各自处理批次大小为 B 的数据,有效全局批次大小为 N × B,在保持单工作进程内存不变的情况下线性扩展吞吐量。对于一个 10 亿参数、BF16 下占 2 GB 的模型,1,024 个工作进程可实现单 GPU 吞吐量的 1,024 倍——直到梯度 AllReduce(环形最优下每工作进程 2(N − 1)/N,每步 2 GB)超过了反向计算时间,从而产生通信瓶颈。

  2. 区别:与将参数分区以致没有单个工作进程持有完整模型的模型并行不同,数据并行要求每个工作进程都有足够的内存容量来存储完整的模型状态——这使得在不结合优化器状态或参数分片的情况下,它不适用于大于单个加速器内存的模型。

  3. 常见陷阱:一个常见的误解是数据并行可以随工作进程数量无限扩展。将有效批次大小 B 扩展到与工作负载相关的临界批次大小之上会降低统计效率,需要更多训练步数才能达到目标损失,从而抵消增加工作进程带来的吞吐量增益 (Shallue et al. 2019)。

每个设备使用其分配的数据子集训练一份完整的模型副本。例如,使用 4 个 GPU 在 100 万张图像上训练图像分类模型时,每个 GPU 处理 25 万张图像,同时维护一份相同的模型架构副本。

当数据集规模大但模型规模可控时,数据并行最为有效,因为每个设备必须在内存中存储一份完整的模型副本。该方法广泛用于图像分类和自然语言处理,在这些场景下数据集可以并行处理,且数据样本之间没有依赖关系。例如在 ImageNet 上训练 ResNet 模型 (He et al. 2016) 时,每个 GPU 可以独立处理其分配的图像部分,因为一张图像的分类不依赖于另一张图像的结果。

数据并行的有效性源于随机梯度下降的一个性质。在不同小批次上计算的梯度可以求平均,同时保持与单设备训练在数学上的等价性。该性质支持跨设备的并行计算,其数学基础直接源于期望的线性性。

考虑一个参数为 θ、在数据集 D 上训练的模型。单个数据点 x[i] 的损失函数为 ℒ(θ, x[i])。在批次大小为 B 的标准 SGD 中,小批次的梯度更新为:

g = \frac{1}{B} \sum_{i=1}^B \nabla_{\theta} \mathcal{L}(\theta, x_i)

在具有 N 个设备的数据并行中,每个设备 k 在其自己的小批次 B[k] 上计算梯度:

g_k = \frac{1}{|B_k|} \sum_{x_i \in B_k} \nabla_{\theta} \mathcal{L}(\theta, x_i)

当所有工作进程使用相同的局部批次大小时,全局更新对这些局部梯度求平均:

g_{\text{global}} = \frac{1}{N} \sum_{k=1}^N g_k

在该等批次假设下,该平均在数学上等价于在合并批次 B_total = \bigcup_{k=1}^N B_k 上计算梯度:

g_{\text{global}} = \frac{1}{|B_{\text{total}}|} \sum_{x_i \in B_{\text{total}}} \nabla_{\theta} \mathcal{L}(\theta, x_i)

对于不等的局部批次大小,合并批次梯度为加权平均 g_global = (1/|B_total|) \sum_k |B_k| g_k。这种等价性说明了 为什么 数据并行保持了 SGD 训练的统计特性:将不同的数据子集分发到各设备,独立计算局部梯度并对其求平均,近似于全批次梯度。该平均步骤本身就是在局部梯度张量上执行 AllReduce;D.2.1 节 推导了环形和树状 AllReduce 的成本模型,这些模型决定了该同步的代价,因此读者可以预测其何时会超过并行节省的本地计算量。

验证您对数据并行如何分发工作的理解:

该方法与梯度累积类似,后者指单个设备在更新参数前跨多次前向传播累积梯度。这两种技术都利用梯度的加法特性来高效处理大批次。然而,将同一思路引入集群会带来单机推导无法捕捉的运维挑战。通信开销、节点故障和成本约束各自施加了二阶效应。

数据并行实现

实现细节至关重要,因为 SGD 等价性仅在每个阶段都保持数据不相交、局部梯度完整且执行单次同步更新时才成立。因此,具体工作流追踪了从分发数据子集到同步计算出的梯度的路径。请参考 图 5.3;它追踪了从数据集拆分到梯度聚合的完整工作流,展示了每个 GPU 如何处理其分配的批次,随后通过同步将所有梯度汇聚在一起以进行参数更新。

图 5.3:数据并行实现流水线:数据并行训练的五阶段工作流:(1) 将输入数据拆分为不重叠的子集,(2) 将批次分配给 GPU,(3) 独立计算前向和后向传播,(4) 通过 AllReduce 同步梯度,(5) 在所有设备上统一更新参数。这种方法与模型并行形成对比,后者是对模型本身进行分区而非复制。

正如 图 5.3 所示,关键同步点在于阶段 4:AllReduce 必须在任何 GPU 更新参数前完成,这使得梯度通信成为设备数量增长时的主要瓶颈。

数据集拆分

数据分割与分布式采样

数据分割是 SGD 等价性可能失效的第一个环节:每个工作进程必须看到一个唯一的确定性 Epoch 切片。以包含 100,000 个训练样本的数据集和 4 块 GPU 为例,每块 GPU 每个 Epoch 接收 25,000 个样本。DistributedSampler 必须确保子集间无重叠,以维持梯度估计的有效性:若两块 GPU 处理相同样本,生成的梯度平均会过度加权该样本,破坏使数据并行在数学上等同于单设备训练的无偏梯度假设。

因此,采样器是训练系统的一部分,而不仅是输入加载的便利工具。现代分布式训练框架通过实现预取和缓存机制的分布式采样器自动处理此分发,在不改变样本归属的前提下保持加速器饱和。采样器利用进程秩(由分布式运行时分配的整数工作进程标识符)跨工作进程协作,以确定性方式分区索引,确保使用相同随机种子时的可复现性。对于跨 32 块 GPU 分发的 120 万样本数据集,每块 GPU 每个 Epoch 约处理 37,500 个样本,采样器会填充最终批次以维持所有工作进程间一致的批次大小。

计算阶段:前向与反向传播

数据并行的决定性特征在于计算阶段(前向和反向传播)是完全并行的。每块 GPU 作为孤立岛屿运行,在唯一的微批次数据上执行模型的同一副本。此处微批次指用于激活核算的每 GPU 本地切片;流水线微批次虽用同一词汇,却指代不同机制——细分全局批次以保持流水线阶段忙碌。以我们的 1750 亿参数参考模型为例,这种隔离至关重要:前向传播期间,每块 GPU 独立为其本地批次(微批次大小 4,序列长度 2048)计算激活。若无优化,为反向传播存储这些激活将消耗约 1.1 TB HBM,超出即便是 H100 GPU 容量的一个数量级。激活检查点技术(反向传播时重新计算激活而非存储)在此场景下成为将占用压缩至约 19.3 GB 的必要手段。

反向传播镜像了这种独立性,却引入了系统的主要瓶颈。当 GPU 逆向遍历计算图时,它计算该副本或分片持有参数的梯度。纯数据并行下,完整的 1750 亿 FP16 副本意味着每工作进程 350 GB 的梯度张量,超出单加速器内存预算;大规模运行因此将数据并行与分片、张量并行或流水线并行结合。计算本身零通信,但生成的梯度仅代表真实损失面的局部视角,仅对本地微批次有效。优化器步骤发生前,对应的局部梯度或梯度分片必须跨数据并行工作进程聚合,以形成有效的全局更新。从孤立的高吞吐计算到同步的过渡,定义了数据并行训练的节奏:长时间的静默高强度算术运算,穿插着突发的大量网络流量。

梯度同步

梯度同步是独立 SGD 估计合为单一更新之处,其开销决定数据并行是表现为扩展策略还是网络基准测试。直接的训练需求很简单:每个数据并行副本必须在进入下一步前应用相同的平均梯度。AllReduce 是为复制张量执行此操作的原语:每个工作进程贡献其局部梯度张量,集群求和张量,每个工作进程接收相同的归约结果。在分片变体中,ReduceScatterAllGather 移动对应片段而非在每个设备上移动完整张量,但同步开销依然存在。小张量受同步延迟主导,因每轮通信有启动成本;大张量受带宽主导,因相关梯度载荷必须跨越链路。第 6 章 后续将推导在真实网络结构上实现此平均操作的环形、树形和分层算法。

当同步性能偏离理论预期时,舰队栈框架提供了结构化方法来隔离瓶颈。

场景

跨 128 个节点(1,024 块 GPU)对 3 GB 梯度张量执行 AllReduce 耗时 100 ms。本章将 AllReduce 视为数据并行创建的同步原语;第 6 章 详细推导环形、树形和分层算法。目前有用的调试举措是将观测值与两个粗略界限对比:悲观的扁平通信路径和区分快速节点内流量与较慢节点间流量的拓扑感知路径。

上下文

  • 拓扑:128 个节点,每节点 8 块 GPU

  • 快速本地路径:GPU 间 NVLink 双向 300 GB/s

  • 较慢网络路径:InfiniBand HDR 每端口 200 Gb/s(25 GB/s)

  • 诊断界限:扁平全 GPU 网络路径预计约 239.8 ms;拓扑感知的本地优先网络路径预测约 48.5 ms,因每节点仅发送其份额跨越较慢链路

分析

  • 诊断:集合通信库使用分层路径,将节点边界视为昂贵通信边界

  • 预期行为:节点间阶段应随每节点的归约份额缩放,而非随每 GPU 的完整张量缩放

  • 诊断核对:覆盖所有 1,024 块 GPU 的扁平单环测得值将接近 239.8 ms,故算法已是分层的

观测

  • 观测延迟:100 ms(远低于 239.8 ms 扁平环上界,但比分层理想值高约 51.5 ms)

  • 带宽利用率:节点间阶段交换机计数器显示受影响上行链路仅达理论 InfiniBand 吞吐量的 60%(端到端效率更低,因节点内阶段增加耗时)

  • 网络计数器:显示特定交换机上行链路拥塞

诊断结论

NCCL 已使用分层算法(100 ms 对比 239.8 ms 扁平环预测证实此点)。观测与建模延迟的剩余差距可能源于少数节点间上行链路的交换机拥塞、有效带宽损耗或实现开销,而非完全扁平算法。

解决方案

监控 InfiniBand 交换机端口利用率以识别热点。考虑 3.5.3 节 的轨道优化拓扑,或进一步调整分层切换以显式分区节点内(NVLink)与节点间(InfiniBand)通信。剩余差距很可能代表可通过更好网络配置和带宽利用实现的优化,而非彻底算法变更。

系统层面的教训

该分析展示了舰队栈层间如何交互:物理约束(带宽)限界运维选择(算法),在服务指标(延迟)中体现。调试需审视所有三层,而非孤立调优单层。

分布式训练中的同步拓扑与模型

跳出具体集群来看设计空间,图 5.4 对比了三种高层同步拓扑:带宽最优的环形 AllReduce、中心化的参数服务器,以及全连接的 All-to-All 网状拓扑。在稠密同步数据并行场景中,环形 AllReduce 通过将流量均匀分布到参与链路来避免单一归约器瓶颈,百度的实现便是例证 (Gibiansky 2017)。

图 5.4:梯度同步拓扑:通信模式的可视化对比。(A) 环形 AllReduce 在所有链路上均匀分配带宽,但延迟随跳数线性增长。(B) 参数服务器使用中心节点,在服务端产生带宽瓶颈。(C) All-to-All 网状拓扑将每个节点与其他所有节点相连,实现常数轮通信,但链路数随节点数呈二次方增长,使其在超过约八个节点时变得不切实际。

图 5.4 中可见的权衡在于带宽延迟之间:环形拓扑均匀分摊带宽但每跳增加一个延迟步长,而 All-to-All 网状拓扑以链路数随节点数二次方增长为代价,将延迟路径压缩为常数轮次。NCCL 等高性能库会根据消息大小和集群拓扑自动选择其中的拓扑。第 6 章 形式化了各拓扑的延迟复杂度,并推导了何时环形优于其他方案。

同步模型

分布式训练系统在显式同步模型下运行,该模型管控工作节点何时观测到彼此的更新。模型的选择决定了系统是保证与单设备训练在数学上等价,还是以一致性换取吞吐量。基线模型批量同步并行 (BSP) (Valiant 1990) 要求所有工作节点完成前向和后向传播的本地计算,通过带有 AllReduce 的屏障同步梯度,然后同时更新参数。

BSP 提供了强保证,即每个工作节点在每一步都看到相同的参数值,确保与单设备训练在数学上等价。代价是最慢的工作节点决定迭代时间,造成拖后腿问题。

陈旧同步并行 (SSP) 通过允许工作节点最多领先最慢工作节点 s 次迭代后再阻塞,放宽了这一约束。这在限制陈旧度的同时减少了同步延迟。SSP 需要仔细调整学习率,因为工作节点在略微不同的参数版本上计算梯度。有界陈旧度保证在 BSP 的强一致性和全异步方法之间提供了一个折中方案 (Ho et al. 2013)。

异步 SGD 完全消除了同步屏障,工作节点独立更新参数。这最大化了硬件利用率,但引入的梯度陈旧度可能损害收敛。操作保证决定了系统接受多少收敛风险;5.4.4 节 将展开各模型所需的收敛率、陈旧惩罚及补偿技术。

各同步模型的关键权衡汇总于表 5.2,图 5.5 直观展示了各策略如何随时间在工作节点间调度工作。

| 模型 | 一致性 | 吞吐量 | 收敛性 | 用例 |

| --- | --- | --- | --- | --- |

| BSP | 强 | 受最慢工作节点限制 | 等价于单 GPU | 最终训练运行、可复现性 |

| SSP | 有界陈旧度 | 高于 BSP | 经调优接近等价 | 超参数搜索 |

| Async | 弱 | 最大 | 退化,需补偿 | 大规模异构集群 |

表 5.2:同步模型权衡:BSP、SSP 和异步 SGD 在一个谱系上以一致性换取吞吐量。BSP 保证可复现性但为最慢工作节点买单;SSP 在有界陈旧度下释放吞吐量;异步以牺牲收敛保证为代价最大化吞吐量。

将调度放在时间线上对比,同样的权衡会更清晰。

图 5.5:分布式同步模型:三种同步策略的时间线对比。(A) 批量同步并行 (BSP) 强制所有工作节点每步在全局屏障处等待。(B) 陈旧同步并行 (SSP) 允许工作节点最多领先最慢工作节点 s 步。(C) 异步 SGD 完全消除屏障,允许最大吞吐量但引入梯度陈旧度。

同步模型的选择直接影响系统吞吐量和模型收敛。训练团队常在最终运行中使用 BSP 以保持可复现性,而在超参数搜索中探索 SSP 或异步方法,因为此时精确复现性不那么关键。

屏障语义与失效模式

AllReduce 操作实现了隐式屏障,所有工作节点贡献梯度前,没有任何工作节点可以继续。这种耦合产生了单设备训练中不存在的失效模式。

AllReduce 期间的工作节点故障会导致所有其他工作节点在等待缺失贡献时无限期阻塞。若无超时机制,整个训练作业会挂起而非干净地失败。部署系统通常实现分钟级的看门狗定时器以检测并终止卡住的作业。

Five worker dots approach a vertical barrier; one red missing worker causes the peer lanes to stop before the barrier.

一个缺失的工作节点导致整个 AllReduce 屏障停滞。

当工作节点因条件计算路径或动态批处理而对同步哪些张量存在分歧时,会发生梯度不匹配。AllReduce 操作可能因等待某些工作节点从未发送的张量而阻塞。这常见于 NLP 模型中的变长序列、动态计算图,以及路由决策不同的混合专家模型。

拖后腿引起的延迟源于迭代时间等于最慢工作节点的时间加上同步开销。单个慢速工作节点(无论因热节流、网络拥塞还是操作系统抖动)会延迟所有工作节点并降低集群利用率。以 1000 GPU 为例,若每 GPU 每步有 1% 的拖后腿概率,则每次迭代约有 10 个 GPU 拖后腿。

部署系统通过超时、心跳监控和弹性训练机制解决这些问题。第 7 章 全面涵盖了故障检测、检查点策略和恢复机制,使训练作业能在不可避免的硬件故障中完成。

参数更新

参数更新闭合了数据并行不变量:聚合后,每个设备必须从相同的梯度值应用相同的优化器更新。每个设备独立使用所选优化算法(如带动量的 SGD 或自适应矩估计 Adam)更新模型参数。这种去中心化更新策略避免了中央协调服务器,因为同步已使本地梯度相同。

在一个使用 8 块 GPU 训练 ResNet 模型的系统中,每块 GPU 基于其数据子集计算局部梯度。通过环形全归约进行梯度平均后,每块 GPU 都拥有相同的全局梯度值。随后,每个设备独立使用优化器的更新规则应用这些梯度。以 SGD 和学习率 0.1 为例,更新公式变为 weights = weights - 0.1 * gradients。该示例说明了为何在不牺牲与单设备训练数学等价性的前提下,更新过程可以保持去中心化。

数据拆分、梯度计算、结果同步和参数更新的循环会针对每个批次重复执行。框架会自动化这一循环,但无法消除使副本保持一致的排序约束:每个工作进程必须在同步梯度完成后才能更新。

权衡:通信墙

数据并行之所以成为常见的起始策略,是有原因的:只要模型能装进内存且通信不是瓶颈,它就能随设备数量线性扩展吞吐量。然而,它会触及由通信计算比定义的硬性上限:一旦梯度交换主导了有效计算,增加工作进程主要只会增加同步工作 (Ben-Nun and Hoefler 2019)。

数据并行提供三大主要优势。对于计算密集型模型,吞吐量近似线性扩展:将 ResNet-50 在 ImageNet 上从 1 块扩展到 256 块 GPU 可获得近线性加速,因为梯度交换相对于计算时间很小。模型架构也保持不变;框架将模型包装在数据并行容器中,该容器拦截反向传播钩子以自动触发梯度同步。利用率保持较高,因为与模型并行不同,不存在流水线气泡:所有 GPU 同时进行前向和反向传播。

三个硬性上限制约了这些优势。内存墙要求每块 GPU 持有模型参数、梯度和优化器状态的完整副本;对于 1750 亿参数的模型,这需要每块 GPU 超过 1 TB 内存,若无 ZeRO 分片则超出单设备 HBM 预算。带宽墙随着 N 增长而显现:AllReduce 成本 \(\frac{2(N-1)}{N} \times \frac{M}{\text{BW}_{\text{net}}}\) 最终会占主导地位,对于大语言模型,梯度同步可能消耗超过 50% 的步骤时间,导致效率崩塌。批量大小陷阱加剧了该问题,因为扩展到数千块 GPU 需要增大全局批量大小 (\(B_{\text{global}} = N \times B_{\text{local}}\)),最终达到临界批量大小,此时每步增加更多数据对收敛的边际收益递减。

一个 GPT-2 扩展场景通过固定模型仅改变通信域,直观地展示了效率损失。第一种情况将所有 8 块 GPU 保持在单个 NVLink 节点内。

单 GPU 基准

  • 批量大小:16(开启梯度检查点,适配 32 GB)

  • 单步耗时:1.8 s

  • 50K 步耗时:25 小时

8 GPU:单节点 NVLink

  • 单 GPU 批量:16,全局批量:128

  • 梯度同步:5.2 GB @ 450 GB/s (NVLink, 单向) ≈ 11.7 ms

性能结果:

  • 计算:1800 ms/步

  • 通信:11.7 ms/步

  • 总计:1811.7 ms/步

  • 加速比(吞吐量):8×

  • 并行效率:99.4%

训练时间:25 小时 ÷ 8 = 3.1 小时

在节点内部,NVLink 使梯度交换相对于计算保持很小,因此效率维持在 99.4% 左右。一旦相同模型必须通过商用网络跨节点同步,情况就会逆转。

第二种情况将相同的 GPT-2 运行扩展到多个节点,用节点间以太网替代节点内 NVLink 跳转,用于 AllReduce 的主要部分。

商用网络配置:跨 4 个节点的 32 GPU

  • 单 GPU 批量:16,全局批量:512

  • 节点内通信:11.7 ms (NVLink)

  • 节点间通信:5.8 GB @ 1.25 GB/s (10GbE) ≈ 4650 ms

性能结果:

  • 计算:1800 ms(占 27.9% 时间)

  • 通信:4661.7 ms(占 72.1% 时间),通信占主导并成为瓶颈。

  • 总计:6461.7 ms/步

  • 加速比(吞吐量):8.9× → 2.8 小时

  • 并行效率:27.9%

梯度累积提供了一种直接补救措施,即将所有通信保持在单个节点的 NVLink 域内,同时仍在等效的大有效批量上训练。

问题:在商用 10G 网络上运行 GPT-2,32 GPU 受通信瓶颈限制,以固定样本数计算成本 $3,021。单个 8 GPU 节点能否利用梯度累积更高效地达到相同有效批量大小?

数学分析

  1. 有效批量大小:8 GPU × 批量 16 × 4 累积步数 = 512。

  2. 通信开销:4 步累积即 AllReduce 一次。

    • 开销 = 5.8 ms / (4 × 1800 ms) ≈ 0.081%。
  3. 训练时长:总时间 3.1 小时。

  4. 总成本:3.1 小时 × $128/小时 = $400。

系统洞察:梯度累积节省 $2,621(86.7%),方法是将计算集中在带宽充裕处(节点内 NVLink)并最小化同步频率。网络较慢时,不要横向扩展——在本地扩大批量大小。

该计算改变了扩展决策:当节点间带宽是约束瓶颈时,带宽丰富节点上的梯度累积可击败朴素的横向扩展,即使墙钟时间略长。可得出四点洞见。NVLink 支持单节点内高效扩展(99.4% 效率),而节点间通信扼杀效率(降至 27.9%)。对于通信受限运行,梯度累积优于朴素横向扩展,因此该 GPT-2 场景的最佳点是单节点 8 GPU 配合梯度累积,而非盲目扩展至 32+ GPU。OpenAI 的 GPT-2 论文报告在 4 个节点的 32 块 V100 上训练,使用了优化通信(可能结合梯度累积与流水线并行),而非纯数据并行。

内存高效数据并行:ZeRO 与 FSDP

数据并行的内存约束催生了一系列技术,它们在保持数据并行训练简单性的同时,将内存状态跨工作进程分片。ZeRO (Zero Redundancy Optimizer)⁶¹ (Rajbhandari et al. 2020) 及其 PyTorch 实现 FSDP (Fully Sharded Data Parallel) (Zhao et al. 2023) 使得原本需要模型并行的模型得以训练。

分片数据并行 是数据并行的变体(实现为 ZeRO 阶段和 FSDP),它将优化器状态、梯度,以及在最深阶段的参数本身分区到各数据并行工作进程中,通过集合通信按需重构每个分片,使单工作进程内存降至完整训练状态的 1/N,而每个工作进程仍处理自己的小批量分片。

  1. 意义:混合精度 Adam 训练每参数携带 16 字节状态,因此 70 亿参数模型需要 112 GB 训练状态,即便模型推理舒适地适配 80 GB 加速器,训练仍 OOM。跨 64 工作进程的 ZeRO Stage 3 将单设备状态降至 1.75 GB(每参数 0.25 字节),将内存墙转化为通信成本:按需重组分片的 AllGather 和 ReduceScatter 流量使每步开销增加 10–25%。

区别

与普通数据并行(在每个工作节点上复制完整训练状态)以及模型并行(对计算本身进行分区)不同,分片数据并行仅对存储进行分区:每个工作节点仍执行完整的前向和后向传播,在需要时即时收集每层参数,用完后立即丢弃。

常见陷阱

一个常见的误解是分片能免费获得内存。按需 AllGather 将参数通信置于每一步每一层的关键路径上;在较慢的互联互通或较小的每工作节点批量大小时,暴露的通信开销侵蚀吞吐量的速度比内存节省带来的帮助更快。容量是用带宽换来的。

内存节省分析:面向 7B Llama-2 模型的 ZeRO

为了理解 ZeRO 提供的内存节省规模,请考虑大语言模型具体的内存预算。

问题

以混合精度训练 7B 参数的 Llama-2 模型,每个参数需要 16 字节用于权重、梯度和优化器状态。完整的训练状态能否放入单张 A100-80 GB?使用 64 张 GPU 的 ZeRO Stage 3 如何改变每设备内存需求?

基准:标准 DDP(复制状态)

每参数内存成本:

  • 权重(FP16:2 字节

  • 梯度(FP16:2 字节

  • 优化器状态(FP32:12 字节(4 字节主权重 + 4 字节动量 + 4 字节方差)

  • 总计:16 字节/参数

7B 模型总内存:

\[C_{\text{state,total}} = 7 \times 10⁹ \times 16 \text{ bytes} \approx 112 \text{ GB} \]

基准结果:在 A100-80 GB 上会出现显存不足(OOM)。

优化:ZeRO-3(完全分片)

使用 \(N = 64\) 块 GPU 时,状态被分区:

  • 权重\(2/N\) 字节

  • 梯度\(2/N\) 字节

  • 优化器\(12/N\) 字节

  • 总计\(16/N = 0.25\) 字节/参数 有效存储!

每 GPU 内存:

\[C_{\text{state,ZeRO3}} = \frac{112 \text{ GB}}{64} \approx 1.75 \text{ GB} \]

结果:轻松放入,剩余约 78 GB 用于激活值(批量大小)。

ZeRO 通过渐进式分片解决了这种冗余,如图 5.6 所示,并由表 5.3 总结。

ZeRO 内存减少

图 5.6:ZeRO 内存减少:标准数据并行(DDP)在每个 GPU 上复制所有模型状态。ZeRO 逐步对这些状态进行分区:ZeRO-1 分片优化器状态,ZeRO-2 增加梯度分片,ZeRO-3 分片参数本身。ZeRO-3 实现了线性内存扩展,使得拥有 100B+ 参数的模型能够放入无法容纳复制状态的加速器上。

表 5.3:ZeRO 分片阶段

| 阶段 | 分片内容 | 内存减少 | 通信开销 |

| :--- | :--- | :--- | :--- |

| ZeRO-1 | 仅优化器状态 | ~4× | 无(同 DDP) |

| ZeRO-2 | + 梯度 | ~8× | ReduceScatter 替代 AllReduce |

| ZeRO-3/FSDP | + 参数 | ~\(N\)(随工作节点线性) | 每层前 AllGather |

ZeRO 分片阶段:每个阶段分片训练状态的不同层级,并权衡不同的通信原语。相对于 DDP,ZeRO-1 无额外开销;ZeRO-2 将 AllReduce 换为 ReduceScatter;ZeRO-3(FSDP)在每层前增加 AllGather,但解锁了近乎线性的内存扩展。

ZeRO 阶段详情

ZeRO-1 跨 GPU 分片优化器状态。每个 GPU 仅存储 \(1/N\) 的 Adam 优化器相关状态。在梯度 AllReduce 之后,每个 GPU 仅更新其分片的参数,然后将更新广播给其他 GPU。按照将 FP32 主权重、动量和方差计为 12 字节的惯例,内存节省使优化器状态从集群总计 \(12 \times N\) 字节/参数降至 12 字节/参数。

ZeRO-2 额外分片梯度。ZeRO-2 使用 ReduceScatter 代替 AllReduce(后者会在每个 GPU 上留下完整梯度),使每个 GPU 接收 \(1/N\) 的归约梯度。按照此处使用的 FP16 梯度惯例,内存节省使梯度从跨 \(N\) 个工作节点复制的 \(2 \times N\) 字节/参数降至总计 2 字节/参数,即每 GPU \(2/N\) 字节/参数。

ZeRO-3 和 FSDP 分片参数本身。每个 GPU 仅存储 \(1/N\) 的模型。在每层前向传播前,通过 AllGather 收集参数;后向传播后,通过 ReduceScatter 归约梯度,随后丢弃参数。这以 FSDP 相对于标准 DDP 引入的额外通信为代价,实现了最大内存效率。

关键路径上的通信

这种分片将通信置于 DDP 避免的关键路径上。前向传播需要 AllGather 重建每层参数(\(M_{\text{layer}}\) 字节);后向传播需要第二次 AllGather 重建参数(因为前向传播后参数被重新分片,\(M_{\text{layer}}\) 字节),随后进行梯度 ReduceScatter\(M_{\text{layer}}\) 字节)。对于拥有 \(N_L\) 层的模型,带重新分片的全分片 FSDP 每训练步约执行 \(3 \times N_L\) 次集合操作,而 DDP 仅需单次 AllReduce,总通信量提升至约 \(3 \times M_{\text{state}}\) 字节,对比 DDP 的 \(2 \times M_{\text{state}}\)。然而,这些集合操作分散在更多操作中并存在重叠机会:当第 \(i\) 层计算时,第 \(i+1\) 层可预取其参数。

FSDP 与 DDP 选择

FSDP 与 DDP 的选择取决于模型大小和内存限制。当模型适配 GPU 内存且留有激活值空间时,DDP 通常更优,因为它避免了重复的 AllGather 工作。随着内存压力增加,ZeRO-2 变得有吸引力,因为它分片梯度和优化器状态同时保留参数复制;一旦参数本身超出单 GPU 内存,即使将 AllGather 置于关键路径,ZeRO-3/FSDP 也成为必要。在 80 GB 上训练 70B+ 模型时,FSDP 通常必须结合张量并行而非替代它。

内存高效的数据并行需要仔细调整分片策略(按层、按 Transformer 块或扁平)和混合精度设置。分片粒度决定了权衡:更细的分片降低每 GPU 内存,但增加通信频率,因为每训练步必须执行更多 AllGatherReduceScatter 操作。

超越内存的扩展:优化格局

通过 ZeRO 和 FSDP 消除内存瓶颈后,将数据并行扩展到数百块 GPU 变得诱人。然而,仅凭通信分析无法预测,这会以改变优化格局的方式进行。大规模全局批量大小改变梯度噪声统计,针对 8 GPU 运行调优的学习率调度在 256 GPU 上可能灾难性发散。一个具有里程碑意义的大规模演示展示了这种失效模式,以及一种广泛产生影响的工程响应,源于单次实验。

背景:Facebook AI Research 致力于使用 256 块 GPU 在一小时内于 ImageNet 上训练 ResNet-50,将数据并行推向远超早期实践视为安全的批量大小(Goyal et al. 2017)。

失效模式:盲目增大全局批量大小使优化不稳定。通信系统能提供更高吞吐,但优化器不再表现得像单机运行。

后果:该团队通过将线性缩放规则与逐渐学习率预热相结合,恢复了准确率,在保持收敛的同时稳定了 8,192 的全局批量大小。

系统教训:分布式训练不仅仅是并行硬件。扩展改变优化机制,因此集群、通信调度、批量大小和学习率调度必须作为一个整体系统进行联合调优。

扩展效率与收敛

数学框架:扩展效率

在研究具体的并行策略之前,我们必须理解决定从单设备扩展到多设备是否值得的指标:扩展效率。如果模型在单设备上训练时间为 T₁,理想(线性)扩展将在 N 个设备上用时 T₁/N。实践中,通信开销、流水线气泡和负载不均衡会降低加速比。扩展效率由公式(5.1)定义:

\[ \eta_{\text{scaling}} = \frac{T_1}{N \times T_N} \qquad(5.1) \]

其中 Tₙ 是 N 个设备上的训练时间。效率 1.0 意味着完美线性扩展;效率 0.5 意味着仅获得预期加速比的一半。

扩展效率 (ηₛ𝒸𝒶𝓁𝒾𝓃𝑔) 是增加 ML 计算设备数量 (N) 时,实际训练吞吐量与理想线性吞吐量的比率。

  1. 意义:这是集群生产力最重要的指标(\(\eta_{\text{scaling}} = \frac{T_1}{N \times T_N}\))。扩展效率 0.50 意味着 10,000 GPU 集群仅提供相当于 5,000 GPU 集群的有效工作量,浪费了 50% 的硬件投资。

  2. 区别:与单节点效率(捕捉带宽等局部瓶颈)不同,扩展效率捕捉通信时间 (T꜀ₒₘₘ(N)) 和同步的集群级开销。

  3. 常见误区:常误认为扩展效率是常数。实际上它是问题规模的函数:随着 N 增加,通信计算比通常恶化(阿姆达尔定律),使得小模型更难维持高效率。

对 175B 模型的数据并行训练,每步通信成本主要由 350 GB 梯度的 AllReduce 主导。

使用 InfiniBand 上的环形 AllReduce(有效带宽 50 GB/s),原始通信时间约为 2 × (N−1)/N × 350/50,大 N 时趋近 2 × 350/50 秒。假设梯度通信与反向传播有 75% 重叠,有效暴露通信时间降至 T꜀ₒₘₘ(N) ≈ 3.5 秒。同等假设下计算项 T꜀ₒₘₚᵤₜₑ/N ≈ 2.1 秒,暴露步时间约 5.6 秒,朴素数据并行扩展效率为 2.1 s / 5.6 s ≈ 37.5%。配置良好的系统可通过结合张量并行、流水线并行、拓扑感知放置和更有效的通信重叠来恢复大部分损失。

扩展效率关键取决于计算与通信比率。三因素主导该比率:模型越大,每训练步计算量越大(每权重更新 FLOPs 更多),相同通信开销占总步时间比例越小;这种模型‑通信比是扩展效率随模型增大而提升的原因。大批量朝同方向推进,提高每步计算量而不成比例增加通信(梯度大小与批量无关);代价是大批量可能损害收敛,需学习率调优和预热调度。网络带宽最直接影响:InfiniBand 带宽翻倍将通信时间减半,按比例提升扩展效率。网络结构因此是集群生产力的一阶决定因素,而非次要考量;其成本(总系统成本 10–15%)若能提升扩展效率哪怕几个百分点就极易证明合理,因为低扩展效率浪费了其余 85–90% 的投资。

悖论在于:大模型比小模型更易扩展。

这三因素有重要交互。大模型配合大批量获得更好扩展效率,意味着最昂贵的训练负载也是最受益于规模的负载。这为大规模基础设施创造了良性循环:证明建设千卡集群合理的负载,也是利用率最高的负载。反之,小模型和小批量扩展差,这也是团队在 64 GPU 上训练 1 B 参数模型时常仅达 40–60% 扩展效率的原因。

然而,即使对大模型,扩展也非无限。存在扩展悬崖,超过后增加 GPU 反而降低成本效率。按 175B 模型假设,成本高效区约为 1,024–4,096 GPU,通信计算比仍有利且扩展效率保持 70% 以上。超 8,192 GPU 后,AllReduce 通信时间开始主导反向传播计算时间,效率跌破 50%。虽墙钟训练时间随 GPU 增加仍略有下降,但每有效 FLOP 成本上升,因为组织为 8,000 GPU 付费却仅完成 4,000 GPU 的工作。这种非线性关系决定了大模型训练的经济可行性受互联延迟物理限制,而非仅受硬件可用性限制。集群规模须在扩展曲线线性区运行,模型架构(批量、序列长度、并行维度)须与集群规模协同设计以维持平衡。

实验设定:175B 模型在 DGX H100 集群上训练(每 GPU 400 Gb/s InfiniBand)

  • 每步计算量(假设批量 2 M token,每参数每 token 6 FLOPs):Oₛₜₑₚ = 6 × 175 × 10⁹ × 2 × 10⁶ ≈ 2.1 × 10¹⁸ FLOPs

  • 1,024 GPU 上的单 GPU 计算时间,每卡 FP8 峰值 1979 TFLOP/s(50% 利用率):T꜀ₒₘₚᵤₜₑ/N ≈ 2.1 秒

  • 350 GB 梯度的 AllReduce 时间(环形 AllReduce + 重叠):T꜀ₒₘₘ(N) ≈ 3.5 秒(原始传输约 14 秒;本例假设与反向传播 75% 重叠)

  • 扩展效率

    \[ \eta_{\text{scaling}} \approx \frac{T_{\text{compute}}/N}{T_{\text{compute}}/N + T_{\text{comm}}(N) + T_{\text{sync}}(N) - T_{\text{overlap}}} = 2.1/5.6 \approx 0.375 \]

    此简化示例中同步已包含在 AllReduce 项内,且已应用 75% 通信重叠。

低效率(37.5%)表明此规模下朴素数据并行不足。层次感知系统通过结合数据并行与张量并行(经 NVLink 通信)及流水线并行(计算通信重叠)恢复大部分损失效率。

并行‑基础设施交互

扩展效率分析揭示了一个更深层的洞见:有效的并行策略不是由模型架构本身决定的,而是由模型的通信需求与基础设施带宽层级之间的交互决定的。每种并行策略与基础设施拓扑的组合都会产生不同的扩展效率曲线,选择错误的组合可能会浪费集群容量的很大一部分。上面的估算已经表明,纯数据并行在该规模下会在 37.5% 效率附近停滞;要恢复这些丢失的容量,意味着将每个并行维度映射到能承载其流量的带宽层级。第 5.7 节通过三个具体配置详细阐述了这一组合,并将其形式化为层级感知并行,这在张量并行和流水线并行开发完成后进行。

因此,选择和组合并行策略直接导向它们产生的流量。第 3 章探讨了拓扑如何与并行映射协同设计以最大化扩展效率;这里的下一个问题是集合操作本身如何塑造步时间。在数据并行系统中,AllReduce 操作可占用总训练时间的 10%–40%,且这种开销随集群规模增长而增加。BERT-Large 在 128 个 GPU 上运行时,通信开销可达到总运行时间的很大一部分,而 GPT-3 规模的模型需要张量、流水线和数据并行以及通信重叠,以避免数据并行梯度同步主导步时间。

AllReduce 复杂度取决于两个组件:延迟(α)和带宽(β)。Ring AllReduce(N−1)/N 的利用率实现带宽高效的通信,而基于树的方法在 𝒪(log N) 步骤中提供更低的延迟。选择取决于消息大小:树算法在延迟主导的小消息场景下胜出,环算法在带宽主导的大梯度场景下胜出。高性能实现(如 NCCL)使用分层算法,结合节点内树延迟和节点间环带宽。第 6 章提供了详细的算法分析,包括复杂度公式、分层变体以及大规模集合操作的拓扑感知优化。

互联选择决定了大规模部署是保持计算受限还是崩溃为通信受限状态,而高效分布式训练的带宽需求巨大,尤其是针对 Transformer 模型。高效系统通常需要每个节点 100–400 GB/s 的聚合带宽以支持 Transformer 架构。BERT-Base(1.1 亿参数)在 FP32 下每次迭代需要约 440 MB 的梯度同步,而 BERT-Large(3.4 亿参数)需要约 1.4 GB。在 64 个 GPU 上,这些同步需求要求 100–200 GB/s 的持续带宽,以实现 50 ms 以下的同步延迟。对于 1750 亿参数的语言模型,确切的带宽需求取决于 3D 并行配置、梯度累积、重叠和互联拓扑,而不是单一的通用数值。

同步频率在通信效率和收敛行为之间存在权衡。累积 4 个微步的梯度可将同步频率降低 75%,但实际的步时间减少量取决于计算/通信组合以及有多少通信可以与反向传播重叠。在标准实现中,梯度累积复用一个驻留的梯度缓冲区并原地累积;内存压力来自保持该缓冲区驻留以及任何更大的微批次或激活选择,而非存储 4 个独立的梯度张量。异步方法完全消除了同步成本,但引入的陈旧性会使大范围学习率下的收敛性降低 15%–30%。

扩展的物理学:含通信的阿姆达尔定律

正如铁律支配单线程执行,分布式训练受显式考虑通信开销的扩展版阿姆达尔定律支配。在 N 个设备上完成一个训练步所需的时间不仅仅是 T[single]/N,而是受同步的顺序性质制约。附录 C.3.2 节推导了舰队规模的加速上限并通过具体示例进行了演示;此处我们确立这样一个概念:固定的同步分量无论添加多少加速器都会限制加速比。

图 5.7 将这种偏离理想线性扩展的现象可视化为扩展税——扩展效率界限的直接后果。它展示了通信开销(r)如何作为性能的阻力,制造出一堵通信墙,在此处添加更多 GPU 会产生边际效益递减。

图 5.7:扩展税:有效加速比 vs. GPU 数量。理想扩展(虚线黑色)呈线性,但现实系统需为通信支付“税”。像 ResNet 这样计算受限的模型(绿色)扩展良好,因其算术强度高。像 GPT-3 这样带宽受限的模型(红色)撞上通信墙,添加更多 GPU 产生边际收益递减(效率 < 50%)。

舰队定律直接映射到铁律的变量上,将计算、带宽和协调分离为相加项:

\[T_{\text{step}}(N) = \underbrace{\frac{T_{\text{compute}}}{N}}_{\text{计算时间项}} + \underbrace{T_{\text{comm}}(N)}_{\text{带宽项}} + \underbrace{T_{\text{sync}}(N)}_{\text{协调项}} - T_{\text{overlap}} \]

舰队定律项将步时间分解为四个不同的成本:

  • 计算时间项 (T[compute]/N):将批次理想地划分到 N 个设备后所需的总计算量。

  • 带宽项 (Tcomm):移动数据所花费的时间。这由铁律的数据移动项 D[vol]/BW 支配。对于 Ring AllReduce,该项为 \(\frac{2(N-1)}{N} \times \frac{M}{\text{BW}_{\text{net}}}\),其中 M 为通信的梯度或模型状态字节数,BW[net] 为网络带宽。

  • 协调项 (Tsync):屏障、排序和等待跟随者的非重叠成本。

  • 重叠 (T[overlap]):隐藏在计算后面的通信部分。

舰队定律导出了固定全局工作量下的扩展效率指标,其中 T[compute] 是该工作量的单设备计算时间:

\[\eta_{\text{scaling}} = \frac{T_{\text{compute}}}{N \times T_{\text{step}}(N)} = \frac{1}{1 + \frac{N(T_{\text{comm}}(N) + T_{\text{sync}}(N) - T_{\text{overlap}})}{T_{\text{compute}}}}} \]

这就是扩展效率界限:完美的线性扩展(η[scaling] = 1.0)是理论极限,而非实践目标。配置良好的系统在中等规模下可达到 η[scaling] = 0.85–0.95,并随 N 增长进一步下降。η[scaling] = 1.0 与实现效率之间的差距即为通信税和协调税:分布式执行的代价。

该方程揭示了扩展墙:随着 N 增加,计算项(T[compute]/N)缩小,但通信和同步项可能保持不变或增长。最终,分母由开销主导,驱使效率趋近于零。除了挂钟时间,这种通信开销还施加了一个随设备间物理距离缩放的能量税。

橙色能量-比特阶梯,具有代表性的中点:InfiniBand 35 pJ,NVLink 7.5 pJ,HBM 1.5 pJ。

通信能耗从 HBM 向外延伸至网络。

分布式训练既是与能耗的赛跑,也是与时间的赛跑。在单个 GPU 内,将一个字节从 HBM 移动到核心的成本约为 1–2 pJ/bit。将同一个字节跨 NVLink 互连移动的成本为 5–10 pJ/bit。将其通过交换机跨 InfiniBand 网络移动的成本为 20–50 pJ/bit。

在 10,000 GPU 的规模下,正是总带宽的乘法效应改变了工程问题的本质。完整的 H100 NVLink 封包在集群内约为 9 PB/s;按 7.5 pJ/bit 计算,该移动在不计冷却和交换机开销的情况下约为 540 kW。聚合 NDR InfiniBand 封包较小(500 TB/s),但按 35 pJ/bit 计算仍约为 140 kW。因此,通信计算重叠对时钟周期效率是必要的,但避免不必要的移动才是直接降低功耗项的途径。

随着 N 的增长,时钟周期效率和能耗代价会一起下降,并在不同 GPU 数量下遵循可预测的模式。作为一个具有代表性的经验法则,处于 2–32 GPU 线性扩展阶段的系统通常能达到 85–95% 的并行效率,因为通信开销仍然很小。在 64–256 GPU 时会出现通信受限阶段,即使互连匹配良好,效率也可能降至 60–80%。超过 512 GPU 后,协调开销可能变为主导因素,由于集合操作延迟,将效率限制在 40–60%。

硬件选择对这些扩展特性至关重要。NVIDIA DGX A100 系统为每个 GPU 提供 600 GB/s 的双向 NVLink 带宽,系统层面的聚合 NVSwitch 带宽使单个 8-GPU 节点内能保持高并行效率。多节点扩展需要具备足够剖分带宽的网络结构;EDR 级 100 Gbps 链路可支持较小规模的多节点作业,而 HDR 级 200 Gbps 链路在拓扑、布局和重叠匹配良好时可支持更大规模的集群。

效率指标直接影响并行策略的选择。数据并行在线性扩展阶段效果良好,但在大规模时变为通信受限。模型并行解决了内存约束,但引入了限制效率的顺序依赖。流水线并行减少了设备空闲时间,但增加了管理微批次的复杂性。有效策略取决于哪个约束——内存、带宽或同步——主导了目标工作负载。

硬件效率指标决定吞吐量,但收敛理论决定分布式训练能否达到与单设备训练相同的解质量。并行性以三种方式影响优化收敛:收敛率随批次大小变化,增加工作线程在超过关键批次大小后收益递减,学习率必须随批次大小缩放。

同步数据并行 SGD 的收敛率

分布式 SGD 的基本收敛结果解释了为何增加工作线程可在不改变优化器含义的情况下减少迭代次数。该定理使用两个标准正则性假设:损失足够平滑,使得梯度不能任意快速变化;每个随机梯度在真实梯度周围具有有界方差。形式化表述为:对于具有 L[s]-Lipschitz 梯度且方差有界的随机梯度 𝔼[∥g[i] − ∇ℒ(θ)∥²] ≤ σ² 的损失函数 ℒ(θ),N 个工作线程的同步数据并行 SGD 达到如下收敛率。

对于具有 N 个工作线程的同步数据并行 SGD,每个工作线程在大小为 b 的本地批次上计算梯度,经过 K 次总迭代后,预期优化误差满足:

$$\mathbb{E}[\mathcal{L}(\theta_K)] - \mathcal{L}^{\star} \leq \underbrace{\frac{L_s \|\theta_0 - \theta^{\star}\|²}{2K}}_{\text{优化误差}} + \underbrace{\frac{\eta L_s \sigma²}{2Nb}}_{\text{方差下限}}$$

其中 η 为学习率,L[s] 为平滑度常数,ℒ^⋆ 为最优损失值,θ^⋆ 为最优参数向量,σ² 为梯度方差。当学习率最优调整为 \(\eta = \mathcal{O}(\sqrt{Nb/K})\) 时,有效收敛率为 \(\mathcal{O}(1/\sqrt{NbK})\)

该定理揭示了几个重要见解。首先,方差下限随工作线程数 N 线性下降,解释了为何分布式训练可在线性扩展阶段用更少迭代达到相同最终损失。在相同本地批次大小 b 下,同步数据并行 SGD 计算出的小批量梯度等同于单工作线程使用批次大小 N**b 进行该步骤计算的梯度;收敛速度仍取决于学习率缩放和关键批次大小效应。其次,收敛率 \(\mathcal{O}(1/\sqrt{NbK})\) 表明,假设带宽无限,工作线程、本地批次大小和迭代次数共同决定统计进度。这是分布式训练的统计效率,区别于硬件效率。

然而,该定理假设完美同步(BSP)。当工作线程以不同速率推进或使用陈旧梯度时,收敛保证会下降,我们将在下一节探讨。

陈旧性影响:BSP vs. SSP vs. ASP

第 5.3.1.4 节 中确立的 BSP、SSP 和 ASP 的操作保证是有收敛代价的,本节通过风格化的平滑目标界来勾勒该代价。陈旧性参数 τ[stale] 量化了梯度应用于参数时可能落后的迭代次数,它是将吞吐量与解质量联系起来的变量(Ho et al. 2013; Dutta et al. 2018)。

梯度陈旧性 (τ[stale])** 是指从梯度计算完成到应用于全局模型状态之间,发生的参数更新次数。

  1. 意义:它代表分布式优化中的同步误差。增加 τ[stale] 可通过减少屏障等待 (Tsync) 提高吞吐量,但通常会降低收敛速度,需要更多操作 (O) 才能达到相同精度。

  2. 区别:与作为物理延迟的网络延迟不同,陈旧性是由同步协议选择(如 ASP、SSP)产生的算法层面的偏移量。

  3. 常见误区:常见的误解认为陈旧性“总是有害的”。实际上,它是吞吐量与收敛的权衡:对于某些大规模工作负载,允许有界陈旧性是让成千上万 GPU 持续工作的一种实用方法。

这些模型的收敛行为存在差异,直接影响训练成本和解质量。在批量同步并行(BSP,τ[stale] = 0)中,所有工作线程在同一参数版本上计算梯度,然后通过屏障同步再更新。这保证了与使用更大批次大小的单设备训练在数学上等价,具有最优收敛率 \(\mathcal{O}(1/\sqrt{NbK})\),且除批次大小缩放外无需调整超参数。

陈旧同步并行(SSP,τ[stale] ≤ s)通过允许工作线程领先最慢工作线程最多 s 次迭代来放宽屏障。在风格化界中,收敛率获得一个加性延迟项:

$$\mathbb{E}[\mathcal{L}(\theta_K)] - \mathcal{L}^{\star} \leq \mathcal{O}\left(\frac{1}{\sqrt{NbK}}\right) + \mathcal{O}\left(\frac{s² \eta² L_s²}{Nb}\right)$$

第二项代表陈旧性惩罚。对于有界陈旧性 s,通过保持延迟较小并重新调整学习率可控制该惩罚。选择有界陈旧性的部署接受收敛退换取在异构集群上的吞吐量提升。

异步 SGD(ASP,τ[stale] = ∞)完全消除了等待:工作进程会立即更新参数。虽然这最大化了吞吐量,但同样的教学模型显示出更大的延迟依赖项:

\[ \mathbb{E}\big[\mathcal{L}(\theta_K)\big] - \mathcal{L}^{\star} \leq \mathcal{O}\Big(\frac{1}{\sqrt{K}}\Big) + \mathcal{O}\big(\bar{\tau}_{\text{stale}}^{2}\,\eta^{2}\,L_s^{2}\big) \]

其中,\bar{\tau}[stale] 表示平均过时程度。过时惩罚现在与平均延迟的平方成正比,并且更重要的是,来自 N 个工作进程的方差减少在主导项中消失了。几种技术用于补偿这一点。学习率衰减(\eta' = \eta / \sqrt{1 + \bar{\tau}_{\text{stale}}})可以降低过时惩罚,但会减慢收敛。动量校正调整动量项以考虑延迟更新。梯度裁剪防止具有大幅度的过时梯度 destabilizing 训练过程。

表 5.4 汇总了每种同步模型的收敛性质。

| 模型 | 过时程度 | 收敛速率 | 方差减少 | 最适用场景 |

|---|---|---|---|---|

| BSP | τ[stale] = 0 | \(\mathcal{O}(1/\sqrt{NbK})\) | 完全(1/N) | 最终训练,可重复性 |

| SSP | τ[stale] ≤ s | \(\mathcal{O}(1/\sqrt{NbK}) + \mathcal{O}(s^{2}\eta^{2}L_s^{2}/Nb)\) | 部分 | 异构集群 |

| ASP | τ[stale] = ∞ | \(\mathcal{O}(1/\sqrt{K}) + \mathcal{O}(\bar{\tau}_{\text{stale}}^{2}\eta^{2}L_s^{2})\) | 无 | 最大吞吐量,早期探索 |

表 5.4: 按同步模型划分的收敛性质:BSP 在同步开销的代价下提供最优的收敛保证。SSP 提供了吞吐量和收敛之间可调节的权衡。ASP 最大化吞吐量,但会失去并行带来的方差减少好处。

学习率缩放规则

当通过数据并行性增加有效批次大小时,必须调整学习率以保持收敛质量。理论和实践中已经出现了两种主要的缩放规则。

线性缩放规则(Goyal et al. 2017)。一个预热期可以让学习率从 η[base] 线性增加到 k ⋅ η[base],在前 W 次迭代中完成,使模型进入一个在该区域内大学习率是稳定的损失景观区域:

\eta_t = \eta_{\text{base}} + \frac{t}{W}\big(k \cdot \eta_{\text{base}} - \eta_{\text{base}}\big)
\quad \text{for } t < W

平方根缩放规则适用于批次大小增长到线性缩放失效的情况:

\eta_{\text{large}} = \sqrt{k} \cdot \eta_{\text{base}}

更保守的平方根规则的动机在于:观察到梯度噪声(而不仅仅是幅度)会影响优化动态。平方根规则能更好地保持梯度更新的信噪比。经验表明,当批次大小超过临界批次大小范围时,需要使用平方根缩放。

对于极端批次大小(32K – 1M),层级自适应学习率缩放(LARS)(You et al. 2017) 会根据权重范数与梯度范数的比率自动调整每层的学习率:

\eta_{\ell} = \eta_{\text{global}} \cdot
\frac{\|w_{\ell}\|}{\|g_{\ell}\| + \lambda_{\text{wd}}\|w_{\ell}\|}

这里的 λ[wd] 是 LARS/LAMB 更新中优化器的权重衰减系数,用于避免与可靠性分析中使用的故障率参数冲突。层级缩放可以防止权重较小的层获得不成比例地大的更新。LAMB 使得在保持收敛质量的前提下,可以使用高达 64K 的批次大小训练 BERT。

临界批次大小和收益递减

分布式训练的一个基本性质是:增加工作进程的帮助会在达到某个阈值后停止。临界批次大小 B* 标志着 transition point,在此点之后,增加批次大小在每看到一个样本上的收敛改善会带来递减回报。

临界批次大小 (B*) 是指分布式训练批次大小,此时梯度噪声规模足够大,以至于进一步增加批次大小在样本效率上会带来递减回报(McCandlish et al. 2018)。

  1. 意义:它标志着并行缩放效率的转折点。在 B* 以下,增加批次大小会线性改善每一步的收敛。在 B* 以上,更大的批次会带来递减回报,需要按比例增加更多样本(D)才能达到相同的损失。

  2. 区分:与由内存容量和激活占用决定的内存受限批次大小不同,临界批次大小是模型和数据集的算法属性。

  3. 常见陷阱:一种常见的误解是认为可以无限通过添加 GPU 来加速训练。实际上,B* 定义了数据并行性的物理上限:在此点之后增加工作进程会浪费能源和计算资源(O),而不会减少总训练时间(T)。

临界批次大小的梯度噪声尺度代理可估算为:

B^{*} \approx \frac{\operatorname{tr}(\Sigma)}{\|\nabla \mathcal{L}(\theta)\|^{2}}

其中,tr(Σ) 是梯度协方差矩阵的迹(总梯度方差),而 ‖∇ℒ(θ)‖² 是梯度范数的平方(信号强度)。直觉上,B* 是使梯度方差通过平均降到真实梯度幅度水平的批次大小。

已发表的大批次训练范例说明了规模而非提供固定阈值。在 ImageNet 上,Goyal 等人通过预热将 ResNet-50 稳定在全局批次 8,192,而 LARS 在某些情况下进一步扩展了 ImageNet 的大批次训练(Goyal et al. 2017)。对于 BERT-Large 预训练,LAMB 报告称在使用层级自适应更新的情况下,批次大小可达 32K – 64K(You et al. 2020)。在其他领域,McCandlish 等人表明,有用的批次大小上限在类似 ImageNet 的场景中从数万到在强化学习工作负载中达到数百万不等,因此新的前沿语言模型阈值应当通过测量得出,而不应简单复制(McCandlish et al. 2018)。

缩放定律区域呈现出三种截然不同的行为。在临界批量大小以下(B* < B*^(*)),线性缩放成立:将批量大小加倍会使达到目标损失所需的迭代次数减半,硬件效率决定吞吐量。在临界点(B* ≈ B*^(*)),每秒样本数效率最大化,达到最优权衡。在此以上(B* > B*^(*)),收益递减,因为将批量大小加倍所需的总样本数增加超过了两倍,因此增加工作进程会提高吞吐量,但不会提高样本效率。

图 5.8 展示了批量大小与训练效率之间的这种关系。

图 5.8:临界批量大小与缩放区域:在临界批量大小 B*^(*) 以下,较大的批量可减少噪声并提高样本效率(线性缩放区域)。在 B*^(*) 以上,较大的批量带来边际收益递减:虽然吞吐量增加,但所需的总样本数也随之增加,从而降低了样本效率。最佳操作点在硬件利用率与收敛效率之间取得平衡。

临界批量大小对分布式训练系统设计具有重要意义。在 B*^(*) / b 之外增加工作进程(其中 b 为每工作进程批量大小)会提高吞吐量但不提高样本效率,如果额外工作进程的边际成本较低,这可能仍因成本原因而值得。学习率调度也很重要,因为在 B*^(*) 以上激进的预热变得至关重要,因为初始化附近的损失地形可能不支持线性缩放所产生的大幅更新。通信权衡也会发生变化,因为在 B*^(*) 以上较大批量的收益减小使得通信开销相对更昂贵,这加强了采用梯度压缩或异步方法的理由。

给定一个分布在 64 块 A100 GPU(每块 80 GB HBM2e)集群上的 7B 参数模型,确定最大有用批量大小。使用带有 tr(Σ) = 1.6 × 10⁶||μ||² = 100、每 GPU 批量 b = 4 以及提议的梯度累积深度 32 步的试运行迹,得到 B[global] = 64 × 4 × 32 = 8192 个样本。

实例解析:8 与 64 工作进程的收敛对比

为了具体说明这些概念,考虑在训练具有基准批量大小 B = 32(每工作进程)的 Transformer 语言模型时,从 8 个工作进程扩展到 64 个工作进程。

设置:具有 1.3B 参数的 Transformer 模型,目标困惑度 15,基准训练:单 GPU 上 100K 次迭代,b = 32

8 工作进程 (BSP)

  • 有效批量大小:B = 8 × 32 = 256

  • 学习率:η = 8 × η[base](带预热的线性缩放)

  • 预期迭代数:100K / 8 = 12.5K 次迭代

  • 收敛:在 12.8K 次迭代中达到目标困惑度(98% 效率)

  • 通信开销:15%(NVLink 节点内)

  • 墙钟加速比:100K × 1 / (12.8K × 1.15×) = 6.8×

64 工作进程 (BSP)

  • 有效批量大小:B = 64 × 32 = 2,048

  • 学习率(N = 64 工作进程):η = N × η[base](若 B < B*^(*))或 η = √N × η_base(若 B > B*^(*)

  • 假设 B*^(*) ≈ 4,000(低于临界值):适用线性缩放

  • 预期迭代数:100K / 64 = 1.56K 次迭代

  • 收敛:在 1.72K 次迭代中达到目标困惑度(91% 效率)

  • 通信开销:45%(InfiniBand 节点间,8 个节点)

  • 墙钟加速比:100K × 1 / (1.72K × 1.45×) = 40.1×

SSP 工作进程:64 工作进程(s = 4

  • 相同的有效批量大小:B = 2,048

  • 学习率:η' = η_BSP / √(1 + s),其中 s = 4,得 η'^≈ 0.45 × η[BSP]

  • 预期迭代数:因陈旧度惩罚而增加

  • 收敛:在 2.1K 次迭代中达到目标困惑度(74% 效率)

  • 通信开销:25%(同步开销降低)

  • 墙钟加速比:100K × 1 / (2.1K × 1.25×) = 38.1×

分析 (表 5.5)

| 配置 | 迭代数 | 通信开销 | 墙钟加速比 | 样本效率 |

| --- | --- | --- | --- | --- |

| 1 GPU (基准) | 100,000 | 0% | 1× | 100% |

| 8 GPU BSP | 12,800 | 15% | 6.8× | 98% |

| 64 GPU BSP | 1,720 | 45% | 40.1× | 91% |

| 64 GPU SSP | 2,100 | 25% | 38.1× | 74% |

表 5.5:跨集群规模缩放 SGD:在 1、8 和 64 GPU 上相同训练任务的迭代次数、通信开销、墙钟加速比和样本效率。

尽管样本效率仅为 91%,64-GPU BSP 配置仍实现了 40× 加速,因为通信开销(45%)被海量并行性抵消。SSP 以较低的通信开销提供了相当的墙钟时间,但需要更多总样本。

成本分析(假设 $3/GPU-小时):

  • 8:12.8K 迭代 × 0.4 秒/迭代 × 8/3600 × $3/GPU-小时 = $34

  • 64 BSP:1.72K 迭代 × 0.58 秒/迭代 × 64/3600 × $3/GPU-小时 = $53

  • 64 SSP:2.1K 迭代 × 0.50 秒/迭代 × 64/3600 × $3/GPU-小时 = $56

尽管并行度更高,但由于通信开销和样本效率降低,64-GPU 训练每次运行成本更高。8-GPU 配置更具成本效率,但需要 6 倍的墙钟时间。选择取决于优先考虑最小化成本还是最小化产出时间。

权衡:通信成本 vs. 收敛速度

分布式训练中的根本权衡在于通信效率与收敛质量之间。图 5.9 可视化了这一权衡空间。

图 5.9:通信-收敛权衡空间:每个点代表一种不同的分布式训练配置。帕累托前沿(虚线)显示了最优配置,在这些配置中,改进一项指标需要牺牲另一项。BSP 位于高收敛质量但较低吞吐量处;ASP 在收敛成本下提供最大吞吐量。梯度压缩和 SSP 处于中间位置。

有几种技术占据了这条权衡曲线上的不同位置。梯度压缩通过量化、稀疏化或低秩近似传输压缩的梯度信息来减少通信量。示例包括基于符号的更新 (Bernstein 等人 2018)、深度梯度压缩 (Lin 等人 2018)、PowerSGD (Vogels 等人 2019)、带收敛分析的 QSGD 量化 (Alistarh 等人 2017) 以及用于神经机器翻译的经验性稀疏梯度丢弃 (Aji 和 Heafield 2017)。

局部 SGD 采取了不同的方法:工作进程在同步前执行 H[local] 次局部更新,将通信频率降低 H[local] 倍。收敛分析表明,对于光滑、强凸目标,局部 SGD 在适当调整学习率下可达到与同步 SGD 相同的渐近速率 (Stich 2019)。

去中心化 SGD 将工作进程限制为仅与通信图中的邻居通信,而不是执行全局 AllReduce。这以较慢的混合为代价降低了带宽需求,使其适用于全局同步代价高昂的地理分布式训练。

这些方法的选择取决于具体的瓶颈。当网络带宽限制吞吐量时,梯度压缩提供最佳权衡。当同步延迟占主导时,首选局部 SGD 或 SSP。当存在网络拓扑约束时,可能需要去中心化方法。

模型并行

当模型状态、激活值、优化器状态或单个张量运算超出一个加速器组的内存或通信预算时,数据并行将彻底失效。这就是运算形式下的内存容量鸿沟(原理 ):模型参数增长速度超过了设备内存增长速度,迫使模型的计算和状态必须被分区。数据并行的内存优化延长了复制的极限,但最终模型本身必须被分区。

模型并行是一种分布式训练策略,它将单个神经网络的参数或运算划分到多个设备上,使每个设备计算模型的一个不同部分。

  1. 意义:如果一个具有参数状态 S[model] 的模型被分割到 N[mp] 个设备上,理想情况下的单设备参数占用将降至 S[model]/N[mp],从而实现训练超出单设备内存的模型。代价是激活值和梯度的额外通信,以及因顺序依赖产生的空闲时间,因此内存收益必须超过分区引入的传输和流水线气泡开销。

  2. 区别:与在每个工作节点上复制完整模型并分片小批量数据的数据并行不同,模型并行分片的是模型本身。流水线并行和张量并行是模型并行的实现形式:前者将层划分为阶段,后者划分单个张量运算。

  3. 常见误区:一个常见的误解是模型并行会自动加速训练。其主要优势在于容量而非吞吐量;如果没有精心的调度与重叠,下游设备将等待上游激活值,导致利用率甚至低于单设备基线。

三个内存台阶的**图,展示 1750 亿参数混合精度 Adam 训练状态:优化器状态 2,100 GB(最长),梯度 350 GB,权重 350 GB。优化器台阶约为其他两者的六倍。

主导单副本训练预算的是优化器状态,而非权重。

即使完全部署 ZeRO-3,将优化器状态、梯度和参数分片到各工作节点,某些架构仍难以处理。张量并行(如图 5.10 所示)通过跨设备分区单个权重矩阵来解决此问题;流水线并行将层划分到各阶段,稍后将在第 5.5.2.2 节介绍。对于 1750 亿参数模型,仅权重就占 350 GB,在 64 张 GPU 上平均每张约 5.5 GB。完整的混合精度 Adam 训练状态大得多:全局 2,800 GB,激活值之前每 GPU 约 43.8 GB。这一区别很重要,因为优化器分片减少了静态状态,但无法消除迫使采用模型并行的激活值和逐层容量约束。

图 5.10:张量并行数据流:张量并行将单个权重矩阵(W[0]…W[3])按列分割到四张 GPU 上,每张 GPU 计算部分输出,再通过 AllGather 重组。配套策略流水线并行则将层垂直分区到各阶段,在层块图中单独展示。这两种策略常结合用于大模型的混合 3D 并行。

对于激活值内存占主导的长上下文 Transformer,2048 token 的序列通过 1750 亿参数,即便在很小的微批次下也会生成约 1 TB 级别的中间激活值(按前文计算,微批次 4 时约 1.1 TB),任何程度的优化器分片都无法解决这一约束。模型并行通过跨设备分割模型架构本身(而非带分片状态的复制)来应对这些限制。

问题:在 NVIDIA A100(80 GB)上训练 1750 亿参数模型(如 GPT-3)。搭配 ZeRO-3 的数据并行能否应对?

数学分析

  1. 参数存储:1750 亿参数 × 2 字节 (FP16) = 350 GB。

  2. 梯度与优化器状态:梯度增加 1750 亿参数 × 2 字节 = 350 GB,而 Adam FP32 主权重、动量和方差增加 1750 亿参数 × 12 字节 = 2,100 GB。

  3. 总静态内存:2,800 GB。

  4. ZeRO-3 分片:使用 64 张 GPU,单 GPU 静态内存 = 2,800 GB / 64 GPU ≈ 43.8 GB。

  5. 激活值内存:序列长度 2048、批次大小 1 时,96 层 Transformer 每 GPU 生成约 50 GB 激活值。

系统洞察:43.8 GB(静态)+ 50 GB(动态)= 93.8 GB。这超过了 A100 的 80 GB 容量。即使完全启用 ZeRO-3 分片,纯数据并行也会失败。唯一出路是张量并行,将层本身进行分割。

模型并行通过跨设备分发模型架构本身来解决这一限制,但关键决策在于在哪里切分模型。基于层的切分将设备分配给顺序层组,如设备 1 负责第 1–4 层,设备 2 负责第 5–8 层。基于通道的切分划分层内的通道,如设备 1 负责 512 个通道,设备 2 负责剩余通道。Transformer 架构增加了注意力头切分,不同设备拥有不同的头。每种切分改变不同的成本:层切分节省内存但串行化流水线,而通道和头切分暴露更多并行工作,但需要更快的层内通信。

这些切分针对不同原因超出单设备极限的大规模场景。GPT-3 拥有 1750 亿参数,依赖模型并行进行训练。处理高分辨率 16k × 16k 像素图像的视觉 Transformer 利用模型并行管理激活值和内存约束。混合专家架构利用此方法将其条件计算路径分发到硬件上(Shazeer et al. 2017;Lepikhin et al. 2021;Fedus et al. 2022)。

设备协调分三个阶段进行。前向传播中,数据按顺序流经不同设备上的模型段。反向传播中,梯度按逆序回传通过这些段。参数更新阶段,每个设备仅修改其分配到的模型部分。协调确保了数学上等同于单设备训练,同时支持训练超出单设备内存容量的模型。代价是阶段间的顺序依赖:当一台设备计算前向传播时,下游设备处于空闲,产生的流水线气泡主导了模型并行系统的利用率分析。

验证你对模型分片的理解:

模型并行实现

一旦模型本身成为被分区的对象,实现问题就变成了放置:每一切分都必须节省足够的内存,以抵消其引入的激活值传输和空闲时间。图 5.11 捕捉了这种双向数据流:输入数据按顺序向前传播通过分配的模型分区,而梯度在更新阶段向后流动以更新参数,中间结果在每个阶段的设备边界间传输。

图 5.11:模型并行数据流:模型分区按顺序分布在三个设备上:输入数据按顺序通过各分区向前流动(上部路径),而梯度在更新阶段向后传播(下部路径)。每个设备处理模型的不同部分,中间激活值和梯度在分区边界传输。这种方法以顺序依赖降低硬件利用率为代价,实现了训练超出单设备内存的模型。

考虑我们的持续示例:1750 亿参数模型在 FP16 下需要 350 GB 内存,超过单张 A100 80 GB 容量的四倍。模型并行通过跨多设备分区模型权重来解决这个“容量墙”问题,有效地将它们拼接成一个超级加速器。与数据并行不同,后者中每个 GPU 持有完整的模型副本并处理全局批次的独特部分,模型并行要求每个 GPU 持有模型的“唯一”部分并“按顺序”处理相同的数据流。在 A100 上采用 8 路分区时,仅 BF16 权重每 GPU 就占用约 44 GB。这只是权重预算:完整训练仍需激活值、梯度和优化器状态,因此纯 8 路模型并行本身是不够的。训练还需要为优化器和梯度状态提供额外的分片或卸载,以及流水线或混合并行来将激活内存控制在容量范围内。

在典型的流水线并行实现中,训练循环像接力赛一样运作。前向传播在 GPU 1 上启动,它计算初始的 Transformer 块,并将生成的中间激活张量通过互联传输到 GPU 2。对于隐藏维度为 12,288、微批次大小为 4 个序列(每个 2,048 个 token)的 1750 亿模型,此交接涉及每步每阶段边界移动约 200 MB 数据。GPU 2 必须等待此负载到达才能开始计算,从而形成一条贯穿所有阶段的严格依赖链。反向传播沿相反路径镜像此过程,将误差梯度从最后一层传播回输入,每个设备仅计算其本地参数的梯度。

与数据并行相比,这种架构从根本上改变了优化动态。每个 GPU 在其特定参数切片上执行局部优化器步骤(Adam(Kingma and Ba 2014)、AdaFactor 或类似算法),而不是通过全局 AllReduce 跨副本平均梯度。持有 Transformer 第 1-12 层的设备仅更新这些层的权重和偏置,优化步骤期间无需跨设备同步。虽然这消除了数据并行中占用大量带宽的梯度同步,但它用一个瓶颈换来了另一个:流水线气泡。如果分配给 GPU 1 的层比 GPU 2 的层计算量更大(当注意力层头数不同或嵌入层大小不均时常见),宝贵的计算周期就会浪费在等待上。因此,主要的工程挑战从最大化算术强度转移到了最小化序列化延迟,并确保分区机群间的负载均衡(Rasley et al. 2020)。

并行变体

分区策略决定了当内存无法容纳单设备时模型付出的代价。逐层分区以空闲时间换取降低单设备内存;流水线并行以调度复杂度换取隐藏空闲时间;张量并行以 NVLink 带宽换取拆分单层。以下方法按使各方法可行的架构和互联将这些选择分开。

逐层分区

逐层分区是侵入性最小的模型切分:它通过将连续层分配给独立设备来节省内存,但保留了模型的顺序依赖。在 Transformer 架构中,这转化为特定设备管理定义好的注意力和前馈块集合。图 5.12 演示了针对 16 层 Transformer 的这种分区:四个连续块驻留在四个设备中的每一个上,前向激活从左向右流动,反向梯度从右向左跨设备边界传播。

图 5.12:逐层模型并行:一个 16 层 Transformer 分布在四张 GPU 上,每个设备分配四个连续的 Transformer 块(L1–L4、L5–L8、L9–L12、L13–L16)。前向激活(黑色箭头)通过设备边界从左向右流动,而反向梯度(红色箭头)在参数更新期间从右向左传播。这种分区将单 GPU 内存从完整模型大小降低到 1/4,但引入了顺序依赖,即下游设备需等待上游计算完成。

顺序处理引入了设备空闲时间,因为每个设备必须等待前一个设备完成计算才能开始工作。当设备 1 处理初始块时,设备 2、3 和 4 处于非活动状态。同样,当设备 2 开始计算时,设备 1 处于空闲。这种等待和空闲的模式相比其他并行化策略降低了硬件利用效率。

流水线并行

流水线并行通过引入微批次来最小化设备空闲时间,扩展了逐层分区。它不等待整个批次顺序通过所有设备,而是将计算划分为称为微批次的较小片段,并跨流水线阶段重叠执行。图 5.13 展示了这种重叠如何工作:当设备 1 处理微批次 i + 1 时,设备 2 计算微批次 i,设备 3 处理 i − 1,设备 4 执行 i − 2,形成持续流使所有设备同时保持活跃。

如图 5.13 所示,图中每行代表的每个设备同时为不同微批次处理其分配的模型层。前向传播涉及设备将激活传递给下一阶段,例如 F[0, 0] 到 F[1, 0]。反向传播将梯度沿流水线回传,例如 B[3, 3] 到 B[2, 3]。这种重叠计算减少了空闲时间并提高了吞吐量,同时保持了跨设备操作的逻辑序列,赋予了该策略其正式形式。

图 5.13:流水线并行调度:一个 4 阶段流水线处理 4 个微批次,展示跨时间步的前向传递(F[i, j])和反向传递(B[i, j])。行代表流水线阶段(GPU),列代表时间步。错峰执行使所有设备保持活跃:当阶段 0 计算 F[0, 1] 时,阶段 1 处理来自前一微批次的 F[1, 0]。所有前向传递完成后,反向传递按逆序传播。“Update”列显示跨所有微批次累积梯度后的同步参数更新。

流水线并行 是一种模型并行技术,它将神经网络的层划分为分配给不同设备的顺序阶段,在阶段间前向传递激活和反向传递梯度,同时利用微批次跨阶段重叠计算以维持吞吐量。

  1. 重要性:阶段间通信仅在每个阶段边界传输激活张量,其在 BF16 下大小为 B[μ] × S × d[model] × 2 字节,其中 B[μ] 是微批次大小。对于隐藏维度 8,192、微批次大小 1 和 2,048 token 序列,这约为 8,192 × 2,048 × 2 ≈ 32 MB 每边界,相比之下数据并行中的梯度 AllReduce 需要千兆字节。这种低通信量使流水线并行成为通过 50 GB/s InfiniBand 连接的节点扩展模型深度的主要技术。流水线气泡浪费约 (p − 1)/(m + p − 1) 的总计算量,其中 p 是阶段数,m 是微批次数——当 p = 8 阶段且 m = 32 微批次时,气泡开销约为 18%。

  2. 区别:与一次只分片单层内的单个矩阵乘法的张量并行不同,张量并行需要在每一层的输出上执行 AllReduce(在每个阶段内需具备 NVLink 级别的带宽),而流水线并行是在层边界处进行分片,仅需在阶段之间进行点对点激活传输——可以容忍节点间的 InfiniBand 带宽。

  3. 常见陷阱:一个常见的误解是,增加流水线阶段总是能提高吞吐量。流水线气泡分数随着 (p − 1)/(m + p − 1) 的增长:当 p = 16 个阶段且 m = 16 个微批次时,有 48% 的计算被闲置浪费——这就需要较大的微批次数量(m ≫ p)才能将气泡降至 10% 以下,而这又会增加峰值激活内存并加大每个阶段的内存压力。

比较流水线气泡税的两个比例堆叠条形图:p8 m32 约有 18% 的闲置时间,而 p16 m16 约有 48% 的闲置时间。

更多阶段需要足够的微批次,否则气泡将主导利用率。

相同的气泡项在下面的示例中变得具体:一个 8 阶段流水线和 32 个微批次仍会留下可测量的空闲时间税。

问题:一个大规模模型训练任务在 8 个阶段上使用流水线并行。为了掩盖顺序延迟,批次被拆分为 32 个微批次。那么“气泡税”——即因空闲等待而损失的 GPU 循环比例——是多少?

数学:在同步流水线(1F1B)中,气泡分数由阶段数与微批次数的比率决定。

  1. 等待时间:在每个批次的开始和结束时,GPU 会因 p − 1 个步骤而空闲。

  2. 生产时间:GPU 用于计算的时间为 m 步。

  3. 气泡分数:(p − 1)/(p − 1 + m)= 7 / (7 + 32) ≈ 17.9%。

系统洞察:流水线并行是一种 利用率-内存权衡。将气泡从 18% 降至 5% 需要更多的微批次(m),这会增加每个 GPU 上的激活内存消耗。在机器学习集群中,深度并非免费可扩展;集群需要承担 17.9% 的容量税,仅仅是为了保持阶段协同。这就是为什么像交错流水线这样的技术至关重要:它们将气泡切分成更小的块,以恢复那 18% 的集群容量损失。

GPipe⁶² (Huang et al. 2019) 引入了带有微批次累积的同步流水线并行,而 PipeDream (Narayanan et al. 2019) 则发展出了带有权重缓存的异步方法。现代系统采用 1F1B(一次前向-一次后向)⁶³ 调度,以更早地释放激活内存;这种较低的内存压力使得更大的微批次数量变得实际可行,而更大的 m 正是压缩气泡项的关键。

在调度层面,1F1B 是一个包含预热、稳态和排空的循环。算法 1 明确规定了:每个阶段在梯度到达后,交替进行较新微批次的前向工作和较旧微批次的后向工作,这样激活就可以在整个批次完成前被释放。

算法 1 一次前向-一次后向流水线调度

要求: p 个有序流水线阶段;m 个微批次;激活链接 j → j + 1 和梯度链接 j + 1 → j

保证: 在所有 m 个微批次的梯度累积后,进行一次优化器更新

  1. 将全局批次拆分为 m 个微批次;将连续的层分配给 p 个阶段

  2. 预热:从阶段 0 启动前向传播,随着激活的产生,将其下游传输

  3. 稳态中的每个阶段 执行

  4. 交替执行一个后向(最旧就绪的微批次)和一个前向(下一个微批次)

  5. 一旦该阶段的后向消耗了激活,立即释放该微批次的激活

  6. 结束循环

  7. 排空:按阶段逆序完成剩余的后向传递

  8. 一旦每个阶段为所有 m 个微批次累积了梯度,则应用优化器更新

预热和排空阶段仍会留下大约 (p − 1)/(m + p − 1) 的气泡,但在稳态中交错后向工作可以在激活被其后向消耗的同时立即释放,因此峰值激活内存将随阶段数增长,而不是随微批次数增长。这正是相比于 GPipe 风格的“全部前向-全部后向”调度所带来的收益,它使得系统能够增加 m,直到激活内存、重新计算或全局批次约束成为下一个限制。在跨四个设备分布的 transformer 模型中,设备 1 会为微批次 i + 1 处理块 1-6,而设备 2 计算微批次 i 的块 7-12。同时,设备 3 执行微批次 i − 1 的块 13-18,而设备 4 处理微批次 i − 2 的块 19-24。每个设备维护其分配的 transformer 块,但在不同的微批次上运行,从而形成连续的计算流。

隐藏状态在设备之间的传输是持续进行的,而不是分阶段进行的。当设备 1 完成一个微批次的处理时,它会立即将形状为 B[μ] × S × d[model] 的输出张量传输到设备 2,并开始处理下一个微批次。这种计算重叠模式在保持模型数学性质的同时,保持了硬件的完全利用率。

一个重要的变体直接针对剩余的气泡。流水线开始阶段的 GPU 在等待从末端返回的梯度时会空闲,反之亦然。这些气泡代表了浪费的计算。算法 1 中的 1F1B 调度在稳态期间让 SM 保持忙碌并减少激活驻留,但除非系统增加 m 或在这些原本空闲的时段内调度有用的工作,否则填充和排空的气泡仍会存在。

零气泡流水线调度 通过将权重梯度计算与激活梯度通信重叠,进一步减少了空闲时间。在标准的后向传递中,GPU 同时计算 ∂ℒ/∂W(权重梯度)和 ∂ℒ/∂X(激活梯度,发送到前一阶段)。零气泡调度将这些拆分为单独的内核:一个 B 内核仅计算激活梯度 ∂ℒ/∂X 并将其发送到前一阶段,以及一个 W 内核仅在本地计算权重梯度 ∂ℒ/∂WB 内核必须及时执行(它处于关键路径上),但 W 内核可以被机会性地调度以填充气泡。

此 B/W 拆分所提供的调度自由度相当大。在一个包含 4 个阶段和 8 个微批次的流水线中,标准的 1F1B 调度的气泡分数约为 (p − 1)/(m + p − 1),其中 p 为阶段数,m 为微批次数。当 p = 4、m = 8 时,这相当于 3/11 ≈ 27% 的空闲时间。零气泡调度可以通过用 W 计算填充启动和收尾气泡,将其降低至接近零。

权衡在于内存:零气泡调度需要更长时间地存储中间激活(因为 W 计算被推迟),从而增加峰值内存使用量。一些实现通过将零气泡调度与激活检查点相结合来解决这一问题,即选择性地重新计算某些激活而不是存储它们。这些技术之间的相互作用创造了一个三元权衡:流水线气泡大小、内存消耗和重新计算开销之间的权衡,这是开销置换(原则 )的一个例子。

张量并行

流水线并行与张量并行在模型并行中的应用

流水线并行

上文中探讨的流水线并行通过在各阶段之间重叠微批处理来解决设备空闲时间问题。每个设备存储完整的层并按顺序处理它们,通信仅在设备之间传输激活值时发生,即在阶段边界处。该方法能够容忍中等互连带宽,因为通信发生得不频繁——每个微批处理每层边界仅发生一次。然而,当单个层本身超出设备内存时,或当层内部的通信模式从不同的粒度获益时(不同于层边界),流水线并行则无济于事。

张量并行

张量并行采取了一种根本不同的方法:它不是将完整的层分配给设备,而是在每一层内部拆分权重矩阵。这种操作级别的并行(也称为层内并行)使得更细粒度的分布成为可能,但需要高带宽互连来支持它引入的频繁层内通信。

张量并行 是一种模型并行技术,它使用列并行或行并行权重拆分将单个张量操作(主要是矩阵乘法)划分到多个设备上,在 Megatron 风格的 Transformer 块中通常需要每个 Transformer 层进行两次 AllReduce 操作来求和所有参与设备的部分结果。

1. 意义

Megatron-LM 风格的张量并行在每个 Transformer 块中放置两个 AllReduce 操作——一个在注意力机制之后,一个在 MLP 之后——每次集合操作减少的激活张量大小为 B[batch] × S × d[model] × 2 字节(BF16)。环形 AllReduce 每个 GPU 大约移动 2(t − 1)/t 倍的有效载荷。在带宽为 900 GB/s 的 NVLink 上,当并行度 t = 8 时,每次 AllReduce 大约耗时 0.1–0.5 ms。同样的操作在 InfiniBand(50 GB/s)上每层耗时 2–10 ms——对于 96 个 Transformer 块来说,这每步会增加 200–960 ms 的纯同步开销,导致 MFU 降至个位数。

2. 区别

与仅在层边界通信并传输小激活切片的流水线并行不同,张量并行通过每一层内的 AllReduce 进行同步,使其对互连的每操作延迟非常敏感,在实际中通常局限于节点内的 NVLink 结构中。

3. 常见误区

一个常见的误解是张量并行可以简单地通过 InfiniBand 在节点之间扩展。NVLink 提供大约 900 GB/s 的双向带宽,或每方向 450 GB/s;InfiniBand NDR 每端口提供 50 GB/s——每方向相差约 9×。在 InfiniBand 上运行张量并行通常会导致 MFU 降至个位数百分比,使得通信开销超过分布式矩阵乘法带来的计算收益。

互连拓扑决定了在集群层次结构的每一级哪种模型并行形式是可行的。张量并行的每层同步需求需要 NVLink 级别的带宽;流水线并行的边界通信则可以容忍 InfiniBand。这种带宽模式产生了混合方案的设计压力:在需要高带宽的层内拆分中使用张量并行,而在粗粒度的层间拆分中使用流水线并行。

Violet bandwidth ladder comparing per-direction NVLink bandwidth with NDR InfiniBand, illustrating the order-of-magnitude bandwidth drop from intra-node to inter-node links.

张量并行需要 NVLink 带宽;流水线并行可以容忍较慢的互连结构。

Megatron 风格的张量并行

Megatron 风格的张量并行⁶⁴ (Shoeybi et al. 2019) 以两种方式分割矩阵乘法。查看 图 5.14:列并行分割沿着列方向划分权重矩阵,用于 QKV 投影,使得各 GPU 可以独立计算;而行并行分割沿着行方向划分,用于输出层,需要在每个块末尾通过 AllReduce 来合并部分和。

列并行线性层

列并行线性层沿着列方向分割权重矩阵。对于输入 X 和被分割到 2 个 GPU 上的权重矩阵 W = [W[1]|W[2]],有:Y = X**W = X[W[1]|W[2]] = [X**W[1]|X**W[2]] 每个 GPU 独立计算其分区,当后续操作是行并行时,输出可以直接拼接而无需通信。

行并行线性层

行并行部分的模式沿着行方向分割权重矩阵,\(W = \\begin{bmatrix} W\_1 \\\\ W\_2 \\end{bmatrix}\)Y = X**W = X[1]W[1] + X[2]W[2] 每个 GPU 计算一个部分和,系统需要支付一次 AllReduce 来合并部分输出。

图 5.14:张量并行 - 矩阵分割:展示了 Megatron-LM 风格的张量并行。第一个线性层(例如,QKV)沿列方向分割 [W[1]|W[2]。第二个层(例如,输出)沿行方向分割 [W[1]; W[2]。这种安排使得第一层的输出可以直接流入第二层,无需同步,仅需在第二层之后进行一次 AllReduce 来求和部分结果。

图 5.14 中所示的列-则-行排列是关键的设计洞察。通过将列并行层与行并行层配对,中间激活值可以在这两者之间直接流动,无需通信,从而将同步限制在该对的末尾。Megatron 在每个 Transformer 块中应用该模式两次:QKV 投影是列并行的,其注意力输出投影是行并行的;然后第一个前馈层是列并行的,第二个前馈层是行并行的。该设计战略性地放置 AllReduce 操作——注意力之后一次,前馈网络之后一次——从而在每个 Transformer 层中实现两次 AllReduce 操作。每个 Transformer 层的通信量取决于序列长度 S、隐藏维度 d[model]、张量并行度 t 和批量大小 B。激活有效载荷为:M[act] = B × S × d[model] × sizeof(dtype)

环形 AllReduce 每个 GPU 传输大约 2(t − 1)/t × M[act] 字节,当 t 很大时,这接近于 2M[act]。当 S = 2048,d[model] = 4096,B = 4,且使用 FP16 时,激活有效载荷为 4 × 2048 × 4096 × 2 ≈ 67 MB。应用环形因子后,每次 AllReduce 大约传输 134 MB。对于一个包含 96 层的模型,每层有两次 AllReduce 操作,前向传播的总通信量约为 25.7 GB;包括对称的反向传播 AllReduce 后,每个训练步的通信总量约为 51 GB,这就需要 NVLink 带宽来避免成为瓶颈。

扩展限制

张量并行的扩展性在超过 8 方式并行时会急剧下降,因为同样的拆分虽然减少了内存占用,但也减少了用于掩盖通信的可用计算。随着张量并行度的增加,每个 GPU 的计算工作量降低,而集体延迟和总流量增加;每个 GPU 的环形 AllReduce 字节数趋近于激活张量大小的 2×,而不是无界线性增长,但更小的局部矩阵乘法留下的有用工作量减少,无法有效覆盖该交换。一旦 NVLink 带宽饱和,增加另一个张量并行分片主要会暴露更多的同步开销,而不是提升吞吐量。已发表的系统(如 Llama 3 405B)在由 H100 组成的 NVLink 互连节点中使用 TP=8 (Dubey et al. 2024),类似的节点本地张量并行是大规模 LLM 训练的常见设计模式 (Z. Jiang et al. 2024)。

同样的分区思路既适用于序列长度,也适用于矩阵维度。标准张量并行在节点内将计算跨 HBM 与 NVLink 的边界进行平铺,而 Ring Attention(环状注意力) 则将注意力状态分散到各个 GPU 上,而不是强迫每个设备都保存完整的前缀(Liu et al. 2023)。这对于超过单块 GPU 显存容量的序列(如百万 Token 的上下文窗口)变得至关重要。每个 GPU 拥有一个查询 Q 的块,并沿环状拓扑循环传递键/值(K/V)块,计算当前驻留在内存中的块的注意力,然后将该块转发给邻居。

该算法在 N1 轮通信中进行(其中 N 为 GPU 数量)。在每一轮中,每个 GPU 计算其本地 Q 块与当前驻留的 K/V 块之间的注意力,将该 K/V 块发送给其环状邻居,并从另一侧邻居接收下一个块。实现将这种交换与计算重叠:当 GPU i 使用 K[j]/V[j] 计算注意力时,它同时从环中接收 K[j+1]/V[j+1]。

如果单个 Tile 的计算时间超过通过 NVLink 传输一个 Tile 的通信时间,通信就会被完全隐藏。在拥有 3.35 TB/s HBM 带宽和 900 GB/s NVLink 带宽的 H100 上,对于典型的 Tile 大小,这种重叠是可以实现的。

Ring Attention 的实际影响体现在上下文长度上。没有它,单块 GPU 的注意力计算受限于 HBM 容量:长度为 S 的序列的键/值状态必须完全放入一块 GPU 的内存中。跨 N 块 GPU 使用 Ring Attention 时,每块 GPU 仅保存该状态的 S/N 个 Token,从而实现 N×S[single] 的上下文长度。第 9 章 后文会将这种序列级分发与 FlashAttention 的 Tile 级 HBM 复用联系起来;此处,分布式训练的启示是序列长度可以成为像层、张量或数据一样的分区维度。

参数服务器与嵌入分片

虽然 AllReduce 主导了稠密模型训练,但 参数服务器(Parameter Server, PS)架构(由 Li 等人于 2014 年形式化)在推荐系统和其他稀疏工作负载中仍很常见。参数服务器架构将 Worker(计算梯度)与 Server(存储参数并应用更新)分离开来。

对于稠密模型(如 ResNet 或 BERT),PS 架构会产生瓶颈:有 N 个 Worker 时,Server 的入站带宽必须同时吸收 N 个梯度流,饱和 Server 的网络带宽,使其成为通信瓶颈。这种“Incast(多对一拥塞)”问题推动了 Ring AllReduce 的采用,在 Ring AllReduce 中,每个 Worker 以自己的链路速率收发数据,通过将负载分散到所有节点,实现了 N 倍的聚合带宽。对于超过 4-8 块 GPU 的稠密模型,去中心化的 AllReduce 具有决定性优势。

然而,对于 推荐系统——特别是 DLRM——模型参数主要由海量嵌入表(通常 10 TB+)主导,无法装入任何单块 GPU。此外,更新是稀疏的:一批用户仅与极小比例(如 0.001%)的物品交互。

在这种稀疏模式下,参数服务器仅移动批次涉及的行,避免了稠密同步:

  1. 嵌入分片:海量表被分区到 PS 集群中(通常是拥有海量 DRAM 的 CPU 节点)。

  2. 稀疏查找:Worker 向 PS 发送 ID 列表。

  3. 稀疏拉取:PS 仅返回请求的嵌入向量,而非完整表。

  4. 稀疏推送:Worker 仅为触及的嵌入行发送梯度。

稀疏拉取/稀疏推送 模式避免了稠密 AllReduce 的带宽瓶颈。如 TorchRec 或 Meta 的分层分片等实现,将“热门”嵌入放在 GPU 上,“冷门”嵌入放在 CPU PS 节点上,为模型参数创建了分级内存层次结构。

专家并行(混合专家模型)

张量并行将稠密层拆分到各设备上,而专家并行则通过条件计算,在不增加计算成本(FLOPs)的情况下扩展模型容量(参数量)。在混合专家(Mixture-of-Experts, MoE)架构(Shazeer et al. 2017)中,每个 Transformer 块的前馈网络(FFN)被一组 E 个“专家”(独立的 FFN)取代。对于每个 Token,门控网络选择一小部分(通常是 top-1 或 top-2)专家来处理它。

在分布式环境下,专家被分区到各个 Worker 上。如果我们有 8 块 GPU 和 8 个专家,每块 GPU 托管一个专家。训练过程引入了一个独特的通信模式,如 图 5.15 所示:

  1. 门控:每个 Token 确定其目标专家。

  2. All-to-All 分发:Token 被路由穿过网络,发送到托管其选定专家的设备。

  3. 计算:专家处理分配给它们的 Token。

  4. All-to-All 合并:处理后的 Token 被路由回其原始设备以恢复序列。

图 5.15:混合专家(MoE)All-to-All 路由:专家并行需要一种独特的“All-to-All”通信模式。本地批次中的 Token 由门控网络路由到分布在集群中分片的专家。这涉及一次全局洗牌,每个 GPU 都要向其他每个 GPU 发送和接收 Token,这对集群的剖分带宽造成压力。

主要优势是将模型规模与计算预算解耦。一个万亿参数的 MoE 模型每个 Token 可能仅使用 100 亿参数,从而能够在可行的硬件预算上进行训练。约束在于 All-to-All 通信,它对带宽要求高且对负载不均衡敏感。

专家并行 是混合专家模型的分发策略,它将不同的专家放置在不同的设备上,并通过 All-to-All 集体通信将每个 Token 的激活路由到托管其选定专家的设备,从而在不增加每个 Token 计算量的情况下,跨集群扩展总参数量。

  1. 意义:它将模型容量与计算预算解耦:万亿参数的 MoE 模型每个 Token 可能仅激活约 100 亿参数,因此每 Token FLOPs 保持在 100 亿稠密模型的水平附近,而容量却增长了约 100 倍。代价是通信结构:每个 MoE 层每次前向/反向传播都需要两次 All-to-All 洗牌(分发和合并),其通信量随移动的 Token 数乘以隐藏维度而扩展,这给集群的剖分带宽而非单个链路带来压力。

  2. 区别:与张量并行不同,张量并行拆分稠密操作,使每个设备计算每个 Token 的一小部分;专家并行是条件计算:每个设备仅计算路由到其专家的 Token,这使得通信模式和负载分布都变得依赖于数据,而非由架构固定。

  3. 常见陷阱:一个常见的误解是专家之间负载均衡。自然的 Token 分布是倾斜的:热门专家可能接收 3-5 倍于其公平份额的流量,导致其设备的激活缓冲区溢出,而同级设备却闲置。这就是为什么生产级 MoE 系统会限制每个专家的 Token 预算,并在训练中惩罚路由偏斜,而不是仅仅信任路由器。

专家并行的核心在于基于动态路由决策在集群间进行通信的 All-to-All 原语。考虑一个配置:在 64 个 GPU 上分布了 E = 64 个专家,处理长度为 S = 2048、B = 4 个序列的批次,隐藏维度为 d[model] = 4096。对于每个 MoE 层,系统必须将 B × S 个 token 调度到其分配的专家上。在 FP16 下,这将移动 B ⋅ S ⋅ d[model] ⋅ 2 字节——大约 67 MB——在单个方向上。由于处理后的嵌入必须返回到其原始设备以用于残差连接,因此每个 transformer 块的网络开销大约为 134 MB。虽然在孤立情况下是可以管理的,但在像 Switch Transformer(Fedus et al. 2022)(多达 2,048 个专家)或 GShard(Lepikhin et al. 2021)这样的深度稀疏架构中,这种延迟会迅速累积。

网络效率依赖于均匀 token 分布的假设,但自然语言本质上是偏斜的:处理常见语法或连接词的特定专家可能会获得 3–5 倍于其公平份额的流量。热点专家既是内存问题也是调度问题。如果一个专家接收了过多的 token,其 GPU 将耗尽激活缓冲区,而其他专家则闲置未用。

因此,MoE 系统通过容量因子 C(通常设置在 1.25 和 1.5 之间)强制执行硬性限制。此参数将一个专家处理的最大 token 数量限制在大约 C ⋅ (B × S)/E(针对路由组)。如果路由门分配的 token 数量超过此缓冲区允许的范围,则超额的 token 将被丢弃,通过残差连接未经处理地传递过去。为了减少这种数据损失,训练目标包括一个辅助负载均衡损失,相对于主要交叉熵损失的权重为 0.01–0.1,用于惩罚路由器偏向特定专家的行为。诸如 Mixtral 8x7B 之类的模型在 8 个专家上使用 top-2 路由,在容量扩展和路由稳定性之间达到有利的平衡。

稀疏通信模式将推荐和 MoE 工作负载(原型 B(DLRM at Scale))与密集 LLM 训练(原型 A(GPT-4/Llama-3))区分开来(Section 1.6.1)。

每种模型通过不同的机制进行扩展,受限于不同的资源:

  • LLMs(密集型):通过张量/流水线并行性进行扩展。限制:计算和互连带宽(NVLink)。

  • DLRM 风格模型(稀疏型):通过 Embedding 分片(参数服务器)进行扩展。限制:内存容量和互连延迟(随机访问)。

这种区分决定了根本不同的集群设计:用于 LLMs 的密集 GPU pod,以及用于 RecSys 的内存丰富的 CPU/GPU 混合体。

权衡:气泡 vs. 带宽困境

模型并行打破了内存墙,但引入了顺序依赖,降低了硬件利用率。工程挑战在于平衡流水线气泡(空闲时间)与 all-to-all 带宽(通信时间)。

模型并行提供了三个主要优势。内存扩展使得能够训练超出单设备容量的模型:通过 8-way 张量并行性,175B 模型的 FP16 权重片可以在考虑优化器状态和激活预算之前放入 A100 HBM 中。拆分模型还能在不出现内存不足错误的情况下实现更大的全局批量大小,因为每个 GPU 处理的是较小的参数片。这种方法自然映射到变压器的物理结构上,其中注意力头通过张量并行性进行分割,而层通过流水线并行性进行分割。

这些优势伴随着三个基本限制的代价。流水线气泡导致 GPU 在填充和排空流水线时保持空闲;气泡分数大约为 (p − 1)/(m + p − 1),其中 p 是流水线阶段,m 是微批次,要实现超过 90% 的效率需要 m ≫ p,这会增加激活内存。张量并行性中的通信强度同样具有约束性:在关键路径上每层执行两次 AllReduce 操作,要求具有极高带宽和低延迟的互连(NVLink),通常在达到带宽墙之前阻止超过单个节点(8 个 GPU)的扩展。实现复杂性则要求对模型定义进行侵入性修改——用列并行和行并行变体替换标准线性层——这与数据并行不同,数据并行是在外部包装模型而无需修改内部。

FP8 用于分布式训练

数值精度本身也是一个分布式杠杆:将每个值的字节数减半会将 AllReduce 和张量并行集体操作所移动的梯度和激活有效载荷减半,从而缩小 Tcomm 项,这一项在规模上主导步骤时间。8 位浮点(FP8)格式通过降低数值精度来利用这一杠杆,同时尽量保留足够的范围以实现稳定的优化。

FP8 在分布式训练中很重要,因为其较窄的格式在不破坏优化的情况下减少了通信。传统的混合精度训练使用 FP32 主权重,并在 FP16 前向和反向传递中进行计算。支持 FP8 的加速器(如 NVIDIA H100)提供了两种针对训练不同阶段优化的格式(Micikevicius et al. 2022)。

两种 FP8 格式之间的选择因此不仅是指令集的细节问题;它决定了哪些张量可以在不破坏优化器的情况下被压缩。E4M3(4 位指数,3 位尾数)提供大约 ±448 的范围,具有中等精度,因此适用于前向传递中的权重和激活(在这些位置,值通常聚集在可预测的分布中)。E5M2(5 位指数,2 位尾数)提供大约 ±57344 的更大范围,但精度较粗,这就是为什么它用于梯度(在反向传播过程中,梯度可能跨越多个数量级)。在梯度中使用 E4M3 会导致频繁的溢出和下溢,而 E5M2 则在略微增加更新噪声的代价下捕获了完整的梯度分布。请参阅 Table 5.6 作为分布约束映射:每一行都会改变范围、精度、内存流量和可训练性之间的平衡。

| Format || Exponent || Mantissa || Range || Precision || Use Case |

| --- | --- | --- | --- | --- | --- | --- |

| FP32 || 8 bits || 23 bits || ±3.4 × 10³⁸ || Very high || Master weights |

| --- | --- | --- | --- | --- | --- | --- |

| FP16 || 5 bits || 10 bits || ±65504 || High || Mixed-precision |

| --- | --- | --- | --- | --- | --- | --- |

| BF16 || 8 bits || 7 bits || ±3.4 × 10³⁸ || Moderate || Training |

| --- | --- | --- | --- | --- | --- | --- |

| E4M3 || 4 bits || 3 bits || ±448 || Low || FP8 forward pass |

| --- | --- | --- | --- | --- | --- | --- |

| E5M2 || 5 bits || 2 bits || ±57344 || Very low || FP8 gradients |

Table 5.6: 训练精度格式:每一行代表训练中使用的浮点格式。FP8 格式(E4M3,E5M2)位于 FP16/BF16 训练稳定性与为缩减分布式有效载荷所需的字节减少之间的通信高效中间地带。

正如 Table 5.6 所总结的,FP8 训练中的关键工程挑战是动态缩放。FP8 的狭窄动态范围意味着固定缩放因子会导致溢出或下溢。逐张量缩放将每个张量在 FP8 转换前乘以一个缩放因子,然后在 FP8 计算后除以该因子。缩放因子是动态调整的,通常通过跟踪每个张量的运行最大绝对值来实现,并选择一个使该最大值映射到 FP8 可表示最大值附近的缩放因子。

三精度方法(FP32 主权重、FP8 通用矩阵乘法 (GEMM) 运算、FP16 累加)可在适用的工作负载上实现接近 FP16 训练质量,同时提升支持 FP8 硬件上的有效吞吐量。其系统层面的影响具体但取决于硬件:相较于 FP16/BF16 张量,FP8 将数据移动字节数减半,而运算吞吐量与能效收益则视加速器实现而定。因此,吞吐量提升是数据移动变窄与硬件支持的性能体现,而非额外的乘性因子。

在分布式训练中,只有当较小的负载减少了通信量、且动态缩放保护了收敛性时,FP8 才发挥作用。它改变的是负载大小,而非并行化轴。若精度降低仍不足以满足需求,系统必须决定如何将数据、张量和流水线并行映射到硬件层级上。

混合并行

当数据并行耗尽内存、模型并行耗尽网络带宽时,训练大模型要求跨三个维度同时编排这两种策略。前述章节揭示了一种根本张力:数据并行虽能扩展吞吐,却需海量内存;模型并行虽能支撑大模型,却让计算资源饥渴难耐。

混合并行通过正交应用两种策略化解此张力:模型并行拆分架构以适配可用内存,数据并行则跨多个模型副本扩展吞吐。在 3000 亿 token 数据集上训练 1750 亿参数语言模型即实践了此法。神经网络层通过模型并行分布于多 GPU,数据并行则让不同 GPU 组处理独立批次。这种双重策略同步解决了模型规模带来的内存约束与数据集规模带来的计算需求,正是这种组合定义了大模型规模下的 Archetype A 训练。

Archetype A(GPT-4/Llama-3)是混合并行的首要驱动力。因模型参数量(P)超出任单一加速器内存(C[mem,device]),且训练数据集(D)要求海量吞吐,我们必须沿三个正交轴拆分问题:

  1. 张量并行:拆分单个层以使 P 适配单节点内存。

  2. 流水线并行:当参数占用超出单节点时,跨节点拆分层。

  3. 数据并行:复制整个拆分后的模型流水线,以在 D 上扩展吞吐。

仅结合三者方能高效训练 Archetype A 系统。

为见为何组合关键,以 1750 亿模型在 1,024 GPU 集群(128 节点 × 8 GPU)上训练为例,考察三种并行配置。从纯数据并行经张量、流水线并行的演进,使效率增益具体化。

配置 A:纯数据并行 (DP-1024)。所有 1,024 GPU 复制完整模型(用 ZeRO 分片优化器状态),各 GPU 处理不同数据分片。梯度 AllReduce 在全 InfiniBand 网络交换 350 GB。正如扩展效率分析所示,因节点间通信主导,效率仅约 37.5%。

配置 B:TP-8, DP-128。节点内 8 GPU 经 NVLink 做张量并行;节点间 128 个数据并行组经 InfiniBand 同步梯度。每个数据并行组仅需 AllReduce 自身 1/8 模型分片的梯度(43.75 GB 而非 350 GB),TP 通信限于高速 NVLink 域。节点间 AllReduce 时间从 14 秒降至约 1.75 秒。若计算时间仍为 2.1 秒,效率升至 2.1/(2.1 + 1.75) ≈ 54.5%,若计算通信重叠,实际效率达 75–85%。

配置 C:TP-8, PP-4, DP-32。节点内 8 GPU 做张量并行;跨 4 节点流水线并行将模型分 4 阶段;剩余 32 个数据并行组经 InfiniBand 同步梯度。流水线并行进一步减少梯度 AllReduce 量(每阶段约 1/4 参数),阶段间通信(传递激活)量小于全量 AllReduce。代价是流水线气泡:每微批次首尾,部分阶段因等待前级激活或后级梯度而空闲。气泡比例约为 (p - 1)/m,故 4 阶段、每步 32 微批次时,约浪费 9% 算力。

三配置对比揭示清晰模式:将高带宽通信(张量并行)限于快速 NVLink 域,仅将低带宽通信(小梯度分片的数据并行 AllReduce)推至较慢 InfiniBand 网络的配置,效率最高。基础设施层级决定并行层级。 由此原则称为层级感知并行,是大规模训练系统的常用方法。

层级感知并行

层级感知并行是将不同并行执行模式映射到集群物理带宽层级的策略。

  1. 意义:确保高频同步(如张量并行)留在最快链路(NVLink),而低频任务(如数据并行)用较慢层级(InfiniBand)。这种对齐通过减少暴露通信时间 (Tcomm) 与延迟 (L[lat]) 最大化扩展效率 (η[scaling])。

  2. 区别:不同于视所有节点间链路等同的统一并行,层级感知策略尊重芯片、节点、机架边界间的带宽断崖。

  3. 常见陷阱:常见误解为任意模型可分片至任意节点数。实则若层级映射错误(如跨慢速机架间链路分片大张量),通信时间将远超计算时间,使扩展无效。

交互亦反向流动:并行策略选择影响基础设施设计。采用 TP-8, PP-4, DP-32 的训练系统产生通信模式:最高带宽流量(张量并行 AllReduce)限于节点内;中等带宽流量(流水线阶段通信)流经 4 邻居节点组;最低带宽流量(缩减梯度分片的数据并行 AllReduce)跨全集群。

这种分层通信模式偏好层级化网络拓扑(邻近节点间带宽更高,“局部感知”拓扑),而非所有节点对等带宽的扁平拓扑(均匀胖树)。轨优化与层级胖树设计利用此局部性,将通信最频繁的节点置于同一交换机或机架,为最高带宽流量最小化交换机跳数。

基础设施采购的实际影响在于,网络拓扑必须与并行策略共同设计,而不能独立选择。一个采购了扁平胖树结构网络(针对任意到任意通信进行了优化)但仅使用层级感知并行进行训练(其中大多数流量为本地流量)的组织,在网络全局带宽方面出现了过度供给,同时可能在本地带宽方面出现供给不足。相反,一个采购了针对局部通信优化的轨道优化网络(优化用于局部通信)但后来需要运行具有 AllToAll 通信的专家混合(Mixture-of-Experts)模型(这需要全局带宽)的组织,会发现该网络结构不足。网络结构占总系统成本的 10–15%,必须与预期的工作负载组合相匹配,而部署后更改网络结构的成本极高且会造成严重干扰。

3D 训练循环

训练一个 1750 亿参数的模型需要3D 并行:数据并行、流水线并行和张量并行在数千个设备上的协同组合。^65 这种方法不仅仅是简单地叠加各个并行策略的优势;它通过几何方式将它们组合起来,以匹配硬件的物理拓扑。

考虑一个训练集群,其配置为张量并行(TP)为 8 个 GPU、流水线并行(PP)为 16、数据并行(DP)为 128。此配置使用了 16,384 个 GPU(8 GPU × 16 × 128),并按照带宽域的层次结构进行组织。

训练步骤从数据并行级别开始。128 个模型副本中的每一个都接收全局批次的一个不同切片。在每个副本内,模型被分割为 16 个流水线阶段(节点),微批次从节点 0 的嵌入层顺序流动到节点 15 的损失计算。在最细的粒度级别,在每个节点内,8 个 GPU 通过张量并行(TP)融合为一个单一的“超级加速器”。前向传播中的每个矩阵乘法在这些设备之间被分割,并且在每次操作之后,必须通过高带宽的NVLink交换部分结果。对于 1750 亿规模的隐藏维度,激活负载大约为 201.3 MB;由此产生的张量并行流量在前向传播中每个流水线阶段每个微批次大约为 4.2 GB,或者在包含反向传播时为 8.5 GB。在完整的 96 层副本中,前向张量并行流量大约为 67.6 GB。此流量是系统中强度最高的,但由于芯片在 600 GB/s–900 GB/s 的局部带宽上进行通信,其暴露的延迟相对于计算可以保持较小。

反向传播反转了此流动并暴露了并行维度之间的关键依赖关系。随着梯度通过流水线逆向流动,节点以点对点的方式交换激活梯度。此流量相对较轻——每个阶段边界大约为 201.3 MB——使其能够在不导致流水线停滞的情况下通过较慢的节点间 InfiniBand 链路。真正的瓶颈出现在步骤的末尾:数据并行同步。完整模型的梯度状态仍然大约为 350 GB,但在 3D 并行布局中,每个数据并行组只同步其张量并行和流水线并行副本对应的参数分片,而不是每个 GPU 移动完整的张量。梯度分桶将就绪的层梯度张量分组为更大的消息,随着反向计算的进行;当分桶时间表和横截面带宽匹配时,同步会在有用计算期间运行,而不是延长关键路径。

架构上的 imperative 是带宽匹配:每个算法的通信量必须与硬件互连的延迟成反比映射。频繁且阻塞的张量并行(TP)通信保持在 NVLink 域内(600 GB/s 或更高)。序列化的、点对点的流水线并行(PP)传输在集群骨干网上以 InfiniBand 速度进行。巨大但罕见的数据并行(DP)同步可以在整个训练步骤中摊销。尝试在机架之间运行张量并行(TP),或在没有梯度累积的情况下运行数据并行(DP),可能会违反这一层次结构,导致 16,000 个 GPU 的集群在等待数据穿越网络线路。这一带宽匹配原则完成了第 5.1.5 节中引入的可行性检查。

混合并行工作示例

将此带宽匹配原则应用于物理基础设施,将集群设计转化为一个放置问题:每个并行维度必须位于能够容纳其流量的网络层级上。张量并行是最对延迟敏感的,因为它在 Transformer 块内部频繁启动 AllReduce 操作,因此它属于节点内的 NVLink 域(600 GB/s–900 GB/s)。流水线并行仅在阶段边界移动激活张量,因此它可以通过 InfiniBand(50 GB/s–100 GB/s)跨越相邻的节点。数据并行产生最大的负载——梯度同步,但它每步只发生一次,并且可以与反向计算重叠。每个设备的内存容量(80 GB–80 GB)为每个放置设定了硬性限制。

对于类似 DGX A100 的部署,这种推理导致了一个具体的布局。将张量并行固定为t = 8,以确保带宽密集的矩阵乘法集体操作停留在单个 8-GPU 节点内。将流水线并行映射到同一高带宽机架或岛屿中的节点上,通常根据内存占用情况选择p = 8 或p = 16。使用数据并行作为跨 Pod 的剩余扩展维度。结果不是一个通用规则,而是将每种流量模式与能够承载它的最便宜的结构层级相匹配的必然结果。

在放置固定后,下一个问题是静态模型状态是否能够适配每个加速器的 HBM 预算。

训练一个 1750 亿参数模型的内存预算主要由模型状态主导,并且需要进行激进的分片,以适应 H100 级加速器的 80 GB HBM 容量。仅 FP16 权重大约消耗 350 GB(175 ×10⁹ × 2 字节)。如果仅依赖 8-way 张量并行进行分片,则每个 GPU 将保存权重的 43.75 GB 切片。然而,优化器状态提出了更大的挑战。在此 3D 并行预算中使用的简化优化器状态惯例下,Adam 的 FP32 动量和方差每参数消耗 8 字节,全球总计约为 1.4 TB;更完整的 12 字节核算(同时包含 FP32 主权重)将达到约 2.1 TB。即使采用 8-way 张量并行,简化后的权重和优化器状态组合每 GPU 也将超过 218.8 GB,这将导致内存溢出(OOM)错误。因此,我们必须采用流水线并行(p)来进一步划分模型层。在张量并行 8 和流水线并行 16 的混合配置下,简化的静态内存占用降至每 GPU13.7 GB(通常在考虑框架缓冲后预算为大约 15 GB),这样剩余的 HBM 就可用于前向传播过程中生成的动态激活内存(A),其大小随着微批大小和序列长度线性增长。在完整的 12 字节 Adam 惯例下,相同的 8-way 张量并行、16-stage 流水线并行分割将导致每 GPU19.1 GB 的内存占用。

一旦配置适合内存,相同的布局就必须满足激发它的通信层次结构。每个并行维度都会对网络施加不同的流量配置文件。张量并行是最“健谈”的,在每个 transformer 块(一个用于注意力投影,一个用于 MLP)在前向和反向传递中都需要进行两个 AllReduce 操作。这些消息相对较小,但每一步会发生数千次,因此它们严格受延迟限制,需要使用 NVLink。与此形成对比的是,管道并行仅在流水线阶段的边界上涉及激活张量(大小为 B[μ] × S × d[model]) 的点对点传输。虽然这些消息大小适中,但发生频率较低,因此可以通过标准 InfiniBand 链路进行管理。数据并行会产生最大的流量突发,需要对整个 350 GB 梯度缓冲区进行全局 AllReduce。然而,此通信仅在每个全局批次更新时发生一次。通过使用梯度分桶(bucketing)将此传输与计算密集型的反向传递重叠,只要集群保持足够的横截面带宽,DP 通信的有效成本通常可以被掩盖。

最后的成本不是带宽,而是调度:只有当有足够的微批次在飞行中以保持所有阶段忙碌时,流水线模型才是高效的。这一剩余成本是第 5 节中发展的流水线气泡(pipeline bubble),通过将其气泡分数 \(\\frac{p-1}{m + p - 1}\) 应用于此混合布局,GPT-175B 配置(p = 16 个阶段,m = 32 个微批次)浪费了 \(\\frac{15}{47} \\approx 31.9\\%\) 的理论计算能力,几乎是三分之一。通过重新激活内存的 1F1B 调度以及它使得更大的微批次数量变得可行,这一比例会渐近地降低。

Blackwell 级别的扩展示例

Blackwell 架构提供了一个具体例子,说明加速器封装如何改变 3D 并行的权衡。首先,NVLink 5 每个 GPU 提供 1.8 TB/s 的双向带宽,是霍珀一代内部节点容量的两倍。这使得更大的张量并行(t)组成为可能(例如,在多个节点上 t= 16 或 t= 32),并且延迟开销较低。其次,Blackwell 添加了 FP4 Tensor Core 支持,提高了受支持工作负载的低精度吞吐量和内存效率;FP4 是否用于训练取决于数值配方和软件栈。

对于 1 万亿参数模型,Blackwell 级别的配置可以使用 t= 16(通过 NVLink Switch 跨越两个 8-GPU 节点)和 p= 32,相较于 16,384-GPU Hopper 级别的舰队参考配置,减少了总的管道阶段数量及其相关的气泡开销。Blackwell GPU 内部的 10 TB/s 晶圆到晶圆互连进一步消除了芯片内和封装内通信之间的区别,使得两个受光掩模限制的芯片能够作为一个单一的高带宽张量并行单元运行。总的系统经验是:提升局部带宽会将压力转向外部:一旦节点内通信的限制减弱,面向铁路优化的拓扑结构(以及全对全优化)对于更大的机器学习舰队配置变得更加重要。

相同的放置约束可以通过设计空间搜索得以实现。

情景:工程团队需要在 8,192 块 H100 GPU 的集群上调度原型 A(一个 175B 参数模型)。手动搜索 3D 并行空间(TP × PP × DP)容易出错:最大化 DP 的拆分可能会超过 80 GB HBM 容量,而最大化 TP 的拆分可能会使 NVLink 互连饱和。

方法:与其试错,不如在我们的物理引擎中实现的三层优化器(ParallelismOptimizer)上调用受限的设计空间搜索,针对 TP、PP 和 DP 的因式分解进行约束,以最大化模型 FLOPs 利用率(MFU)为目标。优化器在一秒内对所有有效的代数因式分解进行受限网格搜索,并在给定目标下选择最佳策略:

  • 张量并行(TP):8 个 GPU

  • 管道并行(PP):8

  • 数据并行(DP):128

系统洞察:优化器正式化了工程师通过经验发现的事实:为了避免遍历较慢的节点间结构,TP 应匹配节点内 GPU 数量(8 个 GPU);在此目标下,一旦张量并行和分片状态使内存预算得以容纳,最佳的管道深度仅为 PP=8。更深的管道虽然会进一步降低每阶段的内存需求,但也会增加管道气泡并降低 MFU。这种自动合成实现了预计的 25.9 MFU,表明经验工程启发式方法往往反映了结构性的数学定律。

调整任意一个维度会将压力转移到其他两个维度上,使得因式分解 TP × PP × DP 成为一个紧密耦合的系统,而不是三个独立的旋钮。

验证您对并行策略如何组合的理解:

MFU 值需要历史利用率基准来展示已发表系统如何接近 50% 的利用率。[ 追溯了从 2020 年到 2024 年已发表训练系统中的模型 FLOPs 利用率(MFU)的演变。从 GPT-3 的 21% MFU 到 PaLM 的 46% MFU 的进展反映了并行策略、通信重叠技术以及本章讨论的调度优化的进步,而不仅仅是原始硬件速度的提升。40–46% 附近的平台表明,即使在高度优化的混合并行系统中,由通信开销、管道气泡和内存管理施加的理论上限仍然十分严峻。值得注意的是,Meta 的 Llama 3 在 16,384 块 H100 GPU 上的训练 MFU(41%)略低于同一模型在 8,192 GPU 上的 MFU(43%),证实了[ 中描述的缩放税不仅是理论上的,而且在大型已发表的运行中是可测量的。

![

图 5.16:MFU 演变:已发表训练系统中的模型 FLOPs 利用率(MFU),按时间顺序排列。MFU 从 2020 年 GPT-3 的 21% 增至 2022 年 PaLM 的 46%,受益于混合并行和通信重叠的进步,随后在大规模运行中趋于 40–45% 附近的平台。Llama 3 的数据点展示了缩放税:随着集群规模增大,通信开销增加导致 MFU 从 8,192 GPU 的 43% 降至 16,384 GPU 的 41%。数据来源:Chowdhery 等人([),Dubey 等人([)。

多模型训练:RLHF 和对齐

已经检查过的并行策略假设在单一目标上训练单一模型。来自人类反馈的强化学习(RLHF)及其变体通过要求多个模型在单一训练循环中协调工作,每个模型具有不同的内存占用、计算特性和梯度需求,从而打破了这一假设。从高层次看,RLHF 生成模型输出,使用源自偏好的奖励信号对其进行评分,针对参考模型正则化更新,然后更新策略;近端策略优化(PPO)通过独立的策略、参考、奖励和价值模型实现了这一在线循环,而直接偏好优化(DPO)风格的方法则去除了大量的 rollout 和价值模型机制。这形成了一个异构的机群管理问题,无法通过任何单一的并行策略解决,并且代表了与标准预训练相比在质量上不同的分布式系统挑战。

多模型协调问题

标准预训练涉及一个模型、一个损失函数和一个梯度流。相比之下,RLHF 对齐在每个训练步骤中协调一套相互交互的模型。在近端策略优化(PPO)(Schulman et al. 2017) 中,正如用于 InstructGPT 风格的 RLHF 系统(Ouyang et al. 2022) 那样使用的,四个不同的模型必须协同运行。表 5.7 按角色、执行模式和内存负担将这些模型分开。

| 模型 | 在 PPO 风格的 RLHF 中的角色 | 执行模式 | 内存负担 |

| --- | --- | --- | --- |

| 策略模型 | 正在对齐和优化的模型。 | 带有梯度、优化器矩和激活的训练模式。 | 对于混合精度下的 Adam 优化器的 70B 参数模型,内存大约为 70B × (2 + 2 + 12) = 1,120 GB,与标准预训练的预算相同。 |

| 参考模型 | 提供 KL 散度惩罚的预训练策略的冻结副本。 | 仅推理模式。 | 不带梯度或优化器状态的 FP16/BF16 参数;70B 模型大约需要 70B × 2 = 140 GB,大约是训练配置的 1/8。 |

| 奖励模型 | 为生成序列产生标量奖励的偏好训练评分器。 | 仅推理模式。 | 通常比策略模型小;13B 的奖励模型在 FP16 下大约需要 26 GB,但它必须处理每个生成的序列。 |

| 价值模型 | 估计预期未来奖励的 PPO 评论家。 | 训练模式,可以是独立的或共享策略骨干。 | 完整大小的独立评论家可以再增加 1,120 GB;较小的模型或共享骨干设计可将其减少到 200–400 GB。 |

表 5.7: RLHF Model Coordination Roles: PPO 风格的 RLHF 协调具有不同执行模式和内存占用的模型,因此放置决策必须将可训练的策略和价值状态与仅用于推理的参考和奖励状态分离。

四模型 PPO 系统的总内存需求远远超过标准预训练。一个将 70B 策略模型、70B 参考模型、13B 奖励模型和 13B 价值模型 naive 地 colocated 在一起的系统,在考虑序列生成过程中产生的 KV 缓存和中间激活之前,大约需要 1,494 GB 的加速器内存。在每块具有 80 GB HBM 的 H100 GPU 上,仅用于参数存储,此系统至少需要 19 块 GPU。一旦包括随着输出序列长度线性增长的生成阶段 KV 缓存和训练阶段的激活,实际最小需求将升至 64–128 块 GPU,以支持单个 RLHF 训练实例。

基础设施不对称:训练 vs. 推理模型

RLHF 的基础设施挑战并非在于 总内存占用,而在于模型的 计算特性不对称。策略模型和价值模型需要完整的反向传播,伴随梯度计算、激活检查点和优化器更新——这些计算密集型操作受益于张量并行和高算术强度。相比之下,参考模型和奖励模型仅执行前向传播:它们是嵌入在训练循环中的推理工作负载,其内存访问模式主要由 KV 缓存管理主导,而非梯度累积。

这种不对称产生了放置困境。将训练模型和推理模型 colocated 在同一块 GPU 上会在生成阶段(此时训练模型处于空闲)浪费计算资源,并在梯度阶段(此时推理模型的参数存储本可用于激活)浪费内存。将它们分离到专用的 GPU 池中可以消除浪费,但会为奖励查询引入网络延迟——每个生成的 token 批次必须遍历互连以到达奖励模型并返回一个标量信号,然后才能计算策略梯度。

生成阶段本身引入了标准预训练中不存在的顺序瓶颈。RLHF 需要策略模型在计算奖励和策略梯度之前自回归地生成完整序列(通常为 256–2,048 个标记)。自回归生成受内存带宽限制,而非计算限制:每个标记需要通过模型的一次完整前向传播以生成一个输出标记。对于类似 Llama-2-70B 的分组查询注意力策略,KV 缓存每生成一个标记每序列大约增加 2 × N[L] × H[KV] × d[head] × 2 字节。在 80 层、8 个 KV 头、128 维度的头、1,024 个生成标记和 256 个提示的批量下,缓存消耗 85.9 GB:2 × 80 × 8 × 128 × 1024 × 256 × 2 字节。使用完整 8,192 宽隐藏状态的 K 和 V 的密集多头注意力模型将产生 687.2 GB。即使采用 GQA,缓存也大约使用相当于一块 H100 GPU 的 HBM 用于在奖励计算后被丢弃的中间注意力状态。

RLHF 系统通过时间复用来解决这种不对称。在生成阶段,集群的行为类似于推理系统:策略模型生成序列,而参考模型计算对数概率,此时主导资源是为标记逐个注意力状态预留的 HBM。在训练阶段,同一机群切换回训练模式:策略模型和价值模型通过本章前文开发的 3D 并行配置计算梯度。服务系统通过请求批处理和 KV 缓存放置来处理该阶段;这里的重点是模式切换。RLHF 不是一个稳定的训练循环,而是介于塑造推理的工作和塑造训练的工作之间的交替,而标准训练框架很少能自行管理这种转变。

DPO:简化机群

直接偏好优化(DPO)(Rafailov et al. 2023) 通过将对齐目标改写为偏好对上的分类损失,完全消除了奖励模型和价值模型。与生成序列、计算奖励和估计优势不同,DPO 直接优化策略,使其为偏好响应分配更高的对数概率,而为不偏好响应分配较低的对数概率,仅使用参考模型来计算 KL 散度正则化项。

基础设施影响是重大的。DPO 将多模型系统从四个模型减少到两个:策略模型(训练 2 小时 16 分)和参考模型(推理模式)。对于以下部分量化的 70B 策略模型、伴随 13B 奖励模型和价值模型,静态参数内存从大约 1,494 GB 降至 1,260 GB,降幅达 16%。如果 PPO 的价值模型与策略模型大小相当,则降幅更大:大约从 2,406 GB 降至 1,260 GB,大约减少 48%。DPO 还完全消除了自回归生成阶段。训练在固定数据集 (prompt, preferred response, dispreferred response) 三元组上进行,恢复了标准的预训练数据管道:固定长度序列、确定性批处理,以及无逐个标记的顺序生成。训练循环变为使用修改后的损失函数的标准监督学习步骤,可以使用与预训练相同的 3D 并行性、梯度累积和通信重叠技术。

权衡在于能力。DPO 在静态偏好数据集上运行,这意味着策略在训练过程中无法探索新响应并接收反馈。PPO 的在线生成使策略能够在其自身输出上迭代改进,可能发现静态数据集中不包含的更好策略。对于偏好数据全面覆盖目标分布的部署场景,DPO 的基础设施简化占优势。对于需要自适应探索的场景(例如,训练模型解决新颖的推理任务),PPO 的在线反馈循环可能证明其 基础设施开销是合理的。

定量分析:PPO vs. DPO 资源需求

为使基础设施权衡具体化,考虑在 256 块 H100 GPU(每块 80 GB HBM)集群上对 70B 策略模型进行对齐。

问题:对齐一个 70B 策略模型需要将其与一个 70B 参考模型共置,并且对于 PPO,还需要一个 13B 奖励模型和一个 13B 价值模型。在 3D 并行性下(TP=8 GPU,PP=4DP=8),每种方法每个 GPU 消耗多少内存?通过消除奖励模型、价值模型和生成阶段的 KV 缓存,DPO 能释放出多少余量?

PPO 内存预算(每个 GPU,TP=8 GPU,PP=4DP=8 用于策略)

TP=8 GPU 和 PP=4 的 3D 并行性下,策略模型将其 1,120 GB 训练状态分布在 32 个 GPU 上,每个 GPU 得到 35 GB。参考模型仅需 140GB 用于推理,可以在另一个 8 GPU 的池中分片,每个 GPU 为 17.5 GB,或者与策略 GPU 共置,每个 GPU 额外增加 4.4 GB140 GB / 32 GPU)。13B奖励模型在其池中增加26 GB 共享内存。13B价值模型在训练模式下在其池中增加大约208 GB13B × 16 bytes/param`)。

共置的策略 + 参考模型在 32 GPU 上的总静态内存:每个 GPU 35 GB + 4.4 GB ≈ 39.4 GB,剩余约 40.6 GB 用于激活和 KV 缓存。对于包含 256 个序列和 1,024 个标记的回滚批次(在一个 TP × PP 策略副本内),生成阶段的 KV 缓存要求每个策略 GPU 额外保留约 85.9 GB / 32 GPU ≈ 2.7 GB 用于其 KV 缓存分片。剩余约 37.9 GB 限制了训练阶段的微批大小。

DPO 内存预算(相同的并行性配置)

策略模型使用相同的每个 GPU 35 GB。参考模型每个 GPU 增加 4.4 GB(共置)。无奖励模型,无价值模型,无生成阶段的 KV 缓存。

在策略 GPU 上,总静态内存为 35 GB + 4.4 GB = 39.4 GB 每个 GPU,与 PPO 的共置策略/参考模型占用相同。然而,集群范围内,DPO 避免了 PPO 奖励模型和价值模型的内存,并且还移除了生成阶段的 KV 缓存负担。每个策略 GPU 上剩余的约 40.6 GB 可用于训练激活,允许在这些共置策略/参考模型内存预算下适度增加微批大小(约 1.07×),或减少 PPO 所需的激活检查点需求。

吞吐量比较

PPO 的两阶段设计(生成然后训练)引入了基本的吞吐量惩罚。如果生成占用了步骤时间的 60%(典型用于长序列的自回归解码),那么在 RLHF 步骤中训练硬件仅能达到 40% 的利用率。DPO 作为标准训练循环运行,实现与预训练相同的 40–55% MFU。表 5.8 侧对比了这两种方案。

| 指标 | PPO(70B 策略) | DPO(70B 策略) |

|---|---|---|

| 所需模型 | 4(策略,参考,奖励,价值) | 2(策略,参考) |

| 参数总内存 | ∼ 1,494 GB | ∼ 1,260 GB |

| 最小 GPU(内存) | 64–128 GPU | 32–64 GPU |

| 生成阶段 | 是(顺序,带宽受限) | 无 |

| 有效训练 MFU | 15–25% | 40–55% |

| 数据管道 | 在线生成 | 静态偏好对 |

系统洞察:在此 13B 奖励/价值模型场景中,DPO 通过消除奖励模型、价值模型和自回归生成阶段,将静态参数内存减少了 15.7%,并使有效计算利用率翻倍。在 PPO 和 DPO 之间的选择不仅是建模偏好,也是基础设施限制:GPU 预算有限的组织可能更倾向于 DPO,因为它将工作负载恢复为固定数据训练循环。如果 PPO 的价值模型是策略大小而不是 13B,则内存减少将接近之前讨论的更大比较 2,406 GB1,260 GB

序列长度方差和批处理挑战

PPO 和 DPO 均面临标准预训练中不存在的数据工程挑战:序列长度的极端方差。预训练通常使用固定长度序列(2,0484,096 个标记,填充或打包以填充每个位置),从而实现统一的批处理形状和可预测的内存消耗。RLHF 训练数据由可变长度的提示(10–500 个标记)与可变长度的完成(50–2,048 个标记)连接而成,在单个批次中产生 10–50× 方差的序列长度。

长度方差会导致两个相互交织的问题。固定大小批处理将所有序列填充到批次中的最大长度,在填充标记上浪费计算。一个包含一个 2,048-标记序列和十五个 128-标记序列的批次在填充标记上浪费了 88% 的计算。动态批处理按相似长度分组序列以最小化填充,但会在数据并行工作器之间引入负载不平衡:一个工作器可能接收一批长序列,消耗 60 GB 的激活内存,而另一个工作器处理短序列仅使用 8 GB,导致短序列工作器在同步屏障处停滞,而长序列工作器仍在完成。

RLHF 系统通过序列打包(将多个短序列连接成一个固定长度输入,并使用注意力掩码防止交叉污染)和自适应微批处理(根据聚合标记数而非序列数动态调整每个微批的序列数)来缓解这一问题。这些技术可以恢复大部分因填充方差而损失的计算,但会增加标准预训练框架未提供的数据管道复杂性。变长数据与分布式同步之间的交互仍然是系统工程研究的活跃领域,因为每个工作器必须在每个步骤处理相同数量的标记以保持梯度一致性。

因此,RLHF 并非并行策略的绕道;它是一个案例,展示了为何单靠任何一种策略都不足够。策略模型可能需要张量并行和流水线并行,参考模型和奖励模型的表现类似于推理服务,价值模型可能需要训练状态,而展开数据管道在每一步都会改变形状。考虑到这一压力测试,最终的对比可以回归到通用的工程问题上:首先识别约束瓶颈,然后选择在移动最少数据的同时能容纳所需状态的并行模式。

并行策略对比

仅当并行策略能移动约束瓶颈且不产生更大的瓶颈时,它才是有用的。第 5.1.5 节 中决策树预览的可行性过滤器给这些约束命名了;现在张量并行、流水线并行和混合并行已经开发完成,同一个过滤器可以表述为工程信号,而非通用的排序。对于一个新的 500 亿参数模型,决定性的问题是:哪种内存状态适用、通信发生的频率如何、流水线气泡出现在何处,以及工程团队能容忍哪种抽象成本。表 5.9 将这些问题转化为工程信号。

| 策略 | 约束瓶颈 | 主要数据移动 | 放置决策 | 验证信号 |

|------|----------|----------------|----------|----------|

| DDP | 模型状态可放入单个加速器 | 每步 AllReduce 梯度 | 在数据并行工作器间复制完整模型 | 步耗时由计算主导,而非同步 |

| FSDP/ZeRO | 模型状态超出单个加速器 | 分片并聚合参数或优化器状态 | 放置分片以最小化聚合和 reduce-scatter 成本 | 峰值内存保持在设备预算之下 |

| 张量并行 | 单层或张量操作超出本地容量 | 层块内 AllReduce 或 AllGather | 将紧密耦合的分区保留在快速的节点内链路上 | 层延迟保持在激活传递成本之下 |

| 流水线并行 | 只有按阶段拆分层堆栈才能放入 | 阶段间传递激活和梯度 | 平衡阶段使气泡小于有效工作 | 随微批次数量增加,气泡比例下降 |

| 混合并行 | 多个约束同时生效 | 在不同轴上执行多种集合操作 | 将数据、张量、流水线和分片组与拓扑对齐 | MFU 提升且不超出内存或网络预算 |

表 5.9:并行训练策略:并行策略选择是一项工程可行性测试:识别约束瓶颈、它产生的通信原语、它强制的放置决策,以及证明该策略有效的测量指标。

图 5.17 将同一个过滤器转化为决策树。它从内存适配性开始,因为容量故障是不可商量的,然后询问数据规模本身是否证明复制的合理性。该树经过有意简化;硬件异构性、通信带宽和工作负载不平衡作为二阶检查项,在识别出约束瓶颈后介入。

![](https://github.com/OpenDocCN/dsai-notes-pt3-zh/tree/master/docs/hav-cs249r-mlsys-vol2/img/file97.svg)

图 5.17:并行策略决策树:基于两个关键问题的系统选择指南:模型状态是否适合单设备内存?所需的数据吞吐量是否证明多副本的合理性?超出设备内存的模型需要模型并行或分片;模型虽能放入但吞吐目标超出单设备的工作负载受益于数据并行;在两个维度均有显著约束的情况则需要混合方法。尽管经过简化,该框架在从业者必须考虑硬件异构性和工作负载不平衡等次要因素之前,捕捉了主要决策点。

该表格和流程图有意保持粗略,但其目的精准:在选择实现栈之前先识别约束瓶颈。

从原理到系统

选择 DDP、FSDP/ZeRO 风格分片、张量并行或流水线并行,是关于运行时拥有哪些约束、工程师必须直接管理哪些约束的决策。本章探讨的并行策略(梯度平均、AllReduce 同步、张量分片、流水线调度)通过这种分层抽象层级转化为部署的训练系统。

在数据并行层,最简单的分布式训练抽象包装模型,使框架自动将其跨可用加速器复制、拆分每个批次,并在反向传播后平均梯度。历史上,参数服务器系统用存储共享参数的服务器节点组织这项工作,而工作节点推送和拉取更新(Li et al. 2014)。正如 5.5.3 节 所确立,该设计将梯度带宽集中在服务器层,因此稠密同步训练倾向于下文的去中心化集合操作。

大规模数据并行通过用去中心化 AllReduce 取代中央服务器消除了这一瓶颈。每个工作器对称参与归约:每个设备以自身链路速率发送和接收梯度块,将带宽负载分散到所有节点,而非集中在一点。框架初始化一个将工作器映射到物理拓扑的进程组,根据检测到的互联选择合适的集合算法(环形、树状或分层),并自动将梯度同步钩子插入反向传播。梯度分桶通过在传输前将小张量分组为大消息进一步提升效率,计算通信重叠允许早期层的 AllReduce 在后期层仍在计算梯度时进行。这些优化能在中等规模达到高并行效率——不是因为 API 简单,而是因为底层运行时做出了拓扑感知决策,否则手动实现需要自行复现。

模型并行和流水线并行需要根本不同的抽象,因为框架必须管理跨设备张量放置和分区间的顺序数据流。核心原则是显式设备分配:工程师指定哪些层驻留在哪些设备上,框架处理前向传播期间设备间的激活传输和反向传播期间的梯度反向流动。显式设备分配使模型并行的顺序依赖可见化,因为每个下游设备必须等待其上游邻居完成。工程师必须在架构层面推理流水线气泡和负载均衡,而非将它们隐藏在不透明的包装器之后。

对于大规模模型并行,关键的架构洞见在于张量切分和流水线调度需要框架提供不同层级的支持。

张量并行将标准线性层替换为列并行和行并行变体,这些变体会在 Transformer 计算图的正确位置自动插入 AllReduce 操作,详见第 5.5.2.3 节。流水线并行增加了微批次调度逻辑,在各阶段间交错前向和后向传播,以最小化气泡开销。内存高效分片通过在每次前向传播前用自动 AllGather 操作、每次后向传播后用 ReduceScatter 操作来包装模型层,集成了 ZeRO-3 风格的参数分区。关键的设计决策在于组合哪些抽象层级:当模型能放入内存时,纯数据并行即可;内存高效分片将数据并行扩展到受内存约束的场景;而当单个层或完整模型深度超出单设备容量时,完整的张量或流水线并行就成为必需。

跨越这些抽象层,分布式训练归根结底是一小组集合通信原语。框架抽象关注的是何时以及为何调用这些原语,而不是探究这些原语如何在物理上跨集群路由数据的网络机制,后者在第 6 章中有详尽论述。

这些原语组合成定义各并行策略的通信模式。数据并行每训练步使用一次 AllReduce。带重分片的全分片 FSDP 每步约使用 3*N*[L] 个集合操作(针对每个 N*[L] 层使用两次 AllGather 和一次 ReduceScatter)。张量并行在关键路径上每个 Transformer 块使用 2 次 AllReduce 操作。原语的选择、其消息大小、及其相对于计算的频率,决定了系统处于计算受限还是通信受限。没有任何框架抽象能改变底层物理规律:分布式训练的效率最终取决于物理互联带宽、内存容量和同步延迟。

谬误与陷阱

许多工程团队将集群容量扩大 4 倍,却发现训练迭代耗时反而更长,这反映出对分布式训练物理规律的根本性误解。以下误解捕捉了那些浪费算力资源、延误研究进度的错误。

谬误:只要工程投入足够,就能实现线性加速。

阿姆达尔定律设定了硬性上限:任何串行组件都会限制最大加速比,无论并行度如何。在分布式训练中,梯度同步本质上是串行的,因为必须收集所有梯度才能进行更新。正如第 5.4 节所示,扩展效率公式 η[scaling] = 1/(1 + N*(Tcomm + Tsync - T[overlap])/T[compute]) 揭示了随着 N 增加,通信和同步开销如何占据主导。即使有完美的重叠和理想算法,通信开销也会随集群规模增长。对于数据并行,随着 GPU 数量增加,树算法下 AllReduce 时间呈对数增长,环算法下的延迟项呈线性增长。1000 GPU 的集群永远不会比单 GPU 快 1000 倍;达到 500 倍加速就已属杰出,对于通信密集型负载,100–200 倍更为典型。假设线性扩展来制定项目预算的组织,必然会错过截止日期并在算力上超支。

陷阱:优化 MFU 却不考虑扩展效率。

某团队在单节点上实现 50% MFU,扩展到 1024 GPU,预期获得 512 GPU 等效的有效工作量。对于通信密集型负载,1024 GPU 下的扩展效率可能接近 50%,因此实际有效吞吐约为 0.50 × 0.50 = 0.25 的峰值,即 256 GPU 等效:仅为预期值的一半。MFU 和扩展效率是相互独立的乘法因子。只优化其中一个而不测量另一个,会导致容量估算偏离 2 倍甚至更多。舰队规模的容量规划要求同时报告这两个指标,并追踪其乘积作为“有效吞吐”。

谬误:在小集群上调优的超参数可直接迁移至大规模训练。

工程师在 8 GPU 工作站上调优超参数,部署到 256 GPU 集群时却期望行为完全一致。大规模下,收敛模式会发生根本性变化。最关键的超参数是学习率:如第 5.3 节所述,数据并行中批量大小随 GPU 数量成比例增加,需要相应调整学习率。“线性缩放规则”⁶⁶ (Goyal 等人, 2017) 建议 η[large] = η[base] × (B[large]/B[base]),但该关系仅在特定范围内成立。随着模型扩大,最终会遇到临界批量大小⁶⁷,此时每步增加数据带来的收敛收益呈边际递减。

超出临界批量大小后,这种关系会以依赖模型、数据和优化器的方式失效。某团队若拿小集群的学习率,机械地按 GPU 数量相乘,直接跳转到大得多的全局批量,即使硬件吞吐提升,也可能看到最终质量下降或样本效率变慢。预热调度、权重衰减调整、逐层优化器(如 LARS/LAMB)以及精细的动量调优,在某些范围内可恢复精度,但需要在目标规模上进行系统实验。跳过这些扩展性研究的组织,会在次优运行上浪费数千 GPU 小时。

陷阱:向数据并行作业添加 GPU 却不建模通信。

工程师假设更多 GPU 总能加速训练。大规模下,统计效率限制会淹没硬件收益。正如第 5.3 节所述,数据并行使有效批量大小随 GPU 数量成比例增加(B[total] = N × B[local]),但梯度质量在模型特定阈值之外呈亚线性增长。10 万样本的批量可能仅提供 1 万样本批量 2 倍的梯度信息,而非 10 倍,因为样本在损失景观中变得冗余。临界批量大小定义了边际收益崩溃的临界点:前文示例将 ResNet-50 置于 8K–16K 左右,BERT-Large 置于 32K–65K 左右,具体阈值取决于优化器、调度和目标质量。超过此阈值后,GPU 数量翻倍成本翻倍,却几乎不加速收敛。在一个典型成本模型中,1024 GPU 运行 18 小时收敛,算力成本 $45,000,可能不如 512 GPU 运行 19 小时收敛,成本 $22,000,这演示了超过临界批量大小如何在无实质性省时的情况下浪费资源。

谬误:仅凭内存容量决定并行策略。

工程师看到 70B 模型超出 80 GB 内存,立即选择张量并行或流水线并行来拆分权重。在实际部署运行中,有效的策略取决于内存压力、计算模式和通信拓扑之间的相互作用。正如第 5.9 节所述,张量并行将每一层拆分到各个设备上,并通过每层的 AllReduce 同步,实现了均匀的内存分布,但将通信置于关键路径上。流水线并行将完整的层分配给各个阶段,阶段间通过点对点传输,减少了每步通信,但引入了流水线气泡开销,浪费了 10%-30% 的周期。对于在 64 块 A100 GPU 上运行的 175B 模型,张量并行度为 8 时可支持训练,而 8 阶段的流水线并行因减少了全对全通信,尽管内存占用相似,却实现了 23% 更高的吞吐量。决策需要对通信模式和气泡开销进行剖析,而不仅仅是检查权重是否能装进内存。

陷阱在未衡量效率权衡的情况下应用 FSDP 或 ZeRO。

工程师在读到 FSDP“减少内存并支持更大模型”后,便普遍采用它。在实际部署运行中,分片引入了 10%-25% 的通信开销,仅当内存压力足以证明其合理性时才划算。FSDP 通过在 GPU 间分片优化器状态、梯度和可选的参数来减少内存占用,但这要求在每次前向传播前进行 AllGather 操作,在反向传播后进行 ReduceScatter。对于在 A100-80 GB 上运行、批次大小为 4 的 7B 模型,标准 DDP 可达 145 samples/second,而 FSDP 仅为 118 samples/second(慢 19%),因为模型无需分片即可舒适放入内存,额外的通信开销未带来任何收益。当模型加优化器状态超出单 GPU 内存、启用更大单 GPU 批次大小足以抵消开销,或通过 ZeRO-Offload 卸载至 CPU 内存扩展容量时,FSDP 才能体现价值。一个无法在 80 GB 上放下的 65B 模型,通过 FSDP ZeRO-3 变得可训练,代价是接受 15% 的吞吐量损失以换取训练的可行性。不加区分地应用 FSDP 而不测量内存压力,会浪费性能。

谬误并行开销与模型大小无关,大致恒定。

工程师在方便的小模型上对并行策略进行基准测试,再将结论应用于大规模训练。在大规模下,计算与通信时间的比率随模型大小发生剧烈变化,甚至反转战略决策。AllReduce 通信时间主要取决于梯度张量大小和网络带宽,大致随参数量线性增长;而前向和反向传播的计算时间因更大的矩阵运算而超线性增长。对于 1B 参数模型,前向/反向传播耗时 50 ms,AllReduce 耗时 25 ms,通信开销占步骤时间的 33%。对于 70B 参数模型,前向/反向传播耗时 2400 ms,AllReduce 耗时 180 ms,尽管梯度大小是前者的 70 倍,通信开销却降至 7%。在小模型上做出的决策(“流水线并行 15% 的气泡开销使其总比数据并行慢”)可能在大规模下反转,此时数据并行的通信开销达到 25%-40%。可靠的策略选择要求在目标规模上进行剖析,或使用考虑计算如何按 𝒪(n²) 到 𝒪(n³) 缩放而通信按 𝒪(n) 缩放的解析模型。

陷阱将扩展效率视为仅属于硬件的属性。

供应商基准测试发布“1024 GPU 下 85% 扩展效率”,采购团队按此规划产能,仿佛这是集群常数。扩展效率是四个因素的联合属性:工作负载的通信计算比、并行策略(数据、流水线、张量、混合)、网络拓扑(Fat-tree、Torus、Dragonfly)和软件栈(NCCL 版本、内核融合、重叠质量)。大语言模型训练基准的参考数字可能高估了具有不规则通信的图神经网络的扩展性,也可能低估了以嵌入为主、计算多为局部的推荐模型的扩展性。产能规划应使用目标工作负载的实测扩展数据,而非供应商的英雄基准。

谬误梯度累积是免费的。

工程师使用梯度累积模拟更大的批次大小,将同步频率从每步一次降低为每 K 步一次。该技术看似零成本,因为它消除了 (K − 1)/K 的同步事件。在真实训练循环中,累积引入了更新延迟、内存驻留和数值精度风险。标准实现将梯度累积到一个常驻梯度缓冲区,而非为每个微步骤分配单独的完整梯度张量;对于 7B 模型,该 FP16 梯度缓冲区仍有 14 GB,且必须在整个累积窗口期间保持驻留。若实现保留计算图、重叠多个未完成的微批次或增大局部微批次大小,激活内存可能进一步增长并耗尽 HBM。有效的优化器步骤延迟随累积步数成比例增加,累积 8 步意味着优化器更新频率降低 8 倍,尽管吞吐量更高,却可能减慢收敛。最关键的是,累积的 FP16 梯度在求和数百个梯度张量时面临溢出风险,尤其是在损失值较大的早期训练阶段。某团队使用 16 步 FP16 梯度累积训练 Transformer 模型,于步骤 1200 出现损失尖峰和发散;改为 4 步累积并更频繁地同步后,尽管通信成本更高,却解决了不稳定性。梯度累积以更新延迟、内存驻留和数值稳定性为代价,换取通信频率的降低。

陷阱无论系统特性如何,均使用固定检查点间隔。

工程师凭直觉而非分析,每隔“1 小时”或“每 1000 步”对分布式训练进行一次检查点。第 7 章 推导了 Young-Daly 检查点定律(原理);此处的训练教训是检查点节奏非常数。它取决于检查点写入成本与故障率之间的数学关系。

公式将最优检查点间隔设为 \(\tau\_{\text{opt}} = \sqrt{2 \cdot T\_{\text{write}} \cdot \text{MTBF}\_{\text{system}}}\),其中 T[write] 为检查点写入时间,MTBF[system] 为系统平均故障间隔时间。对于单 GPU MTBF 为规范 50,000 小时的 1024 GPU 集群,集群级 MTBF[system] 为 48.8 小时 (MTBF[GPU]/N)。若检查点耗时 5 分钟,最优间隔约为 171.2 分钟(~2.9 小时),不可避免的检查点加重做税为 5.8%。每 15 分钟检查点一次“以求稳妥”会将总税提高到 33.6%,而每 8 小时检查点一次则在故障时冒着丢失大量工作的风险。对于更大模型,因模型大小和存储带宽导致检查点时间增加到 15 分钟,最优间隔会再次偏移。猜测的代价随集群规模放大:当 1024 GPU 运行使用 15 分钟间隔而非 Young-Daly 最优值时,每天因过度检查点损失约 $13,638。

标记 15 分钟过于频繁、约 2.9 小时为 Young-Daly 最优值、8 小时过于稀疏的紧凑检查点间隔曲线。

最优检查点节奏在写入开销与重做风险之间取得平衡。

总结

扩展墙与 3D 并行立方体

本章介绍了“扩展墙”:即增加更多 GPU 最终会导致训练变慢而非变快的临界点。分布式训练并非简单的硬件问题;它是一个由模型大小、批次大小和互联带宽之间的相互作用所支配的约束满足问题。

3D 并行立方体(图 5.18)是大模型扩展的一个有用框架。数据并行展开训练的外层循环以扩展吞吐量;张量并行将矩阵乘法的内层循环向量化以适应内存;流水线并行将顺序层分阶段以降低通信频率。这些策略共同将大于任何单个内存库的模型映射到带宽有限的加速器集群上。

图 5.18:3D 并行策略空间:填充 3D 并行坐标网格的具体 8 GPU 配置。每个 GPU 位于一个 (d, p, t) 坐标处,该坐标结合了数据并行(d —— 跨工作节点复制模型)、流水线并行(p —— 垂直跨阶段划分层)和张量并行(t —— 在层内水平拆分单个操作)。单个训练步骤协调 d × p × t 个 GPU,层内进行张量并行通信,微批次阶段边界进行流水线通信,每个优化器步骤进行一次数据并行同步。

最终,并行主义的选择是由集群级编译器应用的循环变换。通过将逻辑通信模式与物理硬件层级相匹配,系统从小型集群的“线性扩展区”过渡到超大规模加速器集群的“通信受限”现实。

本章中,划分策略是直接为密集大模型训练和推荐类工作负载开发的。表 5.10 还展示了相同的约束逻辑如何延伸至联邦边缘场景,其中瓶颈从加速器内存和结构带宽转移到了不可靠的设备、隐私和间歇性连接。

3D 并行立方体中的工作点取决于系统的主要瓶颈而变化:

| 原型 | 主要划分策略 | 逻辑 |

| --- | --- | --- |

| 原型 A (GPT-4/Llama-3) | 混合 3D 并行 | 结合张量(宽度)、流水线(深度)和数据(吞吐量)并行以容纳 3.5 TB 权重。 |

| 原型 B (大规模 DLRM) | Embedding 分片 | 将海量 10 TB+ 表跨参数服务器集群分区;使用稀疏 AllToAll 更新。 |

| 原型 C (联邦 MobileNet) | 联邦学习 | 相同逻辑延伸至跨不可靠、隐私受限的边缘设备协调设备端 MobileNet 更新;第 11 章 深入探讨此模式。 |

表 5.10:按原型划分的策略:3D 并行立方体中的工作点随主导瓶颈的变化而变化。

本章探讨的并行策略(数据、张量、流水线和专家并行)为跨集群划分任何训练工作负载提供了概念工具包。关键洞见在于:这些策略并非互斥的替代方案,而是统一优化空间的互补维度。Megatron-LM 等系统之所以能实现高效扩展,正是因为它们组合了多种策略:节点内使用张量并行,节点组间使用流水线并行,数据并行用于吞吐量,专家并行用于容量扩展。

  • 拆分工作转移瓶颈:分布式训练并不会减少底层计算量;它将内存压力转化为通信量、同步延迟或空闲流水线时间。获胜的策略是将开销发送到车队中约束最不严格的部分。

  • 数据并行止步于收敛:复制仅在更大的全局批次仍能改善优化时才能干净地扩展吞吐量。超过临界批次大小后,AllReduce 成本和梯度噪声降低会使“更多工作节点”变成更慢或不稳定的学习,除非调整调度、预热和累积。

  • 分片以消息换内存ZeROFSDP 通过分区优化器状态、梯度和参数,使 100B+ 参数模型成为可能。代价是更严格的 ReduceScatterAllGather 操作通信调度,这些操作必须被重叠或隐藏。

  • 张量并行属于 NVLink 附近:张量并行在层内拆分矩阵操作,因此需要 A100H100 NVLink 提供的高带宽节点内结构。将该流量拉伸至跨机架会将内存容量解决方案变成通信瓶颈。

  • 流水线并行以气泡换字节:层分阶段降低了通信频率,但填充和排空时段使加速器空闲。m ≫ p 的微批处理是将模型深度转化为吞吐量而非流水线松弛的机制。

  • 3D 立方体是硬件映射:真正的前沿训练结合数据、张量、流水线、专家和分片并行,因为没有单一维度能同时适配模型和集群。逻辑组必须共同映射到 HBMNVLinkInfiniBand 和故障域。

立方体及其三轴之下隐藏着单一权衡。没有任何模型拆分方式能减少底层工作量;每种方式只是转移工作,将内存限制转化为通信成本,或将通信成本转化为空闲时间。数据、张量和流水线并行通过通信换取单个加速器所不具备的内存余量,而协调则是为保持拆分一致性而支付的税费。这就是开销置换定律,也是本卷其余部分的核心:执行规模的税费无法消除,只能在计算、通信和协调之间重新定位。立方体是税费可被送往何处的地图;工程则是选择约束最不严重的目的地。

我们已定义机器学习车队的逻辑流量模式,即拆分数学的方法。然而,这些逻辑模式最终会触及数据中心的物理连线。在第 6 章中,我们将打开实现这种一致性的集体操作(AllReduceAllToAll)的黑盒。我们将从发送什么的逻辑,转向如何通过环、树和轨道路由它的机制。

  • 当添加加速器开始降低扩展效率时,工程师应如何在数据、张量、流水线、专家和分片并行中选择?

  • 即使有更多硬件可用,增加全局批次大小何时不再改善时间-质量比?

  • 逻辑并行组应如何映射到 HBMNVLinkInfiniBand、存储和故障域?

  • 内存分片设计在停止改善训练时间之前,应接受多少额外通信?

集体通信

集体通信拓扑图,梯度片段通过环形和树形归约路径循环,随后广播同步状态。

目的

为什么机器间的通信会成为支配大规模机器学习系统的根本约束?

通过添加处理器来扩展计算;通过在它们之间移动数据来扩展通信

计算通过添加处理器来扩展;通信通过在处理器之间移动数据来扩展。它们的扩展方式不同:添加一个处理器会使总体计算能力线性增加,但将该处理器与所有其他处理器协调起来会移动更多的总数据,并且对于最一般的同步模式,工人之间的逻辑连接会以二次方增长。在足够大的规模下,花费在交换梯度、激活值和参数上的时间会超过花费在计算它们上的时间。这个交叉点不是需要修复的错误,而是分布式系统的一个基本属性,它决定了哪些并行策略有效,哪些模型规模可以训练,以及哪些组织能够在最大规模下运行。光速延迟、带宽限制和数据移动的能源成本的物理限制通信,就像晶体管物理限制计算一样严格,但通信远不如计算直观,因此成为那些只理解计算方面的人设计系统时的隐藏瓶颈。在 术语中,集体通信构成了舰队的指令集:决定计算如何因通信而受到惩罚的物理操作。

  • 应用通信成本模型来限制集体延迟、带宽和消息大小交叉点

  • 将集体通信原语与并行策略和模型原型相匹配

  • 使用依赖于规模的成本模型比较环形、树形、蝶形、双树和层次化归约算法

  • 通过将理论成本与测量的延迟、带宽和拓扑映射进行比较,分析通信库的现实差距

  • NVLinkInfiniBand、轨道优化、环形以及网络内归约结构设计拓扑感知的集体调度

  • 通过平衡通信节省与收敛风险来评估压缩、稀疏化和误差反馈

  • 通过桶融合、异步操作和逐层梯度调度实现重叠策略

从并行性到通信模式

当 10,000 个 GPU 需要应用一次权重更新时,由数据、张量、流水线和专家并行性创建的分割必须再次保持一致。工作进程并不是在发送任意消息;它们正在执行编排好的集体操作。在 图 1.13 所示的舰队栈中,这些操作位于分发层,逻辑并行性在此变为物理流量。

3D 并行立方体 假设副本、分片、阶段和专家能够足够快地交换梯度、激活值、参数和令牌,以至于优化器能够看到一次连贯的训练过程。这个假设是从并行性到通信的过渡:模型的每一种分割方式都会将工作重新分配到网络上。

该假设与实际电线之间的差距揭示了计算和通信在 如何 扩展方面的根本不对称性。计算是局部的:每个 GPU 只处理自己的数据,因此总计算能力随 GPU 数量线性增长。通信是全局的:保持舰队同步需要信息在具有有限延迟、带宽、拓扑和能源的物理链路上传输。核心情况是梯度同步;从这里开始,机制通过 alpha-beta 成本模型、集体原语、AllReduce 调度、拓扑感知路由、压缩和重叠变得具体。

梯度同步

梯度同步 是同步数据并行训练中的集体通信步骤,在此步骤中,每个工作进程将其本地计算的梯度张量贡献给一个总体归约,然后接收相同的归约结果,以便所有模型副本应用完全相同的更新。

  1. 重要性:一个使用 BF16 格式的 70B 参数模型在每个工作进程每步骤中会生成 140 GB 的梯度数据。在每条链路上以 50 GB/s 的速度通过环形 AllReduce 在 1,000 个 GPU 上进行同步,大约需要 2 × 140/50 ≈ 5.6 s 的通信时间每步骤,使得带宽(数据移动)项足够大以至于能够主导铁则定律,除非通过重叠、层次结构或压缩来减少暴露的传输。

  2. 区别:与参数服务器方法不同(在该方法中,所有工作进程将梯度发送到一个中心聚合器,其带宽随工作进程数量呈 O(N) 增长),环形 AllReduce 将通信分布在所有工作进程之间,使得每个工作进程每步骤的通信成本保持恒定为 2 × (N − 1)/N × M,与集群规模无关。

  3. 常见陷阱:一个常见的误解是认为环形 AllReduce 的行为类似于全连接交换。在环形 AllReduce 中,每个工作进程根据调度与邻居通信,且随着集群规模的增长,每个节点的通信量大约保持不变。扩展压力来自延迟步骤、拓扑和每步骤的梯度量,而不是每个工作进程向每个其他工作进程打开独立的数据流。

对于标准的同步数据并行 SGD,梯度同步不仅仅是一个设计便利;它是保持一个共享优化轨迹的机制。如果不同的 GPU 将不同的梯度更新应用到其模型的本地副本上,这些副本将发生偏离。经过足够多的步骤后,不同 GPU 上的模型将代表完全不同的函数,训练过程将不再逼近全局损失上的随机梯度下降。同步确保所有副本在每一步骤中保持相同(在浮点精度范围内),通过 AllReduce[⁶⁸] 原语保持优化算法的理论收敛保证。[⁶⁹]

必须同步的数据量与模型大小成正比。一个以 BF16(每参数 2 字节)存储的包含 P 个参数的模型,在每个训练步骤中每个 GPU 会生成 2P 字节的梯度数据。对于一个 700 亿参数的模型,这相当于每个 GPU 每步骤需要发送和接收 140 GB 的梯度。[⁷⁰]

在大规模情况下(数千个 GPU 上拥有数千亿参数),除非应用激进的优化技术,否则梯度同步可能会主导训练步骤时间。下一步是将这种数据运动建模为物理过程,而不是仅仅视为一个 API 调用。

数据移动的物理学

在设计算法之前,我们必须理解管理数据移动的物理约束。第 3.3 节 建立了这些算法背后的线级物理:光速设定了延迟下限[⁷¹](在光纤中大约为每公里 5 μs),带宽-距离乘积限制了快速链路在需要光学转换之前能传播多远(PAM4 信号和铜线-vs-光学传输距离在第 3.3.1 节中有所发展),而诸如 RDMAGPUDirect RDMA 之类的内核旁路传输将每条消息的软件开销降至几微秒(第 3.4.1 节)。集体通信增加了移动一个比特的能源成本,算法必须像尊重延迟和带宽一样严格地尊重这一成本。

移动数据会消耗能源,且能源消耗随距离增加而增加,如 图 6.1 所示。具体参考值为:本地 SRAM 大约为 0.5 pJ/bit,NVLink(封装内 PCB)为几十 pJ/bit,以及节点间 InfiniBand 为几百到上千 pJ/bit。该图使用了包含链路、交换机和收发器侧功率的较高公开估值;章节其他地方的正文使用了每个范围的下界。两种观点都一致认为:从 SRAM 到节点间织物,能源成本将增加两到三个数量级。

![](https://github.com/OpenDocCN/dsai-notes-pt3-zh/tree/master/docs/hav-cs249r-mlsys-vol2/img/file101.svg)

数据移动能量层级

图 6.1:在系统层级中,移动单个比特的能量成本随层级跃升数个数量级 —— 从 SRAM 经 HBM 到节点内 NVLink 再到节点间 InfiniBand。图中绘制的是公开发布的较高估计值(NVLink 在十几 pJ/比特,InfiniBand 接近 1000 pJ/比特),其中包含了交换机和收发器的开销。本章其他章节的文字论述在信封背面估算时采用的是较低的估计值。无论采用哪组数值,这一梯度都使数据局部性成为可持续且高效的分布式 ML 系统设计的首要驱动因素。

在超算规模(数万张 GPU)下,通信的功率预算与计算本身的功率预算不相上下。一个 1 万 GPU 的集群,若以 30 pJ/比特的能耗每步交换 1 GB 梯度,单次 AllReduce 约消耗 4.8 kJ(已计入数据移动中 2× 的因子),占每步总能量预算的显著比例。

这三个约束存在乘性交互。无论消息大小,延迟设定了每条消息的底线。带宽限制了大规模传输的吞吐量。协议开销给每条消息加征“税”,惩罚细粒度通信,这正是前文回顾的内核旁路传输对集合通信性能至关重要的原因。一次快速的 AllReduce 估算能让它们在真实训练场景下的综合代价具体化。

问题:64 GPU 上 70B 模型的 AllReduce 时间

一个 700 亿参数模型在 64 张通过 InfiniBand NDR(每端口 50 GB/s)互联的 GPU 上采用数据并行训练。每张 GPU 以 BF16 计算梯度(每参数 2 字节,Llama 级训练的常见做法)。一次 AllReduce 需要多长时间?

步骤 1:确定梯度载荷大小

每张 GPU 生成一个完整的梯度张量:7 × 10¹⁰ × 2 字节 = 140 GB。

步骤 2:套用环形 AllReduce 带宽公式

T_bandwidth = 2 * (N-1)/N * n/β

代入数值:T_bandwidth = 1.96875 × 140 GB / 50 GB/s ≈ 5512.5 ms。

步骤 3:加上延迟项

T_latency = 2 * (N-1) * α

代入数值:T_latency = 126 × 1.5 μs = 0.2 ms。

步骤 4:总通信时间

总计:T_AllReduce ≈ 5512.5 ms + 0.2 ms ≈ 5512.7 ms。

系统洞察

单是梯度 AllReduce 就耗时五秒多。这是加在每个训练步上的纯通信开销。在这个规模下,除非与反向传播计算重叠,否则通信将主导步耗时。这就是大规模训练系统将 AllReduce 与反向传播流水线化的原因:在后层仍在计算时,提前为前层发起通信。

两条水平对比条,对比大规模下步耗时的铁律项。通信条宽且紫色阴影,占主导;计算条是窄灰色条。

梯度同步吞噬了大规模下每步 30% 至 70% 的时间。

上述计算揭示了为何数据移动而非计算成为规模化时的决定性约束。在 70B 模型的梯度上执行一次 AllReduce 会消耗数秒墙钟时间,期间所有 GPU 本可忙于计算却被迫闲置。局部计算(可完美并行)与全局协调(需要物理数据移动)之间的这种不对称性,驱动了后续所有集合算法的设计。将这个数秒级的 AllReduce 减半,在典型训练周期中可节省数千 GPU 小时,直接转化为成本降低和更快的部署时间。

成本分析同时解释了为何集合算法的选择比大多数从业者意识到的更为关键。若使用次优算法仅达到理论带宽的 60%(拓扑映射不佳时的常见结果),这五秒的 AllReduce 会膨胀至九秒以上,给每个训练步平增数秒纯浪费。因为这种浪费会在整个机群中传播,通信算法位居机群栈的分发层:其下的基础设施层通过 NVLink、InfiniBand 及网络拓扑提供原始带宽(详见第 3 章),其上的服务层依赖高效的梯度同步完成产出可部署模型的训练运行。当通信算法无法饱和可用带宽,训练拉长,服务模型交付延后,整个机群便低于其经济潜力运行。

图 6.2 量化了这一通信开销如何随机群规模复合增长。在单个 NVLink 互联节点内的 8 GPU 时,通信约占每训练步的 25%,因为 900 GB/s 的互联带宽能跟上梯度数据量。当集群扩展到跨 8 个节点的 64 GPU,转用 InfiniBand(每端口 50 GB/s)后,平衡被打破:通信占步耗时约 50%,另加 5% 耗于同步屏障。在 4096 GPU 规模下,通信与同步开销合计占训练步的 80%,仅余 20% 用于有效计算。这一演进解释了为何集合算物至关重要:没有分层集合通信、梯度压缩和通信计算重叠,大规模训练将把数百万美金的算力预算绝大部分花在等数据到达上。

图 6.2:计算-通信时间线:随着训练从 8 GPU(单节点)扩展到 4096 GPU,通信占比从约 25% 增长至每训练步的 65% 以上。这种从 NVLink 主导向 InfiniBand 受限的通信转变,驱动了所有集合优化策略的设计。比例源自 H100 集群实测,采用 70B 参数模型、数据并行与环形 AllReduce

梯度的“出差清单”

旅程始于反向传播完成的那一刻。在单机上,故事到此结束:权重被更新,神经元完成了学习。但在生产规模下,梯度诞生于孤立之中。它存在于一张 GPU 上,而模型的“真理”分布在数千张 GPU 间。为实现全局收敛,梯度必须找到它的同伴。

这趟旅程的具体“出差清单”由第 5 章所选的并行策略决定。如何拆分数学决定了数据如何移动。对我们的灯塔原型(1.6.1 节),这些清单有本质区别。对原型 A(GPT-4/Llama-3),梯度属于一个巨大的稠密张量,必须与机群中所有其他梯度会合以计算全局平均,因此其主要载具是 AllReduce。对原型 B(大规模 DLRM),DLRM 工作负载下,梯度或激活是稀疏且定向的:它不需见所有人,只需找到持有其嵌入分片的那张特定 GPU。专家混合(MoE)路由创造了相同的定向模式,将每个 token 发送到托管其指定专家的 GPU,因此两者都依赖 AllToAll

理解这种映射至关重要

理解这种映射至关重要:并行性的内容(what)直接决定了通信的方式(how)。在大规模场景下,这些策略并非互斥。大语言模型的单次训练运行通常采用 3D 并行(同时结合数据并行、张量并行和流水线并行),这意味着多个集合通信原语在重叠的 GPU 子集上并发执行。张量并行驱动节点内的 AllReduce 操作(通过 NVLink),流水线并行驱动流水线阶段间的点对点发送(通常跨节点),而数据并行驱动跨节点组的 AllReduce 操作(通过 InfiniBand)。

每个原语在不同的进程组(process group)上运行,即参与该特定集合操作的总 GPU 群体的一个子集。经典的 MPI 使用通信子(communicator)作为命名此类参与集合的对象;其通信子大小即该集合中的秩(rank)数量。MPI 库已根据通信子大小和消息大小选择集合算法(Thakur et al. 2005)。GPU 通信库继承了该算法选择问题,并增加了在不产生并发集合操作争用的情况下协调重叠进程组的进一步挑战。

表 6.1:出行清单

表 6.1 预览了从并行策略到集合通信原语的映射。原语名称将在下一节正式介绍;目前,可将该表视为一张交通地图,指明谁必须与谁交换数据。

| 并行策略 | 梯度的目标 | 主要原语 | 主要约束 |

| --- | --- | --- | --- |

| 数据并行 | 与所有人会面,计算全局平均值 | AllReduce | 带宽(大载荷) |

| FSDP/ZeRO | 寻找分片,重构完整张量 | AllGather + ReduceScatter | 带宽(高频次) |

| 张量并行 | 节点内快速握手 | AllReduce | 延迟(速度即生命) |

| 流水线并行 | 移交给下一个邻居 | 点对点(Send/Recv) | 延迟(顺序依赖) |

| 专家并行 (MoE) | 定向路由至专家 | AllToAll | 延迟 + 争用 |

表 6.1:出行清单:第 5 章 的高层数学如何表现为低层流量模式。

在该表中,专家并行指代混合专家(MoE)[⁷⁴]架构模式。该映射表明,不同的并行策略施加了根本不同的通信模式。数据并行和 FSDP 生成大型的、带宽受限的消息,受益于基于环(ring-based)的算法和分层分解。张量并行和流水线并行生成小型的、延迟受限的消息,受益于基于树(tree-based)的算法和低开销软件栈。专家并行生成全对全(all-to-all)流量模式,考验网络的跨割带宽。要定量推理这些差异,我们需要一个网络性能模型。

绘制地形:网络性能建模

如果不知道两个不同的变量——固定的启动开销和每字节的传输费用——数据中心工程师就无法预测在集群中发送十兆字节需要多长时间。当梯度踏上旅程时,它立即遭遇数据中心网络的物理现实。

α-β 成本模型:启动税与过路费

我们梯度发送的每条消息都遵循线性成本模型 T(n) = α + n/β——一个固定的启动税加上每字节的过路费。这两个项决定了本章中的每一个算法选择。

α-β 模型

α-β 模型 是线性通信成本模型 T(n) = α + n/β,它将消息传输时间分解为固定启动延迟(α,每条消息的开销)和与消息大小相关的带宽项(n/β,与传输字节数成比例),使算法设计者能够预测何时消息融合或梯度压缩能提高吞吐量(Hockney 1994)。

  1. 意义:对于 α ≈ 2 μs 和 β ≈ 50 GB/s 的 InfiniBand NDR,临界消息大小 n^() = αβ ≈ 100 KB。一个 4 KB 的路由消息远低于 n^(),因此启动税占主导;将 100 条此类消息融合为一条 400 KB 的消息,可将通信成本从 100α = 200 μs 降低至一次 α + n/β ≈ 10 μs,实现 20 倍提升。一个 140 GB 的梯度张量远高于 n^(*),因此带宽占主导,减少载荷字节是有效的杠杆。

  2. 区别:与将带宽视为唯一通信成本的理想化吞吐模型不同,α-β 模型揭示了当 n/Nn^() 时,N* 条大小为 n/N 的小消息的成本最高可达一条大小为 n 的大消息的 N 倍,这解释了为什么 NCCL 会融合小的 AllReduce 调用,以及为什么 MoE 路由算法在启动集合操作前缓冲 Token。

  3. 常见陷阱:一个常见的误解是梯度压缩总是有帮助的。如果压缩后的梯度大小仍远大于 n^(*),压缩仅减少带宽项,而延迟项保持不变。即使将 70B 模型的梯度载荷从 140 GB 缩减至 1.4 GB,消息仍牢牢处于带宽受限状态;这并不会将操作转变为低延迟交换。

这两个参数具有不同的物理含义。延迟(α)是发送消息的固定启动成本,与大小无关,涵盖软件开销(内核启动、NCCL 初始化)、PCIe 遍历和网络交换时间。带宽(β)是以字节/秒为单位的持续数据传输速率。临界消息大小 n^() = αβ 标记了分水岭:小于 n^() 的消息受延迟限制;大于者受带宽限制。附录 D.1 节 通过具体的消息大小区间和屋顶线分析详细推演了该模型,想了解临界点背后完整推导的读者可前往查阅;此处我们确立参数并直接应用。

表 6.2:互联性能参数

表 6.2 列出了数据中心互联的典型值,临界大小一列承载了关键模式:节点内 NVLink 在数百 KB 之前保持延迟受限,而节点间互联在约 100 KB 处跨越,因此在节点内足以成为带宽受限的消息,一旦跨节点传输可能仍处于延迟受限状态。

| 互联技术 | 延迟 (α) | 带宽 (β) | 临界大小 (n^(*)) |

| --- | --- | --- | --- |

| NVLink 4.0 (节点内) | 1–2 μs | 450 GB/s | ~0.7 MB |

| InfiniBand NDR 400 Gbps | 1–3 μs | 50 GB/s (每端口) | ~100 KB |

| InfiniBand HDR 200 Gbps | 2–5 μs | 25 GB/s | ~87.5 KB |

| PCIe Gen5 (GPU↔CPU) | 2–5 μs | 64 GB/s | ~224 KB |

| 以太网 100 Gbps (RoCE) | 5–10 μs | 12.5 GB/s | ~93.7 KB |

表 6.2:互联性能参数:有效的每消息启动延迟(α,含软件栈),单向链路带宽(β,进入 T(n) = α + n/β 的单向可达速率;NVLink 数据手册的 900 GB/s 为双向总和,故其 β 为 450 GB/s 单向速率),以及各行 α 中位数下的计算临界大小 n^(*) = αβ,即通信从延迟受限转为带宽受限的分界点。

将临界消息大小公式应用于具体工作负载,可揭示哪种优化策略最关键:

问题

考虑一个使用 InfiniBand NDR 400 Gbps、其中 α = 2 μs 和 β = 50 GB/s 的集群。在何种消息大小下,针对带宽的优化开始比针对延迟的优化更重要?

数学推导

n* = α ⋅ β = 2 × 10^(-6) s × 5 × 10¹⁰ B/s = 100 KB

系统洞察

小于 100 KB 的消息(如 MoE token、流水线激活值)是 延迟敏感型:购买更低延迟的交换机并减少软件开销。大于 100 KB 的消息,如大语言模型(LLM)梯度,则是 带宽敏感型:购买更多带宽并压缩数据。应用错误的优化会浪费金钱却无法提升性能。

两个水平条。顶部(大消息):一个标记为 n 除以 beta 的宽蓝色带宽段,加一个细灰色的延迟薄片。底部(小消息):一个标记为 alpha 的宽紫色延迟段,加一个细灰色薄片。主导项发生翻转。

小消息受延迟限制,大消息受带宽限制;成本主导地位会反转。

关键消息大小将两种截然不同的运行模式分开。在该阈值以下,诸如 MoE 路由 token 或标量归约之类的小消息是延迟敏感型的(n < n*),耗时主要由 α 决定,优化重点在于融合(将小消息打批)、拓扑(减少跳数)和软件栈调优(通过 RDMA 实现内核旁路)。在该阈值以上,诸如 LLM 梯度或优化器状态之类的大消息是带宽敏感型的(n ≫ n*),耗时主要由 n/β 决定,优化重点转向压缩(降低精度或利用稀疏性)、算法选择(环形 vs. 树形)和链路聚合(多轨 NIC)。区分延迟敏感型与带宽敏感型通信,是需要练习的核心诊断技能。

验证你对网络性能模式的理解:

LogP 模型

α-β 模型假设处理器在通信期间处于空闲状态。对于重叠通信与计算的流水线系统,这一假设失效。LogP 模型(Culler 等,1993)引入四个参数对 α-β 模型进行扩展:

  • L_lat(延迟):消息穿越网络所需时间(类似于 α)。

  • o(开销):CPU/GPU 发起或接收传输所花费的时间。在此期间处理器无法计算,因此这是不可重叠的成本。在分布式训练语境下,o 是 PyTorch 或 JAX 调度开销、CUDA 内核启动时间、张量向 RDMA 栈注册内存,以及通信线程与训练循环竞争导致的 Python 全局解释器锁(GIL)争用等开销的总和。这些软件层解释了为什么测量到的 NCCL 开销每个集体操作高达 25–50 μs,即使线级 α 仅为 1–3 μs;本节后面的 NCCL 对比将量化这一差距。

  • g(间隙):连续消息注入的最小时间间隔(消息速率的倒数)。这对链路争用建模。

  • N_rank(秩数量):通信组中的秩数量。

LogP 区分了网络延迟(L_lat,可被隐藏)和处理器开销(o无法被隐藏)。只有当计算内核运行时间超过开销时,系统才能将通信与计算重叠。一个具体的重叠计算可以阐明这一区别:

问题

一个训练流水线试图将梯度 AllReduce 与下一层的反向传播重叠。反向传播耗时 500 μs。AllReduce 的网络延迟为 100 μs,但发起和接收的处理器开销 o 均为 50 μs。通信能否被隐藏?

数学推导

  1. 可重叠部分:网络延迟 L_lat = 100 μs(数据在飞行中而 GPU 计算)。

  2. 不可重叠部分2o = 100 μs(GPU 忙于发起/接收)。

  3. 可用计算时间:500 μs。

  4. 已隐藏:全部 100 μs 的 L_lat 均可与计算重叠。

  5. 暴露部分:100 μs 的 2o 开销无法重叠。

结果:有效时间为 100 μs + max(500 μs, 100 μs) = 600 μs。网络延迟被隐藏了,但处理器开销仍然暴露。图 6.3 直观展示了这一重叠。

系统洞察

α-β 模型捕捉总通信时间。LogP 模型揭示了 其中有多少可以被隐藏。设计流水线训练时,应针对低 o(内核旁路、GPUDirect)优化,而非仅追求高 β

图 6.3:通信-计算重叠:通过将集体操作与反向传播流水线化,系统可在算术执行后隐藏网络延迟。只有将不可重叠开销(o)最小化,批量传输项才能在 GPU 计算下一层梯度时进行,重叠才能成功。

模型选择取决于分析场景。α-β 模式适合回信封计算、算法选择(环形 vs. 树形)以及通信为 阻塞式(同步屏障)的场景。其优势在于简单,因为两个参数可直接通过点对点带宽测试和零字节消息延迟测试获得。其劣势在于假设处理器通信时空闲,当重叠成为可能时会过度悲观。当分析转向 流水线执行、计算-通信重叠,或调试为何理论上快的算法表现不佳(通常是高 o)时,LogP 模型凭借额外参数大显身手。它对网络延迟 L_lat(可隐藏)与处理器开销 o(不可隐藏)的区分,决定了给定的重叠策略是否真能隐藏通信。代价是准确测量 o 需要 NVIDIA Nsight Systems 等剖析工具,因为 o 取决于具体的通信库和 GPU 驱动栈。

实践中,大多数工程计算以 α-β 模型起步进行初步规模估算和算法选择,当通信-计算重叠成为目标优化时,再引入 LogP 分析细化。两个模型共享一个共同局限:假设单链路上的单一流。真实通信模式涉及多个同时流竞争共享带宽,可能导致两个模型均无法捕捉的拥塞。对于对拥塞敏感的工作负载(特别是 MoE 的 AllToAll),在目标集群上进行实测基准测试仍是必要的验证步骤。

实战模型应用:Llama 70B 通信预算

α-β 模型在应用于真实训练配置时最具价值。考虑一个具体场景:跨 16 个节点(每节点 8 GPU,共 128 GPU)使用数据并行训练 Llama 级 70B 参数模型。梯度张量为 BF16 格式下的 140 GB(700 亿参数 × 2 字节,Llama 级训练常见)。每个训练步骤中,整个梯度必须在所有工作节点间同步。

该计算使用了本章后续将正式定义的带宽层级:尽可能在节点内通过快速 NVLink 层级完成归约,仅将归约后的分片通过较慢的 InfiniBand 层级发送,最后在本地重构结果。原语名称将在层级章节中变得精确;工程思路已在带宽预算中显现。

采用 BF16 梯度(一种将通信量减半至 140 GB 的常见做法),对比十分明显。像扁平环形 AllReduce 那样将完整梯度路由到慢速节点间网络,耗时约 5,557 ms。将大部分流量限制在 NVLink,仅通过 InfiniBand 发送归约分片,则将其降至约 1200.8 ms。这绝非边际差异;它决定了通信能否隐藏在计算之后,还是成为关键路径。

这些数字背后的各阶段分解(节点内归约、归约后的节点间交换、节点内重新分发)需要本节尚未介绍的集合通信原语。第 6.5.1 节给出了这些原语的名称并完成了完整的三阶段推导;此处的预算仅旨在确立带宽层级值得尊重,以及具体尊重的程度。

理论与实践:NCCL 的现实差距

带宽-延迟权衡(原则 )提供了有用的一阶预测,但真实的通信库引入了理想化 α-β 模型无法捕捉的开销。NCCL 作为一种广泛使用的 GPU 通信库,增加了协议协商、内存注册和内部流水线,这些会修改有效的 αβ 值(Jeaugey 2017;NVIDIA 2026b)。表 6.3 对比了单消息 α-β 载荷预测与在 8 节点 DGX H100 集群(64 块 GPU,InfiniBand NDR 400G)上测得的常见消息大小的 NCCL 性能。

| 消息大小 | α-β 预测 | 实测 NCCL | 比率 (实测/预测) | 解释 |

| --- | --- | --- | --- | --- |

| 1 KB | ~3.1 μs | ~25 μs | ~8.1× | NCCL 协议建立主导 |

| 64 KB | ~4.4 μs | ~30 μs | ~6.9× | 仍受延迟限制;NCCL 开销 |

| 1 MB | ~23.1 μs | ~40 μs | ~1.7× | 过渡到带宽受限 |

| 64 MB | ~1.3 ms | ~1.6 ms | ~1.2× | NCCL 接近理论带宽 |

| 1 GB | ~20 ms | ~23 ms | ~1.15× | 带宽主导;NCCL 接近最优 |

| 10 GB | ~200 ms | ~215 ms | ~1.07× | 大载荷饱和链路 |

表 6.3:α-β 预测 vs. 实测 NCCL 性能**:对于小消息,NCCL 的协议开销(内存注册、通道建立、内核启动)使有效延迟比裸线 α 膨胀了 7–8 倍。对于大消息,NCCL 达到了理论带宽的 85–92%,验证了单消息载荷模型在带宽受限区域的有效性。数值为 8 节点 DGX H100 集群配合 InfiniBand NDR 400G 时有效 NCCL 载荷传输模型的说明性参考数值;实际集合通信计时随算法、拓扑、NCCL 版本和调优而异。

该表揭示了两个关键教训。首先,α-β 模型低估了小消息延迟 7–8 倍,因为它仅考虑了线路级传播,未计入软件栈开销。对于延迟敏感型操作(张量并行 AllReduce、MoE Token 路由),有效 α 比物理线路延迟高 5–10 倍。其次,对于大消息,模型精度在 85–92% 以内,确认带宽是约束瓶颈,且 NCCL 的内部优化(通道流水线、内核融合)成功饱和了可用链路。

这一现实差距对算法选择有实际影响。Ring 与 Tree AllReduce 之间的交叉点在实践中会上移,因为有效 α 大于线路级数值。使用教科书 α 值的工程师会低估延迟成本,可能在 Tree 表现更好时错误选择 Ring。一种稳健的做法是通过基准测试小消息 AllReduce 延迟,在特定集群上测量有效 α,随后在所有后续计算中使用该实测值。

选择载体:集合操作原语

如果 GPU 只是简单地打开一个套接字并将巨大的梯度发送给另一个 GPU,整个集群将迅速陷入无法管理的死锁和拥塞网络中。有了地形图,梯度现在必须选择其载体:严格编排的组交换,称为集合操作,这是一套由 MPI 标准化并被现代 ML 通信库继承的词汇(Message Passing Interface Forum 2015)。图 6.4 展示了该词汇中的四个核心原语:AllReduce、AllGather、ReduceScatter 和 AllToAll。每个面板都作为进程组的前后对比来阅读,每行对应一个秩:蓝色单元格标记每个秩开始时的输入,绿色单元格标记其最终得到的结果,橙色标记 AllToAll 在每对秩之间洗牌的唯一数据。另外两个基础原语,Broadcast(秩 0 发送给所有)和 Reduce(所有聚合到秩 0),是在下文散文中定义并隐式用于图中复合模式的基础构建块。附录 D.2 节形式化了每个原语的语义及其延迟和带宽复杂度;此处散文通过其工作负载用例介绍它们,希望在继续前先建立正式复杂度界限的读者可先在那里确立。

集合操作 是一种分布式通信模式,其中组内所有进程同时参与以聚合、广播或重新分发数据,并保证无论消息顺序或到达时间如何,每个参与者都能收到集合语义规定的结果。

  1. 重要性:正确的集合算法决定了通信是随集群规模扩展还是保持恒定。Ring AllReduce 实现了带宽最优的 2(N − 1)/N × M/β 每节点(在 N 上为常数),而朴素的归约再广播方法代价为 𝒪(N × M/β),当 N = 1024 时要差 500 倍。算法选择因此直接决定铁律中的 BW 项是否成为训练瓶颈。

  2. 区别:与点对点通信(一个发送方和一个接收方独立交换数据)不同,集合操作协调整个进程组——每个参与者必须调用该集合操作后才能完成,且库保证即使不同进程贡献不同数据也能得到一致结果。

  3. 常见误区:一个常见的误解是单一集合算法适用于所有消息大小。Ring AllReduce 为大梯度张量实现近乎最优的带宽,但具有 𝒪(N) 延迟——对于 MoE 路由决策(~4 KB)等小消息,基于树的算法可将延迟降至 𝒪(log N) 步骤,代价是带宽利用率次优(带宽最优的递归减半-加倍替代方案见 第 6.4.5 节),这要求针对工作负载进行特定算法选择,而非单一默认值。

图 6.4:四大核心集合原语:分布式系统中组通信的标准化模式。(1) AllReduce:全员聚合且全员接收结果。(2) AllGather:全员发送给全员并拼接。(3) ReduceScatter:全员聚合且结果以分片形式分发。(4) AllToAll:每个进程向每个其他进程发送唯一数据。两个基础原语 Broadcast(秩 0 发送给所有)和 Reduce(全员聚合到秩 0)组成了此处展示的模式,并在周围散文中定义。

原语选择是第一个扩展决策,因为它固定了哪些秩必须协调、哪些数据必须移动,以及哪个通信项将占主导。不同模型架构对不同操作施加压力,为工作负载选择错误的原语会产生不必要的瓶颈。

六大核心原语

表 6.4 中的六个原语形成了一个决策阶梯:使用能保留模型数学特性的最窄操作,因为每个更宽泛的模式都会增加参与者、屏障或争用。

六种集合体原语

| 原语 | 作用 | 主要使用场景 | 通信成本 |

| :--- | :--- | :--- | :--- |

| 广播 | 一个发送方向所有接收方传输数据。 | 在启动时从 rank 0 向所有并行策略分发初始模型权重。 | 𝒪(log N) 延迟(树形);𝒪(M) 带宽。 |

| 规约 | 将所有工作节点的数据(求和、最小值、最大值)聚合到单个根节点。 | 聚合验证指标,如损失和准确率,以供日志记录进程使用。 | 𝒪(log N) 延迟(树形);𝒪(M) 带宽。 |

| 全规约 | 从所有工作节点聚合数据,然后将结果分发到所有节点(对所有 i,有 \(y_i = \sum_{j=0}^{N-1} x_j\))。 | 数据并行 ism(第 5.3 节):同步梯度,使每个 GPU 计算相同的更新。 | 环形在带宽上最优,为 \(2\frac{N-1}{N}\frac{M}{\beta}\),但延迟为 𝒪(N);树形延迟为 𝒪(log N),但带宽较差。 |

| 全收集 | 每个工作节点的数据发送到所有节点;结果将所有输入连接起来(x[i] → [x[0], …, x[N−1]])。 | 分片数据并行 ism(FSDP/ZeRO):在前向或反向传递前收集分片参数。 | 环形为 \(\frac{N-1}{N}\frac{M}{\beta}\);驻留数据从 M/N 分片增长到完整的 M。 |

| 规约散布 | 在工作节点间进行规约,但将结果分散,使每个节点保持一个不同的块(工作节点 i 接收 ∑x[j] 的第 i-th 块)。 | 分片数据并行 ism:在减少梯度的同时保持其分片状态以节省内存。 | 环形为 \(\frac{N-1}{N}\frac{M}{\beta}\);是全收集的逆运算。 |

| 全部交换 | 最通用的模式:工作节点 i 向每个工作节点 j 发送一个不同的块,形成分布式矩阵转置。 | 专家混合(MoE)将令牌路由到专家;DLRM 在工作节点间交换嵌入查找。 | 每个工作节点为 \(\frac{N-1}{N}M\),但 𝒪(N²) 的逻辑连接使竞争成为扩展瓶颈。 |

表 6.4:六种集合体原语: 每种原语的计算内容、主要用途及其通信成本,按从最窄模式(广播)到最通用模式(全部交换)排序。

环形实现了全规约的带宽最优成本,而树形则以带宽换取对数延迟;第 D.2.1 节 发展了形式定义及一个比较两者的示例,因此希望获得该推导的读者可以前往查阅完整内容(Patarasuk 和 Yuan 2009;Thakur 等人 2005`)。

一个关键的见解是:全规约 可以分解为 规约散布 后跟 全收集。这不仅仅是一个数学等价;它正是环形 全规约 在内部的工作方式(散布-规约阶段是一个 规约散布,全收集阶段是一个 全收集)。ZeRO 风格的分片在 PyTorch 中实现为 FSDP[⁷⁵],利用这一分解:在模块需要参数之前通过 全收集 参数分片,在反向传播后通过 规约散布 梯度,以使每个秩仅保留其分片用于优化器步骤(Rajbhandari 等人 2020`)。这将参数、梯度和优化器状态分片到各工作节点,使持久模型状态内存大约减少 N 倍,而临时完整参数缓冲区和预取则决定峰值内存。

同样的 全收集/规约散布 配对也出现在序列并行 ism 中:一个 全收集 重构了局部操作所需的激活视图,随后的 规约散布 将结果沿序列维度重新分布。此例后来将变得重要,因为它表明即使被分片的张量维度发生变化,原语的语义仍保持不变。

此梯度之所以重要,是因为原语决定了在任何算法优化开始之前的同步形状。全规约 创建全局一致,全收集规约散布 在重复的分片移动与内存之间进行权衡,全部交换 则用扇出和竞争替换带宽问题。

全部交换全规约 之间的对比凸显了集合操作在集群规模变化时如何不同地扩展的根本差异。

全规约 能高效扩展,因为它可以在环形中流水线化(在此情况下,每个节点仅与其邻居通信);而 全部交换 在根本上更难以扩展。

全部交换 中,每个进程都拥有发送给每个其他进程的唯一数据块。这产生了 𝒪(N²) 的逻辑连接。在硬件层面,这导致 网络竞争:如果 1024 个 GPU 同时尝试向不同目标发送数据,“Fat-Tree” 或 “Spine” 交换机将成为数据中心的瓶颈。

All-to-All quadratic connection growth versus the AllReduce curve.

All-to-All 流量的增长速度快于全规约。

这就是为什么专家并行 ism(MoE)和大规模推荐系统往往比标准数据并行模型更早遇到“通信墙”的原因。算法选择(全规约全部交换)决定了扩展上限。

为了使 全部交换 模式具体化,考虑一个包含 64 个专家且分布在 8 个 GPU 上(每 GPU 8 个专家)的专家混合(MoE)模型。在每次前向传递期间,门控网络将每个输入令牌分配给一个或多个专家。分配给远程 GPU 上专家的令牌必须在计算前被物理移动到这些 GPU,并且结果之后也必须移回。

问题

一个 MoE 模型在 8 个 GPU 上处理包含 4096 个令牌的批次(每 GPU 512 个令牌)。每个令牌是一个 2048 维的隐藏状态,采用 BF16 格式(每个令牌 4.096 KB)。门控网络将每个令牌分配给恰好 1 个 64 个专家中的一个(每 GPU 8 个专家)。假设路由均匀(每个专家接收 64 个令牌),那么每个 GPU 发送和接收多少数据?

数学计算

每个 GPU 持有需要传送到 8 个 GPU 上 64 个不同专家的 512 个令牌。在均匀路由下,每个 GPU 向其他每个 GPU 发送 64 个令牌(保留 64 个令牌在本地)。

  • 每次 GPU-to-GPU 传输的数据量:64 个令牌 × 4.096 KB/令牌 = 262.144 KB

  • 每个 GPU 发送的总数据量:7 × 262.144 KB = 1.84 MB

  • 每个 GPU 接收的总数据量:1.84 MB(对称)

延迟分析(InfiniBand NDR,α = 3 μs,β = 50 GB/s)

这 7 次传输每次都是 262.144 KB 的消息。根据 α-β 模型:T[传输] = 3 μs + 5 μs = 8 μs/次传输。

  • 如果串行执行:7 × 8 μs = 56 μs。

  • 如果所有 7 次传输并行运行(全双工、非阻塞):≈ 8 μs。

系统见解

在此 MoE 示例中,每次传输的大小处于关键消息大小的边界附近,因此启动开销与序列化时间相当。这就是为什么 MoE 的扩展常常受限于 α(延迟)、扇出和竞争,而不是单纯的 β(带宽),以及为什么 MoE 系统受益于低延迟交换机、RDMA、消息融合和拓扑感知路由。全部交换 还会在每层运行两次(一次用于令牌分发,一次用于结果收集),从而将启动开销加倍。

在实际中,MoE 路由很少完全均匀。热门专家获得的令牌多于不受欢迎的专家,由此导致负载不平衡,进而转化为通信不平衡。如果一个专家接收的令牌数量是平均值的 3 倍,那么托管该专家的 GPU 接收到的入站数据量也是平均值的 3 倍,这会形成热点,并可能因 全部交换 是一个屏障操作而使整个集体操作停滞。大规模 MoE 系统通过辅助负载均衡损失来应对这一问题,该损失会惩罚门控网络将过多令牌路由至单个专家;同时通过容量因子限制单个专家可接受的最大令牌数量(溢出令牌将被丢弃)。这些技术会牺牲少量模型质量以换取通信平衡,这是在规模上的正确取舍。

集合操作选择指南

表 6.5 将这些原语映射到第 5 章介绍的并行策略以及第 1.6.1 节介绍的灯塔原型。

| 训练策略 | 主要集合操作 | 瓶颈特征 |

| --- | --- | --- |

| 数据并行 | AllReduce | 带宽受限(大梯度) |

| FSDP/ZeRO-3 | AllGather, ReduceScatter | 带宽受限,高频率 |

| 张量并行 | AllReduce, AllGather | 延迟受限(需要 NVLink) |

| 流水线并行 | 点对点 | 延迟受限(微批次传递) |

| 序列并行 | AllGather, ReduceScatter | 带宽受限(激活值交换) |

| MoE (专家) | AllToAll | 延迟敏感 + 竞争 |

| DLRM (推荐系统) | AllToAll | 延迟与带宽(稀疏查找) |

表 6.5:集合操作选择指南:将训练策略与其主要集合操作相匹配,可实现高效的分布式通信设计。MoE 和 DLRM 作为贯穿全书的典型“灯塔”工作负载,分别代表稀疏专家架构和推荐系统。流水线并行独特地依赖点对点通信而非集合通信。

FSDP 通信模式

两行符号公式:数据并行每步使用一次集合操作,而 FSDP 每步每层使用两次 N 次集合操作。

posted @ 2026-09-06 04:02  绝不原创的飞龙  阅读(4)  评论(0)    收藏  举报