概述
如果没有它们，今天的前沿模型根本无法适应内存或在合理的时间内完成训练。
深入探讨
在一个 GPU 上训练大型模型是不可能的，因为权重、梯度和优化器状态不适合。这些堆栈将工作分配给许多 GPU。 Megatron-LM 开创了张量并行性，在 GPU 上的每一层内切片单独的矩阵乘法，以及管道并行性，将不同的层放在不同的 GPU 上。 DeepSpeed 的标志性贡献是 ZeRO（零冗余优化器），它将优化器状态、梯度和参数跨 GPU 分片而不是复制它们，从而大幅削减每个 GPU 的内存。两者经常结合起来（Megatron-DeepSpeed）来训练 BLOOM-176B 和 Megatron-Turing NLG 等模型。他们还添加了混合精度、激活检查点和卸载到 CPU 或 NVMe，因此大型模型可以在有限的硬件上进行训练。
现实世界的实施
使用跨数百个 GPU 的组合 Megatron-DeepSpeed 堆栈来训练开放式多语言 BLOOM-176B 模型。
Microsoft 和 NVIDIA 使用 3D 并行性训练 5300 亿参数的 Megatron-Turing NLG 模型。
ZeRO-Offload 允许研究人员通过将优化器状态溢出到 CPU RAM 来微调单个工作站 GPU 上的数十亿参数模型。
在这些堆栈中使用激活检查点，通过重新计算激活而不是存储全部激活来适应更长的上下文窗口。
DeepSpeed 和 Megatron 训练堆栈的未来
预计与 PyTorch 的原生 FSDP（完全分片数据并行）进行更紧密的集成，它吸收了许多 ZeRO 思想，模糊了研究堆栈和核心框架之间的界限。编译器驱动的方法和自动并行规划器旨在消除手动调整。随着训练集群增长到数十万个加速器，容错、弹性扩展和与计算的重叠通信成为主要的工程前沿，同时支持 NVIDIA Blackwell 和定制训练芯片等新硬件。
付诸实践
为什么这很重要
- 成本与预算. 多年来，架构决策决定着性能和运营成本。
- 更清晰的判决. 技术教育帮助团队选择正确的堆栈，而不仅仅是最新的堆栈。
- 质量控制. 更好的工程选择可以减少生产中的可靠性事故。
留意
- 优化一项基准测试可以隐藏更广泛的系统弱点。
- 基础设施和维护成本常常被低估。
- 随着系统变得更加复杂，安全性和可观察性差距可能会扩大。
后续步骤
- 在实施之前定义延迟、质量和成本目标。
- 在实际负载和数据条件下进行基准测试。
- 仪器监控错误、漂移和用户影响。
- 在扩展之前准备回滚和事件响应路径。
不断探索
常见问题
什么是 DeepSpeed 和 Megatron 训练堆栈？
DeepSpeed (Microsoft) 和 Megatron-LM (NVIDIA) 是软件堆栈，使跨数千个 GPU 的数十亿参数的训练模型变得切实可行。如果没有它们，今天的前沿模型根本无法适应内存或在合理的时间内完成训练。
DeepSpeed 的 ZeRO 优化器的主要目的是什么？
ZeRO（零冗余优化器）通过跨 GPU 划分优化器状态、梯度和参数而不是在每个设备上复制它们来消除内存冗余。
Megatron-LM 中首创的张量并行如何分割模型？
张量并行性将单个层内的数学（例如大型矩阵乘法）跨多个 GPU 进行划分，这就是层内分割。
哪个 ZeRO 阶段通过对模型参数本身进行分片来提供最大的内存节省？
除了梯度和优化器状态之外，ZeRO Stage 3 还对参数进行分片，按需收集它们，从而最大程度地减少内存。
在训练过程中，“激活检查点”会牺牲什么来节省内存？
激活检查点存储较少的中间激活，并在反向传播期间重新计算它们，用额外的计算换取减少的内存。
为什么这些技术的组合通常被称为“3D 并行性”？
3D并行性堆叠了三种正交策略：数据并行性、张量（层内）并行性和管道（层间）并行性。