加速 GPU 训练：流水线、融合与自定义内核

Corwin 和另一位讲者介绍高效 GPU 训练。CPU 擅长低延迟的顺序工作，GPU 则提供巨大的并行吞吐量。问题是如何组织机器学习工作负载，才能利用这种吞吐量。

Corwin 提出一个刻意引人注意的比较：由于 GPU 执行存在开销，单个浮点运算发送到 GPU 后可能比在 CPU 上慢得多。但这并不是现实中的 GPU 工作负载。相比之下，矩阵乘法可以将大量工作分配给许多线程。

在示例中，一个 2048×2048 的矩阵乘法在 CPU 上约需 28 毫秒，在 GPU 上约需 200 微秒。这些示例测量结果与单次运算的比较相反：并行工作足够多时，GPU 会快得多。

第二位讲者回顾训练过程：将输入送入模型，把输出与目标比较以计算损失，计算损失对参数的梯度，再走一小步以降低损失。这个过程重复数千次或数百万次。

PyTorch 提供 GPU 内核和自动微分。数据加载器提供批次，模型计算输出和损失，backward 计算梯度，优化器更新参数。将模型、输入和目标移动到 GPU 后，开销较大的前向和反向计算就在 GPU 上执行。

Corwin 解释 CPU–GPU 流水线。CPU 准备数据并提交内核启动，包括通过驱动传递参数；GPU 执行线性层、ReLU 等操作。内核大致就是实现 GPU 操作的函数。

简单的时间线会让 GPU 在 CPU 准备每个操作时空闲。实际上内核启动是异步的：GPU 处理较早操作时，CPU 可以将下一个内核加入队列并准备下一批数据。队列让有用的 CPU 和 GPU 工作重叠。

NVIDIA Nsight Systems 等性能分析器会显示 CPU 端 CUDA API 调用和 GPU 硬件活动。Corwin 首先寻找 GPU 活动条中的间隙：空闲时间意味着设备正在等待，而非执行有用工作。

将结果复制回 CPU 可能强制同步。CPU 要等待队列中的 GPU 工作和传输完成，流水线因此被排空；随后准备数据的时间就可能表现为 GPU 空闲。成本取决于同步之后发生什么，而不只是同步本身。

用于记录指标的显式 output.cpu() 是明显的同步点。更隐蔽的例子是 Python 的 if 将 GPU 张量（例如 is-NaN 检查结果）转换为布尔值：CPU 必须知道答案才能选择分支。张量操作本身不一定是同步点。

一位观众提出，优化器步骤可能需要同步以重新加载权重。Corwin 解释说，权重可以留在 GPU 上，优化器操作也可以启动 GPU 内核，因此 CPU 端调用不一定意味着传输或同步。

演示中的 CPU 到 GPU 复制也可能阻塞。GPU 读取时源内存必须保持可用，而示例中的默认传输行为会等待完成。极小的传输可能表现不同，但关键是检查实际流水线。

如果 CPU 在短暂同步后立即继续启动工作，同步不一定代价高昂。如果同步后要进行耗时的下一批数据准备，它就尤其有害，因为这项准备不能再与先前的 GPU 计算重叠。

一种办法是在 GPU 上累积指标，只偶尔复制回 CPU。NaN 检查可以异步处理，例如交给单独的 CPU 线程，让主线程继续启动内核。这会改变行为：采取失败处理措施之前，可能会多执行一些工作。

对于异步主机到设备传输，固定主机内存允许直接内存访问，避免源数据被换出。示例还使用 non_blocking=True。程序员必须管理源数据的生命周期，并避免在传输进行时修改它。

即使移除这些同步，复制仍可能留下间隙；可以让传输与计算重叠，但需要更仔细地编写代码。GPU 活动条填满只表示 GPU 很忙，并不表示内核高效利用了资源。

简化的 GPU 由许多流式多处理器（SM）组成，每个 SM 都有计算资源、寄存器和供协作线程使用的共享内存。示例中的 H100 有 132 个 SM。高效内核将问题分成基本可以独立处理的块，再合并或存储结果。

全局内存提供输入并存储输出，L2 缓存由多个 SM 共享。演讲给出 H100 的近似带宽：全局内存为 3 TB/s，L2 为 7 TB/s，共享内存为 26 TB/s。这种层次结构使在计算附近复用数据很重要，但仍受容量和访问方式限制。

主要瓶颈包括计算、内存传输和内核启动开销。示例中 CPU 启动内核可能需要约 5–10 微秒，因此即使每个操作本身很快，一连串微小的 GPU 内核也可能让设备等待工作。

第二位讲者介绍解决方法。Tensor Core 可以加速矩阵乘法等合适的计算密集型操作。内核融合可以减少内存传输并摊薄启动开销；CUDA 图则可以记录并重放小内核序列，减少反复启动的开销。

考虑层归一化、ReLU 和跳跃连接中的加法。如果独立内核反复将中间结果写入全局内存再读出，内存传输就会占主导。融合后的内核可以读取数据、执行多个操作，再存储最终结果。

为每种组合单独编写融合内核并不现实。torch.compile 会追踪模型、寻找融合机会，并可对不同方案进行基准测试。首次编译很慢，但后续执行可以从生成的内核中受益。

计算图中断会减少融合机会。讲者提到的可能原因包括不受支持的 Python 结构、自定义操作以及意外的 CPU–GPU 同步。使用 fullgraph=True 编译可以暴露中断，而不是悄悄拆分计算图。

作为具体的自定义内核示例，讲者希望对许多张量求和以记录平均梯度。Python 循环会顺序启动许多小型归约。在性能分析结果中，CPU 将一个内核加入队列约需 5 微秒，而内核运行可能不到 1 微秒，GPU 因此等待。

这个例子中 torch.compile 未能自动解决问题，甚至可能更慢。它追踪的是展开后的循环，而没有发现所需的并行工作划分；当融合受到限制时，生成的内核可能不如高度优化的库内核。自定义内核可以直接表达这种并行性。

所提出的划分方式是将一个张量分配给一个线程块，线程块由 SM 动态调度。即使有数千个张量，132 个 SM 仍可处理：一个 SM 完成一个块后可以接收下一个，因此不同大小的张量可以随时间得到负载均衡。

在一个线程块内部，各线程计算部分和，然后一起归约。相邻线程读取相邻内存位置，使访问能够高效合并。最终结果写入全局内存。

CUDA 示例为每个张量启动一个线程块，每块有 1,024 个线程。blockIdx 选择张量，threadIdx 选择线程的起始元素。每个线程以线程块大小为步长处理元素，然后线程块归约各部分和。

线程以 32 个为一组的 warp 执行。代码先在每个 warp 内归约，再使用共享内存存放每个 warp 的一个部分结果，并同步线程块，确保所有 warp 的结果可用。最后一次 warp 归约合并这些结果，由线程零存储答案。

报告中的示例从约 20 毫秒改善到 30 微秒，接近提升一千倍。讲者提醒，初学者编写这样的内核可能需要一周，有经验的开发者可能需要一个下午；是否值得取决于工作负载。

被问到自定义内核该用什么语言时，讲者建议先尝试更清晰或向量化的 PyTorch。接下来通常选择 Triton；当 Triton 无法表达所需内核，或极致性能很重要时，可以使用 CUDA。团队很少用 PTX 编写整个内核，但可能插入单条指令。