乐于分享
好东西不私藏

PyTorch源码篇之源码架构总览【第一篇】

PyTorch源码篇之源码架构总览【第一篇】

PyTorch源码篇之源码架构总览【第一篇】

本文基于pytorch源码分析 PyTorch 的整体架构。PyTorch 是 Meta 开源的动态图深度学习框架,核心特性:

  • 动态计算图(Eager Execution 即时执行): 代码运行时实时构建计算图,调试友好;
  • Python 优先,底层 C++ 高性能内核;
  • 支持 CPU / CUDA / MPS / XPU / ROCm 多硬件后端;
  • 两大核心场景:科研原型实验、工业大模型训练 / 推理;
  • 配套生态:torch.compile 编译栈、TorchServe、TorchRec、TorchVision、TorchAudio 等。

核心顶层能力:

  1. 张量 Tensor: 多维数组,自动微分载体,支持异构设备;
  2. Autograd 自动微分引擎: 自动构建反向传播图;
  3. ATen 算子库: CPU/GPU 底层数学算子集合;
  4. nn.Module 神经网络组件: 层封装、参数管理、优化器;
  5. Compiler 编译栈(Dynamo + Inductor): 动态图捕获、编译优化、算子融合;
  6. 分布式训练(c10d): DDP、FSDP、DTensor 分布式张量;
  7. 推理生态: TorchScript、AOTInductor、Export 模型导出。

两大执行模式:

  • Eager Mode(默认即时执行): 运行一行执行一行,动态图,调试简单;
  • Compiled Mode(torch.compile): Dynamo 拦截 Python 字节码,捕获计算图,编译优化静态子图。

用户视角: PyTorch 的核心入口是 torchtorch.nntorch.autogradtorch.optimtorch.utilstorch.distributedtorch.compile 等模块。从源码视角看,更合适的主干分层是:用户 API 层、C++ 桥接层、张量与算子库层、c10 基础抽象层。在这条主干之外,自动求导、dispatcher、编译栈、设备后端、分布式和代码生成系统作为横切机制或扩展层贯穿其中。

一、整体架构和分层

1.1 分层框图

PyTorch 的源码架构可以先按主调用链分层理解:用户从 Python API 进入,经过 C++ 桥接层,下沉到 ATen 张量与算子库,再依赖 c10 基础抽象表达设备、内存、调度和 Tensor 元数据,最后由具体后端执行数值计算。自动求导、编译栈、分布式、代码生成和扩展机制横跨这些层。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
+--------------------------------------------------------------------+| 用户模型与训练代码                                                   || model(x), loss.backward(), optimizer.step(), torch.compile(model)   |+--------------------------------------------------------------------+                                |                                v+--------------------------------------------------------------------+| 第一层:用户 API 层                                                   || torch / torch.nn / torch.optim / torch.autograd / torch.distributed || torch.utils / torch.export / torch.compile / torch.fx               |+--------------------------------------------------------------------+                                |                                v+--------------------------------------------------------------------+| 第二层:C++ 桥接层                                                     || torch._C / torch/csrc / Python binding / generated bindings          || THPVariable / pybind / C extension / C++ frontend                    |+--------------------------------------------------------------------+                                |                                v+--------------------------------------------------------------------+| 第三层:张量与算子库层                                                  || ATen / at::Tensor / TensorBase / native_functions.yaml               || c10::Dispatcher / OperatorEntry / native kernels / torchgen          |+--------------------------------------------------------------------+                                |                                v+--------------------------------------------------------------------+| 第四层:c10 基础抽象层                                                  || TensorImpl / Storage / DataPtr / Allocator / Device / ScalarType     || DispatchKey / DispatchKeySet / SymInt / intrusive_ptr                |+--------------------------------------------------------------------+                                |                                v+--------------------------------------------------------------------+| 第五层:后端执行层                                                      || CPU / CUDA / HIP / XPU / MPS / Vulkan / Metal / MKL / cuDNN / NCCL   || backend kernels / memory allocator / stream / event / libraries      |+--------------------------------------------------------------------+横切机制层:Autograd / Dispatcher / torch.compile / FX / AOTAutograd / Inductor /Distributed / torchgen / custom op / custom backend / Tensor subclass

