Torch Code

导言

研究PTA的线程设置时，需要了解torch的线程调度设计

后备线程池设计

• 如有需要，torch::parallel_for或者autograd会初始化和CPU核数相同的线程池（export OMP_NUM_THREADS=10来限制），以供复杂的intra-op使用，
  • 感知明显的场景是并行读文件torch.nn.DataParallel/DataLoader或者CPU-GPU数据传输tensor.to(device)或者tensor.cuda()。
• 前向计算下发算子到GPU/NPU是在主线程进行的。如果在CPU计算，一般都会考虑并行，使用PyTorch的CPU的API, at::parallel_for()；
• 反向计算autograd会启动线程池

// In torch/csrc/autograd/engine.cpp
void Engine::thread_init() {
    // 初始化线程池，通常与 CPU 核心数量匹配
    std::thread worker_thread(&Engine::thread_main, this);
    worker_threads_.emplace_back(std::move(worker_thread));
}

反向传播比正向计算更容易并行化

在深度学习中，前向传播和反向传播是两个不同的计算过程：

1. 前向传播 (Forward Pass)

前向传播是将输入数据依次通过网络的各层，从输入层传递到输出层。这个过程中，计算每一层的输出值，也就是网络的推理。所有的运算（如矩阵乘法、激活函数、卷积等）会根据当前的网络参数（权重和偏置）计算，最终得到网络的输出。

• 算子：前向传播中的操作符通常涉及线性变换（如矩阵乘法、卷积）和非线性变换（如激活函数、正则化）。
• 并行化：虽然有一定的并行计算（比如在同一层内可以同时计算多个神经元的激活），但整个前向传播过程是顺序依赖的。必须从第一层开始，逐步计算到最后一层，前面层的输出是后面层的输入。

2. 反向传播 (Backward Pass, Autograd)

反向传播是在前向传播完成后，通过链式法则（链式求导）计算每个参数的梯度，用于优化网络参数。在 PyTorch 中，Autograd 系统通过计算图记录前向传播中的操作，并在反向传播时自动计算导数。

• 算子：每个前向操作符对应的反向操作符并不完全相同。例如，对于矩阵乘法，前向是 (\mathbf{C} = \mathbf{A} \cdot \mathbf{B})，反向传播需要分别计算 (\frac{\partial L}{\partial \mathbf{A}}) 和 (\frac{\partial L}{\partial \mathbf{B}})，这通常是不同的计算。
• 并行化：反向传播中的计算可以更容易并行化，原因如下：
  1. 依赖关系减少：虽然反向传播依赖前向传播的结果，但其梯度计算可以同时对多个参数进行，因为反向传播中涉及的是各个参数的梯度更新，而不是层与层之间的顺序计算。
  2. 梯度累加：当多个参数的梯度可以同时计算时，它们可以并行执行，尤其是在大规模并行硬件（如 GPU 或 NPU）上。
  3. 重用计算图：前向传播生成的计算图允许 Autograd 系统自动跟踪操作和依赖，从而能够在反向传播中调度并行执行。

为什么反向传播可以更容易并行化？

反向传播的梯度计算通常可以并行化，而前向传播则更多的是依赖顺序执行。原因在于：

1. 计算的独立性：反向传播过程中，每一层的梯度计算在一定程度上是独立的，多个参数的梯度可以同时计算。
2. 内存优化：在深度网络中，前向传播过程中所有层的输出必须存储，以供反向传播使用。而反向传播的输出则是梯度，不需要保存完整的中间值，可以即时计算并释放内存，这使得反向传播在大批量训练中更具效率。

• 前向传播主要是逐层依赖顺序的计算，难以完全并行化。
• 反向传播则可以通过分解梯度计算，实现更高效的并行性。
• 反向传播和前向传播一样，很多操作都可以通过 GPU/NPU 加速。

反向传播操作是否下发到 GPU/NPU

反向传播中的算子通常也会下发到 GPU/NPU 上执行。这是因为计算梯度的过程同样非常耗时，尤其在深度神经网络中，涉及大量矩阵运算。GPU/NPU 的并行计算能力能够显著加速这些操作，和前向传播一样，反向传播的很多操作都可以在 GPU/NPU 上执行。

参考文献