PyTorch 源码深度剖析与生产应用教程目录总览:PyTorch 的分层架构与设计哲学 c10 与 ATen:张量

分析对象:pytorch/pytorch master 分支分析时的 HEAD: 2c2148d2 (2026-04-23) 本地路径:d:/opencode/q/pytorch 文件规模:21 703 个跟踪文件,native_functions.yaml 单文件 16 267 行,torchgen/gen.py 3 067 行,torch/nn/modules/module.py 3 053 行,c10/core/TensorImpl.h 3 381 行

本文所有 file:line 引用均来自上述仓库快照,可直接在本地 d:/opencode/q/pytorch/<路径> 下打开。

总览:PyTorch 的分层架构与设计哲学
c10 与 ATen:张量与调度器
算子注册与代码生成:native_functions.yaml 到 Kernel 的完整路径
Autograd 引擎:动态图的构造与反向调度
CUDA 后端:Stream、缓存分配器、CUDA Graphs 与 AMP
torch.compile 栈:Dynamo + AOTAutograd + Inductor
分布式子系统:c10d / DDP / FSDP2 / DTensor
用户层:nn.Module / Optimizer / DataLoader / FX / Profiler
部署与导出:torch.export / AOTI / ONNX / 量化
生产应用教程
二次开发:自定义 C++/CUDA 算子
性能关键路径速查表

1. 总览:PyTorch 的分层架构与设计哲学

PyTorch 是一个 "Python 前端 + C++ 核心 + 动态调度 + JIT/AOT 编译" 的混合架构。从下往上大致 7 层:

┌──────────────────────────────────────────────────────────────┐
│  User API Layer  (torch.nn, torch.optim, torch.utils.data)   │
├──────────────────────────────────────────────────────────────┤
│  torch.compile Stack  (Dynamo → AOTAutograd → Inductor)      │
├──────────────────────────────────────────────────────────────┤
│  Autograd Engine  (torch/csrc/autograd/engine.cpp)           │
├──────────────────────────────────────────────────────────────┤
│  ATen Operators  (aten/src/ATen/native/*, CPU/CUDA/MPS 实现) │
├──────────────────────────────────────────────────────────────┤
│  Dispatcher  (aten/src/ATen/core/dispatch/Dispatcher.h)      │
├──────────────────────────────────────────────────────────────┤
│  c10 Core  (Tensor/Storage/Device/Allocator)                 │
├──────────────────────────────────────────────────────────────┤
│  Hardware Backends  (CUDA/cuDNN/cuBLAS/MKL/MPS/XPU/ROCm)     │
└──────────────────────────────────────────────────────────────┘

核心设计哲学(从源码中能看到的选择):

"急切执行优先 + 可选图化" — 前向计算立即执行,autograd 在前向时动态构建反向图;需要高性能时再用 torch.compile 捕获图。这是 PyTorch 相对于静态图框架(旧 TF 1.x)的根本差异。
"Dispatcher 作为一切可扩展性的枢纽" — 每个算子调用都经过一次按 DispatchKeySet 选择 kernel 的过程。Autograd、AMP、functorch(grad/vmap)、量化、FSDP,全部以 DispatchKey 的形式插入优先级链。新后端(如私有芯片)可以只注册到 PrivateUse1。
"YAML 为唯一事实源" — 16 000+ 行的 native_functions.yaml 是算子语义的唯一声明,C++ 头、Python 绑定、autograd 包装、functionalization 都由 torchgen 自动生成。这是 PyTorch 能维护 2 000+ 算子 × 数十个后端的关键。
"C++ 用 intrusive_ptr 而非 shared_ptr" — 为了在 Tensor 内嵌引用计数,所有核心对象都内嵌 c10::intrusive_ptr_target,避免 shared_ptr 的双间接和控制块开销。
"PEP 523 进程级字节码劫持" — Dynamo 不重写用户代码,而是通过 CPython 的 frame 评估 hook 在运行时符号执行字节码,把 Python 动态性蒸馏成 FX 图 + guard。

顶层目录速览(d:/opencode/q/pytorch/):

目录	作用
`c10/`	Core 张量/设备/分配器抽象,最底层,不依赖 ATen
`aten/src/ATen/`	"A TENsor library" — 张量操作与 kernel 实现
`torch/csrc/`	PyTorch 的 C++ 部分(Python 绑定、autograd、jit、distributed)
`torch/`	Python 前端(nn/optim/distributed/compile 等)
`torchgen/`	从 `native_functions.yaml` 生成代码的工具链
`tools/autograd/`	自动微分包装器的代码生成
`functorch/`	旧目录,大多数代码已移入 `torch/_functorch/`
`third_party/`	子模块(NCCL、cuDNN 头、Gloo、MKL-DNN、fmt、pybind11)
`test/`	测试用例(也是学习算子行为的最佳来源)

2. c10 与 ATen:张量与调度器

2.1 Tensor 的真实面貌

用户眼中 torch.Tensor 是一个整体,源码里它其实是 六层对象 + 一条引用计数链:

Python torch.Tensor
         │  继承自
         ▼
torch._C.TensorBase  (C++ 侧 PyObject 包装)
         │  持有
         ▼
at::Tensor                    ← aten/src/ATen/core/Tensor.h
         │ inherits
         ▼
at::TensorBase                ← aten/src/ATen/core/TensorBase.h:94-127
         │ intrusive_ptr<TensorImpl, UndefinedTensorImpl> impl_  (line 103)
         ▼
c10::TensorImpl               ← c10/core/TensorImpl.h:528
         ├─ c10::Storage        → StorageImpl → DataPtr → 裸内存
         ├─ SizesAndStrides    ← c10/core/impl/SizesAndStrides.h  (≤5D 内联)
         ├─ c10::Device
         ├─ ScalarType (dtype)
         ├─ DispatchKeySet
         └─ AutogradMeta*       ← torch/csrc/autograd/variable.h:102

关键点:

c10::TensorImpl 声明见 c10/core/TensorImpl.h:528:struct C10_API TensorImpl : public c10::intrusive_ptr_target。继承 intrusive_ptr_target 意味着 引用计数内嵌在对象本身(64 位原子字,上 32 位弱引用、下 32 位强引用,见 c10/util/intrusive_ptr.h:19-80)。
SizesAndStrides(c10/core/impl/SizesAndStrides.h)对 ≤5 维张量使用内联栈存储,避免堆分配。绝大多数实际 workload 的张量都能命中内联路径。
Storage 与 Tensor 解耦:多个 Tensor 可以 view 同一个 Storage(这是 tensor.view()、tensor[:, 0] 等零拷贝切片的基础)。StorageImpl 定义见 c10/core/StorageImpl.h:52,它才是真正持有 DataPtr 和 Allocator* 的对象。
stride/size/offset 三元组共同决定一个 view 如何映射到底层 storage,这是 PyTorch 所有 view 操作(transpose、permute、unfold)的核心表达。

2.2 Device 与 DispatchKey

c10::Device(c10/core/Device.h:31)只是 (DeviceType, DeviceIndex) 的二元组。
c10::DispatchKey(c10/core/DispatchKey.h)按 功能(Functionality)× 后端(BackendComponent) 两维建模:
- BackendComponent: CPU、CUDA、XLA、MPS、PrivateUse1/2/3、Meta(约 16 个)
- FunctionalityKeys: Dense、Quantized、Sparse、NestedTensor、AutogradFunctionality、FuncTorchBatched、Functionalize、AutocastCUDA…
c10::DispatchKeySet(c10/core/DispatchKeySet.h)是 64 位位集,表示一个张量"同时属于哪些 key"。例如一个 requires_grad=True 的 CUDA half 张量在 autocast 区域中,它的 DispatchKeySet 至少包含 CUDA | AutogradCUDA | AutocastCUDA。

优先级选择:Dispatcher 按 key 的位置从高到低选择第一个有注册 kernel 的 key。这是一个精巧的设计 — autograd、autocast、functorch 能"透明地"拦截算子,靠的就是这些 key 排在后端 key 之前。

2.3 Dispatcher:一次算子调用的真实路径

入口文件:aten/src/ATen/core/dispatch/Dispatcher.h。核心类:

Dispatcher(单例):维护 OperatorDef 全局注册表。
OperatorHandle:指向某个算子 schema 的轻量句柄。
OperatorEntry(OperatorEntry.h):每个算子的注册信息,包含按 DispatchKey 索引的 KernelFunction 表。
KernelFunction:既可以是 unboxed(静态类型,直接函数指针)也可以是 boxed(栈式通用调度),两种模式互转。

一次 a.add(b) 的调用链:

Python: a + b
  │  Tensor.__add__ → aten::add.Tensor
  ▼
torch._C._TensorBase.add  (pybind11 binding)
  │
  ▼
at::_ops::add_Tensor::call(a, b)           ← 生成代码
  │
  ▼
Dispatcher::singleton().call<add_Tensor>(op_handle, dispatch_key_set, a, b)
  │  计算 dispatch_key_set = a.key_set() | b.key_set() | local_tls
  │  在 OperatorEntry 的 kernel 表里找最高优先级的有效 key
  ▼
(如果 requires_grad) AutogradCUDA kernel:创建 AddBackward0 节点 → 去掉 Autograd key → 递归 call
  │
  ▼
CUDA kernel (aten::add_out_cuda)  ← aten/src/ATen/native/cuda/BinaryAddSubKernel.cu

Boxed vs Unboxed:大多数热路径走 unboxed(直接调用函数指针,无类型擦除开销)。Boxed 路径用于少数需要泛化的场景(Python fallback、JIT、functorch 的某些 rule)。两者的 KernelFunction 互相转换依赖生成的 boxing wrapper(aten/src/ATen/core/boxing/)。

2.4 Allocator

c10::Allocator(c10/core/Allocator.h)是一个虚基类,子类覆盖 allocate(size) -> DataPtr。典型实现:

CPUAllocator(c10/core/CPUAllocator.h:15-59):malloc/jemalloc,可通过 SetCPUAllocator() 替换。
CachingDeviceAllocator(CUDA 实现见 2.5 节)。

DataPtr 是一个 "裸指针 + 自定义 deleter + Device" 三元组,所有设备内存的所有权语义都由它承载。

3. 算子注册与代码生成:从 native_functions.yaml 到 Kernel

3.1 YAML 为唯一事实源

aten/src/ATen/native/native_functions.yaml(16 267 行)声明每个算子的 schema 与 dispatch 表:

- func: add.Tensor(Tensor self, Tensor other, *, Scalar alpha=1) -> Tensor
  structured_delegate: add.out
  variants: function, method
  dispatch:
    SparseCPU, SparseCUDA: add_sparse
    MkldnnCPU: mkldnn_add
    ZeroTensor: add_zerotensor
  tags: [core, pointwise]

这一条 YAML 会驱动生成:

at::_ops::add_Tensor(boxed + unboxed 调用桩)
at::Tensor::add 方法
Python torch.add 绑定
Autograd 包装器(若有对应 derivatives.yaml 条目)
Functionalization pass 的处理规则

3.2 代码生成管道:torchgen

入口:torchgen/gen.py(3 067 行)。管道大致为:

native_functions.yaml + derivatives.yaml + tags.yaml
        │
        ▼
torchgen/model.py        (解析为 NativeFunction 对象)
        │
        ▼
torchgen/api/{native, dispatcher, autograd, python}.py  (签名翻译)
        │
        ▼
torchgen/dest/register_dispatch_key.py  (按后端生成 register_CPU.cpp / register_CUDA.cpp)
        │
        ▼
build/aten/src/ATen/{Operators.cpp, Functions.cpp, NativeFunctions.h, RegisterCPU.cpp, …}

实际触发时机:pip install -e . -v --no-build-isolation(仓库 CLAUDE.md 指定的唯一构建命令)会在 CMake 配置阶段调用 torchgen.gen.main()。

3.3 用户侧注册:TORCH_LIBRARY 宏

第三方算子通过 torch/library.h 的宏注册,不触及 YAML:

// my_ops.cpp
#include <torch/library.h>
torch::Tensor my_fused_gelu_cpu(const torch::Tensor& x) { /* … */ }