这张图表达的是主执行路径。以 torch.add(x, y) 为例,用户 API 层接收 Python 调用,C++ 桥接层把 Python 参数转换为 C++ Tensor 调用,ATen 根据算子 schema 进入 dispatcher,dispatcher 依据 DispatchKeySet 选择 CPU、CUDA、Meta、Autograd 或 Composite 等 kernel,kernel 再依赖 c10 中的 Tensor 元数据、设备、内存和调度抽象完成执行。

1.2 各层功能总表

层次
主要源码位置
核心对象
核心功能
用户 API 层
torch/torch.Tensor
nn.ModuleOptimizerDataLoaderProcessGroupGraphModule
提供用户直接使用的 Python API,组织训练、推理、数据加载、模型定义、分布式和编译入口。
C++ 桥接层
torch/csrc/
torch/_C
THPVariable
、Python binding、generated binding、C extension module
连接 Python 对象和 C++ 核心,处理参数转换、异常转换、GIL、引用计数、Tensor Python 类型和 C++ API 暴露。
张量与算子库层
aten/src/ATen/
aten/src/ATen/native/torchgen/
at::Tensor
TensorBasenative_functions.yamlc10::DispatcherOperatorEntry
定义 Tensor C++ API、算子 schema、算子注册、调度入口、native kernel 和代码生成规则。
c10 基础抽象层
c10/TensorImpl
StorageImplDataPtrAllocatorDeviceDispatchKeySetSymInt
提供 Tensor 内部表示、内存所有权、设备描述、dtype、dispatch key、符号形状、引用计数等基础设施。
后端执行层
aten/src/ATen/cpu
cudahipxpumpsc10/cuda 等
CPU/CUDA/HIP/XPU/MPS kernel、allocator、stream、event、backend library wrapper
执行真实数值计算,接入硬件、运行时、加速库和设备内存管理。
横切机制层
torch/autograd
torch/csrc/autogradtorch/_dynamotorch/_inductortorch/distributedtools/autograd
AutogradMeta
NodeEngine、FX graph、AOTAutograd、Inductor、DDP/FSDP、torchgen
贯穿多层实现自动求导、图捕获编译、分布式训练、代码生成、自定义算子和后端扩展。

1.3 用户 API 层功能表

模块
核心类/对象
功能
torchtorch.Tensor
、顶层算子
提供 Tensor 创建、数学算子、设备迁移、序列化、随机数等基础用户接口。
torch.nnModule
Parameter、各类 layer/loss
组织神经网络结构、参数、buffer、hook、状态字典和训练/推理模式。
torch.optimOptimizer
、参数组、优化器状态
根据梯度更新参数,维护动量、二阶矩、学习率等优化状态。
torch.autogradFunction
、grad mode
提供自动求导用户入口、自定义求导函数、梯度开关和 profiler 能力。
torch.utilsDataset
DataLoader、cpp extension
提供数据加载、扩展构建、checkpoint、benchmark、pytree 等工具。
torch.distributed
DDP、FSDP、ProcessGroup、RPC
支持多进程、多设备、多节点训练和通信。
torch.compile
 / torch.export / torch.fx
GraphModule
、exported program
提供图捕获、程序导出、编译优化和图变换入口。

1.4 C++ 桥接层功能表

模块
核心类/对象
功能
torch._C
C extension module
把 C++ dispatcher、autograd、JIT、profiler、distributed、device API 暴露给 Python。
torch/csrc
Python/C++ binding 源码
管理 Python 类型、C++ 类型转换、异常转换、GIL 和模块初始化。
torch/csrc/autogradTHPVariable
、Python autograd binding
连接 Python Tensor、Python Function 和 C++ autograd engine。
torch/csrc/api
C++ frontend
支持 C++ 侧模型、module、optimizer 等 API。
generated binding
torchgen 生成文件
将 native_functions.yaml 中的算子转换成 Python 可调用入口。
extension binding
C++/CUDA extension
允许用户或第三方库把自定义 C++/CUDA 算子接入 PyTorch。

1.5 张量与算子库层功能表

