PyTorch学习（3）：并行训练DataParallel与DistributedDataParallel

1. 简述

在使用pytorch训练网络时，一般都会使用多GPU进行并行训练，以提高训练速度，一般有单机单卡，单机多卡，多机多卡等训练方式。这就会使用到pytorch提供的DataParallel(DP)和DistributedDataParallel(DDP)这两个函数来实现。

2. DP和DDP的区别

DP是使用一个进程来计算模型参数，然后在每个批处理的数据将分发到每个GPU，然后每个GPU计算各自的梯度，然后汇总到GPU0中进行求平均，由GPU0进行反向传播更新参数，然后再把模型的参数由GPU0传播给其他的GPU，GPU利用率通常很低。

DDP是数据并行的分布式，是同时使用多个进程，每个GPU上一个进程，数据也是被进程数等分，相当于每个GPU上都跑了一份代码，前向之后再经过all reduce的处理，再经过梯度反向传播，更新参数。

DDP即可以做单机多卡，也可以多机多卡，但DP只能是单机多卡，DDP即支持一个进程占一张GPU，也支持一个进程上占多张GPU，而DP只支持一个进程占一个或多个GPU。可以说DDP是支持DP的所有功能的。

3. DP的使用

DataParallel更易于使用，只需简单包装单GPU模型，DP的batchsize是每张GPU上的数目乘GPU数目，这个和DDP是不一样的。

model = MODEL()  # 创建模型，需修改

model = model.cuda()

num_gpus = torch.cuda.device_count()

if num_gpus > 1:

    logger.info('use {} gpus!'.format(num_gpus))

    model = nn.DataParallel(model)

通过watch nvidia-smi查看GPU信息，可以看到使用的每张GPU的PID是一样的，如果是4张卡，4个GPU的PID都是一样的。

4. DDP的使用

DDP的使用，相对来说要复杂一些，启动方式也有多种，本文仅通过torch.distributed.launch启动。此外还有一个spawn启动方式，本文不做解释。

（1）认识torch.distributed.launch

torch.distributed.launch是PyTorch 分布式训练工具的一部分，它用于启动多进程训练。这个工具可以帮助用户在多个进程之间同步数据，并且可以很容易地在多个 GPU 或多个节点上运行分布式训练。总结起来，有如下几个功能：

# 多进程支持：torch.distributed.launch 能够启动多个进程，每个进程可以运行在不同的 GPU 上。

# 环境设置：它可以用来设置训练所需的环境变量，如MASTER_ADDR，MASTER_PORT，WORLD_SIZE，LOCAL_RANK和 RANK。

这些环境变量的解释如下

WORLD_SIZE: 通俗的解释下，就是一共有多少个进程参与训练，WORLD_SIZE = nproc_per_node*nnodes,不同的进程中，WORLD_SIZE是唯一的；

RANK：进程的唯一表示符，不同的进程中，这个值是不同的，上述在AB两台机器上共启动了8个进程，则不同进程的RANK号是不同的(又称为)

LOCAL_RANK：同一节点下，LOCAL_RANK是不同的，常根据LOCAL_RANK来指定GPU，但GPU跟LOCAL_RANK不一定一一对应，因为进程不一定被限制在同一块GPU上。[LOCAL_RANK一般用在torch.cuda.set_devices和torch.nn.parallel.DistributedDataParallel作为参数]

MASTER_ADDR：主设备的IP，我们在执行多机多进程训练的时候，需要有一台机器用作主设备，其他从设备计算得到的梯度都会汇集到当前主设备中进行反向传播计算，并将更新后的权重统一分发各从设备，完成训练参数的更新。

MASTER_PORT：主设备的通信端口。

# 灵活性：用户可以指定每个进程使用的 GPU 设备，以及每个节点运行多少个进程。

（2）启动方式

以两台设备，每台设备有两个GPU为例（一般一个GPU为一个训练进程），我们有一个训练入口脚本叫做train.py，然后带几个参数，典型的启动方式如下：

启动设备A：

python -m torch.distributed.launch

--nproc_per_node=[NUM_GPUS_PER_PC]

--nnodes=[NUM_OF_PCs]

--node_rank=0

--master_addr="192.168.1.1"

--master_port=1234

train.py [arg list]

启动设备B：

python -m torch.distributed.launch

--nproc_per_node=[NUM_GPUS_PER_PC]

--nnodes=[NUM_OF_PCs]

--node_rank=1

--master_addr="192.168.1.100"

--master_port=1234

train.py [arg list]

如上，NUM_GPUS_PER_PC表示服务器有几个GPU，这里为2，NUM_OF_PCs知名有几台训练服务器。

此时，我们查看设备中的环境变量，可以发现新增了几个环境变量。

A设备

environ({..., 'MASTER_ADDR': '10.100.37.21', 'MASTER_PORT': '29500', 'WORLD_SIZE': '8', 'OMP_NUM_THREADS': '1', 'RANK': '0', 'LOCAL_RANK': '0'})

environ({..., 'MASTER_ADDR': '10.100.37.21', 'MASTER_PORT': '29500', 'WORLD_SIZE': '4', 'OMP_NUM_THREADS': '1', 'RANK': '1', 'LOCAL_RANK': '1'})

B设备

