超详细讲解：DP和DDP的区别以及使用方法

前言：最近在跑大模型分类，参数全部冻结，用两张A100，发现48层的qwen2.5带不起来，痛定思痛，原来是我用了DP的原因

浅谈DP与DDP的区别

我们一般谈到DP和DDP，他们之间最大的区别就是

DP只支持单机多卡，而DDP既可以支持单机多卡又能支持多机多卡

一般单机多卡跑模型，数据量不大的情况下，其实二者感觉不出来有什么差别，但是数据量一大你会发现，诶，我明明用了DP来进行数据并行了，怎么还是OOM，这是因为DP在进行数据并行的时候各个GPU分摊是不均衡的，而且，你如果指定多张卡的话，一般需要指明哪个是主卡，所以我们要了解DP和DDP是怎么进行数据并行的

Data Parallelism(DP)

DP的执行过程：

准备工作： 假设有N块GPU，其中第0块作为主GPU，一个batch的数据会先被放到主GPU上，且当前主GPU上有一份完整的网络模型
数据分配： 输入的batch数据会被分割成N份，然后由主GPU分发到另外N-1个GPU上，此时每个GPU都有一份数据
模型复制： 整个模型也会被复制到其他的GPU上，这意味着每张GPU都有一份完整的模型副本
前向传播： 每个GPU独立的进行前向传播（线程并行），计算各自的预测输出
汇总输出、计算损失： 主GPU会收集（gatter）各个GPU的输出，并在主GPU上计算损失
反向传播： 将计算之后的损失再分发（scattered）到对应的GPU上，每个GPU进行独立的反向传播，计算各自的梯度
汇总梯度： 每个GPU计算完梯度后，这些梯度会被汇总（reduce）到主GPU上，并在主GPU求梯度平均值
参数更新： 主GPU使用计算出的平均梯度更新模型的参数
重复上述过程，当执行到下一轮的"第三步：模型复制"时，主GPU上更新后的模型参数会同步给其他所有的GPU，从而保证模型参数的一致性

DP的弊端：

DP在单机器上使用多个线程来处理并行计算，这会导致线程之间争抢GIL (GlobaL Interpreter Lock全局解释器锁)，从而引发性能瓶颈，因为每个线程在执行Python代码时都需要获取GIL，这会导致线程不能高效地并行运行。
DP需要将输入数据分发（scatter ）到不同的设备上进行计算，然后再收集（gather）各个设备上的输出结果，这些scatter 和gather操作会带来额外的开销。
每次迭代，都需要为每个GPU复制一份模型，从而保证介GPU上的模型参致是一致的，这种模型复制操作也会带来额外的开销。

还有很重要的一点就是，你会发现上述基本上所有的操作都是在主GPU上进行的，这样就会导致主GPU已经超负荷了，而其他GPU还没有得到完全利用！！！

虽然DP用起来简单，emmmm用的时候还是三思而后行

DP的用法：

（1）设置主GPU

（2）封装模型

（3）将模型放到对应的device上

python 复制代码

# 当你有多张卡的时候，一定要指明哪一张是主卡，这里指明GPU0是主卡
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
if torch.cuda.device_count() > 1:
	# 用DP封装model
    model = torch.nn.DataParallel(model, device_ids=[0, 1])
# 将模型放到对应的device上
model.to(device)

保存模型的时候也要注意，我们要保存的是主GPU上的模型，主卡上的模型通过 model.module 获取

python 复制代码

# 保存主卡上的模型参数
torch.save(model.module.state_dict(), 'main_model.pth')

Distributed Data Parallelism(DDP)

DDP的执行过程：

模型复制： 整个模型会被复制到每张GPU上，这意味着每张GPU上都有一个完整的模型副本
数据分配： 通过DistributedSampler，每个进程独立的加载自己那部分的数据子集
前向传播： 每个GPU独立的进行前向传播，计算各自的预测输出
反向传播： 每个GPU独立的计算损失，进行反向传播，计算梯度
梯度汇总： 所有GPU通过Ring-Allrecude操作汇总梯度，使得每个GPU上的梯度一致
参数更新： 每个GPU使用汇总后的梯度求平均，更新自己的模型参数，更新后，每个GPU上的模型参数都是一致的
重复上述第二步之后的操作

与DP相比，DDP的优势在于：

DDP使用多进程处理，每个GPPU都有自己的一个专用的进程，避免了Python解释器GIL带来的性能开销
每个进程独立更新模型参数，减少了通信开销，提高了训练效率，且梯度汇总时，使用的是Ring-Allrecude，进一步提升了效率（关于Ring-Allrecude不理解的朋友，大家可以参考，我觉得这位博主说的非常清楚：https://juejin.cn/post/7084135971687497759）
具有更好的扩展性，即支持单机多卡，有支持多机多卡

DDP的用法：

（1）设置环境变量：如果在多节点（多台机器）上进行训练，你需要设置如下的环境变量：