模块
核心类/对象
功能
ATen/coreTensorBase
、operator schema、boxing
定义 Tensor C++ 基础 API、算子 schema、boxed/unboxed 调用和 dispatcher 相关基础。
ATen/native
native kernel、native_functions.yaml
声明和实现大量 PyTorch 内建算子。
ATen/ops
generated operator wrappers
生成算子调用包装,连接 C++ API 和 dispatcher。
dispatcher
c10::Dispatcher
OperatorEntryKernelFunction
按 dispatch key 选择具体 kernel,支持后端、autograd、layout、Meta、Composite 和扩展。
torchgen
schema model、模板、生成器
从 yaml 生成 ATen API、注册代码、binding、autograd wrapper、functionalization 等。
Composite/Meta
composite kernel、meta kernel
用组合算子减少后端重复实现,用 meta kernel 做无数据 shape/dtype 推导。

1.6 c10 基础抽象层功能表

模块
核心类/对象
功能
c10/core/TensorImpl.hTensorImpl
保存 Tensor 的 sizes、strides、storage、dtype、device、layout、dispatch key set 和 autograd meta 指针。
c10/core/Storage.hStorage
StorageImpl
表示底层内存存储,支持多个 Tensor 共享同一块 storage。
c10/core/Allocator.hAllocator
DataPtr
抽象设备内存分配和释放逻辑。
c10/core/Device.hDevice
DeviceType
描述 CPU、CUDA、HIP、XPU、MPS 等设备。
c10/core/DispatchKey.hDispatchKey
DispatchKeySet
表达后端、layout、autograd、Meta、Functionalize 等调度维度。
c10/core/SymInt.hSymInt
SymBool
支持动态 shape 和编译期符号约束。
c10/utilintrusive_ptr
、异常、容器工具
提供引用计数、错误处理、轻量容器和跨平台工具。

1.7 后端执行层功能表

后端/模块
核心对象
功能
CPU
CPU native kernel、parallel/vectorized
执行 CPU Tensor 算子,接入向量化、线程并行、MKL/oneDNN 等优化。
CUDA
CUDA kernel、cuBLAS、cuDNN、NCCL、CUDA allocator
执行 NVIDIA GPU 算子,管理 stream、event、显存和加速库调用。
HIP/ROCm
HIP kernel、MIOpen、rocBLAS
支持 AMD GPU 后端。
XPU/MPS/Metal/Vulkan
对应后端 kernel 和 runtime wrapper
支持 Intel GPU、Apple GPU、移动端/图形 API 等后端。
Meta
Meta kernel
不分配真实数据,只推导输出元信息,服务编译、导出和 FakeTensor。
Quantized/Sparse/Nested
特殊 layout kernel
支持量化、稀疏、嵌套张量等非普通 dense Tensor 计算。

1.8 横切机制层功能表

机制
涉及源码
功能
Autograd
torch/autograd
torch/csrc/autogradtools/autograd
构建动态反向图,执行 backward,生成并注册求导 wrapper。
Dispatcher
ATen/core/dispatch
c10/core/DispatchKey.h
统一设备、layout、autograd、Meta、Composite、Python dispatch、自定义后端的调度。
编译栈
torch/_dynamo
torch/fxtorch/_functorchtorch/_inductor
捕获 Python 程序、生成 FX 图、编译 forward/backward、融合和生成后端代码。
分布式
torch/distributed
torch/csrc/distributed
支持 DDP、FSDP、RPC、checkpoint、elastic、DTensor 和通信后端。
代码生成
torchgen
tools/autograd
生成算子 API、dispatcher 注册、binding、autograd、functionalization 等胶水代码。
扩展机制
torch.library
、cpp extension、PrivateUse backend
支持 custom op、Tensor subclass、Python dispatch mode 和第三方硬件后端。

二、源码目录结构详解

2.1 顶层目录

当前本地源码根目录为 pytorch/,主要目录如下。

目录
功能
torch/
Python 包主体,包含用户 API、Python 侧模块、C++ binding 源码入口 torch/csrc/、编译栈、分布式、设备模块等。
aten/
ATen 张量和算子库,包含算子 schema、native kernel、Tensor API、backend kernel、代码生成模板。
c10/
PyTorch 低层基础库,包含 Device、DispatchKey、Storage、Allocator、intrusive_ptr、Exception、CUDA/XPU/HIP 基础封装等。
torchgen/
代码生成系统,从 native_functions.yaml、autograd 配置和模板生成 ATen、dispatcher、binding、autograd 等代码。
tools/
开发工具、autograd 生成工具、lint、代码覆盖率、构建辅助、统计脚本等。
test/
Python、C++、分布式、移动端、扩展、JIT、ONNX、Inductor 等测试。
caffe2/
历史 Caffe2 组件和仍被部分路径复用的核心、序列化、性能 kernel、工具代码。
cmake/
CMake 模块、第三方依赖查找、构建规则。
third_party/
子模块和第三方依赖。
functorch/
functorch 历史入口和文档/示例;核心能力已大量整合到 torch.func、AOTAutograd 等路径。
docs/
文档源码。
benchmarks/
性能基准,包括 operator benchmark、dynamo、distributed、inference、sparse 等。
android/
Android 构建和移动端相关工程。
binaries/
二进制构建相关内容。
scripts/
发布、构建、维护脚本。
mypy_plugins/
类型检查插件。
.github/
.ci/
CI、工作流、自动化配置。

