数据截至 2026-08-27(上游 commit c187ef3271d5)

PyTorch — 架构与原理

30 秒导读： PyTorch 是当今几乎所有 LLM 训练代码的载体。它对外是「像 NumPy 的张量库 + 自动求导」，对内是一台分层机器：Python API 把调用翻译成 ATen 算子，dispatcher 按 dispatch key 选择执行哪一层（autograd / 设备后端 / 编译器），autograd 引擎则在每次前向时悄悄挂好反向图，等 backward() 一声令下按拓扑顺序执行。本 teardown 从 Tensor 数据结构一路讲到 DDP 与 torch.compile，全部锚定到 c187ef3 的真实源码。

1. 这是什么（零基础也能懂）

一句话定义

PyTorch 是一个 Python 优先的张量计算框架：你用它操作多维数组（Tensor），它会自动记录这些操作构成的计算图，并能一键算出所有参数对某个标量的梯度——这正是训练神经网络需要的全部原料。本 commit 版本为 2.15.0a0（version.txt:1）。

它解决谁的什么问题

设想你要训练一个模型：

你需要在 GPU 上做大规模矩阵乘、卷积、归一化——张量计算。
你需要对每个可学习参数求梯度，但手写反向传播公式极易出错——自动微分。
模型有几亿参数、要切到几百张卡上——分布式训练。
eager 模式逐行执行不够快，你还想要 kernel 融合——编译优化。

PyTorch 把这四件事做成了一套连贯的 API，同一份 nn.Module 代码可以从笔记本一路跑到千卡集群。

它能做什么

能力	入口	对应章节
张量计算	`torch.Tensor`、ATen 约 2500+ 算子（`aten/src/ATen/native/native_functions.yaml`）	第 1、2 章
自动微分	`tensor.backward()`（`torch/_tensor.py:566`）	第 3 章
神经网络组件	`torch.nn.Module`（`torch/nn/modules/module.py:407`）	第 4 章
优化器	`torch.optim.AdamW`（`torch/optim/adamw.py:20`）	第 4 章
分布式训练	`DistributedDataParallel`（`torch/nn/parallel/distributed.py:466`）、FSDP	第 5 章
图编译	`torch.compile`（`torch/__init__.py:3023`）	第 6 章

用起来什么样

一段最小训练循环（每个 API 都会在后面章节解剖）：

import torch
import torch.nn as nn

model = nn.Linear(10, 1)                    # 第 4 章：Module 自动登记参数
opt = torch.optim.AdamW(model.parameters()) # 第 4 章：param_groups + 状态字典

x = torch.randn(32, 10)                     # 第 1 章：TensorImpl + Storage
y = model(x)                                # 第 2 章：dispatcher 选 kernel
loss = y.square().mean()                    # 每次调用都在挂反向图（第 3 章）
loss.backward()                             # 第 3 章：autograd 引擎拓扑执行
opt.step()                                  # 第 4 章：AdamW 更新公式

一句话直觉

把 PyTorch 想成一间「带记账员的流水线工厂」。 每道工序（算子）由调度员（dispatcher）按工单上的标签（dispatch key）派给对应车间（CUDA kernel / autograd 层）；记账员（autograd）站在流水线旁，每出一批货（前向输出）就记一笔「这批货由哪台机器、用什么原料做的」；等你说「结算」（backward()），记账员就倒着把账本走一遍，算出每个原料供应商（叶子参数）该记多少账（梯度）。

2. 顶层全景（它大概怎么转）

2.1 一次算子调用的完整纵切面

你的代码:  torch.add(a, b)  /  a + b
     │
     ▼
① Python 前端          torch.Tensor / torch.add
   (torch/_tensor.py, torch/_ops.py 的 OpOverload)
     │
     ▼
② ATen 算子注册表       add.Tensor 的 schema 来自
   (aten/src/ATen/native/native_functions.yaml:542)
     │
     ▼