TORCH_LIBRARY(myops, m) {
  m.def("fused_gelu(Tensor x) -> Tensor");             // schema
}
TORCH_LIBRARY_IMPL(myops, CPU, m) {
  m.impl("fused_gelu", TORCH_FN(my_fused_gelu_cpu));
}

宏展开见 torch/library.h:1-120:TORCH_LIBRARY 构造静态 Library 对象,其析构时机保证注册在 main 之前完成。CppFunction 通过模板推断类型签名,生成 boxing wrapper,最终调用 Dispatcher::registerImpl()。

3.4 Functionalization

aten/src/ATen/FunctionalTensorWrapper.{h,cpp} 实现一个特殊 DispatchKey Functionalize,它把所有 in-place 和 view 操作改写为纯函数形式(add_ → add + copy_)。torch.export 与 AOTAutograd 都强依赖它来产生纯函数 FX 图。

4. Autograd 引擎:动态图的构造与反向调度

4.1 前向时建图

每个带 requires_grad 的张量都绑定一个 AutogradMeta(torch/csrc/autograd/variable.h:102)。生成的 VariableType 包装器(tools/autograd/gen_variable_type.py 生成到 build/aten/src/ATen/VariableType_*.cpp)在每次 ATen 算子调用时:

取输入张量的 grad_fn / grad_accumulator(即 impl::gradient_edge(), function.h:22)
创建一个 Node 子类实例(如 AddBackward0)
调用 collect_next_edges(inputs) 收集输入的 Edge,设置到 node->next_edges_
调用底层 ATen kernel(通过把 Autograd key 从 dispatch set 中排除,递归 redispatch)
把输出张量的 grad_fn_ 指向该 Node

数据结构:

Node(torch/csrc/autograd/node.h:113-150):图节点基类,纯虚 apply(variable_list&&) 实现反向计算。构造时通过 at::sequence_number::get_and_increment() 分配全局单调序号,用于 ReadyQueue 的稳定排序。
Edge(torch/csrc/autograd/edge.h:14-39):(shared_ptr<Node> function, uint32_t input_nr) 的 pair。
SavedVariable(torch/csrc/autograd/saved_variable.h):前向保存的中间张量,带版本号(用于检测 in-place 冲突)和可选 hooks(用于 offload、checkpoint)。

4.2 反向调度:Engine

入口:torch.autograd.backward(loss) → torch/csrc/autograd/python_engine.cpp → Engine::execute()(torch/csrc/autograd/engine.cpp:1286-1365)。

核心结构:

GraphTask(graph_task.h:18-200):一次 backward 的上下文,持有 dependencies_(每个节点的入度)、not_ready_(部分就绪节点的 InputBuffer 映射)、exec_info_(只计算子集梯度时的过滤)。
ReadyQueue(engine.h:86):优先级队列,按"深度优先于序号"排序(CompareNodeTaskTime,engine.h:99-103) — 这让反向尽量线性化,减少内存峰值。
NodeTask:(Node*, InputBuffer, 重入深度)。