2.2 torch/ 目录

torch/ 是 Python 包主体,也是用户最容易接触的源码目录。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
torch/├── nn/                          # 神经网络层、loss、初始化、functional API、parallel wrapper├── autograd/                    # Python 侧 autograd API├── csrc/                        # C++ extension module 源码,连接 Python 和 C++ 核心│   ├── autograd/                # autograd engine、Node、Variable、SavedVariable、Function、binding│   ├── jit/                     # TorchScript/JIT C++ 实现│   ├── distributed/             # 分布式 C++ 支撑│   ├── dynamo/                  # Dynamo 与 CPython eval frame 相关 C++ 支撑│   └── inductor/                # Inductor C++ 支撑├── _dynamo/                     # Python 字节码捕获和 guard 系统├── fx/                          # FX 图 IR 和 symbolic tracing├── _inductor/                   # 编译器后端、lowering、调度、代码生成、runtime cache├── _functorch/                  # vmap、grad、functional transforms、AOTAutograd├── export/                      # 导出 API├── _export/                     # 导出内部实现和约束处理├── distributed/                 # 分布式训练 Python API├── utils/                       # DataLoader、C++ extension、benchmark、checkpoint、pytree├── cuda/                        # CUDA Python API├── xpu/                         # XPU Python API├── mps/                         # MPS Python API├── _subclasses/                 # FakeTensor、ProxyTensor、FunctionalTensor、Tensor subclass├── _decomp/                     # 算子分解├── _refs/                       # 参考实现└── _prims/                      # primitive 算子体系

2.3 aten/ 目录

aten/ 是算子和 Tensor C++ API 的核心目录。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
aten/└── src/ATen/    ├── core/                    # TensorBase、IValue、operator schema、dispatcher、boxing、TensorOptions    ├── native/                  # 大量 native 算子实现    │   └── native_functions.yaml # 算子声明中心    ├── ops/                     # 生成的算子头文件/包装入口    ├── cpu/                     # CPU 相关支持    ├── cuda/                    # CUDA 相关支持    ├── cudnn/                   # cuDNN 集成    ├── hip/                     # HIP/ROCm 后端    ├── xpu/                     # XPU 后端    ├── mps/                     # MPS 后端    ├── metal/                   # Metal 后端    ├── vulkan/                  # Vulkan 后端    ├── quantized/               # 量化算子    ├── functorch/               # functorch 与 ATen 交互部分    ├── templates/               # 代码生成模板    └── test/                    # ATen C++ 测试

2.4 c10/ 目录

c10/ 是 PyTorch C++ 基础库,许多文件不依赖 Python,也不依赖高层 ATen。

1
2
3
4
5
6
7
8
9
10
11
c10/├── core/                        # Device、DispatchKey、TensorImpl、Storage、Allocator、ScalarType、SymInt、Layout├── util/                        # Exception、ArrayRef、Optional、SmallVector、intrusive_ptr、typeid、数值工具├── cuda/                        # CUDA stream、event、allocator、device guard 等基础封装├── hip/                         # HIP/ROCm 基础封装├── xpu/                         # XPU 基础封装├── metal/                       # Metal 基础封装├── macros/                      # 导出宏、编译器宏、平台宏├── mobile/                      # 移动端相关基础设施├── test/                        # c10 测试└── benchmark/                   # c10 基准测试

2.5 torchgen/ 目录

torchgen/ 是理解 PyTorch 源码生成机制的关键。

