PyTorch中怎么实现多GPU分布式训练_采用DDP框架与Python多进程启动脚本-Python教程-PHP中文网

DDP优于DataParallel，因其采用多进程+NCCL实现GPU间高效梯度同步，负载均衡且接近线性加速；而DataParallel受限于单进程多线程架构，存在主卡瓶颈、通信低效与扩展性差等问题。

pytorch中怎么实现多gpu分布式训练_采用ddp框架与python多进程启动脚本

为什么不能直接用 `torch.nn.DataParallel` 而要切到 `DDP`

DataParallel 是单进程多线程，所有 GPU 都挤在主进程里跑，通信走 Python GIL 和 CPU 内存拷贝，GPU 利用率不均、扩展性差——2 卡可能只快 1.3 倍，8 卡甚至不如 4 卡。而 DDP（DistributedDataParallel）是多进程 + NCCL 后端，每卡一个独立进程，梯度 AllReduce 直接走 GPU 显存间通信，吞吐随卡数线性增长更接近理论值。

关键区别在于启动方式：DataParallel 只需包装模型；DDP 必须显式初始化进程组、分发数据、手动管理进程生命周期。

怎么写一个能被 `torch.distributed.launch` 或 `torchrun` 正确调用的训练脚本

核心是把训练逻辑封装进一个可被多次调用的函数（比如 main_worker），并让每个进程自己初始化本地 rank、world size 和 device。

必须在脚本开头解析 local_rank：用 argparse 读取 --local_rank 参数，或从环境变量 LOCAL_RANK 获取（torchrun 默认设这个）
调用 torch.distributed.init_process_group，后端选 'nccl'（Linux+GPU），init_method 用 'env://'（依赖环境变量自动配地址/端口）
模型要先 .to(device)，再包进 torch.nn.parallel.DistributedDataParallel，注意传 device_ids=[local_rank]
DataLoader 的 sampler 必须换为 DistributedSampler，且设 shuffle=True 和 drop_last=True（避免各进程 batch 数不等导致同步卡死）

示例关键片段：

def main_worker(local_rank, args):
    torch.cuda.set_device(local_rank)
    dist.init_process_group(backend='nccl', init_method='env://')
    
    model = MyModel().cuda(local_rank)
    model = torch.nn.parallel.DistributedDataParallel(
        model, device_ids=[local_rank]
    )
    
    train_sampler = torch.utils.data.DistributedSampler(train_dataset)
    train_loader = DataLoader(train_dataset, sampler=train_sampler, ...)

启动命令该用 `torch.distributed.launch` 还是 `torchrun`

torch.distributed.launch 在 PyTorch 1.10+ 已弃用，官方推荐 torchrun（PyTorch 1.9+ 自带）。两者参数高度兼容，但 torchrun 更健壮：自动检测可用 GPU 数、支持故障重启、默认启用 rdzv_backend=c10d 避免 master 节点单点失败。

Secure Python Code Sandbox

安全 Python 代码沙箱：在隔离的沙箱中执行任意 Python 代码，预装 requests、NumPy、pandas 库。返回 stdout、stderr。

下载

单机多卡最简命令：torchrun --nproc_per_node=4 train.py --your-args
多机训练需加 --nnodes=2 --node_rank=0 --master_addr="192.168.1.10" --master_port=29500
别漏掉 --nproc_per_node —— 它决定启多少个子进程，也对应 WORLD_SIZE 的值；local_rank 由 torchrun 自动注入，无需手动传

常见错误：RuntimeError: Address already in use 通常是因为 --master_port 被占，换一个高位端口（如 29501）即可；NCCL operation failed 多因 NCCL 版本与 CUDA 不匹配，优先用 conda 安装 PyTorch（自带适配 NCCL）。

`DDP` 下保存和加载模型要注意什么

只有 rank 0 进程应该保存模型，否则所有进程同时写同一个文件会冲突或覆盖。加载时所有进程都要从同一份 checkpoint 读，但注意：保存的是 model.module.state_dict()，不是 model.state_dict()（后者含 DDP 封装层的冗余 key）。

保存时加 if dist.get_rank() == 0: 判断
加载时统一用 torch.load(..., map_location=f'cuda:{local_rank}')，再调 model.module.load_state_dict(...)
如果保存了优化器状态，也要同样只由 rank 0 保存，各进程加载时都 load（因为优化器状态不含分布式信息）

容易踩坑：用 torch.save(model.state_dict(), ...) 保存后，加载时报 Missing key(s) in state_dict —— 因为 DDP 模型的 key 带 module. 前缀，而普通模型没有。要么保存前去掉前缀，要么加载时用 model.module.load_state_dict()。

DDP 的正确性高度依赖进程间严格同步，哪怕一个进程早 exit、一个 dataloader 报错没 catch、或者某张卡 out of memory 没 properly cleanup，都可能导致其他进程 hang 在 barrier 上。上线前务必在小数据集上跑通完整生命周期：init → train → save → cleanup。