资讯

历史

科技

环境与自然

成长

游戏

财经

文学与艺术

美食

健康

家居

文化

情感

汽车

三农

军事

旅行

运动

教育

生活

星座命理

大模型分布式训练技术：DP、DDP和FSDP详解

创作时间:

作者:

@小白创作中心

大模型分布式训练技术：DP、DDP和FSDP详解

引用

CSDN

https://blog.csdn.net/yangtuoi/article/details/141186671

随着大模型时代的到来，分布式训练技术成为训练大规模模型的关键。本文将深入探讨三种主要的并行训练技术：数据并行（DP）、分布式数据并行（DDP）和完全分片数据并行（FSDP），并对比它们的优缺点和应用场景。

数据并行（PyTorch DP）

数据并行（torch.nn.DataParallel）是PyTorch最早提供的一种数据并行方式，它基于单进程多线程实现。在每个批处理期间，它将数据分发到每个GPU，并在主GPU上计算模型权重。

计算过程

将输入数据从主GPU分发到所有GPU。
将模型从主GPU分发到所有GPU。
每个GPU分别独立进行前向传播，得到输出结果。
将每个GPU的输出结果发回主GPU。
在主GPU上，通过损失函数计算出损失，并对损失函数求导，得到损失梯度。
将计算得到的梯度分发到所有GPU。
反向传播计算参数梯度。
将所有梯度回传到主GPU，通过梯度更新模型权重。
不断重复上述过程。

使用方法

net = torch.nn.DataParallel(model, device_ids=[0, 1, 2])
output = net(input_var)  # input_var can be on any device, including CPU

缺点

单进程多线程带来的性能开销，速度较慢。
受限于全局解释器锁（GIL），只能在单台服务器（单机多卡）上使用。
不能使用Apex进行混合精度训练。
主卡性能和通信开销容易成为瓶颈，GPU利用率通常很低。
不支持模型并行。

分布式数据并行（PyTorch DDP）

分布式数据并行（torch.nn.DistributedDataParallel）基于多进程实现，每个进程都有独立的优化器，执行自己的更新过程。进程之间只传递梯度，这样网络通信就不再是瓶颈。

具体流程

首先将rank=0进程中的模型参数广播到进程组中的其他进程。
然后，每个DDP进程都会创建一个local Reducer来负责梯度同步。
在训练过程中，每个进程从磁盘加载batch数据，并将它们传递到其GPU。每个GPU都有自己独立的前向过程，完成前向传播后，梯度在各个GPUs间进行All-Reduce，每个GPU都收到其他GPU的梯度，从而可以独自进行反向传播和参数更新。
同时，每一层的梯度不依赖于前一层，所以梯度的All-Reduce和后向过程同时计算，以进一步缓解网络瓶颈。
在后向过程的最后，每个节点都得到了平均梯度，这样各个GPU中的模型参数保持同步。

DDP后端的通信由多种CPP编写的协议支持，不同协议具有不同的通信算子的支持，在开发中可以根据需求选择。

DataParallel是将梯度reduce到主卡，在主卡上更新参数，再将参数broadcast给其他GPU，这样无论是主卡的负载还是通信开销都比DDP大很多。

使用示例

import torch
import torch.distributed as dist
import torch.multiprocessing as mp
import torch.nn as nn
import torch.optim as optim
from torch.nn.parallel import DistributedDataParallel as DDP

def example(rank, world_size):
    # create default process group
    dist.init_process_group("gloo", rank=rank, world_size=world_size)
    # create local model
    model = nn.Linear(10, 10).to(rank)
    # construct DDP model
    ddp_model = DDP(model, device_ids=[rank])
    # define loss function and optimizer
    loss_fn = nn.MSELoss()
    optimizer = optim.SGD(ddp_model.parameters(), lr=0.001)
    # forward pass
    outputs = ddp_model(torch.randn(20, 10).to(rank))
    labels = torch.randn(20, 10).to(rank)
    # backward pass
    loss_fn(outputs, labels).backward()
    # update parameters
    optimizer.step()

def main():
    world_size = 2
    mp.spawn(example,
        args=(world_size,),
        nprocs=world_size,
        join=True)

if __name__=="__main__":
    # Environment variables which need to be
    # set when using c10d's default "env"
    # initialization mode.
    os.environ["MASTER_ADDR"] = "localhost"
    os.environ["MASTER_PORT"] = "29500"
    main()

完全分片数据并行（PyTorch FSDP）

完全分片数据并行（torch.distributed.fsdp.FullyShardedDataParallel）是PyTorch最新的数据并行方案，在1.11版本引入的新特性，目的主要是用于训练大模型。它打破了模型分片的障碍（包括模型参数，梯度，优化器状态），同时保持了数据并行的简单性。

工作原理

FSDP是一种新型数据并行训练方法，与传统的数据并行不同，传统的数据并行维护模型参数、梯度和优化器状态的每个GPU副本，而FSDP将所有这些状态跨数据并行工作线程进行分片，并且可以选择将模型参数分片卸载到CPU。

下图显示了FSDP如何在2个数据并行进程中工作流程：

通常，模型层以嵌套方式用FSDP包装，因此，只有单个FSDP实例中的层需要在前向或后向计算期间将完整参数收集到单个设备。计算完成后，收集到的完整参数将立即释放，释放的内存可用于下一层的计算。通过这种方式，可以节省峰值GPU内存，从而可以扩展训练以使用更大的模型大小或更大的批量大小。为了进一步最大化内存效率，当实例在计算中不活动时，FSDP可以将参数、梯度和优化器状态卸载到CPU。

ZeRO DeepSpeed（零冗余优化器）

现有普遍的数据并行模式下的深度学习训练，每一台机器都需要消耗固定大小的全量内存，这部分内存和并不会随着数据的并行而减小，因而，数据并行模式下机器的内存通常会成为训练的瓶颈。这篇论文开发了一种新颖的解决方案Zero Redundancy Optimizer (ZeRO)，主要用于解决数据并行状态下内存不足的问题。

ZeRO通过跨数据并行进程划分模型状态（参数，梯度和优化器状态），而不是复制它们，从而消除了数据并行进程之间的内存冗余。它在训练期间使用动态通信方式，以在分布式设备之间共享必要的状态，以保持数据粒度的计算粒度和通信量。ZeRO支持的数据并行性可以适应任意大小的模型，只要聚合的设备内存（the aggregated device memory）足够共享模型状态即可。

FSDP和ZeRO的异同

FSDP（Fully Sharded Data Parallel）和ZeRO（Zero Redundancy Optimizer）都是用于大规模深度学习模型训练的优化策略，都采用了内存优化，用于解决大型深度学习模型训练中的内存和计算瓶颈，以便在有限的硬件资源上训练更大规模的模型；但它们的焦点和目标略有不同：