1
2
3
4
5
6
7
8
9
torchgen/├── api/                         # schema 到 C++ API / dispatcher API / native API / structured API 的翻译├── model.py                     # native function、schema、dispatch key、operator name 等数据模型├── gen.py                       # ATen 代码生成主入口之一├── dest/                        # 生成目标文件逻辑├── decompositions/              # 算子分解相关生成├── aoti/                        # AOTInductor 相关生成├── selective_build/             # 选择性构建,用于移动端和裁剪场景└── shape_functions/             # shape function 相关内容

PyTorch 大量 C++ 文件不是手写,而是由 yaml、模板和 torchgen 生成。读源码时需要区分“手写源码”和“生成源码”,否则容易找不到入口。

2.6 tools/ 目录

tools/ 包含开发和生成工具。

1
2
3
4
5
6
7
8
tools/├── autograd/                    # 自动求导相关代码生成,包括 derivatives 配置处理├── codegen/                     # 构建期代码生成工具或相关脚本├── linter/                      # lint 工具├── code_coverage/               # 代码覆盖率工具├── coverage_plugins_package/    # 覆盖率插件包├── stats/                       # 统计工具└── experimental/                # 实验性工具

2.7 test/ 目录

test/ 不只是验证目录,也是理解行为语义的重要入口。

1
2
3
4
5
6
7
8
9
10
test/├── test_*.py                    # 大量 Python API 行为测试├── cpp/                         # C++ API、JIT、distributed、profiler、lazy、custom backend 测试├── distributed/                 # 分布式测试├── inductor/                    # Inductor 测试├── dynamo/                      # Dynamo 测试├── functorch/                   # functorch / 函数变换测试├── cpp_extensions/              # C++/CUDA 扩展测试├── mobile/                      # 移动端测试└── onnx/                        # ONNX 导出相关测试

三、用户 API 层

3.1 层功能说明

用户 API 层是 PyTorch 暴露给模型开发者的入口,主要位于 torch/。这一层负责把底层 Tensor、算子、自动求导、优化器、数据加载、分布式和编译能力组织成自然的 Python 编程体验。它本身不承担所有计算,更多承担 API 组织、对象管理、参数校验和用户语义表达,性能敏感路径会继续下沉到 C++ 桥接层、ATen 或编译栈。

模块
功能
典型类/函数
torch
Tensor 创建、数学算子、设备管理、序列化、随机数
torch.Tensor
torch.addtorch.savetorch.device
torch.nn
神经网络模块、参数、loss、functional API
Module
ParameterLinearConv2dfunctional
torch.optim
参数更新和优化器状态管理
Optimizer
SGDAdamAdamW
torch.autograd
自动求导用户入口和自定义求导
Function
gradno_gradenable_grad
torch.utils
数据加载、C++ 扩展、checkpoint、benchmark
Dataset
DataLoadercpp_extension
torch.distributed
分布式训练和通信
DDP
FSDPProcessGroupRPC
torch.compile
 / torch.export / torch.fx
图捕获、导出、编译和图变换
compile
GraphModuleExportedProgram

3.2 torch 顶层 API

torch 顶层模块是最基础的用户入口,提供 Tensor 构造、顶层算子、dtype/device 管理、随机数、序列化和全局配置。用户调用 torch.addtorch.matmultorch.emptytorch.save 等函数时,Python 层通常只做轻量包装,随后进入 generated binding、ATen dispatcher 或设备相关实现。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
torch/├── __init__.py                  # 顶层命名空间,组织 torch.* API├── _tensor.py                   # Python 侧 Tensor 方法补充├── _C/__init__.pyi.in           # torch._C 类型存根,描述 C++ 扩展暴露的接口├── serialization.py             # torch.save / torch.load├── random.py                    # 随机数相关 Python API├── _prims/                      # primitive 算子└── _refs/                       # 参考实现核心对象:├── torch.Tensor                 # Python 用户看到的 Tensor,底层对应 C++ at::Tensor├── torch.dtype                  # dtype 描述├── torch.device                 # 设备描述└── torch.Generator              # 随机数生成器

顶层 API 的价值是统一用户心智:不管 Tensor 在 CPU、CUDA、MPS、XPU 还是 Meta 上,用户都先面对同一组 torch.* 函数。后端差异由 dispatcher、dispatch key 和 backend kernel 处理。

3.3 torch.nn 模块体系

