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

1,020 阅读24分钟

0x00 摘要

在 PyTorch DataParallel 训练过程中,其会在多个GPU之上复制模型副本,然后才开始训练。笔者在分析过程中,发现如果不把一些GPU相关基础知识整理出来,很难理解DataParallel的这个复制模型的过程,遂有此文。

本系列其他文章如下:

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

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

深度学习利器之自动微分(3) --- 示例解读

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

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

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

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

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

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

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

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

0x01 问题

在 DataParallel 进行前向传播之前,需要在GPU之上分散数据,复制模型,具体可见下图。

由此我们有几个问题:

  1. 移动模型到GPU这个动作的背后究竟做了哪些操作?
  2. 如何在 CPU 之上调用 GPU 操作?
  3. 如何在 CPU,GPU 操作之间无缝切换?
  4. 是否需要把损失函数移动到 GPU 之上?

我们接下来就一一分析。

注,关于CUDA和Dispatcher我们只是大致介绍,目的是可以让读者走通整个流程,有兴趣的读者可以自行深入研究。

0x02 移动模型到GPU

2.1 cuda 操作

CUDA 是NVIDIA公司开发的GPU编程模型,其提供了GPU编程接口,用户可以基于CUDA编程来构建基于GPU计算的应用。

torch.cuda用于设置 cuda 和运行cuda操作。它跟踪当前选定的GPU,默认情况下,用户分配的所有CUDA张量都将在该设备上创建。用户可以使用 torch.cuda.device 来修改所选设备。一旦分配了张量,您可以对其执行操作,而不考虑所选设备,PyTorch 会把运行结果与原始张量放在同一设备上。

默认情况下,除了~torch.Tensor.copy_和其他具有类似复制功能的方法(如~torch.Tensor.to~torch.Tensor.cuda)之外,不允许跨GPU操作,除非启用对等(peer-to-peer)内存访问。

我们从源码之中找出一个具体示例如下,大家可以看到,张量可以在设备上被创建,操作。

    cuda = torch.device('cuda')     # Default CUDA device
    cuda0 = torch.device('cuda:0')
    cuda2 = torch.device('cuda:2')  # GPU 2 (these are 0-indexed)

    x = torch.tensor([1., 2.], device=cuda0)
    # x.device is device(type='cuda', index=0)
    y = torch.tensor([1., 2.]).cuda()
    # y.device is device(type='cuda', index=0)

    with torch.cuda.device(1):
        # allocates a tensor on GPU 1
        a = torch.tensor([1., 2.], device=cuda)

        # transfers a tensor from CPU to GPU 1
        b = torch.tensor([1., 2.]).cuda()
        # a.device and b.device are device(type='cuda', index=1)

        # You can also use ``Tensor.to`` to transfer a tensor:
        b2 = torch.tensor([1., 2.]).to(device=cuda)
        # b.device and b2.device are device(type='cuda', index=1)

        c = a + b
        # c.device is device(type='cuda', index=1)

        z = x + y
        # z.device is device(type='cuda', index=0)

        # even within a context, you can specify the device
        # (or give a GPU index to the .cuda call)
        d = torch.randn(2, device=cuda2)
        e = torch.randn(2).to(cuda2)
        f = torch.randn(2).cuda(cuda2)
        # d.device, e.device, and f.device are all device(type='cuda', index=2)

2.2 Module

深度学习的模型可以看做是一种参数的容器,运行模型其实就是对输入参数做了一些基本的矩阵运算。一般来说,用户定义的模型都是派生自 nn.modules.module 类。而分布式训练涉及到同步更新参数和把模型拷贝到多个worker之上,所以我们首先需要看看Module的状况。从定义中可以看出来,Module的成员变量主要分为状态参数和hooks函数。

class Module:
    dump_patches: bool = False
    _version: int = 1
    training: bool
    _is_full_backward_hook: Optional[bool]
​
    def __init__(self):
        """
        Initializes internal Module state, shared by both nn.Module and ScriptModule.
        """
        torch._C._log_api_usage_once("python.nn_module")
​
        self.training = True
        self._parameters = OrderedDict() # 在训练过程中会随着 BP 而更新的参数
        self._buffers = OrderedDict() # 在训练过程中不会随着 BP 而更新的参数
        self._non_persistent_buffers_set = set()
        self._backward_hooks = OrderedDict()
        self._is_full_backward_hook = None
        self._forward_hooks = OrderedDict()
        self._forward_pre_hooks = OrderedDict()
        self._state_dict_hooks = OrderedDict()
        self._load_state_dict_pre_hooks = OrderedDict()
        self._modules = OrderedDict()

我们主要对状态参数进行说明。状态参数之中,主要有四种:

  • self.training

    • 本网络是否正在训练。
  • self._modules

    • 是本网络下属的子模块,采取迭代的方式进行定义。
  • self._parameters

    • 网络的参数。是在训练过程中会随着 BP 而更新的参数,就是梯度更新的对象。
  • self._buffers

    • 在训练过程中,不会随着BP更新的参数,但需要被保存,比如BatchNorm中的moving mean and variance,其优化不是通过梯度反向传播而是通过其他途径。

从本质上讲,当一个模型的网络结构被定义之后,self._parametersself._buffers的组合是一个模型的具体状态。如果需要拷贝一个模型:

  • self._modules属于网络结构的一部分,当我们拷贝模型到其他workers时,会一起拷贝过来。
  • self._parametersself._buffers 都需要显式拷贝到其他worker,这样才能在不同的Python进程之中维持相同的状态。

那么,这是不是意味着我们只需要拷贝 self._modulesself._parametersself._buffers 这些就可以了?让我们继续往下看。

2.3 移动

2.3.1 示例

前面看到了如何在 GPU 上操作张量,我们接下来看看如何把模型放置到 GPU 之上。

首先我们定义了一个模型。