environ({..., 'MASTER_ADDR': '10.100.37.21', 'MASTER_PORT': '29500', 'WORLD_SIZE': '4', 'OMP_NUM_THREADS': '1', 'RANK': '2', 'LOCAL_RANK': '0'})

environ({..., 'MASTER_ADDR': '10.100.37.21', 'MASTER_PORT': '29500', 'WORLD_SIZE': '4', 'OMP_NUM_THREADS': '1', 'RANK': '3', 'LOCAL_RANK': '1'})

（3）init_process_group

init_process_group 是 PyTorch 分布式训练中的一个关键函数，用于初始化所有参与训练的进程之间的通信。这个函数设置了后端通信库、初始化进程组，并为每个进程分配一个唯一的标识符（rank）。

init_process_group 函数接受多个参数，以下是一些常见的参数：

backend: 指定后端通信库。常用的后端有 "nccl"（NVIDIA Collective Communications Library，用于NVIDIA GPUs），"gloo"（PyTorch的默认后端，适用于CPU和GPU），和 "mpi"（消息传递接口）。

init_method: 指定初始化方法，用于在不同进程间建立连接。例如，可以使用 "env://" 从环境变量中读取初始化信息（需要在启动时指定“--master_addr”和“--master_port”），或者使用 "tcp://<master_ip>:<master_port>" 指定一个主节点的IP和端口。

world_size: 指定进程组中的进程总数。

rank: 指定当前进程的排名（从0开始）。

group_name: 指定进程组的名称，通常在动态分配资源时使用。

需要注意的是，如果再启动训练时，并没有指定“--master_addr”和“--master_port”,那么在init_process_group的参数中需要指定。

dist.init_process_group(backend="nccl",
                        init_method="tcp://{}:{}".format(args.master_addr, args.master_port),
                        rank=rank,
                        world_size=world_size)

（4）训练工程（解析参数，设定）

用launch方式需要注意的位置：

需要添加解析 local_rank，rank，world_size，master_addr和master_port的参数：

dist初始化的方式 int_method取env：

dist.init_process_group("gloo", init_method='env://')

DDP的设备都需要指定local_rank

net=torch.nn.parallel.DistributedDataParallel(net,device_ids=[args.local_rank], output_device=args.local_rank)

简单参考代码，train.py

def main():

    # 步骤一: 获取环境变量

    parser = argparse.ArgumentParser()

    ...

    parser.add_argument("--local_rank", default=os.getenv('LOCAL_RANK', -1), type=int)

    parser.add_argument("--rank", default=os.getenv('RANK', -1), type=int)

    parser.add_argument("--world_size", default=os.getenv('WORLD_SIZE', -1), type=int)

    parser.add_argument("--master_addr", default=os.getenv('MASTER_ADDR', -1), type=int)

    parser.add_argument("--master_port", default=os.getenv('MASTER_PORT', -1), type=int)

    parser.add_argument("--train_batch_size", default=1, type=int) # 此次的batchsize是每个进程上的数目，不是总数目

    parser.add_argument("--num_workers", default=1, type=int)

    parser.add_argument("--epochs", default=1000, type=int)

 

    args = parser.parse_args()


 

    # 步骤二: 初始化

    if args.local_rank != -1:

        torch.cuda.set_device(args.local_rank)

        device = torch.device("cuda", args.local_rank)

        torch.distributed.init_process_group(backend="nccl", init_method='env://')

 

    # 步骤三: 模型分布式处理

    model = MODEL()  # 创建模型，需修改

 

    loss_fn = nn.MSELoss()  # loss函数

    optimizer = optim.SGD(ddp_model.parameters(), lr=0.0001)  # 优化器

 

    #记住要先放在device上再进行DistributedDataParallel, DistributedDataParallel需带参数device_ids和output_device

    model.to(device)

    num_gpus = torch.cuda.device_count()

    if num_gpus > 1:

        logger.info('use {} gpus!'.format(num_gpus))

        model = nn.parallel.DistributedDataParallel(model, device_ids=[args.local_rank], output_device=args.local_rank)

 

    # 步骤四: 定义数据集

    train_datasets = ...  # 自己定义的Dataset子类, 需修改

    train_sampler = DistributedSampler(train_datasets)

    # 使用了sampler, shuffle不能为true

    train_dataloader = DataLoader(train_datasets, sampler=train_sampler, batch_size=args.train_batch_size,num_workers=args.num_workers, pin_memory=True)

 

    # 进行前向后向计算

    for epoch in args.epochs:

        # 步骤五: 打乱顺序, 相当于shuffle=TRUE

        train_sampler.set_epoch(epoch)

        for batch in train_dataloader:

            input, label = batch[:2]

            input = input.cuda()

            label = label.cuda()

 

            optimizer.zero_grad()

 

            output = model(input)

            loss = loss_fn(label, output)

            loss.backward()

            optimizer.step()

 

 

if __name__ == "__main__":

    main()

5. 其他事项

使用DDP时，保存模型或者打印log时，如果不加限制，是会有多份的，为了保证只存一份模型，可以用rank来指定一个主进程保存（一般将rank0作为主进程）。

以保存CKPT为例：

if torch.distributed.get_rank() == 0:
    torch.save(model, "last.pth")