torch.nn 负责组织神经网络结构,核心类是 ModuleModule 提供参数登记、buffer 登记、子模块管理、hook、state_dict、训练/推理模式切换等能力。Parameter 是特殊 Tensor,被赋值到 Module 属性上时会自动进入参数集合。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
torch/nn/├── modules/module.py            # Module 基类:参数、buffer、子模块、hook、state_dict├── parameter.py                 # Parameter / Buffer├── modules/                     # Linear、Conv2d、Norm、Activation、Loss 等模块├── functional.py                # 无状态函数式接口├── init.py                      # 参数初始化└── parallel/                    # DataParallel / DistributedDataParallel 相关封装核心类/方法:├── Module                       # 所有 nn 模块基类├── Parameter                    # 自动注册为参数的 Tensor 子类├── Buffer                       # 非参数状态张量├── Sequential / ModuleList / ModuleDict├── Module.__setattr__           # 参数、buffer、子模块自动登记入口├── register_parameter()├── register_buffer()├── state_dict()└── load_state_dict()

torch.nn 中的 layer 例如 LinearConv2dLayerNorm 通常负责保存权重和超参数,真正计算会调用 torch.nn.functional 或 ATen 算子。也就是说,nn 是模型结构组织层,不是底层 kernel 实现层。

3.4 torch.optim 优化器体系

torch.optim 以 Optimizer 为基类,负责维护参数组、学习率、优化器状态和 step() 逻辑。优化器不直接理解 Module 的层级结构,而是接收一组参数 Tensor,根据每个参数的 .grad 和内部状态进行原地更新。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
torch/optim/├── optimizer.py                 # Optimizer 基类、参数组、状态字典├── sgd.py                       # SGD├── adam.py / adamw.py           # Adam / AdamW├── rmsprop.py / lbfgs.py        # RMSprop / LBFGS 等优化器├── lr_scheduler.py              # 学习率调度└── _functional.py               # 函数式参数更新逻辑核心对象/方法:├── Optimizer├── Optimizer.param_groups       # 参数组├── Optimizer.state              # 参数状态,如 momentum、exp_avg├── zero_grad()├── step()└── state_dict()

SGD、Adam、AdamW 等优化器的差异主要体现在状态变量和更新公式上。优化器层依赖 autograd 产生梯度,依赖 Tensor 原地操作更新参数,因此它连接了用户训练循环、autograd 和 ATen Tensor 操作。

3.5 torch.autograd 用户入口

torch.autograd 提供自动求导的用户可见接口,包括 backwardgradFunction、grad mode 和 profiler 等。用户只需要设置 requires_grad=True 并执行前向计算,PyTorch 会在底层构建反向图。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
torch/autograd/├── __init__.py                  # backward / grad 等用户入口├── function.py                  # Python 自定义 autograd.Function├── grad_mode.py                 # no_grad / enable_grad / inference_mode├── profiler.py                  # profiler 用户接口└── profiler_legacy.pytorch/csrc/autograd/             # C++ autograd engine 和 Python binding核心对象/函数:├── Function├── backward()├── grad()├── no_grad()├── enable_grad()└── inference_mode()

Python 层的 torch.autograd.Function 用于自定义求导逻辑,用户实现 forward 和 backward,桥接层再把它连接到 C++ autograd engine。no_gradenable_gradinference_mode 这类上下文则通过线程局部状态影响 dispatcher 和 autograd 行为。

3.6 torch.utils 工具模块

torch.utils 提供训练工程中常用的辅助能力。Dataset 和 DataLoader 负责数据集抽象、批处理、多进程加载、pin memory 等;cpp_extension 支持构建 C++/CUDA 扩展;checkpoint、benchmark、pytree 等工具分别服务内存优化、性能测量和结构化参数处理。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
torch/utils/├── data/│   ├── dataset.py               # Dataset / IterableDataset│   └── dataloader.py            # DataLoader,多进程加载、pin memory、collate├── cpp_extension.py             # C++/CUDA 扩展构建├── checkpoint.py                # activation checkpoint├── benchmark/                   # benchmark 工具└── _pytree.py                   # 嵌套结构 flatten / unflatten核心类:├── Dataset├── IterableDataset├── DataLoader├── Sampler└── BatchSampler

这一层的特点是贴近用户工程实践。它不定义 PyTorch 的核心 Tensor 语义,但大量训练程序离不开这些工具。

3.7 torch.distributed 分布式入口