调度流程(engine.cpp:1063-1210 的 evaluate_function):

等待所有输入梯度张量的 CUDA 事件完成(line 1073-1089)
调用 call_function(graph_task, func, inputs) → 最终调用 node->apply(inputs) → 用户的 backward 代码
如果 keep_graph=false,释放 SavedVariable(line 1132)
把输出梯度按 next_edges_ 分发,累加到下一节点的 InputBuffer(line 1161-1210)
当某节点的依赖归零,推入 ReadyQueue

工作线程:每个 CUDA 设备一个后台线程 + 主线程。重入(backward-in-backward)超过 MAX_DEPTH=60(engine.h:37)时 fork 新线程,避免栈爆。

4.3 自定义 Function

Python 层(torch/autograd/function.py):

class MyFunc(torch.autograd.Function):
    @staticmethod
    def forward(ctx, x):
        ctx.save_for_backward(x)
        return x * x
    @staticmethod
    def backward(ctx, grad_out):
        (x,) = ctx.saved_tensors
        return 2 * x * grad_out

底层实现走 torch/csrc/autograd/custom_function.{h,cpp}:

PyFunction 继承自 Node,其 apply() 通过 pybind11 回调 Python backward
ctx.save_for_backward() 把张量存入 SavedVariable 列表
前向返回时,通过 _wrap_outputs() 把输出张量的 grad_fn_ 设为这个 PyFunction 实例

4.4 functorch:grad / vmap / jvp

代码:torch/_functorch/(Python)+ aten/src/ATen/functorch/(C++)。

核心思想:每个 transform 在 dispatch 栈上压入一个新的 DispatchKey(FuncTorchGradWrapper、FuncTorchBatched、FuncTorchJvp 等)。

grad(f):给每层增加一个 GradInterpreter,前向时启用 autograd,出 wrapper 时 backward 得到梯度。
vmap(f):每个 ATen 算子都有一个 batching rule(aten/src/ATen/functorch/BatchRules*.cpp)定义"批维应如何传播"。例如 matmul 的规则把第一个参数的第 bdim 维和第二个参数的批维合并处理。
jvp / jacrev / hessian:由 grad/vmap 组合而成。

栈叠加:grad(vmap(f)) 时,vmap 层在内,grad 层在外。内层先处理批维,外层再做 autograd。

4.5 Checkpointing 与 Saved Tensor Hooks

torch.utils.checkpoint.checkpoint(fn, *args):用 CheckpointFunction(继承 autograd.Function)包装,前向不保存激活,backward 时重算。
torch.autograd.graph.saved_tensors_hooks(pack, unpack):全局 hook,用于把激活 offload 到 CPU/磁盘。

5. CUDA 后端:Stream、缓存分配器、CUDA Graphs、AMP

5.1 Stream 管理

c10::cuda::CUDAStream(c10/cuda/CUDAStream.h:10-50):PyTorch 每设备维护 3 个 stream 池(default / low / high),每池 32 条,round-robin 分配。
当前 stream 存在 thread_local 中(c10/cuda/CUDAStream.cpp:168),所以多线程不会互相污染。
c10::cuda::CUDAGuard / CUDAStreamGuard(c10/cuda/CUDAGuard.h):RAII 包装,禁用拷贝/移动,保证 exception-safe 的设备与 stream 切换。

5.2 缓存分配器(面试高频)

文件:c10/cuda/CUDACachingAllocator.cpp。这是 PyTorch 最精巧的基础设施之一。

核心数据结构(c10/cuda/CUDACachingAllocator.cpp:168-223):

Block:一段连续的 GPU 内存,带前后指针(双向链表,用于合并)、size、stream 所有权、是否在 expandable segment 里。
BlockPool:两个 std::set<Block*, Compare>,一个按 size 排序(用于 best-fit 查找),一个按地址排序(用于邻居合并)。
小块池(<1 MB)和大块池(≥1 MB)分离 — 避免小分配污染大块池导致碎片(:85-96)。

分配流程:

在当前 stream 的池里按 best-fit 找空闲块,能切就切(受 max_split_size 限制)
若无,尝试从其他 stream 释放的块里"窃取"(需要先在原 stream 上插一个 event 等待)
若还无,cudaMalloc 新段,加入池
OOM 时释放所有空闲块后重试,仍失败则抛出

Expandable Segments(:314-316):新机制,允许单个 segment 动态扩容,减少 fragmentation。通过 PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True 开启,大模型训练几乎必开。

Snapshot / History(c10/cuda/CUDACachingAllocator.h:77-82):

struct SnapshotInfo {
  std::vector<SegmentInfo> segments;
  std::vector<std::vector<TraceEntry>> device_traces;
  std::vector<AnnotationEntry> external_annotations;
  AllocatorConfigInfo config_metadata;
};

torch.cuda.memory._snapshot() 导出的就是这个结构,喂给 pytorch.org/memory_viz 可可视化 OOM。

cudaMallocAsync 路径:c10/cuda/CUDAMallocAsyncAllocator.cpp 在 CUDA ≥ 11.4 上提供替代实现,靠 driver 侧内存池,有时比 caching allocator 更省。

5.3 cuDNN / cuBLAS 句柄池

aten/src/ATen/cuda/CublasHandlePool.cpp:每设备一个 cuBLAS 句柄,线程间共享;cuBLASLt 在 ROCm 下按 (device, stream) 唯一。
PointerModeGuard(CUDABlas.h:24-39):RAII 切换 host/device 指针模式。
aten/src/ATen/native/cudnn/Conv_v8.cpp:使用 cudnn_frontend API,按 allow_tf32、deterministic 过滤算法候选。BF16/FP16 卷积在内部提升到 FP32 再转回。

5.4 CUDA Graphs

文件:aten/src/ATen/cuda/CUDAGraph.cpp:101-162、c10/cuda/CUDAGraphsC10Utils.h。

g = torch.cuda.CUDAGraph()
with torch.cuda.graph(g, pool=mempool, stream=s):
    static_out = model(static_in)
# replay
g.replay()

capture 过程:

cudaStreamBeginCapture,后续 kernel 被记录但不执行
分配器切换到私有 mempool(CUDACachingAllocator.cpp:125-143),保证地址在 replay 时稳定
cudaStreamEndCapture 生成 cudaGraph_t,实例化为 cudaGraphExec_t

replay 时直接提交整张图到 stream,省掉每个 kernel 的 launch 开销 — 对小 kernel 密集的模型(推荐系统、decoder 推理)提升巨大。

5.5 AMP / Autocast

aten/src/ATen/autocast_mode.cpp:

bool is_autocast_enabled(at::DeviceType device_type) {
  at::DispatchKey dispatch_key = get_autocast_dispatch_key_from_device_type(device_type);
  return !c10::impl::tls_is_dispatch_key_excluded(dispatch_key);
}

原理:autocast 是一个 DispatchKey(AutocastCUDA/AutocastCPU),其 kernel 拦截被列为 fp16-safe 的算子(conv, matmul, …),把输入 cast 为 fp16/bf16 后 redispatch 到真实后端。

Cache(:41-52):flat_hash_map<TensorImpl*, weakref> 记录已 cast 过的 fp32 权重,避免重复 cast。
Nesting:thread_local 计数器,退出最外层 autocast 时清空 cache。

反向是否用低精度?—— 反向自动跟随前向 kernel 的 dtype。GradScaler(torch/amp/grad_scaler.py)负责梯度的动态 loss scale。

