PyTorch 训练框架内核深度解析

贯穿 PyTorch 训练栈：从 Tensor 内存布局一路走到 Inductor 的 Triton 代码生成。

本专栏基于 PyTorch v2.11（torch.compile 已成默认编译路径、FSDP-2 已稳定的里程碑版本）源码，系统剖析当代深度学习训练框架的设计与实现。从一个 tensor.add(other) 的完整旅程出发，由外到内、层层下潜：经过 Python 前端、Dispatcher 多分派、ATen 算子注册、c10 核心抽象、CUDA Caching Allocator、Autograd 反向图、TorchDynamo 帧拦截、AOTAutograd 函数化、Inductor Triton 代码生成、DDP 环形 AllReduce、FSDP ZeRO-3 分片，最终回到 NCCL 通信原语 —— 把每一个架构决策背后的"为什么"翻出来。

本专栏规模

每章 6500+ 中文字 + 至少 4 张 mermaid 图 + 多段真实源码摘录，覆盖 v2.11 源码的核心机制（不是 API 教程、不是入门）。

读完你能做到什么

1. 完整跟完一次 tensor.add(other)。 从 Python 前端到 Dispatcher 多分派、ATen 算子注册、c10 核心抽象，每一跳发生了什么（第1章 架构、第3章 c10、第5章 Dispatcher、第6章 ATen codegen）。
2. 解释显存去哪了。 Tensor 与 Storage 的分离、CUDA Caching Allocator 的分配策略与碎片来源——这是"明明还有显存却 OOM"的答案所在（第2章 Tensor 与 Storage、第4章 CUDA 缓存分配器）。
3. 说清反向图是怎么建起来又怎么跑的。 autograd 的设计、引擎的执行方式、与前向的对应关系（第7章 autograd 设计、第8章 autograd 引擎）。
4. 读懂 torch.compile 的三段式。 TorchDynamo 怎么拦截字节码帧、AOTAutograd 的函数化、Inductor 生成 Triton 代码，以及端到端串起来的完整路径（第12章 Dynamo、第13章 AOTAutograd、第14章 Inductor、第15章 端到端）。
5. 理解分布式训练的两条路。 ProcessGroup 与通信原语、DDP 的环形 AllReduce、FSDP 的 ZeRO-3 分片（第16章 ProcessGroup、第17章 DDP、第18章 FSDP）。
6. 写自定义算子并做性能归因。 自定义算子的注册路径、Profiler 怎么读、量化与混合精度（第22章 自定义算子、第21章 Profiler、第20章 量化与 AMP）。
7. 补齐工程侧。 nn.Module 的机制、优化器、DataLoader、序列化（第9章 nn.Module、第10章 优化器、第11章 DataLoader、第19章 序列化、第23章 设计哲学）。

适合谁读

• AI 工程师：每天在写 loss.backward() 和 torch.compile()，想知道这两行背后各自发生了什么（第7章 Autograd 设计原理与反向图构建、第12章 TorchDynamo）。
• 研究者：想弄清"我写的 Python 代码到底怎么变成 GPU 上跑的 kernel"——这条链路从 Dispatcher 一直走到 Inductor 生成 Triton（第5章 Dispatcher、第14章 TorchInductor）。
• 训练基础设施工程师：DDP 的梯度桶怎么划、FSDP 的分片策略怎么选，都要看到实现才能判断（第17章 DDP、第18章 FSDP）。
• 推理优化与自定义算子的开发者：量化、混合精度、C++ 扩展这三块都要求先知道 ATen 那一层长什么样（第6章 ATen 算子库与代码生成、第20章 量化与混合精度训练、第22章 自定义算子与 C++ 扩展）。
• 不满足于"框架黑盒"的深度学习从业者：把 PyTorch 当成一个可读的系统，而不是一个只能调参数的工具。

与本系列其他专栏的关系

源码版本

本专栏的源码引用以 PyTorch v2.11.0（2026-03-20 发布）为准：

git clone https://github.com/pytorch/pytorch.git
cd pytorch && git checkout v2.11.0

选 v2.11 而不是更早的 v2.0 / v2.4，是因为它是几条脉络汇合后的稳定形态：torch.compile 经过三年打磨、实验性 flag 大多已成默认；FSDP-2 成为新代码的推荐路径；PT2E 量化取代旧的 FX graph mode；Distributed Checkpoint 成为千卡训练的事实格式。正文标注的文件路径与函数名是稳定锚点，行号会随补丁漂移。

目录

开篇

• 前言

第一篇 全景

• 第1章 架构总览：一次张量加法的全旅程

第二篇 张量层与运行时

• 第2章 Tensor、Storage、TensorImpl 三件套
• 第3章 c10 核心抽象：Device、DType、Allocator、intrusive_ptr
• 第4章 CUDA Caching Allocator：碎片化战争

第三篇 Dispatcher 与算子

• 第5章 Dispatcher：多分派的工程实现
• 第6章 ATen 算子库与代码生成

第四篇 Autograd 自动求导

• 第7章 Autograd 设计原理与反向图构建
• 第8章 Engine：多线程后向传播的调度器

第五篇 nn 与训练栈

• 第9章 nn.Module：参数注册、Hooks、序列化
• 第10章 优化器与梯度更新
• 第11章 DataLoader 与多进程数据流水线

第六篇 torch.compile 编译器栈

• 第12章 TorchDynamo：CPython 帧拦截与图捕获
• 第13章 AOTAutograd：函数化与正反向图划分
• 第14章 TorchInductor：从 ATen IR 到 Triton kernel
• 第15章 CUDA Graph 与 torch.compile 端到端

第七篇 分布式训练

• 第16章 进程组与 NCCL 通信
• 第17章 DDP：环形 AllReduce 与梯度桶
• 第18章 FSDP：ZeRO 风格的参数分片

第八篇 工程与生态

• 第19章 序列化：torch.save / torch.load 与权重格式
• 第20章 量化与混合精度训练
• 第21章 Profiler 与性能调优
• 第22章 自定义算子与 C++ 扩展
• 第23章 设计哲学与 PyTorch 演进

版权声明

本专栏内容为 杨艺韬 版权所有，保留一切权利。未经书面许可，不得全文或大段转载、改编、翻译，或用于任何商业用途（含以本专栏内容训练模型、生成衍生课程或商品）。

欢迎分享本专栏的链接。引用少量内容用于评论、教学或研究时，请署名 杨艺韬 并附上原文链接。