Recent content on GPU Notes
📌 Introduction Link to heading CUDA Graphs capture a sequence of GPU operations and replay them as a single unit, eliminating the CPU launch overhead for each individual kernel. This is especially beneficial for workloads with small batches or models with many small kernels, where CPU overhead can become a bottleneck. When you use torch.compile(mode="reduce-overhead") or enable…
前段时间因为工作原因重新浏览了 torch.compile 中几个重要的函数调用栈，现将这些内容分享给有兴趣通过源代码理解 torch.compile 行为的朋友。 本文使用 PyTorch 官方提供的 docker 镜像: pytorch/pytorch:2.4.0-cuda12.4-cudnn9-devel ，PyTorch 版本为 2.4.0。
时至今日，GPU 的珍贵程度无需多言，再加上当前特殊的大环境，手头拥有的 GPU 更显得是一种稀缺资源。在过去的几年里，我们在许多 PyTorch 案例中发现，用户手中的 GPU 并未得到充分的发挥，存在着大量的浪费现象。这不仅是对昂贵硬件资源的一种浪费，也限制了计算任务的效率和速度。因此，我们要尽可能把手中的 GPU 充分利用起来。
前些天在给某个 PyTorch 深度学习模型应用 full iteration CUDA graph 时，碰到了一个随机出现的 bug。CUDA graph 是一个威力很大的性能调优手段，特别是在大模型训练的场景中，但用过的人都知道，能把 CUDA graph 用起来是相对比较困难的，特别是把整个 iteration 捕获为一张 CUDA graph。这个 bug 前前后后花了 3 天时间才解决，分享以下自己的 debug 经过。
进入 PyTorch 2.0 时代，抓取计算图变得越来越容易。虽然 PyTorch 及其生态已经有一些抓取计算图的工具，但这篇文章的主要目的是用最少的代码抓取并保存 正向传播和反向传播计算图 ，通过实践掌握 AOTAutograd 中抓取计算图的核心原理。
简介 Link to heading 在 PyTorch 2.0 以前，用户通过 PyTorch 可以直接捕获到正向传播的计算图，比如 JIT trace 和 TorchFX 的 symbolic trace。虽然 PyTorch 的每个算子都包含正向传播和反向传播的实现，但用户并不能直接在反向传播的计算图上面做优化，也无法把正向传播和反向传播的计算图合并在一张计算图中。PyTorch 2.0 中引入了 AOTAutograd，它的出现解决了这个问题，从而使得一些针对 training 的优化变得可能。
相比读博客、看视频学习一个开源项目，读源代码是比较痛苦的。然而 阅读源代码是成为优秀程序员的起点 ，不读源代码往往是知其然而不知其所以然。 PyTorch 的代码对初学者来说并不好读，原因在于：
简介 Link to heading PyTorch 2.0 的使命是更快、更 Pythonic 以及一如既往地支持动态特性。为了达到这个目的，PyTorch 2.0 引入了 torch.compile ，在解决 PyTorch 固有的性能问题的同时，把部分用 C++ 实现的东西引入 Python 中。PyTorch 2.0 利用了 4 个组件: TorchDynamo，AOTAutograd，PrimTorch 和 TorchInductor。本文以几个简单的案例讲解 TorchDynamo 的使用方法和实现原理。
Guard Link to heading 随着函数调用栈返回到 convert_frame.py#L327 ， TorchDynamo 编译的最后需要为 Guard 生成 Python 代码，从而在后续执行编译好的函数时检查函数的输入信息是否发生了变化，从而决定是否需要重新编译 :
Graph Break Link to heading Offset 28 , POP_JUMP_IF_FALSE : # https://github.com/pytorch/pytorch/blob/fe05266fda4f908130dea7cbac37e9264c0429a2/torch/_dynamo/symbolic_convert.py#L853 POP_JUMP_IF_FALSE = generic_jump ( operator . not_ , False ) POP_JUMP_IF_TRUE = generic_jump ( operator . truth , False ) JUMP_IF_FALSE_OR_POP = generic_jump ( operator . not_ , True ) JUMP_IF_TRUE_OR_POP =…
初始化 InstructionTranslator Link to heading 清理后的字节码指令经由 transformations(instructions, code_options) 开始执行变换， transformations 是之前的 transform() ，其中首先实例化了 InstructionTranslator ，定义于 torch/_dynamo/symbolic_convert.py#L1747 。 InstructionTranslator 中有一个 OutputGraph 的实例， OutputGraph 本身是一个 torch.fx.Tracer ， OutputGraph 用于保存 InstructionTranslator 做 字节码翻译 后的输出，以 torch.fx.Graph 表达。 OutputGraph 中还包含一个…
简介 Link to heading PyTorch 最大的优点是其灵活性，但缺点也很显著，PyTorch 生态对图编译器 (graph compiler) 的支持非常有限，现有的基于 TorchScript 的方案非常难用。为了解决这些问题，PyTorch 2.0 引入了 torch.compile ，它包含几个新的组件: TorchDynamo, AOTAutograd, PrimTorch, TorchInductor。其中，TorchDynamo 用于从用户的 PyTorch 代码中以最小的代价 捕获计算图 ，更多关于 PyTorch 2.0 的简介请参考 PyTorch 2.0 。
I’m a performance engineer focused on squeezing every last cycle out of GPU workloads. My interests span high-performance computing, deep learning optimization, and computer architecture — pushing state-of-the-art models to peak performance across speech recognition, machine translation, image classification, and generative AI. This blog is where I document my learning notes in my spare…
Introduction Link to heading The original claim of OpenAI Triton is: We’re releasing Triton 1.0, an open-source Python-like programming language which enables researchers with no CUDA experience to write highly efficient GPU code — most of the time on par with what an expert would be able to produce. The core ideas of Triton : Program GPU with Python : so that the effort to program GPU…
简介 Link to heading torch.fx 用于对 torch.nn.Module 做图变换。它包含三部分： 符号跟踪器（symbolic tracer）：用于捕获 module 的语义，它以符号的方式执行 Python 代码（symbolic execution），通过给 module 提供虚假值（Proxies）并记录涉及到的运算； 中间表示（intermediate representation）：IR 是在 tracing 期间所记录算子的图，包含一系列节点（Node），节点代表输入（placeholder）、函数（get_attr, call_function, call_module, call_method）、输出（output），IR 是用 torch.fx 进行图变换（transformation）的基石； Python 代码生成（code…
在用GitBook整理笔记的过程中，因GitBook年久失修，碰到了很多问题。 安装 gitbook-cli Link to heading 使用GitBook需要先安装 node.js ，之后安装 gitbook-cli ：