6. torch.compile 栈:Dynamo + AOTAutograd + Inductor

这是 PyTorch 2.x 的核心卖点。整体流水线:

Python frame (CPython)
      │ PEP 523 hook
      ▼
TorchDynamo    (torch/_dynamo/)
      │ 输出 (FX GraphModule, guards)
      ▼
AOTAutograd    (torch/_functorch/_aot_autograd/)
      │ 联合前向-后向追踪 → 切分 → 函数化
      │ 输出 (forward_fx, backward_fx)
      ▼
TorchInductor  (torch/_inductor/)
      │ 调度融合 + Triton/C++ 代码生成
      ▼
编译后的 Python 可调用对象 + 生成的 .triton.py / .cpp

6.1 Dynamo:字节码符号执行

帧 hook(torch/_dynamo/eval_frame.py:61, 156-168):通过 C 扩展 torch._C._dynamo.eval_frame 调用 PyInterpreterState_SetEvalFrameFunc(PEP 523),截获每个 Python 函数的帧,零字节码重写。

入口(torch/_dynamo/convert_frame.py:603):ConvertFrameAssert.__call__ 检查 cache → 未命中则 trace_frame() 创建 InstructionTranslator(torch/_dynamo/symbolic_convert.py:4928)。

符号执行:

维护符号栈(VariableTracker 列表)和符号 locals
每条字节码(LOAD_FAST、BINARY_OP、CALL 等)对应一个 handler,产出新的 VariableTracker
遇到"不可追踪"的操作(如打印、复杂 Python 内置)触发 graph break — 当前 FX 图收尾,编译执行,然后回到 eager 继续

VariableTracker(torch/_dynamo/variables/base.py:302):每个 Python 值(Tensor、int、list、module)都有对应子类(TensorVariable、ConstantVariable、ListVariable、UserDefinedObjectVariable…)。每个 tracker 带 source(来自哪个局部变量/属性),用于生成 guard。

Guard(torch/_dynamo/guards.py:1156):

类型:TYPE_MATCH、ID_MATCH、TENSOR_MATCH(dtype/device/shape/stride/dispatch keys)、DICT_KEYS…
在 C++ 层编译为 RootGuardManager 树(torch/csrc/dynamo/guards.cpp),下次同一帧进入时先用 guard 筛选缓存的编译结果
动态 shape:torch._dynamo.mark_dynamic(x, 0) 或 dynamic=True,生成符号形状 guard

输出(torch/_dynamo/output_graph.py:587, 1821):OutputGraph.compile_subgraph() 构造 FX 图,调 backend(默认 Inductor),把结果包成 Python 字节码继续执行。

6.2 AOTAutograd

文件:torch/_functorch/aot_autograd.py、torch/_functorch/_aot_autograd/。

任务:把 Dynamo 出的 "Python 级 FX 图"(含 nn.Module、可能的 in-place)降到 纯函数化、前向+反向分离的 ATen 级 FX 图。

两阶段(_aot_autograd/graph_compile.py):

aot_stage1_graph_capture(:191):
1. run_functionalized_fw_and_collect_metadata():用 FakeTensor 和 Functionalize key 跑一遍前向,分析输入/输出、view、mutation
2. create_joint()(graph_capture_wrappers.py:294):构造联合前向+反向函数,用 make_fx 追踪成一个大 FX 图
aot_stage2_compile(:361):
1. partition_fn(默认 min_cut_rematerialization_partition)按最小割算法切分联合图,决定哪些中间张量 save、哪些 recompute
2. 分别把前向图和反向图送给 fw_compiler / bw_compiler(默认都是 Inductor)

函数化:所有 x.add_(y) 变成 x_new = x.add(y); copy_(x, x_new),所有 view 被追踪为显式的 aten.view 节点。生成的图无副作用、无 alias,可以被任意下游编译器消费。

6.3 Inductor

入口:torch/_inductor/compile_fx.py:2646 的 compile_fx()。

降级:torch/_inductor/lowering.py 把 ATen 算子映射为 Inductor IR(torch/_inductor/ir.py):

Pointwise(ir.py:1126):逐元素
Reduction(:1276):归约
ComputedBuffer(:4843):已分配存储的中间结果
ExternKernel:无法内联的(如 cuBLAS 矩阵乘、aten fallback)

调度:torch/_inductor/scheduler.py:3123 的 Scheduler:

构造 SchedulerNode 依赖图
识别 MemoryDep(读写同一缓冲)、StarDep
贪心融合相邻的 pointwise / reduction 为 FusedSchedulerNode — 这是 Inductor 的核心性能来源

代码生成:

GPU:torch/_inductor/codegen/triton.py 生成 Triton kernel。自动调优(autotune_process.py)并行尝试 block size、num_warps、num_stages 的组合,结果进 codecache。
CPU:codegen/cpp.py 生成向量化 C++ 循环,调用 MKL-DNN。
Wrapper:codegen/wrapper.py 或 cpp_wrapper_gpu.py 生成调用上面 kernel 的 Python/C++ 驱动代码。

缓存:torch/_inductor/codecache.py:352, 3948:

按 FX 图 + config 的 hash 命名,路径 ${TORCHINDUCTOR_CACHE_DIR}/XX/ (XX = hash 前 2 字符)
PyCodeCache 缓存编译后 Python 模块;FxGraphCache 缓存编译结果;AOTAutogradCache 缓存 AOTAutograd 产物

6.4 Backend 注册

torch/_dynamo/backends/registry.py:74:

@register_backend
def my_backend(gm: torch.fx.GraphModule, example_inputs): ...

内置后端:inductor(默认)、cudagraphs、onnxrt、openxla、aot_eager、eager(仅 Dynamo)。

6.5 torch.export vs torch.compile

维度	torch.compile	torch.export
控制流	graph break 回 eager	strict:拒绝;non-strict:生成 guard
Guard	有,用于缓存命中	strict:无;non-strict:有
形状	支持符号 shape	支持符号 shape,可显式标注 `Dim`
用途	训练/推理加速	序列化、AOTI、外部编译、移动端
实现	Dynamo + AOT + Inductor	`_strict_export()` / `_non_strict_export()`,产 `ExportedProgram`

7. 分布式子系统

7.1 c10d 基础

torch/csrc/distributed/c10d/ + torch/distributed/。

ProcessGroup(ProcessGroup.cpp:99):持有一个 deviceTypeToBackend_ 映射(CPU → Gloo,CUDA → NCCL 等),getBackend(DeviceType) 延迟初始化。
Backend:ProcessGroupNCCL / ProcessGroupGloo / ProcessGroupMPI / ProcessGroupXCCL。
Work(Work.hpp:56):异步操作句柄,isCompleted()、wait()、getFuture() 返回 c10::ivalue::Future。
Store:rendezvous 用 KV store(TCPStore、FileStore、HashStore),在 rank 间同步初始化参数。

7.2 ProcessGroupNCCL

文件:torch/csrc/distributed/c10d/ProcessGroupNCCL.{cpp,hpp}。

关键机制:

异步线程池:集合操作入队到专用线程,不阻塞用户线程
Watchdog(ProcessGroupNCCL.hpp:69-100):独立监控线程,轮询 work 完成。超时由 TORCH_NCCL_HEARTBEAT_TIMEOUT_SEC 控制,超时可触发 TORCH_NCCL_DUMP_ON_TIMEOUT 生成 FlightRecorder dump(所有 rank 当前挂起的集合操作的完整栈),这是大规模训练挂起时救命的工具。
CUDAEventCache:复用 cudaEvent_t,避免 cudaEventCreate/Destroy 的 overhead
优先级流:TORCH_NCCL_HIGH_PRIORITY=1 让集合 kernel 与 compute kernel 竞争更少
PreMul Sum 等 NCCL 高级 op(ProcessGroupNCCL.cpp:76-98):在 allreduce 前乘以缩放因子,用于一些梯度聚合模式

