[源码解析] PyTorch 分布式(3) ----- DataParallel(下)

[源码解析] PyTorch 分布式(3) ----- DataParallel(下)

目录
  • [源码解析] PyTorch 分布式(3) ----- DataParallel(下)
    • 0x00 摘要
    • 0x01 前向操作
      • 1.1 并行
      • 1.2 Gather
        • 1.2.1 Python世界
        • 1.2.2 C++世界
    • 0x02 计算损失
    • 0x03 后向传播
      • 3.1 分发梯度
        • 3.1.1 Gather.backward
        • 3.1.2 Scatter
        • 3.1.3 C++
      • 3.2 并行后向传播
      • 3.3 归并梯度
        • 3.3.1 Broadcast.backward
        • 3.3.2 ReduceAddCoalesced
        • 3.3.3 c++
      • 3.4 更新模型参数
    • 0x04 总结
    • 0xFF 参考

0x00 摘要

本文是 PyTorch 分布式的第三篇,继续上文,介绍 DataPrallel 的并行操作和反向传播。

本系列其他文章如下:

深度学习利器之自动微分(1)

深度学习利器之自动微分(2)

[源码解析]深度学习利器之自动微分(3) --- 示例解读

[源码解析]PyTorch如何实现前向传播(1) --- 基础类(上)

[源码解析]PyTorch如何实现前向传播(2) --- 基础类(下)

[源码解析] PyTorch如何实现前向传播(3) --- 具体实现

[源码解析] Pytorch 如何实现后向传播 (1)---- 调用引擎

[源码解析] Pytorch 如何实现后向传播 (2)---- 引擎静态结构

[源码解析] Pytorch 如何实现后向传播 (3)---- 引擎动态逻辑

[源码解析] PyTorch 如何实现后向传播 (4)---- 具体算法

[源码解析] PyTorch 分布式(1)------历史和概述

[源码解析] PyTorch 如何使用GPU

[源码解析] PyTorch 分布式(2) ----- DataParallel(上)

0x01 前向操作

我们先回忆一下目前的前向图,replicate 调用了Broadcast.forward,同时往其context 存储了input_device和num_inputs。

+----------------------------------------------------------------------------------------+
| DataParallel.forward                                                                   |
|                                                                                        |
|                                                                                        |
|              replicate +--------------->   parallel_apply             gather           |
|                                                                                        |
+----------------------------------------------------------------------------------------+

     +---------------------------+
     | Broadcast                 |
     |                           |
     |                           |
     |                           |
     |          forward()  +----------->
     |                           |
     |                           |
     |  +---------------------+  |
     |  | ctx                 |  |
     |  |       input_device  |  |
     |  |                     |  |
     |  |       num_inputs    |  |
     |  |                     |  |
     |  +---------------------+  |
     |                           |
     |                           |
     |                           |
     |                           |
     |                           |
     |                           |
     +---------------------------+

1.1 并行

目前,我们已经使用 Scatter 函数将数据从 device[0] 分配并复制到不同的卡,用 Replicate 函数将模型从 device[0] 复制到不同的卡,这样各个卡都有了同样的模型和不同的数据,现在就要分别调用 forward 计算损失和梯度。也就是 parallel_apply 部分。

# 分发数据
inputs, kwargs = self.scatter(inputs, kwargs, self.device_ids)      
# 分发模型
replicas = self.replicate(self.module, self.device_ids[:len(inputs)])

# 并行训练
outputs = self.parallel_apply(replicas, inputs, kwargs)

对应我们传播图是:

[源码解析] PyTorch 分布式(3) ----- DataParallel(下)_第1张图片

parallel_apply 是基于threading 实现,用前面准备好的 replica 和输入数据,然后for 循环启动多线程进行前向传播,最后输出传播结果。