class ToyModel(nn.Module):
    def __init__(self):
        super(ToyModel, self).__init__()
        self.net1 = nn.Linear(10, 10)
        self.relu = nn.ReLU()
        self.net2 = nn.Linear(10, 5)

    def forward(self, x):
        return self.net2(self.relu(self.net1(x)))

然后通过如下方式使用模型。

model = ToyModel().cuda(device_ids[0]) # 这里复制模型到 GPU 之上
ddp_model = DDP(model, device_ids)
​
loss_fn = nn.MSELoss() # 接着进行训练
optimizer = optim.SGD(ddp_model.parameters(), lr=0.001)
​
optimizer.zero_grad()
outputs = ddp_model(torch.randn(20, 10))
labels = torch.randn(20, 5).to(device_ids[0])
loss_fn(outputs, labels).backward()
optimizer.step()

2.3.2 操作

示例之中使用了 cuda 方法把模型复制到 GPU 之上,注释中指出了是把模型的 parameters 和 buffers 移动到 GPU 之上。代码中实际就是使用 self._apply 来调用 cuda(device)。

def cuda(self: T, device: Optional[Union[int, device]] = None) -> T:
    r"""Moves all model parameters and buffers to the GPU.
​
    This also makes associated parameters and buffers different objects. So
    it should be called before constructing optimizer if the module will
    live on GPU while being optimized.
​
    .. note::
        This method modifies the module in-place.
​
    Args:
        device (int, optional): if specified, all parameters will be
            copied to that device
​
    Returns:
        Module: self
    """
    return self._apply(lambda t: t.cuda(device))

我们再看大家熟悉的另外一些函数。

首先,to 方法其实本质也是使用 self._apply 来调用 to(device),我们省略了一些检验代码。

def to(self, *args, **kwargs):
    r"""Moves and/or casts the parameters and buffers.
​
    This can be called as
​
    .. function:: to(device=None, dtype=None, non_blocking=False)
    .. function:: to(dtype, non_blocking=False)
    .. function:: to(tensor, non_blocking=False)
    .. function:: to(memory_format=torch.channels_last)
​
    Its signature is similar to :meth:`torch.Tensor.to`, but only accepts
    floating point or complex :attr:`dtype`s. In addition, this method will
    only cast the floating point or complex parameters and buffers to :attr:`dtype`
    (if given). The integral parameters and buffers will be moved
    :attr:`device`, if that is given, but with dtypes unchanged. When
    :attr:`non_blocking` is set, it tries to convert/move asynchronously
    with respect to the host if possible, e.g., moving CPU Tensors with
    pinned memory to CUDA devices.
​
    See below for examples.
​
    .. note::
        This method modifies the module in-place.
​
    Args:
        device (:class:`torch.device`): the desired device of the parameters
            and buffers in this module
        dtype (:class:`torch.dtype`): the desired floating point or complex dtype of
            the parameters and buffers in this module
        tensor (torch.Tensor): Tensor whose dtype and device are the desired
            dtype and device for all parameters and buffers in this module
        memory_format (:class:`torch.memory_format`): the desired memory
            format for 4D parameters and buffers in this module (keyword
            only argument)
​
    Returns:
        Module: self
    """
​
    device, dtype, non_blocking, convert_to_format = torch._C._nn._parse_to(*args, **kwargs)
​
    def convert(t):
        if convert_to_format is not None and t.dim() in (4, 5):
            return t.to(device, dtype if t.is_floating_point() or t.is_complex() else None,
                        non_blocking, memory_format=convert_to_format)
        return t.to(device, dtype if t.is_floating_point() or t.is_complex() else None, non_blocking)
​
    return self._apply(convert)

其次,cpu 方法也是使用 self._apply 来调用 cpu(device)。

def cpu(self: T) -> T:
    r"""Moves all model parameters and buffers to the CPU.
​
    .. note::
        This method modifies the module in-place.
​
    Returns:
        Module: self
    """
    return self._apply(lambda t: t.cpu())

因此,我们需要分析一下 _apply 方法。

2.3.3 _apply 方法

我们可以看到其主要逻辑是:

  • 遍历 _parameters:

    • 对参数调用fn进行处理,得到param_applied。

      • 用 param_applied 重新设置参数。
    • 如果参数有梯度,则:

      • 对参数的grad调用fn进行处理,得到grad_applied。
      • 用 grad_applied 重新设置参数的梯度。
  • 遍历 _buffers:

    • 对buf调用fn进行处理。
def _apply(self, fn):
    for module in self.children():
        module._apply(fn)
​
    def compute_should_use_set_data(tensor, tensor_applied):
        if torch._has_compatible_shallow_copy_type(tensor, tensor_applied):
            # If the new tensor has compatible tensor type as the existing tensor,
            # the current behavior is to change the tensor in-place using `.data =`,
            # and the future behavior is to overwrite the existing tensor. However,
            # changing the current behavior is a BC-breaking change, and we want it
            # to happen in future releases. So for now we introduce the
            # `torch.__future__.get_overwrite_module_params_on_conversion()`
            # global flag to let the user control whether they want the future
            # behavior of overwriting the existing tensor or not.
            return not torch.__future__.get_overwrite_module_params_on_conversion()
        else:
            return False
​
    # 遍历 _parameters
    for key, param in self._parameters.items():
        if param is not None:
            # Tensors stored in modules are graph leaves, and we don't want to
            # track autograd history of `param_applied`, so we have to use
            # `with torch.no_grad():`
            with torch.no_grad():
                param_applied = fn(param) # 对参数调用fn进行处理,得到param_applied
            should_use_set_data = compute_should_use_set_data(param, param_applied)
            if should_use_set_data:
                param.data = param_applied # 用 param_applied 重新设置
            else:
                assert isinstance(param, Parameter)
                assert param.is_leaf
                # # 用 param_applied 重新设置
                self._parameters[key] = Parameter(param_applied, param.requires_grad)