7.3 DDP

torch/nn/parallel/distributed.py + torch/csrc/distributed/c10d/reducer.cpp。

工作原理:

构造时按参数 反向顺序 分桶(默认桶大小 25 MB),第一个桶可以更小(_BucketCapacityConfig,distributed.py:37-116)
给每个参数的 grad_accumulator 注册 autograd hook(reducer.cpp:127)
backward 中,每个参数梯度计算完就在 hook 里标记"就绪"
桶内所有参数就绪 → mark_bucket_ready(bucket_index)(reducer.cpp:1050) → 启动该桶的 allreduce
计算与通信重叠:靠谱的参数反向顺序分桶让"前面的桶在通信,后面的桶在算梯度"

find_unused_parameters=True 会先做一遍遍历确定哪些参数实际参与反向,代价不小,只在必要时开。

7.4 FSDP 与 FSDP2

FSDP v1(torch/distributed/fsdp/):

FlatParameter(_flat_param.py:40)把同一个 FSDP 单元内的所有参数平铺到一块连续张量
前向前 allgather 平铺张量 → 跑前向 → 立刻 reshard 释放
backward 前再 allgather → 跑反向 → reduce_scatter 梯度 → reshard

FSDP v2 / fully_shard(torch/distributed/_composable/fsdp/):

基于 @contract 装饰器,把 FSDP 状态 附加到 模块而不是 wrap
直接使用 DTensor 作为参数表示(每个参数是一个 DTensor(placements=[Shard(0)]))
可组合:fully_shard + checkpoint + activation_memory_budgeter 可叠加
需要 PyTorch ≥ 2.4,新项目推荐直接上 FSDP2

7.5 DTensor & DeviceMesh

torch/distributed/tensor/ + torch/distributed/device_mesh.py。

DeviceMesh:多维设备拓扑(如 (dp=8, tp=2))。
DTensor:分布式张量,placements 描述每维如何分:
- Shard(dim):沿维度切分
- Replicate:完整复制
- Partial:部分和(需要后续 reduce)
redistribute(_redistribute.py):给定源 placement 和目标 placement,求最小通信成本路径(use_min_cost_redistribution_plan() 开启 Dijkstra,否则贪心)

DTensor 是 2D/3D 并行(DP + TP + PP) 的统一表示,也是 FSDP2 的底层。

7.6 Pipeline 与 torchrun

torch.distributed.pipelining(stage.py):PipelineStage 封装子模型,前向输出通过 send/recv 传到下一 stage,支持 1F1B、interleaved 1F1B 等调度。
torchrun(torch/distributed/run.py):控制台脚本,启动 elastic agent,设置 RANK/LOCAL_RANK/WORLD_SIZE/MASTER_ADDR:PORT,支持节点级重启(via torchelastic)。

8. 用户层:nn.Module / Optimizer / DataLoader / FX / Profiler

8.1 nn.Module

文件:torch/nn/modules/module.py(3 053 行)。

三大容器(Module.__init__,line 482):

self._parameters: dict[str, Parameter | None]
self._buffers:    dict[str, Tensor | None]
self._modules:    dict[str, Module | None]

用 super().__setattr__ 绕过 Module.__setattr__ 钩子初始化。

魔法 __setattr__:拦截 self.weight = Parameter(...) / self.layer = nn.Linear(..) / self.running_mean = tensor(..),路由到对应容器,并触发注册 hooks。

_call_impl(line 1782,forward 的真实入口):

def _call_impl(self, *args, **kwargs):
    forward_call = self._slow_forward if torch._C._get_tracing_state() else self.forward
    # 快路径:没 hook 直接调 forward
    if not (self._backward_hooks or self._forward_hooks or ...):
        return forward_call(*args, **kwargs)
    # 慢路径:依次跑 pre_hooks → forward → hooks

注意:tracing 模式(TorchScript 或 Dynamo 的某些路径)会走 _slow_forward,附加 trace 元数据。

_apply(line 930):.cuda() / .to(device) / .half() 的底层 — 递归子模块,对每个参数/buffer 调用转换函数。支持 swap_tensors(保留 grad 与 optimizer state 的引用)。

state_dict / load_state_dict(2176/2530):

递归收集,key 用点号拼接("layer1.weight")
_metadata 字段记版本号,供 _load_from_state_dict 做向后兼容
strict=True 时不匹配的 key 报错,否则返回 (missing_keys, unexpected_keys)

8.2 Optimizer

文件:torch/optim/optimizer.py:339(Optimizer 基类,1 190 行)。

状态管理:

self.state: defaultdict[Tensor, dict]       # 参数 → 状态字典
self.param_groups: list[dict]                # 参数组,每组有独立超参

三套 kernel(以 Adam 为例,torch/optim/adam.py):

_single_tensor_adam(line 347):Python 循环,每参数单独更新 — 最慢,但可被 TorchScript
_multi_tensor_adam(line 553):torch._foreach_* 原语,批处理多参数 — 默认,torch.compile 友好
_fused_adam(line 802):C++/CUDA 单 kernel 融合 m/v/param 更新 — 最快,仅 CUDA/XPU,需 fused=True

选择(adam.py:963):

if fused and not torch.jit.is_scripting():
    func = _fused_adam
elif foreach and not torch.jit.is_scripting():
    func = _multi_tensor_adam
else:
    func = _single_tensor_adam

生产建议:优先用 foreach=True(默认),大规模训练时 fused=True。两者在不同版本上 bug 出现过,注意固定 PyTorch 版本并做 loss 曲线对比。

8.3 DataLoader

文件:torch/utils/data/dataloader.py(1 686 行)。

单进程(_SingleProcessDataLoaderIter,line 750):主线程直接 dataset[idx] → collate → pin_memory。调试首选。

多进程(_MultiProcessingDataLoaderIter,line 784):

num_workers 个 worker 子进程,通过 multiprocessing.Queue 拉 index / 推 batch
预取(line 1271):初始化时塞 prefetch_factor * num_workers 个索引,稳态下每消费一个就补一个
持久化 worker(persistent_workers=True,line 1214):epoch 结束不杀进程,_reset() 重启迭代。强烈推荐开启 — 避免每 epoch 重建 worker 的开销(尤其 dataset __init__ 重的场景)。
pin_memory 线程(line 656):后台线程把 batch cast 到 pinned memory,host→device 复制可 async

坑位:

Windows / macOS 默认用 spawn,worker 启动要重新 import — 顶层大对象(全局模型)会被复制到每个 worker。解决:DataLoader(..., multiprocessing_context='fork')(仅 Linux)或把资源放在 worker_init_fn 里初始化。
数据量大时 num_workers 不是越多越好 — 受 RAM、fd、IO 带宽限制。实测常见最优值 4-8。

8.4 FX

文件:torch/fx/,核心 graph.py(2 465 行)、graph_module.py、proxy.py、symbolic_trace.py。

symbolic_trace(module):

创建 Tracer,用 Proxy 代理张量
执行一次 forward,Proxy.__getattr__ / __call__ 把调用记录成 Node
输出 GraphModule = nn.Module + FX Graph + 动态生成的 forward 源码

Node op 类型:placeholder(输入)、get_attr(参数/buffer)、call_function、call_method、call_module、output。

