当前位置：首页 > news >正文

了解ddp和fsdp

news 2025/9/27 6:40:56

DDP（分布式数据并行）是PyTorch中一种常用的分布式训练技术。它的核心思想是数据并行，即在多个GPU上（可以是单机多卡或多机多卡）分别保留一份完整的模型副本。[3][4]

其工作流程如下：

数据切分：将整个数据集分成多份，每个GPU（进程）拿到一部分数据进行独立的计算。[1][3][5]
独立计算：每个GPU上的模型副本独立进行前向传播和反向传播，计算出梯度。[3][6]
梯度同步：在反向传播之后，通过All-Reduce操作，将所有GPU上的梯度进行聚合和平均，保证所有模型副本最终以相同的梯度来更新参数。[3][6][7]
参数更新：每个GPU上的优化器使用同步后的梯度独立地更新自己的模型参数，从而确保所有副本在下一次迭代开始前保持一致。

简单来说，DDP就是“模型复制，数据切分，梯度同步”。

FSDP（完全分片数据并行）是DDP的演进版本，旨在解决DDP在训练超大规模模型时遇到的内存瓶颈问题。[7][8][9] 它的核心思想是对模型参数、梯度和优化器状态进行分片（Shard），每个GPU只保存模型的一部分，从而极大降低单个GPU的内存占用。[7][10][11][12]

其工作流程要复杂一些：

分片 (Sharding)：在训练开始前，FSDP会将模型的参数 (parameters)、梯度 (gradients) 和优化器状态 (optimizer states) 全部分片，并将这些分片均匀地分配给集群中的所有GPU。[7][10][13] 每个GPU只持有自己负责的那一小部分。
前向传播 (Forward Pass)：
- 当需要计算某一层时，FSDP会使用All-Gather操作，从其他所有GPU收集该层所需的参数分片，临时在当前GPU上重构出完整的层。[7][14][15]
- 完成该层的计算后，立即释放掉从其他GPU收集来的参数分片，以节省内存。[15][16]
反向传播 (Backward Pass)：
- 与前向传播类似，反向传播时也会通过 All-Gather 重新构建完整的层来计算梯度。[6][14]
- 梯度计算完成后，使用Reduce-Scatter操作将梯度进行聚合，并将聚合后的梯度分片发送回其对应的GPU上。这样每个GPU只保留自己负责那部分参数的梯度。[7][13]
参数更新：每个GPU上的优化器只更新其本地持有的那一小部分参数和优化器状态。