③ Dispatcher           从输入 Tensor 的 key_set 算出最高优先级
   (aten/src/ATen/core/dispatch/Dispatcher.h:784)     dispatch key，查表选 kernel
     │
     ▼
④ 分层 kernel           Autograd → 后端(CPU/CUDA) → Composite
   需要梯度的先过 VariableType 挂图，再 redispatch 到真 kernel
     │
     ▼
⑤ 真实计算              CUDA/CPU kernel 读写 Storage 里的内存

怎么读这张图： 从上到下是调用方向；③ 是全部魔法的枢纽——每个 Tensor 自带一个 key_set_（c10/core/TensorImpl.h:3045），dispatcher 取「最高位」决定先走哪一层，处理完再往下 redispatch。autograd、设备后端、Python 子类拦截全都是这条链上的一环。

2.2 仓库与部件职责

PyTorch 仓库极大（约 2 万+ 文件），但可读的主线集中在少数几个目录：

部件	干什么	在哪里
c10	核心数据结构：TensorImpl、Storage、DispatchKey	`c10/core/`
ATen	算子库：schema 注册表 + 各后端 kernel	`aten/src/ATen/`
torch/csrc/autograd	自动微分引擎（C++）	`torch/csrc/autograd/engine.cpp`
torch/autograd	自动微分的 Python 面（`backward`、`Function`）	`torch/autograd/`
torch/nn	Module 体系与标准层	`torch/nn/modules/`
torch/optim	优化器	`torch/optim/`
torch/distributed	c10d 通信 + DDP/FSDP	`torch/distributed/`、`torch/nn/parallel/`
torch/_dynamo 等	编译栈：Dynamo → AOTAutograd → Inductor	`torch/_dynamo/`、`torch/_functorch/`、`torch/_inductor/`
tools/autograd	代码生成：derivatives.yaml → VariableType	`tools/autograd/`

2.3 主线走一遍（一个训练 step 在框架里的旅程）

① 前向  y = model(x)
   Module.__call__ → Linear.forward → F.linear → aten::linear
   → dispatcher: AutogradCUDA 层给输出挂 grad_fn → CUDA kernel 真算

② 建图（随前向免费发生）
   每个中间 Tensor 的 AutogradMeta.grad_fn_ 指向一个 Node，
   Node 的 next_edges_ 指回输入的 grad_fn，组成一张反向 DAG

③ 反向  loss.backward()
   torch.autograd.backward → Engine::execute
   → ReadyQueue 按依赖计数拓扑调度 Node::operator()（即各算子的 backward 公式）

④ 累加  叶子参数的 AccumulateGrad 节点把梯度加进 param.grad
   （DDP 场景下，这一步同时触发梯度桶的 allreduce）

⑤ 更新  opt.step()
   遍历 param_groups，按 AdamW 公式原地改 param.data

3. 阅读地图（建议顺序）

六章由浅入深。只想抓主干的话读 01 → 02 → 03 就够——它们覆盖了「Tensor 是什么、算子怎么被路由、梯度怎么算出来」这条每个框架都绕不开的纵线。

顺序	章节	讲什么	适合谁
1	01-tensor-storage.md	TensorImpl/Storage、strides、view 语义	所有人必读，这是一切的地基
2	02-operator-dispatch.md	native_functions.yaml、DispatchKey 优先级、OpOverload	想写 custom op 或理解「为什么我的 Tensor 子类能拦截一切」的人
3	03-autograd-engine.md	grad_fn 图、Node、Engine 拓扑执行	想真正理解反向传播在工业框架里怎么落地的人（先读 micrograd 会更有感觉）
4	04-nn-and-optim.md	Module 属性拦截、Linear、AdamW 公式	每天写 `nn.Linear` 的人
5	05-distributed-ddp.md	ProcessGroup、Reducer 梯度桶、通信重叠	要多卡训练、想知道 DDP 为什么快的人
6	06-torch-compile.md	Dynamo 截帧、AOTAutograd、Inductor/Triton	关心性能与编译栈的人

4. 巧妙之处（可借鉴的技术）