限制:FX 不处理数据依赖控制流(if x.sum() > 0: 会报错)。对此用 concrete_args 特化,或直接用 torch.export。

典型用途:

图变换(融合 conv+bn+relu、插入量化观察器)
模型可视化
自定义 pass 作为 torch.compile backend

8.5 Profiler

文件:torch/profiler/profiler.py(1 263 行)+ torch/autograd/profiler.py。

基本用法:

with torch.profiler.profile(
    activities=[torch.profiler.ProfilerActivity.CPU, torch.profiler.ProfilerActivity.CUDA],
    schedule=torch.profiler.schedule(wait=1, warmup=1, active=3, repeat=2),
    on_trace_ready=torch.profiler.tensorboard_trace_handler('./log'),
    record_shapes=True,
    profile_memory=True,
    with_stack=True,
) as prof:
    for step, batch in enumerate(loader):
        train_step(batch)
        prof.step()

内部:底层走 Kineto(NVIDIA CUPTI 封装),ProfilerActivity 决定采集什么。schedule 返回 ProfilerAction(NONE/WARMUP/RECORD/RECORD_AND_SAVE),控制何时开启采样、何时导出。

内存 timeline(_memory_profiler.py:1013):export_memory_timeline('mem.html', device='cuda:0') 生成按类别着色的时间线(参数/梯度/激活/优化器状态/临时)。

Chrome trace:prof.export_chrome_trace('trace.json') 在 chrome://tracing 或 perfetto.dev 打开,看每个 kernel 的 GPU 时间。

9. 部署与导出

9.1 路径选择矩阵

场景	推荐路径	起点	中间件	产物
服务器 Python 推理	torch.compile + inductor	原模型	`torch.compile(model)`	Python 对象
服务器 C++ 推理	AOTI	`torch.export`	`torch._inductor.aoti_compile_and_package()`	`.pt2` (含 `.so`)
跨框架 / 加速器	Dynamo-ONNX	`torch.export`	`torch.onnx.export(..., dynamo=True)`	`.onnx`
移动端	ExecuTorch(独立仓库)	`torch.export`	ExecuTorch 编译器	`.pte`
PTQ 量化	FX quant 或 torchao	训练好的 model	`prepare_fx` → calib → `convert_fx`	量化模块
遗留	TorchScript	`torch.jit.script/trace`	-	`.pt` / `.ptl`

重要事实:

TorchScript 自 PyTorch 2.5 起标记 @deprecated(torch/jit/__init__.py:132, 185, 210, 256, 298, 307,出现多次 "TorchScript is deprecated, please use torch.compile instead.")。新项目不再建议。
PT2E 量化已迁出主仓库,见 torch/ao/quantization/fx/prepare.py 中的提示:"Please use torchao (https://github.com/pytorch/ao) for pt2e quantization flow"。主仓库仅保留兼容 shim。
ExecuTorch 在独立仓库 pytorch/executorch。PyTorch 主仓库 .ci/docker/ci_commit_pins/executorch.txt 只钉住协作版本号。

9.2 torch.export

import torch
from torch.export import export, Dim

model = MyModel().eval()
example = (torch.randn(1, 3, 224, 224),)

batch = Dim("batch", min=1, max=32)
ep = export(model, example, dynamic_shapes={"x": {0: batch}})
torch.export.save(ep, "model.pt2")
# 加载
ep2 = torch.export.load("model.pt2")

产物 ExportedProgram(torch/export/exported_program.py:1048)包含:

_graph_module:规范化的 ATen-level FX 图
_state_dict:参数 + buffers
_graph_signature:输入/输出/参数/buffers 的角色标注
符号形状约束、常量表、模块调用图

strict vs non-strict:

strict=True(默认):用 Dynamo 严格追踪,任何数据依赖控制流都必须通过 torch.cond/torch.while_loop 显式表达
strict=False:用 fake tensor + proxy 追踪,允许更多 Python 构造,但可能漏掉某些条件分支

9.3 AOT Inductor(生产推理利器)

ep = torch.export.export(model, example)
path = torch._inductor.aoti_compile_and_package(ep, package_path="model.pt2")
# 加载执行
loaded = torch._inductor.aoti_load_package(path)
out = loaded(x)

产物:.pt2 里含一个编译好的 .so(CUDA + Triton 产物或 CPU C++ 产物)+ 权重。无需 PyTorch 运行时,C++ 侧可以通过 torch/_inductor/codegen/aoti_runtime/interface.cpp 暴露的 C API 加载。

何时用:

延迟敏感的 C++ 服务
想去掉 Python GIL 开销
需要静态编译产物以避免每次冷启动的编译时间

9.4 ONNX(Dynamo 路径)

import torch
program = torch.onnx.export(
    model, (x,), "model.onnx",
    dynamo=True,                     # 强烈推荐
    dynamic_shapes={"x": {0: "batch"}},
    opset_version=18,
)
program.optimize()                   # 可选 ONNX 优化
program.save("model_opt.onnx")

实现:torch/onnx/__init__.py:65-117。dynamo=True 时先走 torch.export.export,再把 ATen 算子翻译为 ONNX op。旧的 TorchScript 路径(dynamo=False)仍可用但不再演进。

9.5 量化现状

Eager mode(torch.ao.quantization.quantize):prepare → 插入 observer → calib → convert 替换为量化算子。代码量大、易错,但无 trace 限制。
FX mode(torch.ao.quantization.quantize_fx:254):prepare_fx(model, qconfig_mapping, example_inputs) 符号追踪后在图上插 observer,convert_fx 转换。
PT2E:用 torchao — from torchao.quantization.pt2e.quantizer import X86InductorQuantizer → 应用到 ExportedProgram → 编译。这是最新推荐。

10. 生产应用教程

10.1 训练骨架(单机单卡)

# train.py
import torch, torch.nn as nn
from torch.utils.data import DataLoader
from torch.amp import autocast, GradScaler

device = torch.device("cuda")
model = MyModel().to(device, memory_format=torch.channels_last)  # CNN 场景
model = torch.compile(model, mode="max-autotune")                # PT2 加速

opt = torch.optim.AdamW(model.parameters(), lr=3e-4, fused=True)
scaler = GradScaler(device="cuda")

train_loader = DataLoader(
    train_ds, batch_size=256, shuffle=True,
    num_workers=8, pin_memory=True, persistent_workers=True,
    prefetch_factor=4,
)

for epoch in range(epochs):
    for batch in train_loader:
        batch = {k: v.to(device, non_blocking=True) for k, v in batch.items()}
        opt.zero_grad(set_to_none=True)                          # 比置 0 省内存带宽
        with autocast(device_type="cuda", dtype=torch.bfloat16):  # A100+ 用 bf16,V100 用 fp16
            out = model(batch["x"])
            loss = criterion(out, batch["y"])
        scaler.scale(loss).backward()
        scaler.unscale_(opt)
        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
        scaler.step(opt)
        scaler.update()

每一行的"为什么":

memory_format=torch.channels_last:NHWC 布局,cuDNN 在 Ampere+ 上对 conv 更快
torch.compile(mode="max-autotune"):让 Inductor 为每个 kernel 尝试多套 block size,产生最快变体
fused=True:触发 _fused_adam,单 kernel 完成 m/v/param 更新
persistent_workers + prefetch_factor:减少 worker 启动开销、提前供料
set_to_none=True:grad 置 None 而不是 zero,跳过一次写 0,Adam 会检测并重新初始化
bf16 on Ampere+:无需 loss scale(但代码上仍用 GradScaler 为了统一,bf16 时它实际是 no-op;更干净的写法是 bf16 完全去掉 scaler)

