all_reduce与GPU数据并行

all reduce方法

先看代码：

def allreduce(data):
    for i in range(1, len(data)):
        data[0][:] += data[i].to(data[0].device)
    for i in range(1, len(data)):
        data[i][:] = data[0].to(data[i].device)

将所有向量相加，并将结果广播给所有GPU。请注意，我们需要将数据复制到累积结果的设备，才能使函数正常工作。

训练过程中：

def train_batch(X, y, device_params, devices, lr):
    # 将一个batch的数据拆分到多个GPU上
    X_shards, y_shards = split_batch(X, y, devices) 
    # 在每个GPU上分别计算损失
    ls = [loss(net(X_shard, device_W), y_shard).sum()
          for X_shard, y_shard, device_W in zip(
              X_shards, y_shards, device_params)]
    for l in ls:  # 反向传播在每个GPU上分别执行
        l.backward()
    # 将每个GPU的所有梯度相加，并将其广播到所有GPU
    with torch.no_grad():
        for i in range(len(device_params[0])):
            allreduce(
                [device_params[c][i].grad for c in range(len(devices))])
    # 在每个GPU上分别更新模型参数
    for param in device_params:
        sgd(param, lr, X.shape[0]) # 在这里，我们使用全尺寸的小批量

总结：

• 前向过程：
  1. 分发mini-batch到每个GPU上；
  2. 将模型复制到每一个GPU上；
  3. 模型对mini-batch进行forward计算结果；
  4. 将每个GPU上的计算结果汇总到第一个GPU上。
• 反向传播过程：
  1. 在第一个GPU上根据计算的loss计算梯度；
  2. 将第一个GPU上计算的梯度值分发到每个GPU上；
  3. 每个GPU上进行梯度更新；
  4. 将每个GPU上更新的梯度值汇总到第一个GPU上

其中foward过程都会把第一个GPU上的模型重新分发到每个每个GPU上。