阿甘的博客 | Arganzheng's Blog

PyTorch 深度实践（10）：PyTorch 的工程体系——一次改动如何安全地到达用户

The Engineering System of PyTorch: How a Change Travels Safely from Commit to Production

前九篇讲的是 PyTorch 是什么、怎么运行：Tensor 怎么存、Autograd 怎么记、算子怎么分发、Kernel 怎么写、编译器怎么融合、性能怎么测、多卡怎么通信。每一篇都在描述一个已经存在、并且正确运行的系统。这一篇换一个问题：它是怎么做到一直正确、一直可用的？1 PyTorch 有两千多个算子、每个算子有十几种 dtype、两个以上后端、无数种 shape 和 stride 组合，还要在 eager、torch.compile、Meta、Autograd、分布式等模式下行为一致。每天有几十个 PR 合入主干，每个 PR 都可能碰到其中任何一层；每三四个月发一个版本，几百万用户的代码、几十个硬件后端、无数已保存的 checkpoint 都要在新版本...

Posted by arganzheng | February 26, 2026

PyTorch 深度实践（09）：分布式 PyTorch

Distributed Training in PyTorch: Collectives, DDP, FSDP, TP, PP, CP and EP

前八篇都在一张卡上。第八篇末尾算过一笔账：Adam 训练下每个参数的静态显存是 16 字节，7B 参数的模型仅参数、梯度和优化器状态就要 112 GB，激活值还没算。一张 80 GB 的卡放不下。即使放得下，第八篇案例里那个 38M 参数的小模型在单卡上跑到 2207 samples/s 之后，GPU 已经饱和——再要快，只能加卡。这一篇回答加卡之后的问题：当一张卡放不下模型或跑不完数据时，PyTorch 如何把计算和状态切分到多个设备，并让通信与计算重叠？1 分布式训练的资料通常按 API 组织：DDP 一章、FSDP 一章、张量并行一章、流水线并行一章。这样读完会记住一堆包装类，却答不出”FSDP 比 DDP 多的那份通信从哪来、什么条件下才多...

Posted by arganzheng | February 25, 2026

PyTorch 深度实践（08）：性能优化与调试

Performance Optimization and Debugging in PyTorch

前面七篇建立了 PyTorch 的执行模型：Tensor 如何存储（第二篇），Autograd 如何记录反向（第三篇），算子如何分发到 Kernel（第五篇），编译器如何把多个算子融合成更少的 Kernel（第七篇）。每一篇都在某处留下一句”性能问题第八篇讨论”。这一篇兑现这些承诺。它围绕一个问题展开：如何判断一个 PyTorch 程序慢，以及如何定位它为什么慢？1 这个问题比”如何让它变快”更基础。AI-Infra 工作中大量的性能优化失败，不是因为不知道优化手段，而是因为没有正确地测量，或者对着错误的瓶颈用力：GPU 利用率只有 30% 时去优化 Kernel，显存碎片导致 OOM 时去减 batch，Python 开销占主导时去换更快的 cu...

Posted by arganzheng | February 24, 2026

PyTorch 深度实践（07）：编译执行与图优化

Compilation and Graph Optimization in PyTorch

前两篇讨论的是单个算子：第五篇解释一次 torch.add 调用如何经过入口、分发、执行；第六篇把一个自定义算子接入了同样的路径。无论原生还是自定义，每个算子都是独立走完这条路的。这一篇把视角从单个算子拉远到一段程序。当我们写下： def f(x, weight, bias): y = x @ weight + bias if x.shape[0] > 64: return torch.relu(y) return torch.tanh(y) compiled_f = torch.compile(f) torch.compile 到底做了什么？1 它不是把 Python 翻译成 CUDA——这个说法既不准确，...

Posted by arganzheng | February 23, 2026

PyTorch 深度实践（06）：C++ 扩展与自定义算子

C++ Extensions and Custom Operators in PyTorch

上一篇把算子系统拆成两个维度：开发者在构建时定义 → 注册 → 实现，用户在运行时入口 → 分发 → 执行，两者通过 Operator Table 交汇。那一篇站在使用者的角度观察原生算子 add。这一篇换到开发者的位置：自己写一个算子，把它接入 PyTorch 的算子系统1 。这是从“阅读框架”走向“扩展框架”的关键一步。AI-Infra 工作中大量的实际需求都落在这里：一个融合 Kernel、一个新硬件的后端适配、一个推理引擎的定制算子，最终都要经过同样的路径。本文用一个刻意简单的算子贯穿全文：scale_shift(x, alpha, beta) = alpha * x + beta。它简单到不会分散注意力，又足够涉及 Tensor 元数据、dty...

