DeepSpeed 和 Megatron 训练堆栈
DeepSpeed (Microsoft) 和 Megatron-LM (NVIDIA) 是软件堆栈，使跨数千个 GPU 的数十亿参数的训练模型变得切实可行。
概述
如果没有它们，今天的前沿模型根本无法适应内存或在合理的时间内完成训练。
深入探讨
在一个 GPU 上训练大型模型是不可能的，因为权重、梯度和优化器状态不适合。这些堆栈将工作分配给许多 GPU。 Megatron-LM 开创了张量并行性，在 GPU 上的每一层内切片单独的矩阵乘法，以及管道并行性，将不同的层放在不同的 GPU 上。 DeepSpeed 的标志性贡献是 ZeRO（零冗余优化器），它将优化器状态、梯度和参数跨 GPU 分片而不是复制它们，从而大幅削减每个 GPU 的内存。两者经常结合起来（Megatron-DeepSpeed）来训练 BLOOM-176B 和 Megatron-Turing NLG 等模型。他们还添加了混合精度、激活检查点和卸载到 CPU 或 NVMe，因此大型模型可以在有限的硬件上进行训练。
技术洞察
ZeRO 具有三个增加内存节省的阶段：第 1 阶段对优化器状态进行分片，第 2 阶段还对梯度进行分片，第 3 阶段对参数本身进行分片，在前向和后向传递过程中按需收集它们。与张量并行性（层内）和管道并行性（层间）相结合，形成“3D并行性”。关键的压力是通信开销：每个分片分割都会增加 GPU 到 GPU 的流量，因此工程师调整分割以保持快速 NVLink 和 InfiniBand 链路饱和。
战略影响
成本与预算
多年来，架构决策决定着性能和运营成本。
更清晰的判决
技术教育帮助团队选择正确的堆栈，而不仅仅是最新的堆栈。
质量控制
更好的工程选择可以减少生产中的可靠性事故。
DeepSpeed 和 Megatron 训练堆栈的未来
预计与 PyTorch 的原生 FSDP（完全分片数据并行）进行更紧密的集成，它吸收了许多 ZeRO 思想，模糊了研究堆栈和核心框架之间的界限。编译器驱动的方法和自动并行规划器旨在消除手动调整。随着训练集群增长到数十万个加速器，容错、弹性扩展和与计算的重叠通信成为主要的工程前沿，同时支持 NVIDIA Blackwell 和定制训练芯片等新硬件。
现实世界的实施
使用跨数百个 GPU 的组合 Megatron-DeepSpeed 堆栈来训练开放式多语言 BLOOM-176B 模型。
Microsoft 和 NVIDIA 使用 3D 并行性训练 5300 亿参数的 Megatron-Turing NLG 模型。
ZeRO-Offload 允许研究人员通过将优化器状态溢出到 CPU RAM 来微调单个工作站 GPU 上的数十亿参数模型。
在这些堆栈中使用激活检查点，通过重新计算激活而不是存储全部激活来适应更长的上下文窗口。
风险与防护栏
优化一项基准测试可以隐藏更广泛的系统弱点。
基础设施和维护成本常常被低估。
随着系统变得更加复杂，安全性和可观察性差距可能会扩大。
实施路线图
在实施之前定义延迟、质量和成本目标。
在实际负载和数据条件下进行基准测试。
仪器监控错误、漂移和用户影响。
在扩展之前准备回滚和事件响应路径。
不断探索
Free newsletter
Get the daily AI briefing
Three verified AI stories every weekday morning, written in plain English. Free forever, no ads.
One email each weekday. Unsubscribe in one click. We never sell or share your address.
Test yourself
Take the DeepSpeed and Megatron Training Stacks quiz
Instant feedback on every answer, and a shareable certificate with a verifiable ID once you pass a course.
Support free AI education. AI Understanding is a 501(c)(3) nonprofit — no ads, no paywall, ever. Make a donation
常见问题
什么是 DeepSpeed 和 Megatron 训练堆栈？
DeepSpeed (Microsoft) 和 Megatron-LM (NVIDIA) 是软件堆栈，使跨数千个 GPU 的数十亿参数的训练模型变得切实可行。如果没有它们，今天的前沿模型根本无法适应内存或在合理的时间内完成训练。
DeepSpeed 的 ZeRO 优化器的主要目的是什么？
ZeRO（零冗余优化器）通过跨 GPU 划分优化器状态、梯度和参数而不是在每个设备上复制它们来消除内存冗余。
Megatron-LM 中首创的张量并行如何分割模型？
张量并行性将单个层内的数学（例如大型矩阵乘法）跨多个 GPU 进行划分，这就是层内分割。
哪个 ZeRO 阶段通过对模型参数本身进行分片来提供最大的内存节省？
除了梯度和优化器状态之外，ZeRO Stage 3 还对参数进行分片，按需收集它们，从而最大程度地减少内存。
在训练过程中，“激活检查点”会牺牲什么来节省内存？
激活检查点存储较少的中间激活，并在反向传播期间重新计算它们，用额外的计算换取减少的内存。
为什么这些技术的组合通常被称为“3D 并行性”？
3D并行性堆叠了三种正交策略：数据并行性、张量（层内）并行性和管道（层间）并行性。