分布式机器学习：数据并行与模型并行详解

一、为什么需要并行训练

深度学习模型的训练本质上是海量数据驱动的反复迭代计算。每一次迭代都包含前向传播、损失计算、反向传播和参数更新，其中涉及大量矩阵乘法、卷积运算与梯度累加。当模型参数量达到亿级甚至更大、训练样本达到千万级时，单张加速卡的算力和显存都难以支撑，训练周期可能被拉长到数周甚至数月。

分布式训练的核心思路很直接：单卡算力不足就增加卡数，单机资源不足就扩展到多机。但“加卡”并不是简单堆叠硬件，关键在于如何合理地划分计算任务、组织通信、控制同步开销。常见的划分维度有两个：数据与模型。由此衍生出数据并行与模型并行两种基本范式。理解它们的差异与联系，是设计高效训练方案的前提。

二、数据并行：切分数据，复制模型

2.1 基本工作方式

数据并行的核心特征是：每个训练进程持有一份完整的模型副本，但处理不同的数据子集。假设有 N 个训练进程（通常称为 worker），全局批次大小为 B，则每个 worker 分到 B/N 条样本。各 worker 独立完成前向传播与反向传播，得到各自基于本地数据的梯度，再通过某种聚合机制将梯度汇总，最终得到一致的参数更新结果。

从计算角度看，数据并行把“样本维度”切开，模型结构和参数在逻辑上保持单副本语义。只要梯度聚合是确定性的（如求平均或求和），其数学结果与单机大批次训练等价。

2.2 参数聚合的两种典型架构

参数服务器架构：系统中存在一组参数服务器（PS）与若干 worker。worker 计算本地梯度后推送到参数服务器，由参数服务器完成梯度求和与参数更新，再将新参数拉回各 worker。这种架构把通信与计算解耦，参数服务器可横向扩展以承载更大的模型，但引入了额外的网络跳数与潜在的带宽瓶颈。

集合通信架构：不设独立参数服务器，各 worker 之间通过 AllReduce 等集合通信原语直接交换梯度。以 Ring AllReduce 为例，梯度被切分成若干块，在环形拓扑上分阶段传递并累加，最终每个 worker 都获得全局归约结果。这种方式通信量与 worker 数量基本无关，在大规模集群中扩展性较好，是当前较为主流的实现方式。

下面给出一个简化的数据并行训练伪代码，帮助理解其循环结构：

# 每个 worker 持有完整模型副本
model = build_model()
for batch in local_data_loader:          # 各 worker 数据互不重叠
    grads = compute_gradients(model, batch)
    all_reduced = all_reduce_sum(grads)  # 集合通信聚合梯度
    grads = all_reduced / world_size     # 求平均
    model = update_parameters(model, grads)

2.3 优势与代价

数据并行的优势在于实现相对简单、对模型结构无侵入，几乎可以套用到任意网络。其代价主要体现在两方面：一是显存冗余，每张卡都要保存完整的模型参数、梯度与优化器状态，当模型本身接近单卡显存上限时便无法容纳；二是通信开销，每轮迭代都要同步全量梯度，梯度规模与参数量成正比，在跨机场景下对网络带宽提出较高要求。因此数据并行更适合“模型能放进单卡、但数据量大或希望缩短训练时间”的场景。

三、模型并行：切分模型，协同计算

3.1 为什么需要切分模型

深度学习的计算主体是矩阵运算，而这些矩阵需要驻留在显存中。当某一层的参数矩阵规模超出单卡显存容量时，仅靠数据并行无法解决问题。例如分类数达到千万级别时，全连接层的权重矩阵会非常庞大，单卡显存难以容纳。此时必须把大矩阵拆开，分别放到不同设备上计算，从网络结构的角度看就是把模型拆成多个部分——这就是模型并行。

3.2 两种切分维度

层间并行（流水线式切分）：按网络深度把不同的层分配到不同设备，例如前若干层放在设备 A，后续层放在设备 B。计算时数据像流水线一样逐段传递。这种方式实现直观，但存在设备空闲问题：同一时刻只有一个设备在计算，其余设备等待，硬件利用率低。改进方案是 microbatch 流水线，将一个小批次再切分为多个微批次，使不同设备能重叠计算，从而提升吞吐。

层内并行（张量切分）：把单个算子内部的矩阵运算切分到多台设备。例如将一个大的矩阵乘法按行或按列拆分，各设备计算部分结果后再拼接或归约。这种方式能真正把单层显存压力分摊到多卡，但通信频繁、实现复杂度高，需要仔细设计切分方式与通信时机，否则通信开销可能抵消并行收益。

3.3 优势与代价

模型并行的核心价值在于突破单卡显存上限，使超大规模模型成为可能。其代价是强烈的设备间依赖：切分后的计算单元之间存在频繁的数据交换，通信模式比数据并行更复杂，且并行度受模型结构限制。此外，不同层的计算量往往不均衡，容易造成负载倾斜，需要合理的切分策略与调度配合。

四、两者的联系与混合使用

数据并行与模型并行并非互斥，实际训练中常被组合使用。一个直观的判断依据是各层的计算特点：

• 卷积层通常计算量大，但参数（权重）相对较少，单卡显存能够容纳，适合用数据并行来分摊计算量；
• 全连接层往往参数量巨大而单次计算量相对有限，容易触碰显存上限，适合用模型并行来分摊参数存储。

因此一种常见做法是：对卷积部分采用数据并行，对全连接部分采用模型并行，形成混合并行方案。更进一步，在大规模训练中还会叠加流水线并行与张量并行，形成“数据并行 + 流水线并行 + 张量并行”的多维并行组合，分别对应样本维度、层维度和算子维度的切分。

设计混合并行方案时，需要重点权衡三件事：显存占用（参数、梯度、优化器状态、激活值各自占多少）、通信量（哪一维切分带来的同步最频繁）、负载均衡（各设备计算量是否接近）。通常的做法是先估算单卡显存能否容纳模型，若不能则优先引入模型侧切分；在显存允许的前提下，再通过数据并行提升吞吐。

五、实践中的选型建议

1. 先评估模型规模：若单卡可容纳完整模型与训练状态，优先考虑数据并行，实现成本最低。
2. 显存受限时引入模型切分：按层或按张量切分，优先选择通信模式简单、负载较均衡的切分方式。
3. 关注通信瓶颈：跨机通信带宽远低于机内互联，尽量把通信密集的切分放在同一节点内。
4. 控制同步频率：可通过梯度累积、更大的批次或异步更新策略降低同步次数，但需注意对收敛性的影响。
5. 逐步调优：从简单方案起步，借助性能分析定位瓶颈，再决定是否引入更复杂的并行组合。

六、小结

数据并行切分的是样本，每张卡保留完整模型，通过梯度聚合保持一致性，适合模型可单卡容纳、追求训练吞吐的场景；模型并行切分的是模型结构与参数矩阵，通过设备间协作突破显存上限，适合超大规模模型。两者在计算划分、通信模式和资源约束上各有侧重，实际工程中往往按层特性组合使用。理解这些基本原理与权衡关系，才能针对具体模型与硬件条件设计出合理、可扩展的分布式训练方案。