​
            if param.grad is not None: # 如果参数有梯度
                with torch.no_grad():
                    grad_applied = fn(param.grad) # 对参数的grad调用fn进行处理
                should_use_set_data = compute_should_use_set_data(param.grad, grad_applied)
                if should_use_set_data:
                    param.grad.data = grad_applied # 用 grad_applied 重新设置
                else:
                    assert param.grad.is_leaf
                    self._parameters[key].grad = grad_applied.requires_grad_(param.grad.requires_grad) # 用 grad_applied 重新设置
​
    # 遍历 _buffers                
    for key, buf in self._buffers.items():
        if buf is not None:
            self._buffers[key] = fn(buf) # 对buf调用fn进行处理
​
    return self

因此我们可以看到,移动模型到GPU,其实就是把模型的self._parametersself._buffers 移动到 GPU,并没有对 self._modules 进行移动。我们对模型进行 .cuda() 处理,是将模型的参数放到显存上去(实际使用的时候也是通过这些参数做运算)。

                                        +
                                        |
+---------------------------------+     |     +----------------------------------+
| CPU                             |     |     | CPU                              |
|  +--------------+               |     |     |       +--------------------+     |
|  |Module        |               |     |     |       | Module             |     |
|  |              |               |     |     |       |                    |     |
|  | _parameters+----> Parameters |     |     |       |     _parameters ------+  |
|  |              |               |     |     |       |                    |  |  |
|  | _buffers +------> Buffers    |     |     |     +-----+ _buffers       |  |  |
|  |              |               |     |     |     | |                    |  |  |
|  | _modules     |               |     |     |     | |     _modules       |  |  |
|  |              |               |     |     |     | |                    |  |  |
|  +--------------+               |     |     |     | +--------------------+  |  |
|                                 |     |     |     |                         |  |
+---------------------------------+     |     +----------------------------------+
                                        |           |                         |
                                        +           |                         |
+------------------------------->  Module.cuda() +---------------------------------> Time
                                        +           |                         |
                                        |           |                         |
+---------------------------------+     |     +----------------------------------+
| GPU                             |     |     | GPU |                         |  |
|                                 |     |     |     |                         |  |
|                                 |     |     |     |       Parameters  <-----+  |
|                                 |     |     |     |                            |
|                                 |     |     |     |                            |
|                                 |     |     |     +---->  Buffers              |
|                                 |     |     |                                  |
|                                 |     |     |                                  |
+---------------------------------+     |     +----------------------------------+
                                        |
                                        +

为什么 self._modules 没有被移动?这是因为没有必要,因为_modules 可以认为是一个list,其主要起到了桥梁作用,对其递归遍历可以被用来获取网络所有的 parameters。而这个功能在后续操作之中不是必须的。

2.4 小结

现在我们可以回答了第一个问题:移动模型到GPU这个动作的背后究竟做了哪些操作?

答案时:调用 cuda 或者 to 方法来移动模型到GPU,其实就是把模型的self._parametersself._buffers 移动到 GPU,并没有对 self._modules 进行移动。这个移动过程是递归调用的,是把模型每个叶子都移动到了 GPU 之上。

0x03 在GPU之上调用函数

3.1 CUDA编程模型基础

我们首先介绍一下CUDA编程模型基础。

3.1.1 异构模型

CUDA编程模型是一个异构模型。程序运行在一个异构系统之上,这个异构系统由CPU和GPU构成,它们之间由总线分开,程序运行时候是由CPU和GPU协同工作。

在CUDA之中,有两个重要概念:host和device。

  • Host :CPU及其内存。
  • Device :GPU及其内存。

因此,CUDA 架构下的一个程序也对应分为两个部份:Host 代码和Device代码,它们分别在CPU和GPU上运行。host与device之间可以通信进行数据拷贝。

  • 主机代码(Host Code):在 CPU 上执行的部份,使用Linux(GNU gcc)和Windows(Microsoft Visual C)编译器来编译。大致可以认为认为C语言工作对象是CPU和内存条。
  • 设备代码(Device Code):在GPU上执行的部份,使用 NVIDIA NVCC 编译器来编译。大致可以认为 CUDA C工作对象是GPU及GPU上内存(也叫设备内存)。
+-------------------+        +--------------------+
|                   |        |                    |
|   +----------+    |        |    +----------+    |
|   |          |    |        |    |          |    |
|   |   RAM    |    |        |    |   RAM    |    |
|   |          |    |        |    |          |    |
|   +----+-----+    |        |    +----+-----+    |
|        |          +--------+         |          |
|        |          |        |         |          |
|   +----+-----+    |        |    +----+-----+    |
|   |          |    |        |    |          |    |
|   |   CPU    |    |        |    |   GPU    |    |
|   |          |    |        |    |          |    |
|   +----------+    |        |    +----------+    |
|                   |        |                    |
+-------------------+        +--------------------+

      Host                           Device


3.1.2 并行思想

CUDA 编程的思路是并行思想,大致如下:

  • 把一个很大的执行任务划分成若干个简单的可以重复的操作,然后使用若干个线程来分别执行这些操作,达到并行的目的。
  • 执行任务处理的数据也要对应分组成多个小数据块。比如一个大数据分成若干个GPU组,每个GPU组要再次分成多个线程组,线程组内的张量可能需要再细分为张量处理器能处理的小组。

因此,一个典型的CUDA程序包括串行代码和并行代码。

  • 串行代码是标准C代码,由host执行。
  • 并行代码是CUDA C代码,在device中执行。

CUDA 主程序由CPU开始,即程序由host执行串行代码开始,当遇到需要数据并行处理的部分,则由device执行并行代码来作为补足。device可以独立于host进行大部分操作。当一个device代码启动之后,控制权会立刻返还给CPU来执行其他任务,所以这是一个异步过程。

图来自 docs.nvidia.com/cuda/cuda-c…

3.1.3 处理流程

