深度学习多卡训练

Kevin 吴嘉文大约 4 分钟

单机单卡

梯度累加

单机单卡下可以通过梯度累加来增大训练时候的理论 batch size。如：

for i, (inputs, labels) in enumerate(training_set):
  loss = model(inputs, labels)                   
  loss = loss / accumulation_steps                
  loss.backward()                                 
  if (i+1) % accumulation_steps == 0: 
      global_step += 1
      optimizer.step()                            
      optimizer.clean_grad()

假设目标复现论文采用 batch size=64，lr=4e-5。复现过程中，在单机单卡上训练时的batch size=16。此时有两种方案达到论文复现效果：

方案一：小batch=16，小学习率 lr=1e-5，高频率 accumulation_steps=1。这种情况下，假设每个样本平均损失为 m ，并且我们对每个 batch 的 loss 取平均，那么 64 个样本训练完后，模型得以更新的损失为 64/16 * 1e-5 * m= m*4e-5
方案二：小 batch=16，大学习率 lr=4e-5，低频率accumulation_steps=4。这种情况下，每个 step 处理 16 个样本后，得到的损失为 loss/accumulation_steps=m/4 * 4e-5；每 64 个样本提供的损失为 4e-5 * m

混合精度训练

混合精度训练可以为计算密集型任务带来效率上的提升，若任务为 IO 密集型（Batch 小），那么提升注定是有限的。

amp 库下采用混合精度训练，大致流程为：

FP32 权重 -> FP16 权重 -> FP16 计算前向 -> FP32 的 loss，扩大 -> 转为 FP16 -> FP16 反向计算梯度 -> 缩放为 FP32 的梯度更新权重

一文搞懂神经网络混合精度训练open in new window

单机多卡

多 GPU 操作总结

以下为单卡多 GPU，数据并行计算总结

启动训练：

python -m paddle.distributed.launch --gpus=0,1,2,3 run_train.py

代码中调用训练函数：

可以通过 spawn 启动，

n_gpu = len(os.getenv("CUDA_VISIBLE_DEVICES", "").split(","))
paddle.distributed.spawn(self.train, nprocs=n_gpu)

在训练中需要初始化平行训练

if paddle.distributed.get_world_size() > 1:
    paddle.distributed.init_parallel_env()

模型加载 model = paddle.DataParellel(model)，在模型加载好权重之后运行。torch 下采用 DDP 加载。

模型加载权重方式不变，同样可以采用 state_dict 方式加载

if paddle.distributed.get_world_size() > 1:
    paddle.distributed.init_parallel_env()

distrubutedbatchsampler 在多进程启动的时候，会将数据集根据进程数量进行划分。

state_dict = emb.state_dict()
paddle.save(state_dict, "paddle_dy.pdparams")

para_state_dict = paddle.load("paddle_dy.pdparams")
emb.set_state_dict(para_state_dict)

在并行训练下，每个 dataloader 的大小都会变为 total_samples/n_gpu 因此总的训练 step 也会相应的减少，对于 lr_scheduler 会有所影响。

lr_scheduler = paddlenlp.transformers.LinearDecayWithWarmup(
    learning_rate=self.learning_rate,
    total_steps=num_train_steps,
    warmup=int(num_train_steps * self.warmup_proportion))

进程间通信

进程间通信包括 all_gather, all_reduce 等方法。

如以下通过 all_gather 方式通信后，每个卡上会有一份所有 GPU 上 loss 的数据，此时仅需要在一个 GPU 上进行操作，如打印日志等。避免重复操作。

# 合并多进程数据，进行 logging loss_list=[] losses=paddle.tensor
loss_list = []
if (paddle.distributed.get_world_size() > 1):
    paddle.distributed.all_gather(loss_list, losses)
    if paddle.distributed.get_rank() == 0:
        losses = paddle.stack(loss_list).sum() / len(
            loss_list)