def parallel_apply(modules, inputs, kwargs_tup=None, devices=None):

    # 确保模型和输入大小一致
    assert len(modules) == len(inputs)

    # 确保每个 GPU 都有相应的元数据,如没有就空白补全
    if kwargs_tup is not None:
        # 在前面已经补全
        assert len(modules) == len(kwargs_tup)
    else:
        kwargs_tup = ({},) * len(modules)

    # 确保模型数目和CPU数目一致    
    if devices is not None:
        assert len(modules) == len(devices)
    else:
        devices = [None] * len(modules)

    devices = [_get_device_index(x, True) for x in devices]

    # 基于threading多线程实现
    lock = threading.Lock()
    results = {}
    grad_enabled, autocast_enabled = torch.is_grad_enabled(), torch.is_autocast_enabled()

    # 定义 worker
    def _worker(i, module, input, kwargs, device=None):
        torch.set_grad_enabled(grad_enabled)
        if device is None:
            device = get_a_var(input).get_device()
        try:
            # 设置当前的设备
            with torch.cuda.device(device), autocast(enabled=autocast_enabled):
                # this also avoids accidental slicing of `input` if it is a Tensor
                if not isinstance(input, (list, tuple)):
                    input = (input,)
                output = module(*input, **kwargs) # 前向操作
            with lock:
                # 并行计算得到输出
                results[i] = output
        except Exception:
            with lock:
                results[i] = ExceptionWrapper(
                    where="in replica {} on device {}".format(i, device))

    if len(modules) > 1:
        # 如有一个进程控制多个 GPU ,起多个线程
        # 注意,这里就是每个 worker 调用了 modules 数组中的一个模型copy
        threads = [threading.Thread(target=_worker,
                                    args=(i, module, input, kwargs, device))
                   for i, (module, input, kwargs, device) in
                   enumerate(zip(modules, inputs, kwargs_tup, devices))]

        for thread in threads:
            thread.start()
        for thread in threads:
            thread.join()
    else:
        # 一个GPU对应一个进程
        _worker(0, modules[0], inputs[0], kwargs_tup[0], devices[0])

    outputs = []
    for i in range(len(inputs)):
        output = results[i]

        # error handle
        if isinstance(output, ExceptionWrapper):
            output.reraise()
        outputs.append(output)
        
    # 输出 n 个计算结果
    return outputs

此时前向传播具体对应如下图,现在并行操作调用了 module 的forward方法。

+----------------------------------------------------------------------------------------+
| DataParallel.forward                                                                   |
|                                                                                        |
|                  1                               2                      3              |
|              replicate +--------------->   parallel_apply             gather           |
|                                                                                        |
+----------------------------------------------------------------------------------------+

     +---------------------------+       +-------------------+
     | Broadcast                 |       | module            |
     |                           |       |                   |
     |                           |       |                   |
     |              1            |       |         2         |
     |          forward()  +-----------> |      forward() +--------->
     |                           |       |                   |
     |                           |       |                   |
     |  +---------------------+  |       |                   |
     |  | ctx                 |  |       |                   |
     |  |       input_device  |  |       |                   |
     |  |                     |  |       |                   |
     |  |       num_inputs    |  |       |                   |
     |  |                     |  |       |                   |
     |  +---------------------+  |       |                   |
     |                           |       |                   |
     |                           |       |                   |
     |                           |       |                   |
     |                           |       |                   |
     |                           |       |                   |
     |                           |       |                   |
     +---------------------------+       +-------------------+

1.2 Gather

目前,我们已经使用 Scatter 函数将数据从 device[0] 分配并复制到不同的卡,用 Replicate 函数将模型从 device[0] 复制到不同的卡,这样各个卡都有了同样的模型和不同的数据,然后分别调用 forward 计算损失和梯度。也就是 parallel_apply 部分。

现在要做的就是把分布式计算的梯度合并到 device[0],就是 self.output_device。

# 分发数据
inputs, kwargs = self.scatter(inputs, kwargs, self.device_ids)      
# 分发模型
replicas = self.replicate(self.module, self.device_ids[:len(inputs)])
# 并行训练
outputs = self.parallel_apply(replicas, inputs, kwargs)
# 收集到 devices[0]
return self.gather(outputs, self.output_device)

对应我们传播图是:

[源码解析] PyTorch 分布式(3) ----- DataParallel(下)_第2张图片

我们看看如何把结果收集到 device[0],以及device[0]如何作为参数服务器。

1.2.1 Python世界

gather 主要是调用 Gather.apply(target_device, dim, *outputs) 完成收集工作。

def gather(outputs, target_device, dim=0): # target_device 就是 device[0]
    r"""
    Gathers tensors from different GPUs on a specified device
      (-1 means the CPU).
    """
    def gather_map(outputs):
        out = outputs[0]
        if isinstance(out, torch.Tensor):
            return Gather.apply(target_device, dim, *outputs) # 调用下面的 Gather
        if out is None:
            return None
        if isinstance(out, dict):
            return type(out)(((k, gather_map([d[k] for d in outputs]))
                              for k in out))
        return type(out)(map(gather_map, zip(*outputs)))

    # Recursive function calls like this create reference cycles.
    # Setting the function to None clears the refcycle.
    try:
        res = gather_map(outputs)
    finally:
        gather_map = None
    return res

Gather 则调用了 comm.gather 完成工作,而 comm.gather 则会带领我们进入到 C++世界。

我们省略一些校验代码。