Posted by arganzheng | February 22, 2026

PyTorch 深度实践（05）：Dispatcher 与算子系统

The Dispatcher and Operator System in PyTorch

上一篇讨论了 nn.Module 与训练系统：Module 如何组织子模块、Parameter 和 Buffer，Optimizer 如何更新参数，DataLoader 如何把数据送入训练循环。但当我们写下： z = x + y 到底是谁决定了这个操作对应哪个算子、x 和 y 位于哪个设备、当前是否需要 Autograd、最终应该调用哪个 Kernel？1 这些问题属于 PyTorch 的算子系统（Operator System）。本文不会一开始分析复杂的 Attention，而是选择最简单的 add 算子，从两个视角把算子系统拆开：算子开发者在构建时做了什么，以及算子使用者在调用时发生了什么。 Dispatcher 解决的不是“调用哪个 Pyt...

Posted by arganzheng | February 21, 2026

PyTorch 深度实践（04）：nn.Module 与训练系统

nn.Module and Training Systems in PyTorch

上一篇讨论了 Autograd：Tensor 运算如何形成动态计算图，backward() 如何沿图传播梯度，以及梯度状态和计算图生命周期之间有什么关系。但一个真实的模型不会只是几个散落的 Tensor 运算。它通常包含：多层嵌套的网络结构；需要训练的参数；不参与梯度更新但属于模型状态的 Buffer；优化器和优化器状态；训练、验证和推理三种不同阶段；数据集、采样器和 DataLoader； CPU 到 GPU 的数据搬运； checkpoint、恢复训练和混合精度。 PyTorch 用 nn.Module 把这些对象组织起来，再通过 Optimizer、Dataset 和 DataLoader 将它们连接成训...

Posted by arganzheng | February 20, 2026

PyTorch 深度实践（03）：自动求导与动态计算图

Autograd and Dynamic Computation Graphs in PyTorch

上一篇介绍了 Tensor 的核心模型：它不是一组孤立的数字，而是由 Storage、Shape、Stride、Storage Offset、dtype、device 和 layout 共同描述的一种数据抽象。但 Tensor 只有数据和布局，还不能完成模型训练。训练还需要回答一个问题：模型输出发生变化时，参数应该沿着什么方向、以多大的幅度变化？1 这需要计算梯度。PyTorch 通过 Autograd 把数学上的求导过程变成了一个可以执行的运行时系统： Tensor 运算动态记录运算关系形成计算图从结果反向遍历计算并累积梯度 Optimizer 更新参数初学者通常只需要记住： loss.backward...

Posted by arganzheng | February 19, 2026

PyTorch 深度实践（02）：Tensor 与内存布局

Tensor Abstraction and Memory Layout in PyTorch

上一篇从整体上介绍了 PyTorch：它不是只有 Python API 的库，而是连接模型代码、Tensor 编程模型、算子运行时、设备后端、Kernel 和硬件的一套计算平台。这一次进入这张地图的核心数据抽象： Tensor 到底是什么？1 初学 PyTorch 时，Tensor 很容易被理解成“支持 GPU 的 NumPy 数组”。这个理解可以帮助开始使用，但不足以解释真实工程中的许多现象：为什么 transpose() 通常不会复制数据？2 为什么 view() 有时成功，有时会报错？3 为什么 reshape() 有时是零拷贝，有时会产生一份新数据？4 为什么一个 Tensor 的 shape 相同，性能却可能完全不同？...

Posted by arganzheng | February 18, 2026

PyTorch 深度实践（01）：PyTorch 整体介绍

PyTorch Overall Introduction

PyTorch 经常被介绍成一个“深度学习框架”，也经常被使用成一个 Python 库：导入 torch，创建 Tensor，定义 nn.Module，然后训练模型。这种理解对于开始使用 PyTorch 已经足够，但对于训练平台、推理引擎、算子开发和 AI-Infra 来说还不够。真正需要理解的是： PyTorch 如何把 Python 中表达的张量计算，转化为可以自动求导、跨设备执行、编译优化和分布式协作的运行时系统？1 一行看起来很普通的代码： z = torch.add(x, y) 背后可能涉及： Python API Python Binding Operator Schema Dispatcher ATen O...

Posted by arganzheng | February 17, 2026

Arganzheng's Blog