阿甘的博客 | Arganzheng's Blog

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

PyTorch 深度实践：从 Tensor 到深度学习运行时（总纲）

Deep Dive into PyTorch, from Tensor to Deep Learning Runtime

内容简介《PyTorch 深度实践：从 Tensor 到深度学习运行时》是一组共十篇的系列文章，面向有后端工程经验、尤其是 Java 背景、准备进入 AI-Infra 方向的工程师，系统梳理 PyTorch 从 Python API 到 C++、CUDA、编译器、性能分析和分布式运行时的主要机制。它同时是《AI 算法工程师学习地图》 L1 工具箱的深入篇：算法侧的《算法工程师的工具箱》讲 PyTorch 的”用”——五个对象、二十行训练循环、显存的账——本系列讲”改”：Dispatcher、Autograd 引擎、编译、分布式。两张地图共享本系列，算法方向的读者按需读。这个系列不是 PyTorch API 速查表，也不是机器学习算法教程，而是试图回答一个...

Posted by arganzheng | February 16, 2026

C++ 在 AI-Infra（09）：系列总结与通关自测

C++ for AI-Infra: Series Recap and Final Self-Test

八篇正文回答了一个问题：PyTorch 和 vLLM 的 C++ 源码里，这段代码为什么这样写。第一篇讲一个 .cpp 怎么变成 .so、import torch 拉起哪几个库；第二篇讲 at::Tensor 为什么是一个 8 字节的句柄、数据什么时候释放；第三篇讲 AT_DISPATCH 里的 scalar_t 从哪里来；第四篇讲 Dispatcher 用什么机制调到 kernel；第五篇讲算子怎么在没有任何显式调用的情况下登记进 torch.ops；第六篇讲 no_grad 为什么对其他线程不生效；第七篇讲一个 Tensor 跨过 Python/C++ 边界经过了什么；第八篇讲一个改动从写完到能提 PR 要跑哪些东西。八篇合起来，是总纲开篇那段 scale_sh...

Posted by arganzheng | February 15, 2026

Arganzheng's Blog