avatar

Chen Kunpeng

Tongji University. Focusing on Trustworthy AI & Federated Learning.

深度学习框架简要对比：Megatron、PyTorch和DeepSpeed

从技术的角度解析大模型训练基础设施中的三个核心组件：PyTorch、DeepSpeed 以及 Megatron。

PyTorch

PyTorch 是目前绝大多数 LLM 训练的底层计算环境。它的核心职责是提供张量（Tensor）的数学运算抽象、构建前向传播（Forward Pass）的动态计算图，并通过自动微分（Autograd）机制计算反向传播（Backward Pass）的梯度。

在大模型训练的语境下，PyTorch 原生提供了基于数据并行（Data Parallelism, DP）和分布式数据并行（Distributed Data Parallel, DDP）的功能。在 DDP 模式下，模型被完整地复制到集群中的每一个 GPU 上，每个 GPU 处理不同的数据微批次（Micro-batch），随后通过 All-Reduce 通信原语对计算出的梯度进行同步。

当模型参数量超过单卡显存上限时，PyTorch 原生的 DDP 由于要求单卡必须能够容纳完整的模型参数、梯度和优化器状态（即所谓的 1+1+2 或更多倍的模型体积），导致无法进行训练（Out of Memory, OOM）。虽然 PyTorch 后来引入了 FSDP（Fully Sharded Data Parallel），但在极致的显存优化和算子融合层面，工业界通常会寻求更专业的解决方案。

DeepSpeed

DeepSpeed 是微软开源的深度学习优化库，它在软件层面上对 PyTorch 进行了深度的 Hook 和封装。DeepSpeed 的核心技术贡献是 ZeRO（Zero Redundancy Optimizer，零冗余优化器） 算法。

ZeRO 算法旨在打破 DDP 中每个 GPU 都保存完整模型状态的冗余。它将训练过程中的显存占用分为三个阶段（Stages）进行切分：

  • ZeRO-1： 仅对优化器状态（Optimizer States，如 Adam 的动量和方差）进行分片存储在不同 GPU 上。
  • ZeRO-2： 在 ZeRO-1 的基础上，进一步对梯度（Gradients）进行分片。
  • ZeRO-3： 在前两者的基础上，对模型参数（Parameters）本身也进行分片。

在 ZeRO-3 模式下，单张 GPU 仅保存整个模型的极小一部分。只有在前向传播或反向传播需要用到某一层参数时，DeepSpeed 才会通过 All-Gather 通信从其他 GPU 动态获取这些参数，并在计算完成后立即释放（销毁）这部分参数的显存占用。此外，DeepSpeed 还提供了 ZeRO-Offload 技术，能够将部分显存数据卸载到 CPU 内存或 NVMe 固态硬盘上，进一步突破 GPU 显存限制。

Megatron (Megatron-LM)

Megatron 是由 NVIDIA 开发的，专门针对大规模 Transformer 架构进行极致优化的分布式训练框架。不同于 DeepSpeed 侧重于显存状态的管理，Megatron 侧重于计算维度的拆解，即如何在多张物理 GPU 之间最高效地分配矩阵乘法（GEMM）运算。

Megatron 实现了 并行架构：

  1. 张量并行 (Tensor Parallelism, TP)： 在单个矩阵乘法层面进行垂直或水平切分。例如，将 Transformer 层中的多头注意力机制（Multi-Head Attention）的查询、键、值（Q, K, V）权重矩阵切分到多张同一节点内的 GPU 上同时计算，最后通过 All-Reduce 合并结果。TP 的通信极其密集，因此通常被限制在具有超高带宽（如 NVLink）的单台服务器内部。
  2. 流水线并行 (Pipeline Parallelism, PP)： 按层（Layers）对模型进行切分。将模型的前 NN 层放置在一个 GPU 节点，接下来的 MM 层放置在下一个节点。通过按微批次（Micro-batch）注入数据，形成计算流水线。
  3. 数据并行 (Data Parallelism, DP)： 在 TP 和 PP 之外的维度进行常规的数据切分。

除此之外，Megatron 包含了大量针对 NVIDIA 硬件底层深度优化的定制 CUDA 算子（如 FlashAttention 的早期集成、LayerNorm 的算子融合等），以最大化 GPU 的硬件浮点运算效率（MFU）。

四、 架构总结与对比

简而言之：

  • PyTorch 提供了定义神经网络的语言和基础数学运算；
  • DeepSpeed 解决的是显存不够的问题，通过 ZeRO 切分技术让小显存集群能够训练大模型，代价是引入了更高的通信开销；
  • Megatron 解决的是计算效率的问题，通过 3D 并行和算子优化，让巨型模型在超大型集群（万卡级别）上跑出极高的吞吐量。

在实际的工业界大型预训练中，这三者并非互斥关系。最常见的组合是：底层基于 PyTorch，使用 Megatron 的张量并行和流水线并行处理计算，并在数据并行维度上套用 DeepSpeed 的 ZeRO-1 进行优化器分片，这种架构通常被称为 Megatron-DeepSpeed。