torch.distributed 是分布式训练的用户入口,覆盖通信原语、进程组、DDP、FSDP、RPC、elastic、checkpoint 和 DTensor 等能力。它把底层通信后端包装成用户可用的训练抽象。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
torch/distributed/├── distributed_c10d.py          # ProcessGroup 和集合通信 Python API├── device_mesh.py               # DeviceMesh├── fsdp/│   ├── fully_sharded_data_parallel.py  # FSDP 主类│   └── _flat_param.py           # 参数 flatten / shard├── rpc/                         # RPC├── elastic/                     # 弹性训练├── checkpoint/                  # 分布式 checkpoint└── tensor/                      # DTensor 相关能力torch/csrc/distributed/          # 底层 C++ 通信实现核心对象:├── ProcessGroup├── DistributedDataParallel├── FullyShardedDataParallel├── DeviceMesh├── RPC agent└── checkpoint planner

DDP 主要围绕梯度同步工作,FSDP 主要围绕参数、梯度和优化器状态分片工作。它们都与 Module、autograd hook、optimizer 和通信后端深度交互。

3.8 torch.compiletorch.export 与 torch.fx

torch.compile 提供动态图编译入口,目标是在尽量不改变用户代码的情况下捕获可优化区域。torch.export 更偏稳定导出语义,适合保存和下游消费。torch.fx 提供 Python 级图表示和图变换能力。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
torch/_dynamo/├── eval_frame.py                # torch.compile 入口之一├── convert_frame.py             # frame 转换├── symbolic_convert.py          # 字节码符号执行└── guards.py                    # guard 生成和检查torch/fx/├── graph.py                     # Graph IR├── node.py                      # Node└── graph_module.py              # GraphModuletorch/export/ and torch/_export/ # 导出 API 和内部实现torch/_inductor/├── compile_fx.py                # FX 图编译入口├── graph.py                     # GraphLowering├── ir.py                        # Inductor IR├── scheduler.py                 # 调度和融合└── codegen/                     # 后端代码生成核心对象:├── Graph / Node / GraphModule├── ExportedProgram├── Dynamo guard├── Inductor GraphLowering└── Inductor IR node

这些模块让用户 API 层不只支持 eager execution,也支持图捕获、编译优化和部署导出。底层会进一步进入 Dynamo、FX、AOTAutograd、Inductor 和后端代码生成。

四、C++ 桥接层

4.1 层功能说明

C++ 桥接层主要位于 torch/csrc/,编译后以 torch._C 的形式被 Python 侧导入。它是 Python API 和 C++ 核心之间的边界层,负责对象表示、参数转换、异常转换、GIL、引用计数、模块初始化和 generated binding 接入。

模块
功能
典型对象
torch._C
暴露 C++ 核心能力给 Python
C extension module
torch/csrc
Python/C++ 类型绑定和模块初始化
binding 源码、初始化函数
torch/csrc/autograd
Python Tensor 与 C++ autograd 连接
THPVariable
、autograd binding
torch/csrc/api
C++ frontend
C++ Module、optimizer 等
torch/csrc/distributed
分布式 Python API 的 C++ 支撑
process group binding
torch/csrc/dynamo
 / inductor
编译栈 C++ 支撑
eval frame、runtime 支撑

4.2 torch._C 扩展模块

torch._C 是 Python 侧看到的 C extension module。它不是普通 Python 文件,而是编译产物,集中暴露 dispatcher、Tensor 类型、autograd、JIT、profiler、distributed 和设备管理等 C++ 能力。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
torch/_C/└── __init__.pyi.in              # torch._C Python 类型存根torch/csrc/├── Module.cpp                   # 主扩展模块初始化,大量 C++ 能力注册├── Exceptions.cpp               # C++ 异常到 Python 异常转换├── Device.cpp                   # Device Python binding├── Dtype.cpp                    # dtype Python binding├── Layout.cpp                   # layout Python binding└── Stream.cpp                   # stream Python binding核心对象:├── THPModule 初始化逻辑├── THPVariableClass├── torch._C._VariableFunctions├── Python type object└── pybind 注册对象

很多 Python API 表面上位于 torch/,实际执行会进入 torch._C。例如 Tensor 方法、部分设备 API、Dynamo eval frame hook、C++ profiler 和底层 distributed binding 都依赖这一模块。

4.3 torch/csrc 绑定源码