典型的CUDA程序的执行流程如下:

  • 分配host内存空间并且初始化数据。
  • 分配device显存空间。
  • 将要计算的数据从Host内存之上复制到device显存之上。
  • 调用CUDA核函数在device上完成用户指定的运算。
  • 将计算后GPU内存上的结果复制到Host内存上。
  • 释放device和host上分配的内存。

具体可以参见下图。

010da16d222a960934288b03c67ad6dd.png

3.2 函数

3.2.1 核函数

核函数是在device线程中并行执行的函数。在 CUDA 程序中,主程序在调用GPU内核之前需要对核进行执行配置,以确定线程块数,每个线程块中线程数和共享内存大小。比如在调用时需要用<<参数1,参数2>>来指定核函数需要的线程数量以及线程是如何组织,这样在GPU之中就会启动若干个线程来并行执行这个核函数,每个线程被分配一个唯一的线程号。

限定符执行调用备注
__global__设备端执行可以从主机调用也可以从某些特定设备调用异步操作,host 将并行计算任务发射到GPU的任务调用单之后,不会等待kernel执行完就执行下一步
__device__设备端执行设备端调用不可以和__global__同时用
__host__主机端执行主机调用可省略,不可和__global__同时用,可和__device__同时用,此时函数在device和host都编译。

具体如下:

+------------------------+   +------------------------+
|                        |   |                        |
|                        |   |                        |
| __host__   __global__  |   |          __device__    |
|    +           +       |   |                        |
|    |           |       |   |              +         |
|    |           |       |   |              |         |
|    |           v--------------->          |         |
|    |                   |   |   |          |         |
|    |                   |   |   |          |         |
|    |                   |   |   |          |         |
|    |                   |   |   |          |         |
|    |                   |   |   |          |         |
|    |                   |   |   |          |         |
|    |                   |   |   |          |         |
|    |                   |   |   |          |         |
|    |                   |   |   |          |         |
|    |           +<--------------v          |         |
|    |           |       |   |              |         |
|    |           |       |   |              |         |
|    |           |       |   |              |         |
|    v           v       |   |              v         |
|                        |   |                        |
+------------------------+   +------------------------+

      Host                           Device


这三个限定词其实也是 CUDA 中常见的三种运行场景。其中,device 函数和global函数因为需要在GPU上运行,因此不能调用常见的一些 C/C++ 函数(因为这些函数没有对应的 GPU 实现)。

如下代码是 NVIDIA 的例子,使用内置的 threadIdx 变量,把 A 和 B 两个张量进行相加,得到 C。因此,N 个线程之中每个都会执行 VecAdd() 。

// Kernel definition
__global__ void VecAdd(float* A, float* B, float* C)
{
    int i = threadIdx.x;
    C[i] = A[i] + B[i];
}
​
int main()
{
    ...
    // Kernel invocation with N threads
    VecAdd<<<1, N>>>(A, B, C);
    ...
}

3.2.2 PyTorch 样例

我们从 third_party/cub/cub/device/dispatch/dispatch_reduce.cuh 找一个核函数例子来看看。

/**
 * Reduce region kernel entry point (multi-block).  Computes privatized reductions, one per thread block.
 */
template <
    typename                ChainedPolicyT,             ///< Chained tuning policy
    typename                InputIteratorT,             ///< Random-access input iterator type for reading input items \iterator
    typename                OutputIteratorT,            ///< Output iterator type for recording the reduced aggregate \iterator
    typename                OffsetT,                    ///< Signed integer type for global offsets
    typename                ReductionOpT>               ///< Binary reduction functor type having member <tt>T operator()(const T &a, const T &b)</tt>
__launch_bounds__ (int(ChainedPolicyT::ActivePolicy::ReducePolicy::BLOCK_THREADS))
__global__ void DeviceReduceKernel(
    InputIteratorT          d_in,                       ///< [in] Pointer to the input sequence of data items
    OutputIteratorT         d_out,                      ///< [out] Pointer to the output aggregate
    OffsetT                 num_items,                  ///< [in] Total number of input data items
    GridEvenShare<OffsetT>  even_share,                 ///< [in] Even-share descriptor for mapping an equal number of tiles onto each thread block
    ReductionOpT            reduction_op)               ///< [in] Binary reduction functor
{
    // The output value type
    typedef typename If<(Equals<typename std::iterator_traits<OutputIteratorT>::value_type, void>::VALUE),  // OutputT =  (if output iterator's value type is void) ?
        typename std::iterator_traits<InputIteratorT>::value_type,                                          // ... then the input iterator's value type,
        typename std::iterator_traits<OutputIteratorT>::value_type>::Type OutputT;                          // ... else the output iterator's value type
​
    // Thread block type for reducing input tiles
    typedef AgentReduce<
            typename ChainedPolicyT::ActivePolicy::ReducePolicy,
            InputIteratorT,
            OutputIteratorT,
            OffsetT,
            ReductionOpT>
        AgentReduceT;
​
    // Shared memory storage
    __shared__ typename AgentReduceT::TempStorage temp_storage;
​
    // Consume input tiles
    OutputT block_aggregate = AgentReduceT(temp_storage, d_in, reduction_op).ConsumeTiles(even_share);
​
    // Output result
    if (threadIdx.x == 0)
        d_out[blockIdx.x] = block_aggregate;
}

3.3 小结

目前我们知道了,PyTorch 其实可以通过调用 __global__ 方法来在GPU之上执行并行操作。这回答了我们的第二个问题:如何在 CPU 之上调用 GPU 操作?

0x04 在GPU/CPU之间切换

我们接下来分析如何在GPU/CPU之间切换。

由示例代码可以知道,只要调用了 cuda 函数把模型移动到 GPU 之上,我们就可以使用 CUDA global 核函数在GPU上进行并行运算。

model = ToyModel().cuda(device_ids[0]) # 这里复制模型到 GPU 之上
ddp_model = DDP(model, device_ids)
​
loss_fn = nn.MSELoss()
optimizer = optim.SGD(ddp_model.parameters(), lr=0.001)
​
optimizer.zero_grad()
outputs = ddp_model(torch.randn(20, 10))

