DDP、pytorch的分布式 torch.distributed.launch 训练说明

torch.distributed.launch 是PyTorch提供的一个工具，用于启动多个Python进程以运行分布式训练。这里提供一个简单的使用示例：

假设你有一个名为 train.py 的训练脚本，你想用4个GPU进行分布式训练。

首先，在命令行中使用以下命令来启动分布式训练：

python -m torch.distributed.launch --nproc_per_node=4 train.py

在 train.py 文件中，你需要初始化进程组，并在进程组内部的每个进程上设置适当的区域，如下所示：

import torch
import torch.distributed as dist
import torch.multiprocessing as mp
 
def main():
    # 启动多进程
    mp.spawn(train_fn, nprocs=4, args=(args,))
 
def train_fn(gpu, args):
    # 每个进程的设备ID
    rank = gpu
    # 初始化进程组
    dist.init_process_group(
        "nccl", 
        init_method='tcp://localhost:23456', 
        rank=rank, 
        world_size=4)
    torch.cuda.set_device(rank)
    
    # 模型和训练代码
    model = Net().cuda(rank)
    # ... 训练代码 ...
 
if __name__ == "__main__":
    main()

在这个例子中，mp.spawn 会为每个GPU启动一个进程，dist.init_process_group 会初始化分布式进程组，并且每个进程都会使用 torch.cuda.set_device 来设置它的CUDA设备。然后，你可以在模型和训练循环中使用分布式数据并行来分配数据和计算。

注意：init_method='tcp://localhost:23456' 指定了初始化通信所使用的TCP地址和端口。在实际应用中，这个地址需要根据实际网络环境进行调整。