分布式训练
分布式训练是一种利用多台机器或多个计算设备（如GPU、TPU）协同工作来加速模型训练的技术，主要用于处理大规模数据或参数量极大的模型（如深度学习）。其核心目标是通过并行计算解决单设备内存不足、训练速度慢的问题。
一、为什么需要分布式训练？
- 数据/模型规模大：现代模型参数量可达千亿（如GPT-3），单设备无法存储。
- 训练速度需求：单GPU训练大型数据集耗时过长（如几周甚至几个月）。
- 资源利用：充分利用集群计算资源，降低成本。
二、分布式训练的并行策略
1. 数据并行（Data Parallelism）
- 原理：将训练数据拆分到多个设备，每个设备保存完整的模型副本，独立计算梯度后同步更新。
- 实现：
- 同步更新：所有设备计算完梯度后，聚合（如取平均）再更新模型（PyTorch
DistributedDataParallel）。 - 异步更新：设备独立更新参数，无需等待（更快但可能收敛不稳定）。
- 同步更新：所有设备计算完梯度后，聚合（如取平均）再更新模型（PyTorch
- 适用场景：数据量大、模型可单卡存放（如ResNet）。
2. 模型并行（Model Parallelism）
2.1流水线并行（Pipeline Parallelism）
- 原理：将模型按层分段，不同设备处理不同段，通过流水线提高设备利用率。
- 示例：GPU1处理第1-5层，GPU2处理6-10层，数据像流水线依次传递。
- 优化技术：梯度累积、微批次（Micro-batching）减少气泡（Bubble）时间。
2.2 张量并行（Tensor Parallelism）
将单个张量操作（如矩阵乘法）的计算拆分到多个设备上，每个设备处理张量的一个子块，通过通信协作完成整体计算。
实现方式
- 行拆分（Row-wise Split）：将权重矩阵按行拆分到不同设备。
- 列拆分（Column-wise Split）：将权重矩阵按列拆分到不同设备。
- 示例：反向传播时梯度矩阵按列聚合。
- 混合拆分
- 结合行和列拆分，适用于复杂张量操作（如多头注意力机制）。
具体案例：前馈网络（FFN）：将两层线性变换的权重矩阵按行或列拆分，每个设备计算局部结果，通过 All-Reduce 同步全局输出。
4. 混合并行
- 结合数据、张量、流水线并行（如Megatron-LM、DeepSpeed的ZeRO优化器）。
三、关键技术
- 通信机制
- 参数服务器（Parameter Server）：中心化节点聚合梯度（易成瓶颈）。
- All-Reduce：去中心化通信，所有节点共同参与梯度聚合（如NCCL库）。
- Ring-AllReduce：设备组成环状拓扑，高效传输数据（带宽利用率高）。
- 框架支持
- PyTorch：
DistributedDataParallel（DDP）、RPC（模型并行）。 - TensorFlow：
tf.distribute.Strategy（如MirroredStrategy、MultiWorkerMirroredStrategy）。 - 专用库：DeepSpeed（支持ZeRO、3D并行）、Horovod（All-Reduce优化）。
- PyTorch：
- 显存优化
- 梯度检查点（Gradient Checkpointing）：用计算换显存，只保留部分中间结果。
- 混合精度训练：FP16/FP32混合计算（节省显存，加速训练）。
- ZeRO（Zero Redundancy Optimizer）：消除数据并行中的显存冗余（DeepSpeed）。