# Gather 源码
class Gather(Function):

    @staticmethod
    def forward(ctx, target_device, dim, *inputs): # target_device 就是 device[0]

        # 下面会往 context 内部存放几个变量,后续会用到
        target_device = _get_device_index(target_device, True)
        ctx.target_device = target_device
        ctx.dim = dim
        ctx.input_gpus = tuple(i.get_device() for i in inputs)

        if all(t.dim() == 0 for t in inputs) and dim == 0:
            inputs = tuple(t.view(1) for t in inputs)
            ctx.unsqueezed_scalar = True
        else:
            ctx.unsqueezed_scalar = False
            
        ctx.input_sizes = tuple(i.size(ctx.dim) for i in inputs)
        return comm.gather(inputs, ctx.dim, ctx.target_device) # 这里会进入C++世界

    @staticmethod
    def backward(ctx, grad_output): # 注意,这里后续会用到
        scattered_grads = Scatter.apply(ctx.input_gpus, ctx.input_sizes, ctx.dim, grad_output)
        if ctx.unsqueezed_scalar:
            scattered_grads = tuple(g[0] for g in scattered_grads)
        return (None, None) + scattered_grads

现在前向计算如图:

gather 调用到了Gather的forward 函数,forward 方法在 ctx 存储了 input_gpus, input_sizes, dim 这三个变量,这些变量后续会用到。

+-----------------------------------------------------------------------------------------+
| DataParallel.forward                                                                    |
|                                                                                         |
|                  1                               2                           3          |
|              replicate +--------------->   parallel_apply +--------------> gather       |
|                                                                                         |
+-----------------------------------------------------------------------------------------+

     +---------------------------+       +-------------------+       +--------------------+
     | Broadcast                 |       | module            |       |Gather              |
     |                           |       |                   |       |                    |
     |                           |       |                   |       |                    |
     |              1            |       |         2         |       |         3          |
     |          forward()  +-----------> |      forward() +--------> |      forward()     |
     |                           |       |                   |       |                    |
     |                           |       |                   |       |                    |
     |  +---------------------+  |       |                   |       | +----------------+ |
     |  | ctx                 |  |       |                   |       | |ctx             | |
     |  |       input_device  |  |       |                   |       | |     input_gpus | |
     |  |                     |  |       |                   |       | |                | |
     |  |       num_inputs    |  |       |                   |       | |     input_sizes| |
     |  |                     |  |       |                   |       | |                | |
     |  +---------------------+  |       |                   |       | |     dim        | |
     |                           |       |                   |       | +----------------+ |
     |                           |       |                   |       |                    |
     |                           |       |                   |       |                    |
     |                           |       |                   |       |                    |
     |                           |       |                   |       |                    |
     |                           |       |                   |       |                    |
     +---------------------------+       +-------------------+       +--------------------+

1.2.2 C++世界

gather 函数调用了 _gather_out_impl 来完成拷贝操作。

at::Tensor gather(
    at::TensorList tensors,
    int64_t dim,
    c10::optional destination_index) { // destination_index 就是 device[0] 的index

  int64_t total_size = 0;
  auto& first = tensors.front();
  const auto first_size = first.sizes();
  dim = at::maybe_wrap_dim(dim, first);
  std::vector expected_size(first_size.begin(), first_size.end());
  auto memory_format = first.suggest_memory_format();
  for (size_t i = 0; i < tensors.size(); i++) {
    const auto& tensor = tensors[i];
    expected_size[dim] = tensor.size(dim);
    total_size += tensor.size(dim);
    if (memory_format != MemoryFormat::Contiguous &&
        tensor.suggest_memory_format() != memory_format) {
      memory_format = MemoryFormat::Contiguous;
    }
  }
  expected_size[dim] = total_size;
  at::Device device(DeviceType::CPU);
  // 根据 index 得到输出的目标设备
  if (!destination_index || *destination_index != -1) {
    // device 就是 GPU 0 这个设备
    device = at::Device(
        DeviceType::CUDA, destination_index ? *destination_index : -1);
  }

  //首先,构建一个空的目标tensor建立在目标设备之上,命名为result
  at::Tensor result =
      at::empty(expected_size, first.options().device(device), memory_format);
  
  return _gather_out_impl(tensors, result, dim); // 然后对result进行gather
}

_gather_out_impl 执行了具体的gather 操作,就是把 输入的tensors 拷贝到 目标 tensor 之上,即拷贝到 GPU0 之上。

// ***************** Gather *******************
//
// Gather a list of CUDA tensors on one or more devices to a target tensor or
// device, either CPU or CUDA.

