PyTorch 分布式训练底层原理与 DDP 实战指南

作者：郝培强

2025-11-29
广东
本文字数：18352 字
阅读完需：约 60 分钟

深度学习模型参数量和训练数据集的爆炸式增长，以 Llama 3.1 为例：4050 亿参数、15.6 万亿 token 的训练量，如果仅靠单 GPU 可能需要数百年才能跑完，或者根本无法加载模型。

并行计算（Parallelism）通过将训练任务分发到多个 GPU（单机多卡或多机多卡），并利用通信原语同步状态，能让训练过程变得可控且高效。

本文讲详细探讨 Pytorch 的数据并行（Data Parallelism）

扩展算力的两种路径

扩展训练规模无非两种方式：纵向扩展（Vertical Scaling） 和 横向扩展（Horizontal Scaling）。

纵向扩展：

简单粗暴地升级硬件。比如把 10GB 显存的显卡换成 30GB 的。这种方式不需要改动代码，原本跑不起来的模型现在能跑了或者可以调大 batch size 加快训练速度。

横向扩展：

增加机器数量。比如增加 7 台同配置（10GB）的机器，通过网络互联每台机器可以挂载单卡或多卡。这种方式需要代码层面的适配，利用 PyTorch 的分布式模块进行通信。

数据并行 vs 模型并行

数据并行 (Data Parallelism):

当模型本身能塞进单张 GPU，但数据量太大时，我们可以将模型复制到所有 GPU 上，将数据切分（Split），每个 GPU 并行处理一部分数据，在反向传播时同步梯度。

模型并行 (Model Parallelism):

当模型大到单张 GPU 放不下时使用。将模型切分成不同部分，分配给不同 GPU。每个 GPU 负责计算前向/后向传播中的一部分层。

实际超大模型训练中，通常是两者的混合。

前置概念：梯度累积

在讲 DDP 之前，先回顾一个基础技巧：梯度累积（Gradient Accumulation）。PyTorch 的设计中，

loss.backward()

复制代码

计算出的梯度默认是累加的，而非覆盖。

 import torch  
# Let us define a tensor with requires_grad = True  x = torch.tensor(4.0, requires_grad=True)  # Creating a function y=x^2  y = x*x  # Calculating dy/dx  y.backward(retain_graph=True)  # retain_graph = True because pytorch automatically removes the computation  # graph and intermediate tensors once backward is called to save memory  # If we want to call backward again, we need to tell pytorch not to delete  # the computation graph and intermediate tensors  print(f"Gradient of y w.r.t x after the first backward: {x.grad}")  # Output: 8.0 as dy/dx = 2*x = 2*4  # Now let us try to call backward again  y.backward()  print(f"Gradient of y w.r.t x after the second backward: {x.grad}")   # Output = 16 because 8 from the previous backward is added up here

复制代码

利用这个特性，当大 Batch Size 导致 OOM 时，可以将其切分为多个 Mini-batch，连续执行多次

backward()

复制代码

累积梯度，最后执行一次

optimizer.step()

复制代码

。这是单卡解决显存瓶颈的常用手段。

分布式数据并行 (DDP) 工作流

PyTorch 的

DistributedDataParallel

复制代码

(DDP) 是实现数据并行的核心模块，基于

c10d

复制代码

的 ProcessGroup 进行通信，每个进程（Process）通常控制一个 GPU 并持有一个模型副本。

DDP 的标准执行流程如下：

初始化 ProcessGroup：建立进程间的通信握手。
广播权重（Broadcast）：训练开始时，Rank 0 节点的模型权重被广播到所有其他节点，确保初始状态一致。
前向反向传播：每个节点在自己的数据子集上独立计算。
梯度归约（All-Reduce）：各节点的梯度被汇聚（求和或平均），然后同步回所有节点。
参数更新：各节点使用同步后的梯度独立更新权重。

集合通信原语 (Collective Communication Primitives)

分布式训练中，点对点（Point-to-Point）通信效率太低。假设要把 5MB 参数发给 5 个节点，逐个发送会导致耗时线性增长。集合通信（Collective Communication） 利用拓扑结构（如树状、环状）并行传输，是高性能计算的基石。

创作场景

PyTorch 分布式训练底层原理与 DDP 实战指南

扩展算力的两种路径

数据并行 vs 模型并行

前置概念：梯度累积

分布式数据并行 (DDP) 工作流

郝培强

评论