但是我们忽略了一个问题,就是 PyTorch 怎么知道此时应该调用GPU对应的 global 核函数?为什么 PyTorch 就不调用 CPU 函数或者其他设备的函数了?这就是我们接下来需要分析的。

4.1 Dispatcher 机制

此处我们主要借鉴 blog.ezyang.com/2020/09/let…

4.1.1 问题

在PyTorch中,operator 所表现出预期行为是由很多机制共同作用导致的,比如:

  • 做实际工作的kernel。
  • 是否支持反向自动微分,例如,使 loss.backward() 正常工作的标记位。
  • 是否启用了torch.jit.trace。
  • 如果你正在vmap调用中,所运行operator将会表现出不同的批处理行为。

因此,我们知道有太多不同的方式可以对PyTorch operator进行不同的解释,如果我们试图在一个名为add的单一函数里面处理所有的行为,我们的实现代码会很快演变成一个不可维护的混乱局面。

所以我们需要有一个机制来解决这个问题,这个机制不仅仅是一个if语句这么简单,而是PyTorch内部一个非常重要的抽象,而且它必须在尽可能不降低PyTorch性能的情况下做到这一点。这个机制就是 Dispatcher。

4.1.2 什么是 Dispatcher

什么是dispatcher?dispatcher对于每个operator都会维护一个函数指针表,这些函数为每个dispatch key提供了对应的实现,这套机制大致对应于PyTorch中的一个横切关注点。在上图中,你可以看到在这个表中有针对不同后端(CPU、CUDA、XLA)以及更高级概念(例如 autograd 和跟踪)的dispatch条目。dispatcher的工作是根据输入的tensor和其他一些东西来计算出一个dispatch key,然后跳转到函数指针表所指向的函数。

熟悉 C++ 的人可能会注意到,这个函数指针表与C++中的虚表非常相似。在C++中,对象的虚函数是通过将每个对象与一个虚表的指针相关联来实现的,该虚表包含了有关对象上每个虚函数的实现。在PyTorch中,我们基本上重新实现了虚拟表,但有一些区别。

  • dispatch表是按operator分配的,而虚表是按类分配的。这意味着我们可以通过分配一个新的dispatch表来扩展所支持的operator集。与其不同的是,对于一个C++对象,你可以通过继承子类来扩展类型,但你不能轻易添加虚函数。与普通的面向对象系统不同,PyTorch大部分的可扩展性在于定义新的operator(而不是新的子类),所以这种权衡是合理的。此外,dispatch key的种类不是公开可扩展的,我们希望那些想添加新dispatch key的使用者通过向PyTorch核心团队提交一个补丁来添加他们的dispatch key。
  • 我们的dispatch key的计算考虑了operator的所有参数(multiple dispatch)以及线程本地状态(TLS)。这与虚表不同,在虚表中只有第一个对象(this指针)很重要。
  • 最后,dispatcher支持boxing和unboxing作为op的调用约定的一部分。在文章的最后部分会有更多关于这个的内容。

有趣的历史笔记:我们曾经使用虚函数来实现动态dispatch,当我们意识到需要比虚表更多的能力时,我们重新实现了动态dispatch。

4.1.3 如何计算key

那么,我们究竟是如何计算dispatch key的呢?我们是基于dispatch key set来完成的,dispatch key set是一个基本抽象,它是dispatch key的一个bitset。大致来讲,我们综合来自不同来源的dispatch key sets(在某些情况下屏蔽一些key)来得到一个最终的dispatch key set。然后我们在这个set中挑选优先级最高的key(dispatch keys按某些优先级隐式排序),这就是我们这次应该调用的结果。那么,这些dispatch key sets的来源是什么?

  • 每个张量输入都有一个由该张量上的所有dispatch key组成的dispatch key set(直观地说,这些dispatch key的值会是类似 “CPU”字符串这样的东西,这告诉我们该张量是一个CPU张量,所以应该由dispatch表中的CPU handler来处理)。
  • 我们还有一个local include set,用于 "模态(modal) "功能,例如tracing,它不与任何张量关联,而是某种线程的本地模态,用户可以在某些范围内打开或关闭。
  • 最后,我们有一个global set,它包含了始终应该被考虑的dispatch key(自从写下这张PPT以来,Autograd已经从global set转移到了张量之上。然而系统的高级结构并没有改变)。

除了这些,还有一个local exclude set,其用从dispatch排除某些dispatch key。一个常见的场景是一个handler负责处理一个key,然后通过local exclude set将自己屏蔽掉,这样我们以后就不会尝试重新处理这个key。

4.1.4 注册

我们接下来看看如何注册这个dispatch key 到 dispatch 表之中。这个过程通过operator registration API来实现。操作符注册 API 有三种主要方式:

  • 为operator定义模式。
  • 然后在对应的key上注册实现。
  • 最后,有一个 fallback 方法,用户可以使用它为某个key对应的所有运算符定义同一个处理程序。

为了可视化 operator registration的工作,让我们想象一下,所有op的dispatch表共同形成一个二维网格,像这样:

  • 纵轴上是PyTorch中支持的每个op。
  • 横轴上是系统支持的每个dispatch key。

operator registration 行为就是在这两个轴定义出的单元格中填写对应的实现。

在一个特定的dispatch key上为一个operator注册kernel函数时,我们会填写一个单元格(下面的蓝色)的内容。

4.2 Dispatcher 代码

我们接下来通过源码来看看。

4.2.1 虚函数表

4.2.1.1 例子

我们可以从 aten/src/ATen/native/native_functions.yaml 之中找到一些虚函数的例子。

# zero 操作对应的虚函数表
- func: zero_(Tensor(a!) self) -> Tensor(a!) 
  device_check: NoCheck   # TensorIterator
  variants: method, function
  dispatch:
    CPU, CUDA: zero_
    Meta: zero_meta_
    SparseCPU, SparseCUDA: zero_sparse_
    MkldnnCPU: mkldnn_zero_

