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 模式代码解析
  • x = torch.tensor([1, 2, 3], device="cuda"): 分配 GPU 内存。
  • x = x + 1: 使用 CUDA 广播数字 1 并执行向量加法。
  • print(x): 同步 GPU 并读取结果。这使得调试和交互式编程变得简单。
  • x = torch.cos(x): 启动一个用于计算余弦的 CUDA 核函数。
  • x = torch.relu(x): 启动一个用于 ReLU 的 CUDA 核函数。

硬件瓶颈的演变

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

Page 5 - GPU 算力与内存带宽的差距
Page 5 - GPU 算力与内存带宽的差距
  • CPU 和 GPU 之间的数据传输，在实践中受限于铜线中电子的速度，成为了一个显著的瓶颈。
  • 更多信息可参考 Stephen Jones 在 GTC 2022 上的演讲《How CUDA Programming Works》。

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

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

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

  • 示例：

    • torch.compile(torch.square): 编译成 1 个核函数。
    • torch.compile(torch.square(torch.square)): 同样只编译成 1 个核函数。

  • 融合的优势（以 Double Square 为例）：

    • 不使用融合: 需要 2 次内存读取和 2 次内存写入。
    • 使用融合: 仅需 1 次内存读取和 1 次内存写入。

  • 警告: 融合并非总是有益的，因为它可能导致寄存器溢出（spilling registers），因此需要启发式方法来决定何时进行融合。

为何仍需编写 CUDA？

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

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

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

<blockquote>

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

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

</blockquote>

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

案例研究：量化

  • 相关博文：https://pytorch.org/blog/pytorch-native-architecture-optimization/
  • 对训练和推理过程中的激活值、权重、梯度、优化器进行量化。
  • 我们的大部分功能实现是通过 torch.compile() 完成的，但我们最快的 int4 推理内核是使用 CUTLASS 编写的。
  • 这对于减少内存带宽瓶颈非常重要。

自定义操作 (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！

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

  • Python 版本（最近由 Jane Xu 修复）
  • CUDA 版本
  • Torch 版本
  • 操作系统

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

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

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

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

融合到用户定义代码中

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

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

Page 20
Page 20

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

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

  • 低比特矩阵乘法 (Low-bit matmuls)
  • 稀疏注意力 (Sparse attention)
  • KV 缓存管理器 (KV cache managers)
  • 关注技术栈更高层级的性能问题
  • 如果您编写的许多内核看起来很相似，尝试更多地使用元编程 (metaprogramming)

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

TL;DR (内容摘要)

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

附录 (Appendix)

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

反馈 (Feedback)

  • 来自 Fred 的反馈：

    • 我们是否应该提及使用云服务提供商？
    • 深度学习容器（不像 CPU 那样可共享）。
    • 像性能计数器这样的东西在云供应商那里很少见，只有少数提供（例如 Lightning, AWS Bare Metal）。
    • 提供一些更清晰的性能数据。

  • 来自 Vikram 的反馈：

    • 请准确提及 cutlass 后端是如何工作的，它会生成一段 cutlass 代码。并解释其发生过程。
    • 更详细地介绍模板元编程及其挑战。
    • 将动机部分的时间从7分钟减少到3分钟，并扩展关于代码生成等内容。
    • 主要结论应该是我们可以代码生成“难以编写”的语言，使其更易于使用。
    • 针对特定形状提供更多性能数据：例如选择 llama 70B 的形状。

  • 来自 Tami 的反馈：

    • CUDA 技术简报会与此会议时间冲突。
    • 选择上午的会议时段。