每条先白话点出「妙在哪」，再给锚点；细节见对应章节。

妙处	一句话	锚点
元数据与数据分离	Tensor 只是「一段内存 + 怎么读它」的说明书，view/transpose 零拷贝	`c10/core/TensorImpl.h:2888`（`storage_`）
小数组内联	sizes/strides ≤5 维时存在 TensorImpl 体内，不走堆分配	`c10/core/impl/SizesAndStrides.h:10`
惰性 autograd 元数据	不需要梯度的 Tensor 连 `AutogradMeta` 都不分配（nullptr 即默认）	`c10/core/TensorImpl.h:2890-2908`
一张 yaml 生成一切	native_functions.yaml 一份 schema 生成 C++ API、Python 绑定、autograd 挂图、 boxing	`aten/src/ATen/native/native_functions.yaml:542`
dispatch key 即位掩码	每个 Tensor 的 `key_set_` 是 64 位掩码，dispatch 就是「数前导零」取最高位	`c10/core/DispatchKey.h:104-112`
反向不用显式拓扑排序	前向时顺手给每个 Node 记「入度」，反向时倒数到 0 即入队	`torch/csrc/autograd/engine.cpp:1224-1280`
DDP 梯度桶	把参数梯度按桶聚合，桶一装满立刻 allreduce——通信与反向计算重叠	`torch/csrc/distributed/c10d/reducer.cpp:668`
`zero_grad` 默认置 None	省一次清零 kernel，也让「没收到梯度的参数」可被跳过	`torch/optim/optimizer.py:1058`
Dynamo 不解析源码	直接拦截 CPython 字节码帧逐条解释执行，eval 出的图自带 guards	`torch/_dynamo/symbolic_convert.py:5587`

5. 边界与局限

诚实清单（哪些是从代码里能看到的，哪些标注推断）：

eager 的调度开销是真实存在的。 每个 Python 算子调用都要走 dispatcher 查表；对逐元素小算子这是纯开销。torch.compile（第 6 章）就是为消灭它而生。
反向图是一次性的。 默认 backward() 会释放图，再次调用报错；retain_graph=True 才能保留——create_graph=True 场景官方还警告会造成引用环内存泄漏（torch/csrc/autograd/engine.cpp:1375-1382 的 TORCH_WARN_ONCE）。
Dispatcher 不保证跨 key 的组合顺序之外的语义。 自定义 dispatch key 的注册顺序有坑，社区为此写了大量 note（c10/core/DispatchKey.h 的 Note [DispatchKey Classification]，c10/core/DispatchKey.h:114）。
DDP 对没收到梯度的参数直接报错：反向结束时 Reducer 检查每个参数是否都收到了梯度，没收到的以 "Parameter indices which did not receive grad for rank …" 报错（torch/csrc/distributed/c10d/reducer.cpp:2166）；static_graph=True 时第一轮就会记录未用参数集合（torch/csrc/distributed/c10d/reducer.cpp:625-644）。
torch.compile 遇到图断点会静默回落 eager（fullgraph=False 默认），性能收益取决于你的代码里有多少「Dynamo 解释不了的字节码」。
in-place 操作与 view 的组合是 autograd 的经典坑：version counter 检出冲突时报「a view ... is being used in-place」一类错误，机制见第 3 章。
看不出：各 CUDA kernel 的具体性能调优参数（本 teardown 不深入 kernel 层）。

6. 横向对比

维度	PyTorch（本文）	micrograd	tinygrad	verl
定位	工业级训练框架	教学用反向传播最小实现	小而全的框架，主打自己的编译栈	建在 PyTorch 之上的 RL 训练库
autograd	C++ 引擎 + ReadyQueue 拓扑调度（第 3 章）	纯 Python `Value._backward` 递归 + 显式 topo 排序	惰性计算图，backward 在调度前展开成 kernel	直接用 PyTorch autograd，自己管分布式编排
算子路由	dispatch key 位掩码 + 每层 redispatch（第 2 章）	无（只有十几个手写算子）	lazy buffer + shapetracker，延迟到 kernel 生成	无算子层，调度的是「方法调用」
编译	Dynamo/AOT/Inductor 外挂栈（第 6 章）	无	编译就是主路径（uop 重写 → 渲染 kernel）	复用 torch.compile / 后端引擎
分布式	DDP/FSDP 内建（第 5 章）	无	`Shard`/`multi.py` 简版多卡	单控制器 + WorkerGroup 抽象（见 verl 第 1 章）

