一、全局鸟瞰torch/ 目录到底有什么打开 PyTorch 的torch/目录你会看到超过 80 个子目录和数十个独立 Python 文件。初次面对这个庞然大物很容易迷失方向。但如果你从职责分层的角度去看它可以被清晰地划分为五大层次┌─────────────────────────────────────────────────────┐ │ 第 1 层Python API 门面层 │ │ torch/__init__.py, torch/nn, torch/optim, ... │ ├─────────────────────────────────────────────────────┤ │ 第 2 层编译器与图系统层 │ │ torch/_dynamo, torch/_inductor, torch/fx, │ │ torch/export, torch/compiler │ ├─────────────────────────────────────────────────────┤ │ 第 3 层自动微分与函数变换层 │ │ torch/autograd, torch/_functorch, torch/func │ ├─────────────────────────────────────────────────────┤ │ 第 4 层算子原语与分解层 │ │ torch/_prims, torch/_refs, torch/_decomp, │ │ torch/_ops, torch/_dispatch │ ├─────────────────────────────────────────────────────┤ │ 第 5 层Python-C 桥接层 │ │ torch/csrc, torch/_C, torch/include, torch/lib │ └─────────────────────────────────────────────────────┘下面我们从最贴近用户的顶层开始逐层深入。二、第 1 层Python API 门面层2.1 入口文件 torch/init.py这是整个 PyTorch 的总入口。文件超过 2000 行其核心职责包括加载 C 扩展torch._C这是所有底层能力的来源。暴露公共 API将compile、export、no_grad、Tensor等关键符号注入torch命名空间。平台适配如 Windows 下的 DLL 加载逻辑、ROCm 运行时初始化。懒加载子模块如torch.cuda、torch.distributed等只在首次访问时初始化。# torch/__init__.py 中的关键导出摘录__all__[Tensor,compile,export,no_grad,enable_grad,inference_mode,autocast,vmap,...]这里有一个重要细节torch.compile并不直接定义在__init__.py中而是通过torch._compile模块做了一层懒加载桥接目的是避免torch._dynamo的循环导入问题。2.2 torch/nn — 神经网络模块torch.nn是大多数用户最常接触的子包其结构如下torch/nn/ ├── modules/ # 所有内置层的实现 │ ├── module.py # nn.Module 基类 │ ├── linear.py # Linear 层 │ ├── conv.py # 卷积层族 │ ├── transformer.py# Transformer │ ├── loss.py # 损失函数 │ ├── activation.py # 激活函数 │ ├── container.py # Sequential, ModuleList, ModuleDict │ └── ... ├── functional.py # 函数式接口 F.relu, F.cross_entropy 等 ├── parameter.py # Parameter 与 Buffer 定义 ├── init.py # 权重初始化方法 ├── utils.py # clip_grad_norm_ 等工具 └── parallel/ # DataParallel 等核心设计nn.Module是一切神经网络组件的基类。它维护_parameters、_buffers、_modules三个有序字典通过__setattr__魔术方法实现属性自动注册。当你在__init__中写self.linear nn.Linear(10, 5)时Module.__setattr__会检测到右侧是Module类型自动将其注册到_modules中。2.3 torch/optim — 优化器所有优化器继承自torch.optim.Optimizer基类内置实现包括 SGD、Adam、AdamW、LBFGS、RAdam 等经典算法以及较新的 Adafactor 和 Muon。每个优化器的step()方法实现了对应的参数更新规则。2.4 torch/distributed — 分布式训练这是 PyTorch 最庞大的子包之一涵盖通信后端distributed_c10d.py封装 NCCL、Gloo、MPI 等后端。DDP / FSDPnn/parallel/,_composable/fsdp/数据并行和全分片数据并行。DTensortensor/分布式张量抽象。Pipeline Parallelpipelining/流水线并行。Checkpointcheckpoint/分布式模型检查点保存与加载。Elastic Launchelastic/弹性训练启动器。其初始化逻辑依赖 C 侧的_c10d_init()底层通过 ProcessGroup 抽象统一了多种通信后端。2.5 其他 API 子包子包职责torch/cudaCUDA 设备管理、Stream、Event、CUDAGraphtorch/cpuCPU 设备及 AMP 自动混合精度torch/xpuIntel GPU (XPU) 支持torch/mpsApple Metal Performance Shaders 后端torch/linalg线性代数操作SVD、QR、Cholesky 等torch/fft快速傅里叶变换torch/sparse稀疏张量操作torch/distributions概率分布Normal、Bernoulli 等 40 种torch/profiler性能分析工具torch/onnxONNX 格式导出torch/quantization/torch/ao量化与架构优化三、第 2 层编译器与图系统层torch.compile 核心路径这是 PyTorch 2.x 最具革命性的部分。从用户调用torch.compile(model)开始数据流经过以下子系统用户代码 │ ▼ torch.compiler ← 对外命名空间 │ ▼ torch._dynamo ← 前端字节码捕获 → FX Graph │ ▼ torch.fx ← 中间表示 (IR) │ ▼ torch._inductor ← 后端代码生成 → Triton / C │ ▼ 执行优化后的代码3.1 torch/compiler — 编译器公共命名空间torch/compiler/__init__.py统一暴露了面向用户的编译 APItorch.compiler.compile()— 编译入口等同于torch.compiletorch.compiler.reset()— 清除所有编译缓存torch.compiler.allow_in_graph()— 告诉 Dynamo 跳过函数追踪torch.compiler.substitute_in_graph()— 注册 polyfill 替代函数torch.compiler.list_backends()— 列出所有可用后端torch.compiler.disable()— 禁用编译torch.compiler.set_stance()— 设置编译立场如 eager fallbacktorch/compiler/config.py则是跨组件配置中枢通过Config(alias...)机制将torch.compiler.config.dynamic_shapes映射到torch._dynamo.config.dynamic_shapes实现了统一配置入口。3.2 torch/_dynamo — TorchDynamo 前端TorchDynamo 是 Python 级别的 JIT 编译器也是整个编译器栈最复杂的子系统之一。核心原理利用 CPython 的 PEP 523 Frame Evaluation API在 Python 字节码执行前拦截并修改它将 PyTorch 操作序列提取为 FX Graph。关键源文件及职责torch/_dynamo/ ├── eval_frame.py # ★ 核心入口注册 frame evaluation 回调 │ # optimize()、export()、explain() 在此定义 ├── convert_frame.py # 将 Python frame 转换为编译后的代码 ├── symbolic_convert.py # 符号化字节码分析构建图 ├── guards.py # Guard 系统跟踪编译假设条件 ├── output_graph.py # 管理生成的 FX Graph ├── codegen.py # 从 FX Graph 生成 Python 字节码 ├── config.py # Dynamo 配置项150 个开关 ├── side_effects.py # 副作用追踪 ├── compiled_autograd.py # 编译模式下的自动微分 ├── variables/ # ★ 变量系统30 个文件 │ ├── tensor.py # Tensor 变量的符号追踪 │ ├── nn_module.py # Module 变量追踪 │ ├── builtin.py # Python 内置函数追踪 │ └── ... ├── backends/ # 后端注册表 │ └── registry.py # register_backend / lookup_backend ├── source.py # Source 对象追踪值的来源 ├── bytecode_transformation.py # 字节码变换工具 └── trace_rules.py # 追踪规则哪些函数要追踪/跳过Guard 机制是 Dynamo 的灵魂。每次编译一段代码时Dynamo 都会记录一组卫士条件如x.shape (32, 784)、x.dtype float32。后续调用时先检查 Guard若全部满足则直接执行已编译的代码若任一 Guard 失败则触发重新编译。Graph Break图中断是另一个核心概念当 Dynamo 遇到无法符号化追踪的代码如print(tensor.item())它会将已捕获的子图编译执行回退到 Python 解释器执行不支持的代码继续追踪后续代码。3.3 torch/fx — FX 中间表示FX 是 PyTorch 的Python-to-Python 变换工具包三大核心对象对象职责fx.Graph有序操作列表节点间形成 DAGfx.Node单个操作节点placeholder/call_function/call_module/output等 6 种 opfx.GraphModule持有 Graph 的nn.Module子类可生成有效的forward()代码FX Graph 是 Dynamo 和 Inductor 之间的契约Dynamo 产出 FX Graph各种 pass 在上面做变换最终 Inductor 消费它生成机器码。torch/fx/ ├── graph.py # Graph 数据结构 ├── node.py # Node 数据结构 ├── graph_module.py # GraphModule 实现 ├── proxy.py # Proxy 对象追踪时的占位符 ├── interpreter.py # 图解释器 ├── subgraph_rewriter.py # 子图模式替换 ├── passes/ # 内置变换 pass │ ├── infra/ # Pass 基础设施 │ └── ... └── experimental/ # 实验性功能 ├── symbolic_shapes.py # ★ 符号形状推理ShapeEnv ├── proxy_tensor.py # make_fx通过 Proxy Tensor 追踪 └── ...特别值得关注的是fx/experimental/symbolic_shapes.py它实现了整个符号形状推理系统ShapeEnv是动态形状支持的核心。3.4 torch/_inductor — TorchInductor 后端TorchInductor 是默认的代码生成后端职责是将 FX Graph降低lower为高性能机器码。torch/_inductor/ ├── compile_fx.py # ★ 主编译入口串联 AOTAutograd → FX Pass → 代码生成 ├── graph.py # GraphLowering将 FX 节点降低为 IR 节点 ├── ir.py # ★ Inductor IRBuffer、ComputedBuffer、LoopBody 等 ├── lowering.py # 算子降低规则aten op → IR op ├── scheduler.py # 内核调度与融合策略 ├── codegen/ # ★ 代码生成器 │ ├── triton.py # Triton GPU 内核生成 │ ├── cpp.py # C/OpenMP CPU 内核生成 │ ├── cuda/ # CUDA C 模板内核 │ └── wrapper.py # 包装器代码生成 ├── fx_passes/ # FX 级别的优化 pass │ ├── pre_grad.py # 梯度计算前的优化 │ ├── post_grad.py # 梯度计算后的优化 │ └── joint_graph.py# 联合图优化 ├── cudagraph_trees.py # CUDA Graph 树管理 ├── codecache.py # 编译产物缓存 ├── pattern_matcher.py # 模式匹配引擎融合规则 ├── select_algorithm.py# 自动调优如 GEMM 算法选择 ├── config.py # Inductor 配置项 ├── package/ # AOTI 打包工具 └── runtime/ # 运行时支持compile_fx.py 主流程简化compile_fx(gm, example_inputs) │ ├─ 1. pre_grad_passes() # 图级预优化 │ ├─ 2. AOTAutograd # 分解前向/反向图 │ └─ min_cut_partition() # 最小切割减少梯度保存 │ ├─ 3. post_grad_passes() # 图级后优化 │ ├─ 4. GraphLowering # FX → Inductor IR │ └─ lowering.py # 逐节点降低 │ ├─ 5. Scheduler # 融合决策 │ └─ pattern_matcher.py # 模式匹配驱动融合 │ ├─ 6. Codegen # IR → Triton / C 源码 │ └─ 7. codecache # 编译 缓存3.5 torch/export — AOT 导出系统torch.export提供提前时刻AOT图捕获能力与torch.compile的 JIT 模式互补fromtorch.exportimportexport,save,load exportedexport(model,args(x,),dynamic_shapes{x:{0:Dim(batch)}})save(exported,model.pt2)loadedload(model.pt2)关键特性产出ExportedProgram包含规范化的 ATen 算子图。不含 Python 语义子模块完全内联展平。可序列化、跨平台执行。支持动态形状约束声明。四、第 3 层自动微分与函数变换层4.1 torch/autograd — 自动微分引擎torch.autograd是 PyTorch 的核心竞争力之一采用反向模式自动微分reverse-mode AD。Python 层面torch/autograd/提供用户 APIbackward()— 触发反向传播grad()— 计算梯度Function— 自定义自动微分算子的基类grad_mode—no_grad、enable_grad、inference_mode上下文管理器C 层面torch/csrc/autograd/是性能热点所在。PyTorch 采用双实现模型概念C 类型Python 类型变量Variablevariable.hTHPVariablepython_variable.h微分节点Nodefunction.hTHPFunctionpython_function.hPython 调用桥PyNodeC 类—PyNode是一个关键的胶水层当用户用 Python 定义torch.autograd.Function时C 侧会创建PyNode实例在反向传播时通过它回调 Python 的backward()方法。4.2 torch/_functorch — 函数变换_functorch实现了vmap向量化映射、grad函数式梯度、jacrev/jacfwd雅可比矩阵等组合式变换。它也是AOTAutograd的宿主模块torch/_functorch/ ├── aot_autograd.py # ★ AOTAutograd 核心实现 ├── _aot_autograd/ # AOTAutograd 子模块 │ ├── schemas.py # 图签名定义 │ └── subclass_parametrization.py ├── compile_utils.py ├── pyfunctorch.py # 函数变换的分发机制 └── config.pyAOTAutograd 的核心思路是在编译期就追踪出前向和反向两张图然后用min_cut_rematerialization_partition做最小切割决定哪些中间值需要保存checkpoint、哪些可以重计算从而平衡内存和计算。五、第 4 层算子原语与分解层这一层回答的问题是PyTorch 的 1000 个算子之间是什么关系如何保证新后端只需实现少量原语就能支持全部算子5.1 torch/_prims — 原语算子_prims定义了一套最小化的原语操作约 100 个包括基本的逐元素运算abs、cos、sin、归约运算sum、prod和张量操作reshape、slice。这些是编译器栈中不可再分解的原子操作。它使用torch.libraryAPI 注册到调度器中以prims::xxx命名空间存在。5.2 torch/_refs — 参考实现_refs为 PyTorch 中绝大多数aten算子提供了纯 Python 的参考实现这些实现仅使用_prims中的原语。例如torch._refs.relu用prims.maximum和prims.full_like来实现。这层的价值在于新后端只要实现_prims就自动获得所有_refs算子的支持。为编译器测试提供了黄金标准。5.3 torch/_decomp — 算子分解表_decomp维护了一张全局分解表global_decomposition_table按阶段分为表名时机用途pre_autograd自动微分之前简化反向传播图post_autograd自动微分之后便于后端代码生成meta元信息推断时形状/dtype 推断当编译器遇到一个高级算子如F.batch_norm时它会查分解表将其拆解为更低级的算子组合最终降低到 Inductor IR 可以处理的粒度。5.4 torch/_ops — 算子分发基础设施_ops.py定义了算子系统的 Python 侧基础OperatorBase— 所有算子的基类维护了一个_dispatch_cache。OpOverload— 表示一个具体的 C ATen 算子重载。HigherOrderOperator— Python 级别的高阶算子如torch.cond、torch.vmap。与之配合的torch/_dispatch/提供了 Python 侧的分发器支持。六、第 5 层Python-C 桥接层6.1 torch/csrc — C 源码torch/csrc/是连接 Python 解释器和 C 运行时的桥梁其内部结构映射了主要功能域torch/csrc/ ├── Module.cpp/.h # Python 扩展初始化入口 ├── autograd/ # ★ 自动微分引擎 C 实现 ├── jit/ # TorchScript JIT 编译器 │ ├── passes/ # 优化 pass │ ├── runtime/ # 运行时含 static_runtime │ └── codegen/ # 代码生成CUDA/oneDNN fuser ├── dynamo/ # Dynamo 的 C 加速部分eval_frame ├── inductor/ # Inductor 的 C 运行时 ├── distributed/ # 分布式通信的 C 绑定 ├── cuda/ # CUDA 相关绑定 ├── tensor/ # Tensor 的 Python 绑定 ├── api/ # LibTorch C API ├── functorch/ # functorch C 后端 └── profiler/ # 性能分析器 C 实现开发者须知来自 torch/csrc/README.md永远先#include Python.h否则会得到_XOPEN_SOURCEredefined 错误。调用 Python API 前必须持有 GIL使用pybind11::gil_scoped_acquire。异常桥接使用HANDLE_TH_ERRORS/END_HANDLE_TH_ERRORS宏将 C 异常转换为 Python 异常。智能指针THPObjectPtr等同于THPPointerPyObject管理 Python 对象引用计数。6.2 torch/_C — C 扩展模块torch._C是编译后的 C 扩展模块.so/.pyd由torch/csrc/Module.cpp驱动初始化。几乎所有核心操作Tensor 运算、自动微分、分发器最终都通过_C下沉到 C 实现。七、关键横切系统7.1 torch/_subclasses — Tensor 子类PyTorch 允许用户自定义 Tensor 子类来拦截和修改操作行为。关键实现FakeTensor—不持有真实数据的 Tensor仅做形状/dtype 推断。它是编译器栈的基石编译期间所有 Tensor 操作都通过 FakeTensor 推断输出形状而不需要真正执行计算。FunctionalTensor— 将 in-place 操作功能化为 out-of-place 操作。7.2 torch/_higher_order_ops — 高阶操作高阶操作Higher-Order Ops是 PyTorch 2.x 的重要扩展它们接受函数作为参数高阶 Op用途torch.cond条件分支取代动态 iftorch.vmap/map向量化映射while_loop循环控制流flex_attentionFlashAttention 灵活变体associative_scan关联扫描triton_kernel_wrap内联 Triton 内核它们在 FX Graph 中表示为特殊的call_function节点编译器对其有专门的处理逻辑。7.3 torch/_logging — 结构化日志PyTorch 编译器栈有自己的结构化日志系统fromtorch._loggingimporttrace_structured trace_structured(artifact,metadata_fnlambda:{name:debug_graph,encoding:string},payload_fnlambda:graph_str,)日志可以通过tlparse工具从结构化 trace 中提取便于生产环境调试。7.4 torch/overrides.py —torch_function协议这个文件实现了__torch_function__的 Python 侧逻辑允许自定义类型拦截torch命名空间下的函数调用。它是 FX Proxy 追踪机制的基础——Proxy 对象正是通过__torch_function__将操作记录到 Graph 中。八、核心数据流从 torch.compile(model) 到执行用一个完整的例子串联所有层次torch.compile(backendinductor)deffn(x):returntorch.sin(torch.cos(x))fn(torch.randn(1000,devicecuda))执行全流程1. torch.compile 装饰器 → torch._dynamo.optimize(inductor) 包装 fn → 返回 OptimizedModule 2. 首次调用 fn(x) → eval_frame.py 截获 Python frame → symbolic_convert.py 分析字节码 → 识别 torch.cos, torch.sin 操作 → 构建 FX Graphplaceholder(x) → cos → sin → output → guards.py 记录 Guard: {shape(1000,), dtypefloat32, devicecuda:0} 3. Dynamo 将 FX Graph 交给 Inductor 后端 → compile_fx.py 接收 (GraphModule, example_inputs) → pre_grad_passes: 图级优化 → AOTAutograd: 无反向推理模式直接传递 → post_grad_passes: 融合优化 4. GraphLowering → lowering.py 将 aten.cos, aten.sin 降低为 Inductor IR → scheduler.py 决定融合cos sin → 单个 fused kernel 5. Codegen → triton.py 生成 Triton 内核 triton.jit def fused_kernel(in_ptr0, out_ptr0, xnumel, XBLOCK): tmp0 tl.load(in_ptr0 xindex) tmp1 tl.cos(tmp0) tmp2 tl.sin(tmp1) tl.store(out_ptr0 xindex, tmp2) 6. 编译 缓存 → codecache.py 检查缓存 → Triton 编译为 GPU 二进制 → 返回可调用的 compiled function 7. 后续调用 → Guard 检查通过 → 直接执行编译后代码 → Guard 失败 → 回到步骤 2 重新编译关键优化点融合原本需要 2 次读写内存cos 读写 sin 读写融合后只需 1 读 1 写中间结果驻留寄存器。CUDA Graphs当 Inductor 检测到图可以整体录制为 CUDA Graph 时消除逐 kernel 的 Python 启动开销。九、Tensor 对象的完整生命周期torch/_tensor.py定义了 Python 侧的Tensor类它继承自 C 实现的torch._C.TensorBase。Tensor 的操作通过三重分发机制路由Python Tensor 方法调用 │ ├─ 1. __torch_function__ ← 用户子类拦截点 │ ├─ 2. Python Dispatcher ← TorchDispatchMode 拦截 │ └─ 3. C Dispatcher ← 按 DispatchKey 路由 ├─ Autograd key → 创建反向图节点 ├─ FuncTorch key → vmap/grad 变换 ├─ Backend key → CPU/CUDA/XPU 内核 └─ ...十、目录速查表最后整理一份可以收藏的速查表目录/文件一句话定位__init__.py总入口暴露公共 API_tensor.pyTensor Python 类定义nn/神经网络层、损失函数、容器optim/SGD/Adam/AdamW 等优化器autograd/自动微分 Python APIcompiler/torch.compile 公共配置与入口_dynamo/Dynamo 前端字节码 → FX Graphfx/FX 中间表示Graph/Node/GraphModule_inductor/Inductor 后端FX Graph → Triton/Cexport/AOT 导出ExportedProgram_functorch/AOTAutograd vmap/grad 变换_prims/原语算子不可分解单元_refs/参考实现用 prims 实现 aten_decomp/算子分解表_ops.pyPython 算子分发基础设施_subclasses/FakeTensor、FunctionalTensor_higher_order_ops/cond/vmap/while_loop/flex_attentioncsrc/C 绑定与运行时_C/编译后的 C 扩展模块cuda/CUDA 设备管理与 Streamdistributed/DDP/FSDP/DTensor/Pipeline/RPCprofiler/性能分析工具_logging/结构化日志系统overrides.py__torch_function__协议实现总结PyTorch 的torch/目录是一个精心设计的分层架构顶层 APInn/optim/distributed面向模型开发者。编译器栈Dynamo → FX → Inductor / Export面向性能优化。算子系统prims → refs → decomp保证新后端的可扩展性。C 运行时承载性能关键路径。横切关注点FakeTensor、Guard、torch_function将各层有机串联。理解这个架构你就能在 PyTorch 源码中按图索骥遇到性能问题看_inductor遇到图中断看_dynamo遇到梯度问题看autograd_functorch遇到新算子需求看_prims_decomp。声明本文基于 PyTorch 开源代码分析代码遵循 BSD-style 许可证。如果本文对你有帮助欢迎点赞收藏。关注作者获取更多 PyTorch 源码解析系列文章。