MASTER_ADDR: 主节点的 IP 地址
MASTER_PORT: 主节点的端口号
WORLD_SIZE: 参与训练的进程总数（通常是 GPU 总数）
RANK: 当前进程的排名（进程编号，从 0 开始）

如果只是在单机多卡上使用 DDP，PyTorch 会自动处理这些。

（2）初始化进程组 ：在使用 DDP 前，通过torch.distributed.init_process_group()来初始化一个进程组（process group），用于在进程间进行通信。初始化进程组之后使用torch.cuda.set_device(rank)将进程与GPU进行绑定，即指定当前GPU应该使用哪个进程，设置之后，rank=0被绑定到GPU0，rank=1被绑定到GPU1

（3）封装模型 ：将模型封装进 torch.nn.parallel.DistributedDataParallel。在多进程中，每个进程需要独立初始化一个模型副本。

（4）使用 DistributedSampler 由于不同 GPU 处理不同的数据子集，必须使用 torch.utils.data.DistributedSampler 来确保每个 GPU 处理的数据是不同的

（5）用完之后要记得销毁进行组torch.distributed.destroy_process_group()

示例：

python 复制代码

import os
import torch
import torch.distributed as dist
import torch.multiprocessing as mp
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.utils.data import DataLoader, DistributedSampler
from torchvision import datasets, transforms

# 定义模型（示例用简单的线性模型）
class MyModel(torch.nn.Module):
    def __init__(self):
        super(MyModel, self).__init__()
        self.fc = torch.nn.Linear(10, 10)

    def forward(self, x):
        return self.fc(x)

# 训练函数
def train(rank, world_size):
    # 初始化进程组
    # 常见的通信后端包括 gloo 和 nccl，其中 nccl 是在 GPU 上最常用的后端
    dist.init_process_group(backend='nccl', init_method='env://', world_size=world_size, rank=rank)

    # 设置当前 GPU
    torch.cuda.set_device(rank)  # 若是在多机上，此处设置为local_rank

    # 创建模型并转移到当前 GPU
    model = MyModel().cuda(rank)  # 若是在多机上，此处设置为local_rank

    # 将模型包装为 DDP 模型
    model = DDP(model, device_ids=[rank])	# 若是在多机上，此处设置为local_rank

    # 定义损失函数和优化器
    criterion = torch.nn.CrossEntropyLoss().cuda(rank)	# 若是在多机上，此处设置为local_rank
    optimizer = torch.optim.SGD(model.parameters(), lr=0.01)

    # 加载数据集（示例使用 CIFAR10）
    transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5,), (0.5,))])
    train_dataset = datasets.CIFAR10(root='./data', train=True, download=True, transform=transform)

    # 使用 DistributedSampler 确保每个 GPU 得到不同的数据
    train_sampler = DistributedSampler(train_dataset, num_replicas=world_size, rank=rank)	# 若是在多机上，此处设置为local_rank
    train_loader = DataLoader(dataset=train_dataset, batch_size=32, sampler=train_sampler)

    # 开始训练
    for epoch in range(10):
        model.train()
        train_sampler.set_epoch(epoch)  # 每个 epoch 设置随机种子以保证不同进程之间的数据一致
        for batch_idx, (data, target) in enumerate(train_loader):
            data, target = data.cuda(rank), target.cuda(rank)	# 若是在多机上，此处设置为local_rank
            optimizer.zero_grad()
            output = model(data)
            loss = criterion(output, target)
            loss.backward()
            optimizer.step()

    # 销毁进程组
    dist.destroy_process_group()

# 主函数
def main():
    world_size = 2  # 使用 2 张 GPU

    # 使用 multiprocessing 启动多进程，每个 GPU 对应一个进程
    mp.spawn(train, args=(world_size,), nprocs=world_size, join=True)

if __name__ == "__main__":
    main()

checkpoint的保存与加载：

如果需要先保存checkpoint再执行加载操作，要使用torch.distributed.barrier()，等到保存完成之后再进行加载
在 DDP 中，每个进程都拥有一个模型副本，通常只需在主进程（rank == 0）上保存模型即可，避免每个进程都进行保存

python 复制代码

if rank == 0:  # 只在主进程保存模型，因为每个GPU上都是一样的，我们保存主进程的就可以了
        torch.save(model.module.state_dict(), filepath)
        print(f"Model saved to {filepath}")
        
torch.distributed.barrier() # 如果保存之后就进行加载的话，这里要使用torch.distributed.barrier()，确保都保存完毕再进行加载

model.load_state_dict(torch.load(filepath))  # 加载模型

启动训练：

--nproc_per_node指定GPU数量

python 复制代码

torchrun --nproc_per_node=4 script.py

参考

https://www.bilibili.com/video/BV1cJ4m1g7sG/?spm_id_from=333.337.search-card.all.click\&vd_source=9b5bcbc75b2a02599faa4bac703679a1