torch/csrc 是桥接层源码主体。它负责 Python 模块初始化、Python 类型定义、C++ 类型转换、异常转换和生命周期管理。这里的代码需要同时理解 CPython C API、C++ Tensor 语义和 PyTorch 自身的 intrusive pointer 机制。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
torch/csrc/├── python_headers.h             # Python 头文件包含和隔离├── utils/│   ├── python_arg_parser.h      # Python 参数解析│   └── python_arg_parser.cpp├── utils.cpp / utils.h          # Tensor 包装、类型检查、通用转换工具├── tensor/python_tensor.cpp     # Tensor Python 类型相关逻辑├── Storage.cpp                  # Storage Python binding└── StorageMethods.cpp           # Storage 方法绑定核心函数:├── THPVariable_Wrap()├── THPVariable_Unpack()├── PythonArgParser└── THP*initModule()

典型调用路径是:

1
2
3
4
5
Python torch.Tensor method  -> torch._C 暴露的 binding  -> 参数解析和类型转换  -> C++ at::Tensor / dispatcher / autograd  -> C++ 返回值包装为 Python 对象

4.4 torch/csrc/autograd 绑定

torch/csrc/autograd 连接 Python Tensor 和 C++ autograd engine。THPVariable 等结构负责让 Python 侧 Tensor 持有 C++ at::Tensor,同时暴露 .grad.grad_fnrequires_grad 等求导属性。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
torch/csrc/autograd/├── python_variable.h            # THPVariable / Wrap / Unpack├── python_variable.cpp          # Python Tensor 类型行为├── python_variable_indexing.cpp # Tensor 索引├── python_function.cpp          # Python autograd.Function 接入├── python_engine.cpp            # C++ Engine 的 Python binding├── engine.h / engine.cpp        # backward 执行引擎├── node.h                       # 反向图 Node├── function.h                   # autograd function 基础结构├── variable.h                   # Variable = at::Tensor 相关接口└── autograd_meta.cpp            # AutogradMeta核心类:├── THPVariable                  # Python Tensor 壳├── torch::autograd::Engine├── Node / Edge├── AutogradMeta├── SavedVariable└── GraphTask

Python 自定义 autograd.Function 也需要通过这一层接入 C++ autograd graph。用户在 Python 中定义 backward,底层仍然需要以 Node/Edge 的形式进入 autograd engine 调度。

4.5 torch/csrc/api C++ frontend

torch/csrc/api 提供 C++ frontend,使用户可以在 C++ 中使用类似 Python nn.Module、optimizer 和 Tensor API 的模型构建方式。它不是 Python API 的简单翻译,而是面向 C++ 用户的高层封装。

1
2
3
4
5
6
7
8
9
10
11
12
13
torch/csrc/api/├── include/torch/               # C++ frontend 公开头文件│   ├── torch.h                  # 常用总入口│   ├── nn/                      # C++ 版模块体系│   ├── optim/                   # C++ 版优化器│   └── data/                    # C++ 版数据 API└── src/                         # C++ frontend 实现核心类:├── torch::nn::Module├── torch::nn::Linear├── torch::optim::Optimizer└── torch::data::Dataset

C++ frontend 依赖 ATen Tensor 和 dispatcher,但提供更接近深度学习框架的对象模型,例如 module、sequential、optimizer 和 serialization。

4.6 分布式和编译栈 C++ 支撑

torch/csrc/distributed 提供分布式 Python API 背后的 C++ 支撑,包括通信后端绑定、process group 对象和部分 runtime 能力。torch/csrc/dynamo 和 torch/csrc/inductor 则支撑现代编译栈,例如 CPython eval frame hook、编译 runtime 和低开销执行路径。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
torch/csrc/distributed/├── c10d/                        # C++ 通信抽象和后端绑定│   ├── nccl/                    # NCCL 后端│   ├── gloo/                    # Gloo 后端│   ├── ucc/                     # UCC 后端│   └── store/                   # Store / rendezvous 支撑├── rpc/                         # RPC C++ 支撑└── autograd/                    # 分布式 autogradtorch/csrc/dynamo/               # Dynamo C++ 支撑,如 guards、compiled autogradtorch/csrc/inductor/             # Inductor runtime / launcher 支撑核心对象:├── ProcessGroup├── Store├── RPC agent├── distributed autograd context├── Dynamo guard manager└── Inductor launcher

这些模块说明桥接层不只是 Tensor binding,而是所有需要从 Python 高层进入 C++ runtime 的系统边界。