10.2 单机多卡 DDP

# ddp_train.py
import os, torch, torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP

def main():
    dist.init_process_group("nccl")
    local_rank = int(os.environ["LOCAL_RANK"])
    torch.cuda.set_device(local_rank)

    model = MyModel().to(local_rank)
    model = torch.compile(model)
    model = DDP(model, device_ids=[local_rank], gradient_as_bucket_view=True)
    # gradient_as_bucket_view=True: 让 param.grad 直接是 bucket 的 view,省一倍梯度内存

    # sampler 保证不同 rank 看不同样本
    sampler = torch.utils.data.distributed.DistributedSampler(ds, shuffle=True)
    loader = DataLoader(ds, batch_size=bs_per_gpu, sampler=sampler,
                        num_workers=4, pin_memory=True, persistent_workers=True)

    opt = torch.optim.AdamW(model.parameters(), lr=3e-4, fused=True)
    for epoch in range(epochs):
        sampler.set_epoch(epoch)   # 关键:保证每 epoch shuffle 种子不同
        for batch in loader:
            # ... forward/backward/step
            pass
    dist.destroy_process_group()

if __name__ == "__main__":
    main()

启动:torchrun --nproc_per_node=8 ddp_train.py。

调优点:

DDP(static_graph=True) 若模型结构完全静态(没有 dynamic control flow),开启后 DDP 会跳过每 step 的 used-param 检测
多机:torchrun --nnodes=N --nproc_per_node=8 --rdzv_id=xxx --rdzv_backend=c10d --rdzv_endpoint=master:29500 train.py
NCCL 调优环境变量:
- NCCL_ASYNC_ERROR_HANDLING=1 + TORCH_NCCL_BLOCKING_WAIT=1:任一 rank 超时,全组失败 — 防止死锁
- TORCH_NCCL_HEARTBEAT_TIMEOUT_SEC=300
- TORCH_NCCL_DUMP_ON_TIMEOUT=1 + TORCH_NCCL_DEBUG_INFO_TEMP_FILE=/tmp/nccl_trace:挂起时救命

10.3 FSDP2(大模型训练)

from torch.distributed._composable.fsdp import fully_shard, MixedPrecisionPolicy
from torch.distributed.device_mesh import init_device_mesh

mesh = init_device_mesh("cuda", (world_size,))  # 纯 DP;或 (dp, tp) 2D mesh

mp = MixedPrecisionPolicy(
    param_dtype=torch.bfloat16,
    reduce_dtype=torch.float32,   # 梯度在 fp32 里 allreduce 保精度
)

for layer in model.layers:       # 每个 transformer block 一个 FSDP 单元
    fully_shard(layer, mesh=mesh, mp_policy=mp)
fully_shard(model, mesh=mesh, mp_policy=mp)

model = torch.compile(model)     # FSDP2 可组合 torch.compile

opt = torch.optim.AdamW(model.parameters(), lr=3e-4, fused=True)
# 训练循环与单卡完全相同,参数/梯度/优化器状态自动分片

为什么比 FSDP v1 更好:无 FlatParameter、状态附加而非 wrap、原生 DTensor、与 torch.compile / fully_shard / checkpoint 可组合。

activation checkpoint:

from torch.distributed._composable.checkpoint_activation import checkpoint
for layer in model.layers:
    checkpoint(layer)
    fully_shard(layer, mesh=mesh, mp_policy=mp)

10.4 推理部署(AOTI + C++)

Python 侧编译:

import torch
from torch.export import export

model = MyModel().eval().cuda()
example = (torch.randn(1, 3, 224, 224, device="cuda"),)
ep = export(model, example)

torch._inductor.aoti_compile_and_package(ep, package_path="model.pt2")

C++ 侧加载(伪代码,实际接口见 torch/csrc/inductor/aoti_runner/):

#include <torch/csrc/inductor/aoti_package/model_package_loader.h>
auto runner = torch::inductor::AOTIModelPackageLoader("model.pt2");
std::vector<at::Tensor> inputs = {input_tensor};
auto outputs = runner.run(inputs);

10.5 监控、诊断与调优

训练挂起:

export TORCH_NCCL_DUMP_ON_TIMEOUT=1
export TORCH_NCCL_TRACE_BUFFER_SIZE=2000
export TORCH_NCCL_DEBUG_INFO_TEMP_FILE=/tmp/nccl_trace_rank

挂起时所有 rank 会 dump 当前 NCCL 工作栈,用 torch.distributed.flight_recorder 工具分析。

OOM 诊断:

torch.cuda.memory._record_memory_history(max_entries=100_000)
# ... 跑几步 ...
torch.cuda.memory._dump_snapshot("mem.pickle")
# 上传到 https://pytorch.org/memory_viz 可视化

可看到每次分配的栈、peak、碎片分布。

性能分析:10.1 节的 Profiler 用法 + torch.compile(..., fullgraph=True, dynamic=False) 看是否有 graph break。用 TORCH_LOGS="dynamo,guards,recompiles" python train.py 输出重编译原因。

Inductor 产物:TORCHINDUCTOR_CACHE_DIR=/tmp/ind torch._inductor.config.debug=True 会把每个编译出的 Triton kernel 写到 /tmp/torchinductor_<user>/,可直接打开 .triton.py 看融合效果。

11. 二次开发:自定义 C++/CUDA 算子

11.1 简单 Python fallback(最快上手)