# sub.out 对应的虚函数表
- func: sub.out(Tensor self, Tensor other, *, Scalar alpha=1, Tensor(a!) out) -> Tensor(a!)
  device_check: NoCheck   # TensorIterator
  structured: True
  structured_inherits: TensorIteratorBase
  dispatch:
    CPU, CUDA: sub_out
    SparseCPU, SparseCUDA: sub_out_sparse

# sub.Tensor 对应的虚函数表
- func: sub.Tensor(Tensor self, Tensor other, *, Scalar alpha=1) -> Tensor
  device_check: NoCheck   # TensorIterator
  variants: function, method
  structured_delegate: sub.out
  dispatch:
    SparseCPU, SparseCUDA: sub_sparse
4.2.1.2 Operator的实现

我们可以看看 zero 的两个实现,下面是MkldnnCPU的实现。

Tensor& mkldnn_zero_(Tensor& self) {
  using Vec = vec::Vectorized<float>;
​
  ideep::tensor& x = itensor_from_mkldnn(self);
​
  auto n = x.get_nelems();
  auto* x_ = static_cast<float*>(x.get_data_handle());
  parallel_for(0, n, 2048, [x_](int64_t begin, int64_t end) {
    vec::map(
        [](Vec /* unused */) { return 0.0; },
        x_ + begin,
        x_ + begin,
        end - begin);
  });
​
  return self;
}

又比如下面是SparseCPU, SparseCUDA 的对应实现:

// --------------------------------------------------------------------
// zero_(SparseTensor)
// --------------------------------------------------------------------
// hummu hummu
SparseTensor& zero_sparse_(SparseTensor& self) {
  AT_ASSERT(self.is_sparse());
  at::zeros_out(self, get_sparse_impl(self)->sizes());
  return self._coalesced_(true);
}

4.2.2 Dispatcher 定义

我们接下来看看Dispatcher的定义,这里只给出部分成员变量。

class TORCH_API Dispatcher final {
private:
  // For direct access to backend fallback information
  friend class impl::OperatorEntry;
​
  struct OperatorDef final {
    explicit OperatorDef(OperatorName&& op_name)
    : op(std::move(op_name)) {}
    impl::OperatorEntry op;
    size_t def_count = 0;
    size_t def_and_impl_count = 0;
  };
  friend class OperatorHandle;
  template<class> friend class TypedOperatorHandle;
​
public:
​
  static Dispatcher& realSingleton();
​
  //存储所有的算子,并在其成员变量中存储了每个算子的不同版本,比如cpu,cuda,autograd....
  std::list<OperatorDef> operators_;
  //注册算子时会将算子名称和方法也存储在这个里面, 这样就可以快速的通过名字查找到算子方法(其中包含了成员OperatorDef) 
  LeftRight<ska::flat_hash_map<OperatorName, OperatorHandle>> operatorLookupTable_;
  // Map from namespace to debug string (saying, e.g., where the library was defined)
  ska::flat_hash_map<std::string, std::string> libraries_;
  std::array<impl::AnnotatedKernel, static_cast<uint8_t>(DispatchKey::NumDispatchKeys)> backendFallbackKernels_;
  std::unique_ptr<detail::RegistrationListenerList> listeners_;
  std::mutex mutex_;
};

4.2.3 注册

我们接下来给出注册虚函数表的方法。

RegistrationHandleRAII Dispatcher::registerImpl(
  OperatorName op_name,
  c10::optional<DispatchKey> dispatch_key,
  KernelFunction kernel,
  c10::optional<impl::CppSignature> cpp_signature,
  std::unique_ptr<FunctionSchema> inferred_function_schema,
  std::string debug
) {
  std::lock_guard<std::mutex> lock(mutex_);
  auto op = findOrRegisterName_(op_name);
  auto handle = op.operatorDef_->op.registerKernel( // 进行注册
    *this,
    dispatch_key,
    std::move(kernel),
    std::move(cpp_signature),
    std::move(inferred_function_schema),
    std::move(debug)
  );
​
  ++op.operatorDef_->def_and_impl_count;
​
  return RegistrationHandleRAII([this, op, op_name, dispatch_key, handle] {
    deregisterImpl_(op, op_name, dispatch_key, handle);
  });
}
4.2.3.1 注册表

OperatorEntry代表了一个算子,以及该算子的dispatch table,这里只给出成员变量。

class TORCH_API OperatorEntry final { //代表了一个算子,以及该算子的dispatch table
public:
  OperatorName name_;
  c10::optional<AnnotatedSchema> schema_;
  //存储了不同key对应的算子实现版本,比如cpu,cuda,autograd 等等,所有的算子版本都会在这个table里面
  std::array<KernelFunction, static_cast<uint8_t>(DispatchKey::NumDispatchKeys)> dispatchTable_;
  DispatchKeyExtractor dispatchKeyExtractor_;
  //不同 DispatchKey对应了不同的版本的kernel算子实现版本
  ska::flat_hash_map<DispatchKey, std::list<AnnotatedKernel>> kernels_;
};
4.2.3.2 注册行为

最终注册行为就是往 dispatchTable_ 之中设置。

void OperatorEntry::updateDispatchTableEntry_(const c10::Dispatcher& dispatcher, DispatchKey dispatch_key) {
  auto dispatch_ix = static_cast<uint8_t>(dispatch_key);
  dispatchTable_[dispatch_ix] = computeDispatchTableEntry(dispatcher, dispatch_key);
  dispatchKeyExtractor_.setOperatorHasFallthroughForKey(dispatch_key, dispatchTable_[dispatch_ix].isFallthrough());
}

4.2.4 如何dispatch

4.2.4.1 调度依据