// no checks
static inline at::Tensor& _gather_out_impl(
    at::TensorList tensors,
    at::Tensor& out_tensor,
    int64_t dim) {
  std::vector chunk_sizes;
  chunk_sizes.reserve(tensors.size());
  for (auto& tensor : tensors) {
    chunk_sizes.push_back(tensor.size(dim));
  }
  auto chunks =
      out_tensor.split_with_sizes(/*split_sizes=*/chunk_sizes, /*dim=*/dim);
  for (size_t i = 0; i < tensors.size(); i++) { // 拷贝到GPU 0 之上
    chunks[i].copy_(tensors[i], /*non_blocking=*/out_tensor.is_cuda());
  }
  return out_tensor;
}

0x02 计算损失

[源码解析] PyTorch 分布式(3) ----- DataParallel(下)_第3张图片

现在,我们已经把梯度收集到 device[0] 之上,现在我们开始进行反向传播,其整体逻辑如上图所示。首先是在 device[0] 计算损失。其实这步计算损失算是前向计算和后向传播的中间环节,这里把它算成是反向传播的开端,如下图。

[源码解析] PyTorch 分布式(3) ----- DataParallel(下)_第4张图片

我们找出来示例代码看看,里面关键的几点:

  1. 数据已经放到了默认GPU,即GPU 0上。
  2. prediction 是gather到 GPU 0 的前向计算输出
  3. 使用 loss = criterion(prediction,target_var) 在默认GPU之上计算loss。
  4. 使用 loss.backward() 开始反向传播。
for batch_idx, (data, label) in pbar:   
    if args.cuda:
        data,label= data.cuda(),label.cuda(); # 1. 数据已经放到了默认GPU上
    data_v = Variable(data)
    target_var = Variable(label)
    prediction= model(data_v,target_var,args) # 2. prediction 是gather到 GPU 0 的前向计算输出
    
    # 到目前为止,我们完成了DataParallel.forward()
    #这里的prediction 预测结果是由两个gpu合并过的,并行计算只存在于前向传播里
    #前向传播每个gpu计算量为 batch_size/len(device_ids),等前向传播完了将结果聚合到主gpu里

    criterion = nn.CrossEntropyLoss()
    loss = criterion(prediction,target_var)  # 3. 在默认GPU之上计算loss
    optimizer.zero_grad()
    loss.backward()   # 4. 开始反向传播
    optimizer.step()

0x03 后向传播

我们前面运行的是上面的 Forward 部分,计算损失,接下来就运行上面代码中 loss.backward() 部分。

3.1 分发梯度

我们首先来到分发梯度部分,这部分作用是:把损失在 GPUs 之间 scatter,这样后续才可以在每个GPU之上独立进行后向传播。对应下图:

[源码解析] PyTorch 分布式(3) ----- DataParallel(下)_第5张图片

3.1.1 Gather.backward

前面有提到,prediction 是gather到 GPU 0 的前向计算输出。而 loss 又是根据 prediction 计算出来,所以从 loss.backward() 开始反向传播,从后向前的第一个步骤就来到了 gather 的传播操作,对应的就是 Gather 的 backward 函数,其中的核心代码是 Scatter.apply。

class Gather(Function):

    # 这里前向传播用到了,为了对照,我们依然贴出来
    @staticmethod
    def forward(ctx, target_device, dim, *inputs): # target_device 就是 device[0]

        # 下面会往 context 内部存放几个变量,后续会用到
        target_device = _get_device_index(target_device, True)
        ctx.target_device = target_device
        ctx.dim = dim
        ctx.input_gpus = tuple(i.get_device() for i in inputs)

        if all(t.dim() == 0 for t in inputs) and dim == 0:
            inputs = tuple(t.view(1) for t in inputs)
            ctx.unsqueezed_scalar = True
        else:
            ctx.unsqueezed_scalar = False
            
        ctx.input_sizes = tuple(i.size(ctx.dim) for i in inputs)
        # 这里会进入C++世界,把输出聚集到 GPU 0。
        return comm.gather(inputs, ctx.dim, ctx.target_device) 

    @staticmethod
    def backward(ctx, grad_output): # 这里现在后向传播用到了!
        # 把前向传播在 context 之中存放的变量取出,作为 Scatter 的输入 
        scattered_grads = Scatter.apply(ctx.input_gpus, ctx.input_sizes, ctx.dim, grad_output)
        if ctx.unsqueezed_scalar:
            scattered_grads = tuple(g[0] for g in scattered_grads)
        return (None, None) + scattered_grads

具体如下,可以看到,backward 使用了之前前向传播时候存储的 ctx.input_gpus, ctx.input_sizes, ctx.dim, grad_output,以此调用 Scatter.apply。

图中,最上面是前向传播过程,最下面是反向传播过程,中间是某些在前后传播中都用到的代码模块。

