数据并行｜显存 “减负” 从 DDP 到 FSDP

数据并行简单易用，是大模型训练基础方式。DDP在每个GPU上拷贝一份模型状态，大量的显存被冗余状态占用。FSDP借鉴零冗余优化器的思想对模型状态进行分片后，由不同的GPU负责维护，使用数据并行训练超大规模的模型成为可能。

1）数据并行中，训练设备配置和数据分发机制介绍。 2）反向计算梯度后怎样进行梯度同步，使计算和通信耗时重叠。 3）FSDP在数据并行基础上，分片模型状态优化内存占用。

1，基本概念DDP

1.1，数据并行方式

每个计算的设备都有整个模型的副本，根据整个数据集的一个子集进行前向计算。假设训练样本为D，使用N个加速卡进行计算，每个计算设备会分配到 个样本，并使用这些样本分批计算梯度 ，并将梯度进行广播，所有计算设备聚合到其他N-1个梯度后，使用平均梯度 进行模型参数更新。在实际模型训练中一般采用分布式数据并行 DDP (Distributed Data Parallel)。

1.2，数据切片

假设训练设备有：2个主机 (node)，每个主机上有3张GPU计算设备，共6张卡。名词解释：

• • Host：可以理解为一台主机，每个主机有自己的IP地址，用于通信。
• • Local Rank：每个主机上，对不同GPU设备的编号。
• • Global Rank：全局的GPU设备编号，Global Rank = Host * num GPUs per host + Local Rank。
• • Worldsize：总的GPU个数。num Hosts * num GPUs per host

数据分配代码：

indices = indices[self.rank:self.total_size:self.num_replicas]
"""
self.rank 即 Global Rank，
self.total_size，总的数据量。
self.num_replicas，即worldsize。
"""

在每个host上运行启动命令，建立DDP组，等待每个GPU设备测试通信正常后，即可开始训练。

图1，数据并行设备介绍

2，DDP分布式数据并行

2.1，DDP参数更新过程

分布式训练需要同步每个GPU设备上的梯度，获得一致的参数更新。具体步骤如下：

• • 加载数据到内存，并为每个设备分配不同的数据片。
• • 数据加载到GPU的内存中。
• • 每个GPU设备开始独立的进行前向计算，获得损失。
• • 反向传播计算梯度。并且需要进行梯度的 all-reduce 通信，使得每个GPU设备获得平均梯度。
• • 更新参数，所有的GPU设备上的模型副本获得一致性的参数更新。

图2，分布式数据并行训练过程。

在第四步时进行梯度的all-reduce 操作，将所有GPU设备上的梯度进行求和取平均。使用ring-all-reduce算法，通信量为两倍的参数量。

图3，集合通信all-reduce操作示意图

2.2，并行梯度

梯度依赖关系：反向传播的梯度计算，天然具有层间独立性（如第L层梯度不依赖第L-1层），使得流水线设计可行。

目标：通过梯度流水线化，将梯度全归约的通信时间隐藏在计算梯度的时间内，减少整体训练耗时。理想情况下，通信时间被完全覆盖，训练速度仅由计算时间决定。

• • 梯度归约串行 设备1：计算L层梯度 → 计算L-1层梯度 →...→ 全归约L层梯度 → 全归约L-1层梯度→.… 设备2：计算L层梯度 → 计算L-1层梯度 →...→ 全归约L层梯度 → 全归约L-1层梯度→.…
• • 梯度归约并行 全归约操作在梯度计算完成后立即触发，与其他层的梯度计算并行执行。 设备1：计算L层梯度→[全归约梯度并行 L层 ]→计算L-1层梯度→[全归约梯度 L-1层]→... 设备2：计算L层梯度→[全归约梯度并行 L层]→计算L-1层梯度→[全归约梯度并行 L-1层]→...

小结：DDP通过分布式多进程设计、去中心化梯度同步、计算与通信重叠等优化，显著解决了数据并行的冗余拷贝、线程开销、主GPU瓶颈等问题，适合大规模分布式训练场景。

3，DDP更进一步！FSDP

DDP数据并行中每个GPU设备都拷贝了一份模型状态副本，冗余的拷贝使得显存开销巨大。FSDP (Fully Sharded Data Parallelism) 是一种分布式训练技术，通过分片(Sharding) 模型参数、梯度和优化器状态，将大型模型的训练负载分散到多个GPU或计算节点上，从而解决传统数据并行在训练超大模型时的内存瓶颈问题。

其核心思想源自 ZeRO (Zero Redundancy Optimizer) 优化器，通过分片模型状态，消除内存冗余实现内存高效利用。

图4，数据并行中FSDP训练流程。

4，总结

参考： [1] blog: Distributed data parallel training using Pytorch on AWS Telesens