The Performance of CUDA with the Flexibility of PyTorch

Mark Saroufim

目录

  1. 摘要
  2. llm.c：性能与开发的权衡
  3. PyTorch 的灵活性：Eager 模式
  4. 硬件瓶颈的演变
  5. torch.compile：应对内存带宽瓶颈
  6. 为何仍需编写 CUDA
  7. 案例研究：量化
  8. 自定义操作 (Custom ops
  9. 如何支持自定义 CUDA 扩展
  10. 问题：软件包非纯 Python
  11. 融合到用户定义代码中
  12. 你应该编写什么样的内核
  13. TL;DR (内容摘要
  14. 附录 (Appendix
  15. 反馈 (Feedback

摘要

本次演讲的目标是展示 PyTorch 用户如何从机器学习研究者转变为系统研究者，并说服听众也进行同样的转变。

llm.c：性能与开发的权衡

llm.c 是一个用原生 CUDA 实现的完整 GPT-2 训练循环。
- GitHub 仓库: https://github.com/karpathy/llm.c
- 其性能优于 torch.compile。
- 然而，该项目由 4 位顶尖工程师耗时约 6 个月完成。
- 这引出了一个核心问题：什么是更好的权衡（tradeoff）？

PyTorch 的灵活性：Eager 模式

PyTorch 的 Eager 编程模式因其直观易用而广受欢迎。用户执行的每一行代码都会立即在后端执行相应的操作。

Page 4 - PyTorch Eager 模式代码解析
Page 4 - PyTorch Eager 模式代码解析

硬件瓶颈的演变

现代 GPU 的计算能力增长速度远超内存带宽的增长速度。这导致许多操作受限于内存带宽，而非计算能力。

Page 5 - GPU 算力与内存带宽的差距
Page 5 - GPU 算力与内存带宽的差距

torch.compile：应对内存带宽瓶颈

torch.compile 是一个融合编译器（fusion compiler），它是应对内存带宽瓶颈的主要方法之一。通过将多个操作融合成一个单一的 CUDA 核函数，可以显著减少内存的读写次数。

Page 6 - torch.compile 的核函数融合示例
Page 6 - torch.compile 的核函数融合示例

为何仍需编写 CUDA？

尽管 torch.compile 等编译器功能强大，但在某些情况下，手动编写 CUDA 仍然是必要的。

编译器在处理需要对表达式进行数学重写（同时保持数值稳定性）的优化时，通常会遇到困难。一个典型的例子是 Flash Attention 2。

在关于 Flash Attention 2 论文的 OpenReview 讨论中，一位审稿人提出了编译器是否能自动生成类似 FA-2 的优化的问题。作者回答：

<blockquote>

“编译器通常可以执行融合。然而，对于那些需要对同一表达式进行数学重写（同时保持数值稳定性）的优化，编译器通常更难处理。”

来源: https://openreview.net/forum?id=mZn2Xyh9Ec

</blockquote>

这表明，对于需要深度领域知识和复杂算法重构的顶尖性能优化，手动编写 CUDA 仍然具有不可替代的价值。

案例研究：量化

自定义操作 (Custom ops)

PyTorch 拥有许多如 torch.add、torch.sum 这样的操作符。

如果没有自定义操作的注册，torch.compile() 和 torch.autograd 将无法理解如何处理它们。

以下代码展示了如何注册一个“伪”（fake）实现：

@torch.library.register_fake("extension_cpp::mymuladd")
def _(a, b, c):
    torch.check(a.shape == b.shape)
    torch.check(a.dtype == torch.float)
    torch.check(b.dtype == torch.float)
    torch.check(a.device == b.device)
    return torch.empty_like(a)

这种方法不允许输入发生变异（mutation）。

更多信息请参考：https://pytorch.org/tutorials/advanced/custom_ops_landing_page.html

如何支持自定义 CUDA 扩展

Page 18
Page 18

为了支持自定义 CUDA 扩展，您需要在两个地方进行设置：

  1. 在 setup.py 文件中：

    • 如左侧代码所示，通过 get_extensions() 函数配置编译选项。
    • 使用 torch.utils.cpp_extension.CUDAExtension 来定义扩展，指定源文件（.cpp 和 .cu），并添加额外的编译和链接参数。

  2. 在 C++ 代码中注册操作：

    • 如右侧代码所示，使用 TORCH_LIBRARY_FRAGMENT 宏在 C++ 中将自定义操作注册到 PyTorch 的调度器中。
    • m.impl_abstract_pystub 用于链接到 Python 端定义的抽象实现（pystub）。

更多详情可参考：https://github.com/pytorch/ao/pull/135

问题：软件包非纯 Python！

这意味着您现在需要为所有可能的组合构建您的软件包：

对于 torchao 项目，我们在持续集成（CI）中启动了许多机器来完成这项工作。

许多项目会要求您从源码安装，或者将支持限制在带有最新稳定版 CUDA 和 Torch 的 Linux 系统上。

如果您不支持某个版本，pip 会在用户机器上静默地降级软件包。🤯

长期来看，我们需要转向一个类似 JIT 的工作流，但其代价是较长的冷启动时间。

融合到用户定义代码中

通常情况下，torch.compile 将内核视为黑盒，并且没有好的工具来解决黑盒融合问题。

下面的方法对于简单的逐点操作（pointwise ops）是可行的，但通常不具有普适性。

Page 20
Page 20

如上图所示，对于两个由 triton.jit 定义的函数 func1 和 func2，一种简单的融合思路是直接通过字符串拼接来创建一个新函数 func3，该函数按顺序调用前两者。但这并非一个通用的解决方案。

你应该编写什么样的内核？

如果您需要同行交流和发布工作的平台，可以考虑加入 Discord 社区：https://discord.gg/gpumode

TL;DR (内容摘要)

  1. 世界追求性能。
  2. 像 PyTorch 这样的灵活框架可以解决融合问题，但无法解决算法系统层面的问题。
  3. 自定义内核将长期存在并持续发展，PyTorch 正在使用它们，您也应该这样做！
  4. 得益于 CUTLASS、Triton 和各种元编程技术，自定义内核的入门门槛已经降低。字符串（代码）易于内省（introspect），但难以调试。
  5. 借助我们的自定义操作支持，集成和发布自定义内核变得很容易。JIT（即时编译）与 AOT（提前编译）代表了不同的权衡。

附录 (Appendix)

（此页为空白，作为附录部分的开始）

反馈 (Feedback)