+--------------------------------------------------------------------------------------+
| DataParallel.forward                                                                 |
|                                                                                      |
|               1                               2                           3          |
|           replicate +--------------->   parallel_apply +--------------> gather       |
|                                                                                      |
+--------------------------------------------------------------------------------------+

  +---------------------------+       +-------------------+       +--------------------+
  | Broadcast                 |       | module            |       |Gather              |
  |                           |       |                   |       |                    |
  |                           |       |                   |       |                    |
  |              1            |       |         2         |       |         3          |
  |          forward()  +-----------> |      forward() +--------> |      forward()     |
  |                           |       |                   |       |                    |
  |                           |       |                   |       |                    |
  |  +---------------------+  |       |                   |       | +----------------+ |
  |  | ctx                 |  |       |                   |       | |ctx             | |
  |  |       input_device  |  |       |                   |       | |     input_gpus | |
  |  |                     |  |       |                   |       | |                | |
  |  |       num_inputs    |  |       |                   |       | |     input_sizes| |
  |  |                     |  |       |                   |       | |                | |
  |  +---------------------+  |       |                   |       | |     dim        | |
  |                           |       |                   |       | +----------------+ |
  |                           |       |                   |       |                    |
  |                           |       |                   |       |                    |
  |                           |       |                   | <---------+ backward()     |
  |                           |       |                   |       |         3          |
  |                           |       |                   |       |                    |
  +---------------------------+       +-------------------+       +--------------------+

+--------------------------------------------------------------------------------------+
| loss.backward()                                                                      |
|                                                                           3          |
|                                                           <--------------------+     |
|                                                                                      |
|                                                                                      |
+--------------------------------------------------------------------------------------+

3.1.2 Scatter

Scatter.apply 实际上调用到了其 forward 方法。

  • 首先从上下文之中提取之前存储的变量,这里主要是输入设备 input_device(源设备)和 target_gpus(目标设备)。
  • 获取到目标设备的流。
  • 调用 comm.scatter 把梯度分发到目标设备。
class Scatter(Function):

    @staticmethod
    def forward(ctx, target_gpus, chunk_sizes, dim, input):
        target_gpus = [_get_device_index(x, True) for x in target_gpus]
        ctx.dim = dim
        ctx.input_device = input.get_device() if input.device.type != "cpu" else -1
        streams = None
        if torch.cuda.is_available() and ctx.input_device == -1:
            # Perform CPU to GPU copies in a background stream
            streams = [_get_stream(device) for device in target_gpus]
         
        # 分发到其他GPU
        outputs = comm.scatter(input, target_gpus, chunk_sizes, ctx.dim, streams)
        
        # Synchronize with the copy stream
        if streams is not None:
            for i, output in enumerate(outputs):
                with torch.cuda.device(target_gpus[i]):
                    main_stream = torch.cuda.current_stream()
                    main_stream.wait_stream(streams[i])
                    output.record_stream(main_stream)
        return outputs

    @staticmethod
    def backward(ctx, *grad_output):
        return None, None, None, Gather.apply(ctx.input_device, ctx.dim, *grad_output)

3.1.3 C++

上面python代码 outputs = comm.scatter(input, target_gpus, chunk_sizes, ctx.dim, streams) 会直接进入到C++世界。具体代码位于 torch/csrc/cuda/comm.cpp。

scatter 的作用就是把tensor进行split,然后分发给各个设备的流。

std::vector scatter(
    const at::Tensor& tensor,
    at::IntArrayRef devices,
    const c10::optional>& chunk_sizes,
    int64_t dim,
    const c10::optional>>&
        streams) {
  dim = at::maybe_wrap_dim(dim, tensor);
  
  // 把tensor进行split
  std::vector chunks = chunk_sizes
      ? tensor.split_with_sizes(/*split_sizes=*/*chunk_sizes, /*dim=*/dim)
      : tensor.chunk(/*chunks=*/devices.size(), /*dim=*/dim);
  
  at::cuda::OptionalCUDAStreamGuard cuda_guard;
  for (size_t i = 0; i < chunks.size(); ++i) {
    const auto device_index = static_cast(devices[i]);
    if (device_index != tensor.get_device()) {
      if (i < (streams ? streams->size() : 0U) && (*streams)[i]) {
        cuda_guard.reset_stream(*(*streams)[i]);
      }
      // 发送给各个设备的流
      chunks[i] = chunks[i].to( 
          {DeviceType::CUDA, device_index},
          /*non_blocking=*/true,
          /*copy=*/false,
          /*memory_format=*/at::MemoryFormat::Preserve);
    }
  }
  return chunks;
}

3.2 并行后向传播

现在梯度已经分发到各个 GPU,接下来正式进入并行后向传播,这部分作用是:在各个GPU之上并行运行后向传播,计算参数梯度。对应下图:

