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

十篇怎么连起来

%%{init: {"flowchart": {"wrappingWidth": 170}}}%%
flowchart LR
    P1["01 整体介绍"] --> T2["02 Tensor<br/>与内存"] --> A3["03 Autograd<br/>动态图"] --> M4["04 nn.Module"] --> NX["05–10 系统与工程<br/>（下图）"]

05–10：系统与工程

%%{init: {"flowchart": {"wrappingWidth": 170}}}%%
flowchart TB
    D5["05 Dispatcher"] --> X6["06 C++ 扩展"] & C7["07 编译与图优化"]
    X6 & C7 --> P8["08 性能与调试"] --> D9["09 分布式"] --> E10["10 工程体系"]

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 执行

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

view 只改元数据，contiguous / .to() / clone 才复制

三维 stride 与 expand 的 stride = 0

三个误区

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

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

保存一个输出等于保存整张图

%%{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

要点

VJP 只算 (v^\top J)——不物化 Jacobian

自定义 Function：apply 立即执行，引擎回调 backward

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

Adam 每参数 16 字节——7B = 112 GB 放不下的根源

从模型对象到训练循环

%%{init: {"flowchart": {"wrappingWidth": 210}}}%%
flowchart LR
    A["nn.Module 树"] --> B["__setattr__<br/>注册 P/M/B"] --> C["state_dict<br/>路径 → Tensor"] --> E["Optimizer<br/>.grad 更新"]
    F["Dataset<br/>→ DataLoader"] --> G["循环 fwd→loss<br/>→bwd→step"]
    E --> G
    G --> H["autocast · Hook<br/>· ckpt"]

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

一次 add 两次经过 Dispatcher：包装 Key 做完事再次分发

开发态填表、运行态查表：交汇于 Operator Table

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

要点

多维可多次分发——不是按 device 的 if/else

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

Dynamo 捕获 → AOTAutograd 变换 → Inductor 生成，共享 FX Graph；运行时 Guard 定复用，捕获不了就 graph break 退 Eager。

%%{init: {"flowchart": {"wrappingWidth": 170}}}%%
flowchart LR
    CALL["调用"] --> GD{"Guard 成立?"} -->|是| RUN["执行产物"]
    GD -->|否| CA{"缓存命中?"} -->|是| RUN
    CA -->|否| TRIG["触发编译"] --> DY["编译器链<br/>（下图）"]
    DY --> RUN

编译器：触发时跑一次

Dynamo 定边界 → AOTAutograd 展开反向 → Inductor 融合出 Triton/C++，共享 FX Graph

%%{init: {"flowchart": {"wrappingWidth": 150}}}%%
flowchart LR
    DY["Dynamo 捕获<br/>Graph Break 定边界"] --> FX1[("FX Graph<br/>torch 级")]
    FX1 --> AOT["AOTAutograd<br/>反向展开 · 分解"] --> FX2[("FX × 2<br/>ATen 级")]
    FX2 --> IND["Inductor<br/>融合 · 内存规划"] --> CODE["Triton / C++"]

要点

热路径 3 分发 → 0、3 launch → 2

常见误区

• 「换 batch size 变慢是编译器 bug」——静态 shape 下每个新 shape 是一次 Guard 失败 + 重编译；dynamic=True / mark_dynamic

Inductor 的内存规划：三个逻辑值共用一次分配

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

CPU 提交与 GPU 执行是两条异步时间线；优化是把瓶颈从一类推到另一类

按 GPU 泳道形态定瓶颈类别

%%{init: {"flowchart": {"wrappingWidth": 160}}}%%
flowchart LR
    tl["profiler 时间线"] --> q1{"GPU 泳道？"}
    q1 -->|"稀疏，大段空闲"| q2{"CPU 在忙什么？"}
    q1 -->|"与 CPU 交替空洞"| sync["Sync-bound<br/>.item() 过多"]
    q1 -->|"密集"| q3{"时间被谁占？"}
    q2 -->|"launch 占大头"| launch["Launch-bound<br/>增大 batch"]
    q2 -->|"框架逻辑"| py["Python-bound<br/>向量化"]
    q3 -->|"逐元素"| mem["Memory-bound<br/>融合 / bf16"]
    q3 -->|"mm / 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

09 · 分布式：五类状态各做一个决定

五类状态各选复制还是分片——ring all_reduce 每 rank 收发 → 2n

常见误区

• 「FSDP 比 DDP 省显存也省通信」——通信 3P 对 2P 多 50%；省的是 16P → 16P/N
• 「加卡就该线性加速」——带宽项通信量不随 N 减少，per-rank 计算随 N 缩小；藏不住就换策略