pytorch中的分布式训练

1.简介

Pytorch的数据分为两种，torch.nn.parallel.DataParallel(DP)和torch.nn.parallel.DistributedDataParallel(DDP).

使用场景：

• DP模式用于单机多卡
• DDP模式可以用于单机多卡、多机多卡以及模型并行。

2. DP模式

DP模式用于单机多卡的情况。实现起来较为简单，只需要用torch.nn.parallel.DataParallel包装模型即可，其余与单机单卡的情况相同。

1.1 流程

1. master(一般是GPU0)从磁盘或者合页内存中取数据。
2. master将数据分到其他GPU上
3. master将模型复制到其他GPU上
4. 每块GPU单独进行前向计算，得到输出
5. master收集每块GPU上的输出，计算损失
6. master将损失分到其他卡上，每块卡单独进行反向传播，计算梯度
7. master收集每块GPU上的梯度，汇总以后，进行reduce操作，结果分发到每块卡上。
8. 每块GPU根据梯度，单独更新模型

1.2 图示流程

可以看出DP模型比较多的操作是在0号卡上进行的。分数据、将模型的副本传到其他模型，收集每块卡的输出、计算loss,将loss传到每块卡上，在每块卡上进行反向传播得到梯度后，收集每块卡上的梯度，进行reduce上操作后，传到其他卡上。这样使得0号卡所占的内存比较多，使得内存使用不均衡，而且会经常出现其他卡等待0好卡计算的情形。

2. Distributed Data Parallel

2.1 处理流程

1. 从master(一般是GPU0)从磁盘或者合页内存中取数据。
2. 所有GPU同时去取数据，不需要GPU0去分
3. 每块GPU单独进行前向计算
4. 每块GPU单计算loss
5. 每块GPU单独进行反向过程，计算出参数的梯度，并进行参数之间的传递(计算和参数传递间存在交叉过程)
6. 在GPU0上进行梯度的allreduce操作，然后将梯度传递到其他GPU上。
7. 每个GPU单独地进行参数更新
   
   可以看出DDP只有在梯度收集的时候，存在GPU间的通信，其余的操作都是在各自的GPU上进行的，这样可以均衡负载，也可以节省时间。

1.3 实现

1. 进程初始化
   backend 可选nccl，gloo和mpi模式，表示不同的collective communication方式。使用英伟达显卡，运行比较快的是nccl方式

local_rank= os.getenv('LOCAL_RANK', -1)
    if local_rank != -1:
        torch.cuda.set_device(local_rank)
        device = torch.device("cuda", local_rank)
        torch.distributed.init_process_group(backend="nccl", init_method='env://')

2. 数据集创建

def create_dataloader(...):

    train_datasets= torchvision.datasets.MNIST(
        root='./mnist',
        train=True,
        transform=torchvision.transforms.ToTensor(),  # (0,255) -->(0,1)
        download=False
    )

    train_sampler = DistributedSampler(train_datasets) ## 防止不同进程交叉取数据
    train_dataloader = DataLoader(train_datasets, sampler=train_sampler, batch_size=int(args.test_batch_size/args.world_size),
                                  num_workers=args.num_workers, pin_memory=True)
    ...
    ....

train_sampler,train_dataloader,test_sampler,test_dataloader = create_dataloader(args,main_process)

3. 使用DDP模型对模型进行包装

model = ...
model.to(device)
num_gpus = torch.cuda.device_count()
if num_gpus > 1:
   print('use {} gpus!'.format(num_gpus))
   model = nn.parallel.DistributedDataParallel(model, device_ids=[args.local_rank],
                                                    output_device=args.local_rank)

4. 模型训练

    for epoch in range(EPOCH):
        # 步骤五：打乱数据顺序
    train_sampler.set_epoch(epoch)
    model.train()
    start= time()        
    ### train
    for step, (b_x, b_y) in enumerate(train_dataloader):
        predit_y = model(b_x)
        loss = loss_func(predit_y, b_y)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

5. 启动方式

• 单机多卡 python -m torch.distributed.run --nproc_per_node 2 --master_port 2393 ddp_train.py
• 多机多卡
  在机器0上运行python -m torch.distributed.launch --nproc_per_node=4 --nnodes=2 --node_rank=0 --master_addr="192.0.0.1" --master_port=1234 ddp.py
  在机器1上运行python -m torch.distributed.launch --nproc_per_node=4 --nnodes=2 --node_rank=1 --master_addr="192.0.0.1" --master_port=1234 ddp.py
  不同机器上的node_rank不同，其余相同

TODO

参考

[1] pytorch(分布式)数据并行个人实践总结——DataParallel/DistributedDataParallel
[2] PyTorch Distributed: Experiences on Accelerating Data Parallel Training
[3] 多机多卡