[源码解析] PyTorch 分布式(3) ----- DataParallel(下)_第6张图片

这部分调用到了原始模型的 backward,具体如下图中的数值 4:

+--------------------------------------------------------------------------------------+
| DataParallel.forward                                                                 |
|                                                                                      |
|               1                               2                           3          |
|           replicate +--------------->   parallel_apply +--------------> gather       |
|                                                                                      |
+--------------------------------------------------------------------------------------+

  +---------------------------+       +-------------------+       +--------------------+
  | Broadcast                 |       | module            |       |Gather              |
  |                           |       |                   |       |                    |
  |                           |       |                   |       |                    |
  |              1            |       |         2         |       |         3          |
  |          forward()  +-----------> |      forward() +--------> |      forward()     |
  |                           |       |                   |       |                    |
  |                           |       |                   |       |                    |
  |  +---------------------+  |       |                   |       | +----------------+ |
  |  | ctx                 |  |       |                   |       | |ctx             | |
  |  |       input_device  |  |       |                   |       | |     input_gpus | |
  |  |                     |  |       |                   |       | |                | |
  |  |       num_inputs    |  |       |                   |       | |     input_sizes| |
  |  |                     |  |       |                   |       | |                | |
  |  +---------------------+  |       |                   |       | |     dim        | |
  |                           |       |                   |       | +----------------+ |
  |                           |       |                   |       |                    |
  |                           |       |                   |       |                    |
  |                           | <---------+  backward()   | <---------+ backward()     |
  |                           |       |          4        |       |         3          |
  |                           |       |                   |       |                    |
  +---------------------------+       +-------------------+       +--------------------+

+--------------------------------------------------------------------------------------+
| loss.backward()                                                                      |
|                                                4                          3          |
|                                     <------------------+  <--------------------+     |
|                                                                                      |
|                                                                                      |
+--------------------------------------------------------------------------------------+

3.3 归并梯度

这部分作用是 在 GPU 0 之上归并梯度,总体流程拓展对应下图:

[源码解析] PyTorch 分布式(3) ----- DataParallel(下)_第7张图片

3.3.1 Broadcast.backward

这部分对应了 Broadcast 的 反向传播。

class Broadcast(Function):

    @staticmethod
    def forward(ctx, target_gpus, *inputs):
        target_gpus = [_get_device_index(x, True) for x in target_gpus]
        
        # 前向传播时候,向上下文存入了一些变量
        ctx.target_gpus = target_gpus
        if len(inputs) == 0:
            return tuple()
        ctx.num_inputs = len(inputs)
        # input 放在 device[0],所以 input_device 就是 GPU 0
        ctx.input_device = inputs[0].get_device()
        # 和 detach 的情形一样
        outputs = comm.broadcast_coalesced(inputs, ctx.target_gpus)
        non_differentiables = []
        
        # 在上下文中设置哪些不需要梯度
        for idx, input_requires_grad in enumerate(ctx.needs_input_grad[1:]):
            if not input_requires_grad:
                for output in outputs:
                    non_differentiables.append(output[idx])
        ctx.mark_non_differentiable(*non_differentiables)
        return tuple([t for tensors in outputs for t in tensors])

    @staticmethod
    def backward(ctx, *grad_outputs):
        # 反向传播来到这里,取出之前在上下文存放的变量作为输入。ctx.input_device 就是之前存储的 GPU 0。
        return (None,) + ReduceAddCoalesced.apply(ctx.input_device, ctx.num_inputs, *grad_outputs)

因此,我们可以拓展流程图:

+--------------------------------------------------------------------------------------+
| DataParallel.forward                                                                 |
|                                                                                      |
|               1                               2                           3          |
|           replicate +--------------->   parallel_apply +--------------> gather       |
|                                                                                      |
+--------------------------------------------------------------------------------------+

  +---------------------------+       +-------------------+       +--------------------+
  | Broadcast                 |       | module            |       |Gather              |
  |                           |       |                   |       |                    |
  |                           |       |                   |       |                    |
  |              1            |       |         2         |       |         3          |
  |          forward()  +-----------> |      forward() +--------> |      forward()     |
  |                           |       |                   |       |                    |
  |                           |       |                   |       |                    |
  |  +---------------------+  |       |                   |       | +----------------+ |
  |  | ctx                 |  |       |                   |       | |ctx             | |
  |  |       input_device  |  |       |                   |       | |     input_gpus | |
  |  |                     |  |       |                   |       | |                | |
  |  |       num_inputs    |  |       |                   |       | |     input_sizes| |
  |  |                     |  |       |                   |       | |                | |
  |  +---------------------+  |       |                   |       | |     dim        | |
  |                           |       |                   |       | +----------------+ |
  |                           |       |                   |       |                    |
  |                           |       |                   |       |                    |
  |          backward()       | <---------+  backward()   | <---------+ backward()     |
  |              5            |       |          4        |       |         3          |
  |                           |       |                   |       |                    |
  +---------------------------+       +-------------------+       +--------------------+