@torch.library.custom_op("myops::swiglu", mutates_args=())
def swiglu(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
    return x * torch.nn.functional.silu(y)

@swiglu.register_fake
def _(x, y):                                    # Meta kernel,供 Dynamo/export 用 FakeTensor 形状推断
    return torch.empty_like(x)

这就能:在 eager / torch.compile / torch.export 中都工作,进 FX 图成为一个 torch.ops.myops.swiglu 节点。

11.2 C++/CUDA 扩展

目录结构:

myproj/
├─ setup.py
├─ myops/__init__.py
└─ csrc/
   ├─ swiglu.cpp           # CPU + 绑定
   └─ swiglu_cuda.cu       # CUDA kernel

csrc/swiglu.cpp:

#include <torch/library.h>
#include <torch/extension.h>

at::Tensor swiglu_cpu(const at::Tensor& x, const at::Tensor& y) {
  TORCH_CHECK(x.sizes() == y.sizes(), "shape mismatch");
  return x * at::silu(y);
}

at::Tensor swiglu_cuda(const at::Tensor& x, const at::Tensor& y);  // 在 .cu 里实现

TORCH_LIBRARY(myops, m) {
  m.def("swiglu(Tensor x, Tensor y) -> Tensor");
}
TORCH_LIBRARY_IMPL(myops, CPU, m)  { m.impl("swiglu", swiglu_cpu); }
TORCH_LIBRARY_IMPL(myops, CUDA, m) { m.impl("swiglu", swiglu_cuda); }
TORCH_LIBRARY_IMPL(myops, Meta, m) {                     // fake kernel for torch.compile
  m.impl("swiglu", [](const at::Tensor& x, const at::Tensor& y) {
    return at::empty_like(x);
  });
}

csrc/swiglu_cuda.cu:写 CUDA kernel + host launch,参考 aten/src/ATen/native/cuda/ 下的实现模式(用 TensorIterator 或自写 grid-stride loop)。

setup.py:

from setuptools import setup
from torch.utils.cpp_extension import CUDAExtension, BuildExtension

setup(
    name="myops",
    ext_modules=[CUDAExtension(
        name="myops._C",
        sources=["csrc/swiglu.cpp", "csrc/swiglu_cuda.cu"],
    )],
    cmdclass={"build_ext": BuildExtension},
)

myops/init.py:

import torch
from . import _C   # 副作用:触发 TORCH_LIBRARY 注册
swiglu = torch.ops.myops.swiglu

接入 autograd:

@torch.library.register_autograd("myops::swiglu")
def _backward(ctx, grad):
    x, y = ctx.saved_tensors
    # ... 计算 dx, dy
    return dx, dy

@torch.library.register_autograd("myops::swiglu", setup_context=True)
def _setup_ctx(ctx, inputs, output):
    x, y = inputs
    ctx.save_for_backward(x, y)

11.3 一些原则(来自仓库 `CLAUDE.md` 与实战)

数值测试使用 instantiate_device_type_tests,一套用例跑 CPU/CUDA/Meta/ROCm
Tensor 等式用 self.assertEqual,它会处理 dtype、NaN、tolerance
不要 mock dispatcher — 写集成测试跑真实 dispatch 路径,曾有 mock 通过但生产 redispatch 失败的事故

12. 性能关键路径速查表

症状	排查工具	常见原因
GPU 利用率低(<70%)	`nvidia-smi dmon` + Profiler	DataLoader 瓶颈(加 `num_workers`、`persistent_workers`);host-device 同步过多(`loss.item()` 在训练循环里)
OOM	`torch.cuda.memory._dump_snapshot`	激活过大 → checkpoint;碎片 → 开 `expandable_segments:True`;梯度累积 stale
单机多卡加速比差	Profiler 看 NCCL 时间占比	bucket 配置、`gradient_as_bucket_view`、网络拓扑(NVLink vs PCIe);allreduce 与 compute 没重叠 → 看 reducer hook 触发时机
torch.compile 慢	`TORCH_LOGS=recompiles`	形状/dtype 抖动导致重编译;开 `dynamic=True` 或 `mark_dynamic`
第一次 step 慢	Inductor 冷编译	用 `TORCHINDUCTOR_CACHE_DIR` 持久化缓存,CI 里 warm 过再部署
推理延迟尾部毛刺	Profiler memory timeline	caching allocator 在发生 `cudaMalloc`(缓存不足) → 用 CUDA Graphs 或 AOTI 固化

关键环境变量汇总:

# 分配器
export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True,max_split_size_mb:512

# NCCL 诊断
export TORCH_NCCL_ASYNC_ERROR_HANDLING=1
export TORCH_NCCL_DUMP_ON_TIMEOUT=1
export TORCH_NCCL_HEARTBEAT_TIMEOUT_SEC=300
export TORCH_NCCL_TRACE_BUFFER_SIZE=2000

# Dynamo/Inductor
export TORCH_LOGS="dynamo,guards,recompiles"
export TORCHINDUCTOR_CACHE_DIR=/persistent/inductor_cache
export TORCHINDUCTOR_CPP_WRAPPER=1           # C++ 包装器,减少 Python 开销
export TORCHINDUCTOR_MAX_AUTOTUNE=1          # 等价 mode="max-autotune"

# CUDA 编译器
export CUDA_LAUNCH_BLOCKING=1                 # 调试时打开,定位崩溃 kernel

附录:关键文件快速索引

主题	路径
Tensor/TensorImpl	`c10/core/TensorImpl.h:528`、`aten/src/ATen/core/TensorBase.h:94`
Storage	`c10/core/StorageImpl.h:52`
Allocator	`c10/core/Allocator.h`
Dispatcher	`aten/src/ATen/core/dispatch/Dispatcher.h:41`
DispatchKey	`c10/core/DispatchKey.h`、`DispatchKeySet.h`
算子声明	`aten/src/ATen/native/native_functions.yaml`
代码生成	`torchgen/gen.py:1`、`torchgen/model.py`
TORCH_LIBRARY	`torch/library.h:1`
Autograd Engine	`torch/csrc/autograd/engine.cpp:1063, 1286`
Node/Edge	`torch/csrc/autograd/{node.h:113, edge.h:14}`
SavedVariable	`torch/csrc/autograd/saved_variable.h:22`
VariableType 生成	`tools/autograd/gen_variable_type.py`
CUDA Stream	`c10/cuda/CUDAStream.{h,cpp}`
缓存分配器	`c10/cuda/CUDACachingAllocator.cpp:168`
CUDA Graph	`aten/src/ATen/cuda/CUDAGraph.cpp:101`
Autocast	`aten/src/ATen/autocast_mode.cpp:11`、`torch/amp/autocast_mode.py:52`
Dynamo 入口	`torch/_dynamo/convert_frame.py:603`、`eval_frame.py:61`
Dynamo 符号执行	`torch/_dynamo/symbolic_convert.py:4928`
Dynamo Guard	`torch/_dynamo/guards.py:1156`
AOTAutograd	`torch/_functorch/aot_autograd.py:698`、`_aot_autograd/graph_compile.py:191, 361`
Inductor 入口	`torch/_inductor/compile_fx.py:2646`
Inductor IR	`torch/_inductor/ir.py:1126, 1276, 4843`
Inductor 调度	`torch/_inductor/scheduler.py:3123`
代码缓存	`torch/_inductor/codecache.py:352, 3948`
c10d ProcessGroup	`torch/csrc/distributed/c10d/ProcessGroup.cpp:99`
NCCL 后端	`torch/csrc/distributed/c10d/ProcessGroupNCCL.{hpp,cpp}`
DDP Reducer	`torch/csrc/distributed/c10d/reducer.cpp:86, 1050`
FSDP2	`torch/distributed/_composable/fsdp/fully_shard.py`
DTensor	`torch/distributed/tensor/`
DeviceMesh	`torch/distributed/device_mesh.py:44`
torchrun	`torch/distributed/run.py`
nn.Module	`torch/nn/modules/module.py:407, 1782, 930, 2176, 2530`
Optimizer	`torch/optim/optimizer.py:339, 1094`、`torch/optim/adam.py:347, 553, 802`
DataLoader	`torch/utils/data/dataloader.py:142, 750, 784`
FX	`torch/fx/{graph.py, graph_module.py, proxy.py}`
Profiler	`torch/profiler/profiler.py:651`、`_memory_profiler.py:1013`
torch.export	`torch/export/__init__.py:59`、`exported_program.py:1048`
AOTI	`torch/_inductor/__init__.py:56, 273`
ONNX Dynamo	`torch/onnx/__init__.py:65`
量化 FX	`torch/ao/quantization/quantize_fx.py:254`

读完本文后的推荐学习路径:

跑一遍 10.1 的单卡训练,用 Profiler 看 timeline,找到最耗时的 kernel
开 TORCH_LOGS="dynamo,recompiles",跑 torch.compile 的模型,看 guard 触发什么
读 c10/cuda/CUDACachingAllocator.cpp 前 300 行 + torch/csrc/autograd/engine.cpp:1063-1210
仿照第 11 节写一个自定义算子,完整走 def/impl/autograd/meta 四步
把一个真实模型用 10.4 的流程导出成 .pt2,在 C++ 侧加载运行,对比 Python 侧延迟

PyTorch 的"难"不在 API,而在于 从 Tensor 到 Triton kernel 的 7 层抽象每层都可以被扩展和替换。一旦读通本文引用的核心文件,几乎所有高级特性(新硬件后端、新量化方案、新并行范式)都能找到对应的扩展点。

PyTorch 源码深度剖析与生产应用教程

目录