PyTorch 之中会依据dtype、device和layout的不同来调度不同的operator。

  • 大多数类型(比如int32)可以使用模版方式直接进行映射,但是某些operator 不支持模版功能,就需要dispatcher这样的动态调度器。
  • PyTorch的tensor不仅可以运行在CPU上,还可以跑在GPU,mkldnn和xla等设备,这也需要动态调度。
  • layout是指tensor中元素的排布,这就有strided layout和sparse layout的区别,所以也需要动态调度。
4.2.4.2 调度代码

我们这里这是给出部分代码,有兴趣的读者继续继续深入。

template<class Return, class... Args>
C10_DISPATCHER_INLINE_UNLESS_MOBILE Return Dispatcher::call(const TypedOperatorHandle<Return(Args...)>& op, Args... args) const {
  detail::unused_arg_(args...);  
 
  // 得到key set
  auto dispatchKeySet = op.operatorDef_->op.dispatchKeyExtractor()
    .template getDispatchKeySetUnboxed<Args...>(args...);
  TORCH_INTERNAL_ASSERT_DEBUG_ONLY(!c10::isAliasDispatchKey(dispatchKeySet.highestPriorityTypeId()));
  
  // 得到算子
  const KernelFunction& kernel = op.operatorDef_->op.lookup(dispatchKeySet.highestPriorityTypeId());
  
  // 进行调度
#ifndef PYTORCH_DISABLE_PER_OP_PROFILING
  bool pre_sampled = false;
  if (C10_UNLIKELY(at::shouldRunRecordFunction(&pre_sampled))) {
    return callWithDispatchKeySlowPath<Return, Args...>(op, pre_sampled, dispatchKeySet, kernel, std::forward<Args>(args)...);
  }
#endif  // PYTORCH_DISABLE_PER_OP_PROFILING
  return kernel.template call<Return, Args...>(op, dispatchKeySet, std::forward<Args>(args)...);
}
4.2.4.3 key

我们接下来看看key的定义,因为太多,所以我们只给出部分数值。

enum class DispatchKey : uint8_t {
  CPU, // registered at build/aten/src/ATen/RegisterCPU.cpp
  CUDA, // registered at build/aten/src/ATen/RegisterCUDA.cpp
  HIP, // NB: I think this is not actually used, due to Note [Masquerading as
  // CUDA]
  FPGA, // Xilinx support lives out of tree at
  // https://gitlab.com/pytorch-complex/vitis_kernels
  MSNPU, // unused externally, but tested at
  // test/cpp_extensions/msnpu_extension.cpp
  XLA, // lives out of tree at https://github.com/pytorch/xla
  MLC, // lives out of tree at https://github.com/pytorch/MLCompute
  Vulkan,
  Metal,
  XPU, // For out of tree Intel's heterogeneous computing plug-in
  HPU, // For out of tree & closed source integration of HPU / Habana
  VE, // For out of tree & closed source integration of SX-Aurora / NEC
  Lazy, // For lazy tensor backends
  // A meta tensor is a tensor without any data associated with it.  (They
  // have also colloquially been referred to as tensors on the "null" device).
  // A meta tensor can be used to dry run operators without actually doing any
  // computation, e.g., add on two meta tensors would give you another meta
  // tensor with the output shape and dtype, but wouldn't actually add anything.
  Meta,
  // Here are backends which specify more specialized operators
  // based on the dtype of the tensor.
  QuantizedCPU, // registered at build/aten/src/ATen/RegisterQuantizedCPU.cpp
  QuantizedCUDA, // registered at build/aten/src/ATen/RegisterQuantizedCUDA.cpp
  QuantizedXPU, // For out of tree Intel's heterogeneous computing plug-in
  // This backend is to support custom RNGs; it lets you go
  // to a different kernel if you pass in a generator that is not a
  // traditional CPUGeneratorImpl/CUDAGeneratorImpl.  To make use of this
  // key:
  //  1) set it as a second parameter of at::Generator constructor call in
  //     the user-defined PRNG class.
  //  2) use it as a dispatch key while registering custom kernels
  //     (templatized kernels specialized for user-defined PRNG class)
  // intended for out of tree use; tested by aten/src/ATen/test/rng_test.cpp
  CustomRNGKeyId,
​
  // Here are backends which specify more specialized operators
  // based on the layout of the tensor.  Note that the sparse backends
  // are one case where ordering matters: sparse multi-dispatches with
  // the corresponding dense tensors, and must be handled before them.
  MkldnnCPU, // registered at build/aten/src/ATen/RegisterMkldnnCPU.cpp
  // NB: not to be confused with MKLDNN, which is Caffe2 only
  SparseCPU, // registered at build/aten/src/ATen/RegisterSparseCPU.cpp
  SparseCUDA, // registered at build/aten/src/ATen/RegisterSparseCUDA.cpp
  SparseHIP, // TODO: I think this is not actually used, due to Note
  // [Masquerading as CUDA]
  SparseXPU, // For out of tree Intel's heterogeneous computing plug-in
  SparseVE, // For out of tree & closed source integration of SX-Aurora / NEC
  SparseCsrCPU,
  SparseCsrCUDA,
​
  AutogradOther,
  AutogradCPU,
  AutogradCUDA,
  AutogradXLA,
  AutogradLazy,
  AutogradXPU,
  AutogradMLC,
  AutogradHPU,
​
  ......
};
4.2.4.4 key的使用

因为篇幅所限,我们无法深入分析,这里只给出从 DeviceType 出发的情景。我们从下面函数可以看到,如何从 DeviceType 映射到 DispatchKey 类型。