+--------------------------------------------------------------------------------------+
| loss.backward()                                                                      |
|                5                               4                          3          |
|         <------------------------+  <------------------+  <--------------------+     |
|                                                                                      |
|                                                                                      |
+--------------------------------------------------------------------------------------+

3.3.2 ReduceAddCoalesced

Broadcast.backward 调用了 ReduceAddCoalesced.apply,其对应了 ReduceAddCoalesced 的 forward 方法,目的是把梯度归并到目标设备 destination,就是GPU 0。

class ReduceAddCoalesced(Function):

    @staticmethod
    # 会调用到这里,destination 是GPU 0
    def forward(ctx, destination, num_inputs, *grads): 
        # 从梯度之中提取所在的设备
        ctx.target_gpus = [grads[i].get_device() for i in range(0, len(grads), num_inputs)]

        grads_ = [grads[i:i + num_inputs]
                  for i in range(0, len(grads), num_inputs)]
        # 把梯度归并到目标设备 destination,就是GPU 0
        return comm.reduce_add_coalesced(grads_, destination)

    @staticmethod
    def backward(ctx, *grad_outputs):
        return (None, None,) + Broadcast.apply(ctx.target_gpus, *grad_outputs)

3.3.3 c++

看注释就是:从多个 GPU 来相加梯度,代码之中就是归并相加。

def reduce_add_coalesced(inputs, destination=None, buffer_size=10485760):
    """Sums tensors from multiple GPUs.

    Small tensors are first coalesced into a buffer to reduce the number
    of synchronizations.

    Args:
        inputs (Iterable[Iterable[Tensor]]): iterable of iterables that
            contain tensors from a single device.
        destination (int, optional): a device on which the output will be
            placed (default: current device).
        buffer_size (int): maximum size of the buffer used for coalescing

    Returns:
        A tuple of tensors containing an elementwise sum of each group of
        inputs, placed on the ``destination`` device.
    """
    dense_tensors: List[List] = [[] for _ in inputs]  # shape (num_gpus, num_tensors)
    output = []
    ref_order = []
    # process sparse ones first since they may have different sizes on different gpus
    for tensor_at_gpus in zip(*inputs):
        if all(t.is_sparse for t in tensor_at_gpus):
            # 进行归并
            result = reduce_add(tensor_at_gpus, destination)  # this will be sparse too
            output.append(result)
            ref_order.append(tensor_at_gpus[0])
        else:
            for coll, t in zip(dense_tensors, tensor_at_gpus):
                coll.append(t.to_dense() if t.is_sparse else t)
            ref_order.append(dense_tensors[0][-1])
    itrs = [_take_tensors(tensors, buffer_size) for tensors in dense_tensors]
    # now the dense ones, which have consistent sizes
    for chunks in zip(*itrs):
        flat_tensors = [_flatten_dense_tensors(chunk) for chunk in chunks]  # (num_gpus,)
        # 进行归并
        flat_result = reduce_add(flat_tensors, destination)
        for t in _unflatten_dense_tensors(flat_result, chunks[0]):
            # The unflattened tensors do not share storage, and we don't expose
            # base flat tensor anyways, so give them different version counters.
            # See NOTE [ Version Counter in comm.*_coalesced ]
            output.append(t.data)
    return tuple(_reorder_tensors_as(output, ref_order))

3.4 更新模型参数

这部分功能是:更新梯度参数。进行梯度下降,并更新主GPU上的模型参数。

另外,由于模型参数仅在主GPU上更新,而其他从属GPU此时并不是同步更新的,所以需要将更新后的模型参数复制到剩余的从属 GPU 中,以此来实现并行。这就是在下一次for循环之中进行,以此循环反复。

对应示例代码是:

for batch_idx, (data, label) in pbar:   # 6. 下一次迭代会继续从分发开始
    if args.cuda:
        data,label= data.cuda(),label.cuda(); # 1. 数据已经放到了默认GPU上
    data_v = Variable(data)
    target_var = Variable(label)
    prediction= model(data_v,target_var,args) # 2. prediction 是gather到 GPU 0 的前向计算输出
    
    # 到目前为止,我们完成了DataParallel.forward()
    #这里的prediction 预测结果是由两个gpu合并过的,并行计算只存在在前向传播里
    #前向传播每个gpu计算量为 batch_size/len(device_ids),等前向传播完了将结果和到主gpu里

    criterion = nn.CrossEntropyLoss()
    loss = criterion(prediction,target_var)  # 3. 在默认GPU之上计算loss
    optimizer.zero_grad()
    loss.backward()   # 4. 开始反向传播
    optimizer.step() # 5. 更新模型

