大模型训练方法详解
随着大语言模型（LLM）参数量级的不断攀升，训练数据规模日益庞大，单卡训练往往面临显存不足或计算效率低下的瓶颈。当拥有多张 GPU（单机多卡或多机多卡）时，通过并行训练技术可以有效解决这些问题。常见的并行策略主要包括数据并行、模型张量并行、流水线并行以及 ZeRO 优化。
1. 并行训练基础概念
在大规模分布式训练中，核心目标是平衡计算负载与通信开销，同时确保模型参数能够被有效利用。
1.1 数据并行（Data Parallelism, DP）
数据并行是最基础的并行方式。每个 GPU 都加载完整的模型副本，将输入数据分割成多个批次（Batch），分别输入到不同的 GPU 中进行前向传播和反向传播。在计算 Loss 和梯度后，各 GPU 之间需要同步梯度，通常使用 AllReduce 操作来聚合梯度，然后更新本地模型参数。
DP 与 DDP 的区别：
- 实现机制：DP 通常基于多线程实现，受全局解释器锁（GIL）限制；DDP（Distributed Data Parallel）基于多进程实现，每个 GPU 由独立进程控制，避免了 GIL 锁的干扰。
- 适用范围：DP 仅适用于单机环境；DDP 支持单机及多机集群。
- 性能表现：DDP 通常比 DP 更快，因为它减少了进程间通信的开销。但在网络带宽受限的场景下，DDP 的通信效率可能成为瓶颈。
- 通信机制：DP 可能存在多次数据交换，而 DDP 通常在每个训练步骤中只进行一次梯度同步，并通过 GPU 间直接通信融合数据。
1.2 模型张量并行（Tensor Parallelism, TP）
当单个 GPU 无法存储整个模型的某一层权重时，需要使用张量并行。TP 将一个大的 Tensor 按行或列切分，存储在不同的 GPU 上。每个 GPU 负责计算部分结果，最后通过通信汇总。
以 Transformer 架构中的线性层为例：
- 按列拆分权重：输入不需要拆分，输出通过 Concat 组合。
- 按行拆分权重：输入需要拆分，输出通过 Sum 组合。
这种方式的优点是可以显著降低单卡的显存占用，但缺点是增加了层内通信的频率和开销。TP 通常用于处理极大规模的模型层，如 Megatron-LM 的实现。
1.3 流水线并行（Pipeline Parallelism, PP）
流水线并行将模型按照层进行切分，不同的层分配给不同的 GPU。数据像流水线一样，从前面的 GPU 传递到后面的 GPU。
工作流程：
- 前向传播：GPU0 计算第 1-4 层，输出传递给 GPU1；GPU1 计算第 5-8 层。
- 反向传播：梯度从后向前传递。
气泡问题（Bubble）：由于流水线特性，后面的 GPU 必须等待前面的 GPU 完成计算才能开始，导致部分 GPU 闲置。为了解决这个问题，通常引入微批次（Micro-batch）技术，将一个大 Batch 拆分为更小的 Micro-batch，交错执行不同层的计算，从而减少空闲时间。
1.4 ZeRO（Zero Redundancy Optimizer）
ZeRO 是 DeepSpeed 提出的一种优化技术，旨在进一步减少显存占用。它属于数据并行的范畴，但通过分片优化器状态、梯度和参数来实现。
ZeRO 的三个阶段：
- Stage 1：仅对优化器状态进行分片。
- Stage 2：对优化器状态和梯度进行分片。
- Stage 3：对优化器状态、梯度和模型参数全部进行分片。
ZeRO-Offload：可以将计算量小且低频使用的参数（如优化器状态、FP32 参数）卸载到 CPU 上，从而释放 GPU 显存。这在混合精度训练中尤为有效，能在不影响计算精度的前提下显著提升可训练的模型规模。
2. 适用场景与策略选择
选择合适的并行策略取决于硬件资源、网络环境和模型大小。