template <typename Func>
inline CppFunction dispatch(c10::DeviceType type, Func&& raw_f) {
  auto deviceTypeToDispatchKey = [](c10::DeviceType t){
    switch (t) {
      // This list is synchronized with the k-constants in c10/core/DeviceType.h
      case c10::DeviceType::CPU:
        return c10::DispatchKey::CPU;
      case c10::DeviceType::CUDA:
        return c10::DispatchKey::CUDA;
      case c10::DeviceType::XLA:
        return c10::DispatchKey::XLA;
      case c10::DeviceType::Lazy:
        return c10::DispatchKey::Lazy;
      case c10::DeviceType::MLC:
        return c10::DispatchKey::MLC;
      case c10::DeviceType::Meta:
        return c10::DispatchKey::Meta;
      case c10::DeviceType::HIP:
        return c10::DispatchKey::HIP;
      case c10::DeviceType::MSNPU:
        return c10::DispatchKey::MSNPU;
      case c10::DeviceType::HPU:
        return c10::DispatchKey::HPU;
      default:
        TORCH_CHECK(false,
          "Device type ", t, " cannot be overloaded at dispatch time, "
          "please file a bug report explaining what you were trying to do.");
    }
  };
  return dispatch(deviceTypeToDispatchKey(type), std::forward<Func>(raw_f));
}

4.3 小结

至此,我们知道,通过 Dispatcher 机制,PyTorch 可以依据dtype、device和layout的不同来调度不同的operator。这就解答了我们第三个问题:如何在 CPU,GPU 操作之间无缝切换?

关于第四个问题:是否需要把损失函数移动到 GPU 之上?,我们也有了解答:

损失函数的参数是前向传播的outputs和label,outputs已经在GPU之上(因为训练数据已经在GPU之上),label 也被用户手动设置到GPU之上。所以损失函数的参数都已经在GPU之上,这样 Dispather 就依据device会调用到GPU对应的operator,所以不需要把损失函数移动到GPU之上。

我们整理一个总体逻辑如下,序列是:

  1. 把训练数据 inputs 移动到GPU。
  2. 进行前向操作,假设只有一个operator,就是 op1,使用 device='GPU' 这个 dispatch key 去 Dispatcher 查找。
  3. 找到了 op1-gpu 这个operator,进行计算,得出 outputs。
  4. outputs 就自动存在于 GPU 之上。
  5. 把 Labels 也放到 GPU 之上。
  6. 进行损失函数运算,假设只有一个 operator,就是 op2,此时损失函数的参数都在GPU之上,所以使用 device= 'GPU' 这个 dispatch key 去 Dispatcher 查找。
  7. 找到了 op2-gpu 这个operator,进行计算,得出 loss。
                           +--------------------+
         +-----------+     | Forward            |      +------------+     +------------------+
         | GPU       |     |                    |      | GPU        |     | Loss Function    |
         |           +---> |    op1   op1-gpu() +----> |            +---> |                  |   +--------+
         |   Inputs  | 1   |                    |  4   |   Outputs  |     |                  |   | GPU    |
         |           |     |     +        ^     |      |            |     |                  |   |        |
         +-----------+     |     |        |     |      +------------+     |  op2   op2-gpu() +-->+  loss  |
                           |     |        |     |                         |                  |   |        |
                           +--------------------+      +------------+     |   +        ^     |   |        |
                                 |        |            | GPU        | 5   |   |        |     |   +--------+
                                 |        |            |            +---> |   | 6      | 7   |
                               2 |        | 3          |   Labels   |     |   |        |     |
                                 |        |            |            |     |   |        |     |
                                 |        |            +------------+     +------------------+
    +----------------------------+        +--------------------------------+  |        |
    |                                                                      |  |        |
+-----------------------------------------------------------------------------+        |
|   |                                                                      |           |
|   |          +-------------------------------------------------------+   |           |
|   |          | Dispather                                             |   |           |
|   |          |       +          +          +            +            |   |           |
|   |          |       |   XLA    |   CPU    |    Metal   |    GPU     |   |           |
|   |          | +---------------------------------------------------+ |   |           |
|   |          |       |          |          |            |            |   |           |
|   +--------> |   OP1 | op1-xla  | op1-cpu  |  op1-metal |  op1-gpu   +---+           |
| 'device=GPU' |       |          |          |            |  +------+  |               |
|              | +---------------------------------------------------+ |               |
|              |       |          |          |            |            |               |
+------------> |   OP2 | op2-xla  | op2-cpu  |  op2-metal |  op2-gpu   +---------------+
  'device=GPU' |       |          |          |            |  +------+  |
               | +---------------------------------------------------+ |
               |       |          |          |            |            |
               |   OP3 | op3-xla  | op3-cpu  |  op3-metal |  op3-gpu   |
               |       |          |          |            |            |
               | +---------------------------------------------------+ |
               +-------------------------------------------------------+


手机如下:

至此,GPU相关分析结束,下一篇我们开始分析DataParallel,敬请期待。

0xEE 个人信息

★★★★★★关于生活和技术的思考★★★★★★

微信公众账号:罗西的思考

0xFF 参考

blog.ezyang.com/2020/09/let…

pytorch.org/tutorials/a…

GPU多卡并行训练总结(以pytorch为例)

当代研究生应当掌握的并行训练方法(单机多卡)

分布式训练从入门到放弃

再谈PyTorch的初始化(上)

pytorch中的dispatcher

【译】聊聊Pytorch Dispatcher

扩展Pytorch:利用CUDA实现算子(二)

PyTorch ATen代码的动态生成

blog.csdn.net/qq_23858785…

CUDA 函数前缀

CUDA C编程入门

CPU—GPU并行处理—CUDA编程从想入门到放弃

blog.csdn.net/weixin_4223…

blog.csdn.net/crazy_sunsh…

CPU、GPU、CUDA,CuDNN 介绍

CUDA编程(三): GPU架构了解一下!

CUDA编程入门极简教程

写CUDA到底难在哪?

深入浅出PyTorch(算子篇)

深入浅出全连接层(fully connected layer)

Pytorch拓展进阶(二):Pytorch结合C++以及Cuda拓展

Pytorch拓展进阶(一):Pytorch结合C以及Cuda语言

PyTorch 源码解读之 cpp_extension:揭秘 C++/CUDA 算子实现和调用全流程

pytorch中的dispatcher