Skip to content跳到主要内容
f.FENG/Knowledge Lab知识实验室
LEARNING SERIES学习专题

Inside PyTorchPyTorch 源码之旅

Read Inside PyTorch in order, following the questions developed in each installment.按期阅读《PyTorch 源码之旅》,沿着每一期的问题逐步展开。

For whom适合谁
For readers who use tensor operations and want to understand what happens underneath a Python call.适合已经用过张量运算,希望理解一次 Python 调用背后发生了什么的读者。
Before you begin已有基础
Basic Python, matrix multiplication and the chain rule. C++ snippets are explained in context; a GPU is not required to read the series or run its CPU examples.需要基础 Python、矩阵乘法与链式法则。C++ 摘录随正文解释;阅读专题及运行 CPU 示例不要求配备 GPU。
What you can explain读完能解释什么
Trace a tensor operation from storage and dispatch through gradients and execution, distinguish copying, synchronization and recompilation costs, and read the schema that defines an operator's contract.沿着存储、调度、梯度和执行追踪一次张量运算,分清数据复制、同步与重新编译的开销,并读懂定义算子接口契约的 Schema。

Follow one operation跟踪一次运算

Start with matmul, locate tensor elements in storage, then see how the dispatcher chooses an implementation.从 matmul 出发,找到张量元素的存储位置,再看调度器怎样选择实现。

  1. Part 1第 1 期
  2. Part 2第 2 期
  3. Part 3第 3 期

Gradients and GPU execution梯度与 GPU 执行

After the runtime path, follow gradient accumulation and CUDA dependencies to choose a meaningful timing boundary.了解运行时后,继续看梯度怎样汇合、CUDA 工作怎样建立依赖,以及计时边界如何选择。

  1. Part 4第 4 期
  2. Part 5第 5 期

Distinguish two kinds of code generation分清两种代码生成

Follow graph capture, cached-code reuse and kernel fusion in torch.compile, then compare them with torchgen's build-time generation of operator interfaces and registrations from schemas.先理解 torch.compile 的图捕获、缓存代码复用与内核融合,再与 torchgen 在构建时根据 Schema 生成算子接口及注册代码的过程对照。

  1. Part 6第 6 期
  2. Part 7第 7 期
  1. PART 01第 1 期 ·

    What is PyTorch? Inside torch.matmulPyTorch 到底是什么?跟踪一次 torch.matmul

    For dense 2-D inputs in eager mode, torch.matmul calls mm and dispatches to a CPU or CUDA implementation. With gradient recording enabled and inputs requiring gradients, autograd records dependencies for backward to apply the chain rule and accumulate leaf gradients.在即时执行模式下,稠密二维张量的 torch.matmul 调用 mm,并分派到 CPU 或 CUDA 实现。开启梯度记录且输入需要梯度时,autograd 记录依赖关系,供反向传播应用链式法则并累积叶子张量的梯度。

  2. PART 02第 2 期 ·

    Tensor, strides, and storage: why a transpose can be freeTensor、步幅与存储:转置为什么可以不搬数据

    For a dense strided tensor, strides and storage_offset map logical indices to storage locations. Transpose shares the original storage by swapping sizes and strides; making a noncontiguous result contiguous requires copying its data.对稠密步幅张量,strides 与 storage_offset 决定逻辑索引对应的存储位置。转置通过交换形状和步幅共享原存储;将非连续结果转为连续布局则需要复制数据。

  3. PART 03第 3 期 ·

    Operator registration and dispatch: how PyTorch selects a kernel算子注册与调度:PyTorch 如何选择内核

    PyTorch separates operator schemas from implementations. The dispatcher combines tensor dispatch keys with thread-local state, selects an operator-specific table entry, and lets wrappers redispatch to lower layers. Autograd dispatch and gradient recording are separate decisions.PyTorch 将算子接口与实现分开注册。调度器合并张量的调度键与线程局部状态,从算子专属的调度表中选择入口;包装层再通过 redispatch 进入下层实现。进入 Autograd 调度层与记录反向图是两个不同的判断。

  4. PART 04第 4 期 ·

    Autograd nodes and the backward engine: how gradients meetAutograd 节点与反向引擎:梯度如何汇合

    Backward nodes apply vector–Jacobian products. The autograd engine accumulates incoming contributions in input buffers and schedules a node when its dependencies are satisfied. AccumulateGrad then writes or adds the resulting gradient to a leaf tensor’s .grad.反向节点计算向量–雅可比积。Autograd 引擎用输入缓冲区汇总梯度贡献,在依赖满足后调度节点;AccumulateGrad 随后将结果写入或累加到叶子张量的 .grad。

  5. PART 05第 5 期 ·

    CUDA streams, kernels and timing: when is the GPU actually done?CUDA 流、内核与计时:GPU 到底什么时候算完?

    CUDA streams order submitted work; events connect dependencies across streams. Host return, event completion and tensor storage lifetime are different boundaries. Timing a matrix multiplication requires choosing which boundary to measure and accounting for warm-up, precision and synchronization.CUDA 流为提交的操作建立顺序,事件连接跨流依赖。主机调用返回、事件完成和张量存储生命周期是不同的边界。测量矩阵乘法前,需要明确计时边界,并控制预热、计算精度与同步方式。

  6. PART 06第 6 期 ·

    From torch.compile to generated kernels: capture, guards and fusion从 torch.compile 到生成内核:图捕获、守卫与融合

    TorchDynamo captures tensor operations into FX graphs and guards the assumptions used to specialize them. AOTAutograd prepares differentiation graphs; TorchInductor lowers and schedules computation, fusing compatible operations to avoid materialized intermediates. Cached code is reused only when its guards hold.TorchDynamo 将张量操作捕获为 FX 图,并用守卫检查特化所依赖的条件。AOTAutograd 准备求导图,TorchInductor 将图转换为低层中间表示并调度计算,将兼容操作融合以减少中间张量的物化。只有守卫满足时,缓存代码才能复用。

  7. PART 07第 7 期 ·

    The operator contract: schemas and code generation算子的接口契约:Schema 与代码生成

    An operator schema specifies arguments, returns, aliasing and mutation. Following the add family through native_functions.yaml and torchgen shows how PyTorch generates interfaces and registrations, which numerical and derivative rules still need definitions, and why this build-time work differs from torch.compile.算子 schema 描述参数、返回值、别名与修改约定。沿着 add 家族追踪 native_functions.yaml 与 torchgen,可以看清 PyTorch 怎样生成接口和注册代码、哪些数值与求导规则仍需定义,以及这项构建期工作与 torch.compile 的区别。

Browse all notes浏览全部笔记

Search the notebook搜索知识库

Try “matrix”, “Rust”, or “spikes”试试「矩阵」「Rust」或「脉冲」

Full-text search ·全文检索 · ESC to closeESC 关闭