Skip to main content

Efficiency

David Liu4/3/26About 5 min

Efficiency

训练大模型（LLM）的核心挑战在于显存限制（Memory）**和**计算效率（Compute Efficiency）。单张 GPU 甚至单台服务器往往无法装下千亿参数的模型，因此必须依赖分布式训练技术和显存优化技术。

以下是对大模型训练核心技术栈的详细介绍，重点解析 DeepSpeed、Megatron-LM 等主流方案。

1. 并行策略 (Parallelism Strategies)

这是大模型训练的基础，目的是将模型或数据“拆解”到成百上千张 GPU 上。

1.1 数据并行 (Data Parallelism, DP)

这是最基础的模式。

  • 原理： 每张 GPU 复制一份完整的模型参数。数据被分割成小批次（Micro-batch），每张卡处理一部分数据，计算梯度，然后通过 All-Reduce 同步梯度，更新权重。
  • 局限： 当模型大到单张 GPU 放不下时，纯 DP 就失效了。
  • 进阶： FSDP (Fully Sharded Data Parallel) 和 DeepSpeed ZeRO 是 DP 的进化版（见下文）。

1.2 张量并行 (Tensor Parallelism, TP)

由 NVIDIA 的 Megatron-LM 推广。

  • 原理： 将模型中的具体运算层（如 Transformer 的 MLP 和 Attention 矩阵）切分到多张 GPU 上。
    • 例如，一个 4096×40964096 \times 4096 的矩阵乘法，被切分成两个 2048×40962048 \times 4096 的矩阵在两张卡上并行计算。
  • 适用场景： 节点内通信（Intra-node），因为 TP 需要极为频繁的 GPU 间通信，通常依赖 NVLink 这种高带宽互连。

1.3 流水线并行 (Pipeline Parallelism, PP)

  • 原理： 将模型的不同层（Layer）按顺序切分。例如，前 10 层在 GPU 0，中间 10 层在 GPU 1。数据像流水线一样流过。
  • 挑战： 存在“气泡”（Bubble），即在等待数据依赖时，某些 GPU 会空闲。
  • 优化： 1F1B (One-Forward-One-Backward) 策略，通过交替前向和后向传播来填充气泡，提高利用率。

2. DeepSpeed 与 ZeRO 技术详解

DeepSpeed 是微软开源的深度学习优化库，它的核心杀手锏是 ZeRO (Zero Redundancy Optimizer)。

传统的分布式训练（DDP）中，每张卡都保存完整的参数（Parameters）、梯度（Gradients）和优化器状态（Optimizer States）。这造成了巨大的显存冗余。

ZeRO 的核心思想是：既然数据是并行的，为什么不把模型状态也切分到各个 GPU 上，需要时再通信获取？

ZeRO 的三个阶段 (Stages)

  1. ZeRO Stage 1 (切分优化器状态)：
    • 将 Optimizer States（如 Adam 的动量和方差，通常占显存的大头）切分。
    • 效果： 显存节省约 4 倍，通信量不变。
  2. ZeRO Stage 2 (切分梯度 + 优化器状态)：
    • 进一步将 Gradients 进行切分。
    • 效果： 显存进一步节省约 2 倍，通信量不变。
  3. ZeRO Stage 3 (全切分：参数 + 梯度 + 优化器状态)：
    • 将模型权重（Parameters）也切分了。每张卡只存一部分权重。在前向/后向传播计算某一层时，通过 All-Gather 临时从其他卡拉取权重，计算完立即释放。
    • 效果： 显存占用与 GPU 数量成反比（线性节省），可以训练万亿参数模型。
    • 代价： 通信开销增加约 50%。

ZeRO-Offload / ZeRO-Infinity

  • 原理： 当 GPU 显存实在不够时，DeepSpeed 可以将优化器状态和梯度“卸载”到 CPU 内存（Offload）甚至 NVMe SSD 上。
  • 意义： 允许在有限的 GPU 资源上（甚至单卡）微调超大模型，虽然速度会变慢，但打破了“显存墙”。

3. 显存优化技术 (Memory Optimization)

除了切分模型，还需要其他手段降低单卡显存压力。

3.1 激活重计算 (Activation Checkpointing / Gradient Checkpointing)

  • 背景： 在前向传播时，通常需要保存每一层的激活值（Activation）用于后向传播计算梯度。这占用大量显存。
  • 技术： 不保存中间层的激活值。在后向传播需要时，利用保存的少量检查点（Checkpoint）重新计算前向传播。
  • 代价： 用“时间换空间”，增加了约 33% 的计算量，但显著降低了显存峰值。

3.2 混合精度训练 (Mixed Precision Training)

  • FP16 / BF16： 使用 16 位浮点数（Float16 或 BFloat16）代替传统的 32 位（FP32）进行计算。
  • 优势： 显存占用减半，计算速度在 Tensor Core 上快数倍。
  • 注意： 通常需要维护一份 FP32 的主权重副本（Master Weights）来保证数值稳定性（DeepSpeed 会自动处理）。

3.3 FlashAttention

  • 原理： 这是一个底层 CUDA 算子优化。传统的 Attention 计算需要构建 N×NN \times N 的巨大矩阵，显存消耗是序列长度的平方级 O(N2)O(N^2)。FlashAttention 通过分块计算（Tiling）和重计算，避免将大矩阵写入显存（HBM），主要在高速缓存（SRAM）中完成。
  • 效果： 训练速度大幅提升，且极大地节省显存，是训练长文本（Long Context）模型的必备技术。

4. 主流框架对比

框架核心技术优势适用场景
DeepSpeedZeRO (1/2/3), Offload易用性强，与 HuggingFace 集成好，节省显存效果最强。资源受限（显存不足）或超大规模模型全量微调。
Megatron-LMTensor Parallel (TP) + Pipeline Parallel (PP)计算效率极高，针对 NVIDIA 硬件极致优化。硬核预训练，拥有由 NVLink 连接的高端集群（如 A100/H100 集群）。
PyTorch FSDP类似于 ZeRO-3PyTorch 原生支持，无需第三方库，稳定性好。通用分布式训练，Llama 官方训练代码常采用此方案。
Colossal-AI综合了 ZeRO、TP、PP提供了很多自动并行的策略，试图做全能王。想要结合多种并行策略但不想自己手写复杂代码。

总结：如何选择？

  1. 如果你在做全量预训练（Pre-training）： 通常采用 3D 并行（即 Megatron-LM 的 TP/PP + DeepSpeed 的 ZeRO/DP）。
  2. 如果你在做微调（SFT）且显存不足： 首选 DeepSpeed ZeRO-3 配合 Offload。
  3. 如果你要训练超长上下文（Context Window > 32k）： 必须开启 FlashAttention-2 和 Sequence Parallelism。

Would you like me to explain how to write a configuration file (ds_config.json) for DeepSpeed to start a training task?

David's Blog
Copyright © 2026 David Liu