PyTorch 深度实践：从 Tensor 到深度学习运行时

01 · 整体介绍：torch.add(x, y) 经过了哪几层

结论：六个职责层、一条动态调用路径、四层源码；Autograd 只是 Operator Table 上优先级更高的一个 Key。

%%{init: {"flowchart": {"wrappingWidth": 420}}}%%
flowchart TB
    L1["① 编程接口<br/>Tensor · Autograd 接口 · nn.Module<br/>Optimizer · DataLoader"] --> L2["② 图表示与编译<br/>Dynamo · FX · AOTAutograd · Inductor"] --> L3["③ 算子与分发运行时<br/>Schema · Dispatcher · ATen"] --> L4["④ 设备与通信<br/>内存分配 · Stream · 进程组"] --> L5["⑤ 计算内核<br/>CUDA kernel · cuBLAS · cuDNN"]

要点

• Python 只是表达和控制层；C++ 是运行时和抽象；CUDA 是设备执行——「PyTorch 是一个 Python 库」是第一个误区
• 源码四层 torch/ → torch/csrc/ → aten/src/ATen/ → c10/，依赖只能向下

02 · Tensor 与内存布局：view 只改元数据

结论：Tensor = 数据 + 形状 + 布局 + 类型 + 设备 + 生命周期；(\text{offset} = \text{storage_offset} + \sum_k i_k \cdot \text{stride}_k)；view 只改元数据，contiguous / .to() / clone 才复制。

三维 stride 与 expand 的 stride = 0

三个误区

• reshape 总是零拷贝——不连续时先 contiguous() 复制再 view；返回值是否共享内存不确定
• del x 后显存就释放了——小 view 仍持有同一 Storage；释放的块只回缓存分配器，reserved 不降是缓存不是泄漏
• allocated 与 reserved 是一个东西——前者是 Tensor 实际占用，后者是分配器向 CUDA 要来的；差值是碎片与缓存

03 · Autograd：链式法则变成一次沿图的反向遍历

结论：前向就地建图，每个 Node.apply 算一次 VJP（不物化 Jacobian），梯度 += 到叶子的 .grad；保存一个输出等于保存整张图。

%%{init: {"flowchart": {"wrappingWidth": 200}}}%%
flowchart TB
    subgraph fwd["forward：数据流"]
        direction LR
        X["x（leaf）"] --> M[["mul"]] --> Y["y = x·x"] --> A[["add 3"]] --> Z["z"] --> S[["sum"]] --> L["loss"]
    end
    subgraph bwd["Autograd 记下的反向节点（next_functions）"]
        direction LR
        sm(["SumBackward0"]) -.-> ad(["AddBackward0"]) -.-> mu(["MulBackward0<br/>saved: x"]) -.-> ac(["AccumulateGrad<br/>→ x.grad"])
    end
    fwd -. "loss.grad_fn" .-> bwd

要点

• [B, 4096] 的一层若物化 Jacobian 每样本 (4096^2) 个数；VJP 只算 (v^\top J)
• .grad 是累加，所以要 zero_grad()；version counter 抓 in-place；gradcheck 必须 float64

04 · nn.Module 与训练系统：一切状态都在 state_dict 里

结论：__setattr__ 的注册机制决定框架能否发现对象（list 不是 Module，用 ModuleList）；Adam 训练每参数 16 字节——7B 模型 112 GB，一张卡放不下的根源。

从模型对象到训练循环

%%{init: {"flowchart": {"wrappingWidth": 210}}}%%
flowchart LR
    A["nn.Module 树<br/>__init__ 定结构、forward 定数据流"] --> B["注册机制<br/>__setattr__ 登记 Parameter / 子 Module / Buffer"] --> C["state_dict<br/>属性路径 → Tensor"]
    C --> E["Optimizer<br/>拿 parameters() 用 .grad 更新"]
    F["Dataset → Sampler → DataLoader<br/>pin_memory / non_blocking"] --> G["循环：取 batch → forward → loss → backward → step"]
    E --> G
    G --> H["autocast / GradScaler · Hook · checkpoint"]

• model.eval() 只切换 dropout / BN，不关梯度——评估要 eval() + inference_mode() 一起用

05 · Dispatcher：开发态填表、运行态查表