一句话：micrograd 教你「反向传播是什么」，tinygrad 展示「一个小框架可以把编译做成主路径」，PyTorch 回答「同一件事在十亿参数、千张卡、十年兼容性的约束下怎么做」，verl 则是在 PyTorch 这台机器上再盖一层 RL 编排。

7. 代码地图（导航索引）

主题	文件路径	符号名
Tensor 元数据结构	`c10/core/TensorImpl.h`	`TensorImpl`、`sizes_and_strides_`、`key_set_`
数据缓冲	`c10/core/StorageImpl.h`	`StorageImpl`、`data_ptr_`
算子 schema 注册表	`aten/src/ATen/native/native_functions.yaml`	`- func: add.Tensor` 等条目
分发器	`aten/src/ATen/core/dispatch/Dispatcher.h`	`Dispatcher::call`
dispatch 键定义	`c10/core/DispatchKey.h`	`DispatchKey`、`BackendSelect`
每算子 kernel 表	`aten/src/ATen/core/dispatch/OperatorEntry.h`	`OperatorEntry::lookup`、`dispatchTable_`
Python 算子对象	`torch/_ops.py`	`OpOverload`、`OpOverloadPacket`
autograd 元数据	`torch/csrc/autograd/variable.h`	`AutogradMeta`
反向图节点	`torch/csrc/autograd/node.h`	`Node`、`Node::operator()`
反向引擎	`torch/csrc/autograd/engine.cpp`	`Engine::execute`、`Engine::evaluate_function`
求导公式表	`tools/autograd/derivatives.yaml`	`- name: add.Tensor` 等条目
自定义 autograd	`torch/autograd/function.py`	`Function`、`FunctionCtx.save_for_backward`
Module 基类	`torch/nn/modules/module.py`	`Module.__setattr__`、`_call_impl`
标准层示例	`torch/nn/modules/linear.py`	`Linear.forward`、`reset_parameters`
优化器基类	`torch/optim/optimizer.py`	`Optimizer`、`zero_grad`
AdamW	`torch/optim/adamw.py`、`torch/optim/adam.py`	`AdamW`、`_single_tensor_adam`
DDP	`torch/nn/parallel/distributed.py`	`DistributedDataParallel._ddp_init_helper`
梯度归约器	`torch/csrc/distributed/c10d/reducer.cpp`	`Reducer::autograd_hook`、`all_reduce_bucket`
通信组	`torch/csrc/distributed/c10d/ProcessGroup.hpp`	`ProcessGroup::allreduce`
编译入口	`torch/__init__.py`	`torch.compile`
字节码解释器	`torch/_dynamo/symbolic_convert.py`	`InstructionTranslator`
AOT 抓图	`torch/_functorch/aot_autograd.py`	`aot_function`
代码生成	`torch/_inductor/scheduler.py`、`torch/_inductor/codegen/triton.py`	`Scheduler`、`TritonKernel`

1. 这是什么（零基础也能懂）​

一句话定义​

它解决谁的什么问题​

它能做什么​

用起来什么样​

一句话直觉​

2. 顶层全景（它大概怎么转）​

2.1 一次算子调用的完整纵切面​

2.2 仓库与部件职责​

2.3 主线走一遍（一个训练 step 在框架里的旅程）​

3. 阅读地图（建议顺序）​

4. 巧妙之处（可借鉴的技术）​

5. 边界与局限​

6. 横向对比​

7. 代码地图（导航索引）​