0x04 总结

我们总结一下流程,起初数据和模型被放入到默认GPU,就是 GPU 0,然后迭代如下:

  1. scatter 会把数据分发到其他 GPU。
  2. replicate 会把模型分发到其他 GPU。
  3. parallel_apply 会启动多个线程进行前向计算。
  4. gather 会把计算输出收集到 GPU 0。
  5. GPU 0 会计算损失。
  6. 把梯度 scatter 到其他 GPU。
  7. 模型调用 backward 计算。
  8. 把梯度归并到 GPU 0。
  9. optimizer.step 更新模型。

具体对应下图之中的数字。

                     +-----+                   +-------+
                     |GPU1 |                   | GPU1  |
main thread          +-----+                   +-------+
 +-----> Forward----> scatter +--------------> replicate------->  parallel_apply  +-------->  gather +---------+
                        +                           +                     +                                    |
                      1 |                         2 |                   3 |                                    |
                        |                           |                     |                                    |
                        |  +---------+----------+---+                     |                                    |
                        |  |         |          |                         |                                    |
                        +---------+----------+  |               +--------------------+                         |
                        |  |      |  |       |  |               |         |          |                         |
                        |  | 2    |  | 2     |  | 2       thread|1     thread 2    thread 3                    |
                      1 |  |    1 |  |     1 |  |               |         |          |                         |
                        |  v      |  v       |  v               |         |          |                         |
                        v         v          v                  v         v          v                         |
                     +--+---+  +--+---+   +--+---+           +--+---+  +--+---+   +--+---+    +-------+        |
                     | GPU1 |  | GPU2 |   | GPU3 |           | GPU1 |  | GPU2 |   | GPU3 |    | GPU1  |        |
                     +------+  +------+   +------+           +--+---+  +-+----+   +---+--+    +-+-+--++        |
                                                                |        |            |         ^ ^  ^         |
                                                                |        |            |   4     | |  |         |
                                                                |        |            ----------^ |  |         |
                                                                |        |                4       |  |         |
                                                                |        +------------------------+  |         |
                                                                |                                    |         |
                                                                +------------------------------------+         |
        +------------------------------------------------------------------------------------------------------+
        |                               +------+
        |                               | GPU1 |
        |                               +------+                                                                     main thread
        +-> loss = criterion(...)+-----> scatter   +-------------->  model.backward() +---------->  reduce gradient +-------> optimizer.step
                     +                      +                               +                          +------+         9
                     | 5                    | 6                             | 7                        | GPU1 |
                     |                      |                               |                          +--+---+
                     |              v---------------v             +--------------------+                  ^
                     |              |       |       |             |         |          |                  | 8
                     |              |       |       |         thread 1    thread 2   thread 3             |
                     |              |       |       |             +         |          |           +-------------+
                     |              |       |       |             |         |          |           |      |      |
                     v              v       v       v             v         v          v           |      |      |
                  +--+---+      +---+-+  +--+--+  +-+---+      +--+--+  +---+--+    +--+--+     +--+--+ +-+--+ +-+---+
                  | GPU1 |      | GPU1|  | GPU2|  |GPU3 |      | GPU1|  | GPU2 |    |GPU3 |     | GPU1| |GPU2| | GPU3|
                  +------+      +-----+  +-----+  +-----+      +-----+  +------+    +-----+     +-----+ +----+ +-----+

手机如下:

[源码解析] PyTorch 分布式(3) ----- DataParallel(下)_第8张图片

至此,DP 分析完毕,我们下一篇要介绍 DDP 的一些相关知识。

0xFF 参考

PyTorch 源码解读之 torch.optim:优化算法接口详解

pytorch(分布式)数据并行个人实践总结——DataParallel/DistributedDataParallel

Pytorch的nn.DataParallel

PyTorch 源码解读之分布式训练了解一下?

https://discuss.pytorch.org/t/dataparallel-imbalanced-memory-usage/22551/20

[原创][深度][PyTorch] DDP系列第二篇:实现原理与源代码解析

Pytorch-CUDA从入门到放弃(二)

Pytorch踩坑记:赋值、浅拷贝、深拷贝三者的区别以及model.state_dict()和model.load_state_dict()的坑点

PyTorch 源码解读之 DP & DDP:模型并行和分布式训练解析

你可能感兴趣的:([源码解析] PyTorch 分布式(3) ----- DataParallel(下))