结论：开发态（定义 → 注册 → 实现）与运行态（入口 → 分发 → 执行）交汇于 Operator Table；DispatchKeySet = 各输入 OR + TLS include − exclude，取最高优先级；包装 Key 做完事再次分发——一次 add 两次经过 Dispatcher。

%%{init: {"flowchart": {"wrappingWidth": 200}}}%%
flowchart LR
    subgraph DEV["开发态"]
        direction TB
        A1["定义 Schema<br/>native_functions.yaml / torch.library"] --> A2["注册到各 DispatchKey"] --> A3["编写实现"]
    end
    OT[("Operator Table<br/>算子 → Key → Kernel")]
    subgraph RUN["运行态"]
        direction TB
        B1["入口 torch.add → at::add"] --> B2["分发：查表<br/>KeySet / 包装 Key 再分发"] --> B3["执行：TensorIterator → Kernel"]
    end
    A2 -->|填表| OT -->|查表| B2

要点

• AutocastCUDA → AutogradCUDA → CUDA：autocast 先 cast，autograd 记图，最后 kernel；「Dispatcher 是按 device 的 if/else」是误区

07 · 编译执行：torch.compile 到底做了什么

结论：前端 Dynamo 捕获、中端 AOTAutograd 变换、后端 Inductor 生成，共享 FX Graph；运行时靠 Guard 决定复用——能确定的固化进代码，不能确定的变成运行时检查，捕获不了就 graph break 退回 Eager。

%%{init: {"flowchart": {"wrappingWidth": 170}}}%%
flowchart LR
    subgraph RT["运行时：每次调用"]
        direction LR
        CALL["调用"] --> GD{"Guard 成立?"} -->|是| RUN["执行产物"]
        GD -->|否| CA{"缓存命中?"} -->|否| TRIG["触发编译"]
    end
    subgraph CP["编译器：触发时跑一次"]
        direction LR
        DY["Dynamo 捕获<br/>Graph Break 定边界"] --> FX1[("FX Graph<br/>torch 级")] --> AOT["AOTAutograd<br/>反向展开 · 函数化 · 分解"] --> FX2[("FX × 2<br/>ATen 级")] --> IND["Inductor<br/>融合 · 内存规划"] --> CODE["Triton / C++"]
    end
    TRIG --> DY
    CODE --> RUN

要点

• 热路径：3 次分发 → 0 次、3 次 launch → 2 次；Guard size[0] == 128 → s0 > 64；缓存条目上限 8
• 「换 batch size 变慢是编译器 bug」——静态 shape 下每个新 shape 是一次 Guard 失败 + 重编译；dynamic=True / mark_dynamic

08 · 性能优化：CPU 与 GPU 是两条异步时间线

结论：五类瓶颈——launch / Python / sync-bound（CPU 侧）、memory / compute-bound（GPU 侧）——优化是把瓶颈从一类推到另一类。

%%{init: {"flowchart": {"wrappingWidth": 160}}}%%
flowchart LR
    tl["时间线<br/>torch.profiler / nsys"] --> q1{"GPU 泳道形态？"}
    q1 -->|"稀疏，大段空闲"| q2{"CPU 在忙什么？"}
    q1 -->|"与 CPU 交替空洞"| sync["Sync-bound<br/>批量 .item()、pinned + non_blocking"]
    q1 -->|"密集，首尾相接"| q3{"CUDA 时间被谁占？"}
    q2 -->|"cudaLaunchKernel 占大头"| launch["Launch-bound<br/>增大 batch、融合、CUDA Graphs"]
    q2 -->|"Python / 框架逻辑"| py["Python-bound<br/>向量化、compile"]
    q3 -->|"逐元素与归约"| mem["Memory-bound<br/>融合 / SDPA、bf16"]
    q3 -->|"mm / bmm / conv"| comp["Compute-bound<br/>Tensor Core、减少计算量"]
    classDef cpu fill:#fde9d9,stroke:#c0392b
    classDef gpu fill:#e8f5e9,stroke:#1e8449
    classDef mid fill:#fff4d6,stroke:#b9770e
    class launch,py cpu
    class mem,comp gpu
    class sync mid

案例：五步 166 → 2207 samples/s（13 倍）

• 「低精度总能加速」——launch-bound 下 Kernel 数不变，autocast 还插入 cast Kernel；先解决 CPU 侧问题，再上 bf16

Roofline：A100

• 融合把 5N 访存变 2N；ridge FP32 ≈ 10、BF16 ≈ 156 FLOP/Byte