分布式训练详解
分布式训练是指利用多台计算机（节点）或单个节点上的多个加速卡（如 GPU、TPU）共同协作，来训练大规模深度学习模型的技术，其核心目的是通过并行计算来缩短训练时间，并突破单机内存对模型规模的限制。
为什么需要分布式训练
- 模型规模巨大：现代大语言模型（LLM）拥有数千亿个参数，单张 GPU 的显存无法装下整个模型及其梯度和优化器状态。
- 数据量海量：面对 TB 级甚至 PB 级的数据集，单机训练可能需要数年时间，通过分布式训练可以实现线性加速。
- 计算资源优化：通过集群化部署，可以更高效地利用计算资源，提高硬件利用率。
核心并行策略
根据切分对象（数据或模型）的不同，分布式训练主要分为以下几种策略：
数据并行 (Data Parallelism, DP)
这是最常用的并行方式,每个计算节点都拥有一份完整的模型副本，但处理不同的数据分片。
- 工作流程：
- 将一个 Batch 的数据均匀分给所有 GPU。
- 每个 GPU 独立进行前向传播计算梯度。
- 通过 All-Reduce
操作同步所有 GPU 的梯度，计算平均值。
- 所有 GPU 使用相同的更新梯度来更新各自的模型参数。
- 常见实现：PyTorch 的
DDP(Distributed Data Parallel)。
模型并行 (Model Parallelism, MP)
当模型太大无法放入单张 GPU 时，需要将模型本身拆分到不同设备上。
- 张量并行 (Tensor Parallelism, TP)：
- 将模型中的某个层（如线性层、注意力层）的权重矩阵切分到多个 GPU 上。
- 每个 GPU 计算矩阵乘法的一部分，最后通过通信合并结果。
- 适用于极宽的层,通信频率极高。
- 流水线并行 (Pipeline Parallelism, PP)：
- 将模型的不同层（Layer）分配到不同的 GPU 上。
- 数据像流水线一样流动：GPU 0 处理前几层 $rightarrow$ 传递给 GPU 1 $rightarrow$ 传递给 GPU 2。
- 为了解决 GPU 空闲（Bubble）问题，通常采用 Micro-batch 策略。
混合并行与 ZeRO 技术
为了进一步优化,业界引入了更先进的内存优化技术，如 ZeRO (Zero Redundancy Optimizer)：
- ZeRO-1：仅对优化器状态进行分片。
- ZeRO-2：对优化器状态和梯度进行分片。
- ZeRO-3：对优化器状态、梯度和模型参数全部进行分片。
- 效果：在保持数据并行训练效果的同时，极大地降低了单卡显存占用，使得在有限资源下训练超大规模模型成为可能。
关键通信原语
分布式训练依赖于高效的底层通信,最常见的操作包括：
- Broadcast：将一个节点的数据发送到所有其他节点。
- Scatter：将数据切分并分发到各个节点。
- Gather：将所有节点的数据收集到单个节点。
- All-Reduce：所有节点交换数据并进行求和/平均，最终每个节点都获得相同的汇总结果（分布式训练的核心）。
主流分布式训练框架
- PyTorch DDP/FSDP：PyTorch 原生支持的分布式数据并行和完全分片数据并行。
- DeepSpeed：微软开发的开源库，集成了 ZeRO 技术，是目前训练大模型的主流选择。
- Megatron-LM：NVIDIA 开发的针对 Transformer 架构的高度优化框架，侧重于 TP 和 PP。
- Horovod：由 Uber 开发，基于 MPI 的通用分布式训练框架，支持多种深度学习后端。
总结对比
| 策略 | 切分对象 | 解决的问题 | 通信开销 | 复杂度 |
|---|---|---|---|---|
| 数据并行 (DP) | 数据 | 训练速度慢 | 中 | 低 |
| 张量并行 (TP) | 权重矩阵 | 单层太大 | 极高 | 高 |
| 流水线并行 (PP) | 网络层 | 模型太深 | 低 | 中 |
| ZeRO | 状态/参数 | 显存不足 | 中 | 中 |
首发原创文章，作者：王坚，如若转载，请注明出处：https://idctop.com/article/490504.html
