【AI Infra小白入门⑧】PyTorch框架Infra实战指南：Tensor、autograd与性能调试

本文是「AI Infra 小白入门」系列的第八篇。PyTorch是AI Infra工程师的"主战场"------训练用它，模型从它导出，自定义算子接入它。本文不讲怎么用PyTorch写模型（那是算法工程师的事），而是从Infra视角讲：Tensor、autograd、Module、checkpoint、profiler------这些才是我们天天打交道的东西。

一、写在前面：Infra工程师怎么看PyTorch？

算法工程师用PyTorch，关注的是"怎么搭模型、怎么训出好效果"。

AI Infra工程师用PyTorch，关注的是：

显存：这个模型吃多少显存？哪里可以省？
性能：哪个算子最慢？瓶颈在哪？
导出部署：怎么把模型导出去部署？
自定义算子：怎么把我写的CUDA kernel接进PyTorch？
分布式训练：DDP/FSDP怎么配置？通信瓶颈怎么优化？

一句话：算法工程师关心"效果"，Infra工程师关心"性能与工程"。

先上一张全景图：

二、Tensor：一切的基础

Tensor（张量）是PyTorch最基本的数据结构。对于Infra工程师，有三个属性最重要：

2.1 dtype：精度与显存的起点

dtype	占用字节	特点	典型用途
float32 (FP32)	4	精度高、速度一般	训练（某些场景）、计算精度要求高的
float16 (FP16)	2	精度稍低、速度快（Tensor Core）	混合精度训练、推理
bfloat16 (BF16)	2	动态范围大、精度中等	大模型训练（防止溢出）
int8	1	精度低、速度最快、显存最少	量化推理
int4	0.5	精度更低、显存最少	大模型量化推理（AWQ/GPTQ）

💡 Infra视角：显存与精度的一切权衡，从选择dtype就开始了。FP32变FP16，显存直接减半，速度还可能翻倍------这就是最直接的优化。

2.2 device：数据在哪

cpu：在CPU内存里
cuda：在GPU显存里

python 复制代码

x = x.to('cuda')  # 把数据搬到GPU

性能杀手：频繁的CPU↔GPU数据搬运。

.to()操作是同步的------CPU要等数据搬完才能继续。搬数据走PCIe，带宽比显存低一个数量级还多。频繁搬运会严重拖慢性能。

🔍 优化思路：尽量让数据一直在GPU上，减少搬运。如果必须搬，尽量批量搬。

2.3 视图 vs 拷贝：改形状不一定复制数据

操作	是否共享内存	说明
`view()` / `reshape()`	是（大部分情况）	只是换个"看法"，数据没变
`transpose()` / `permute()`	是（转置是视图）	数据不变，只是 stride 变了
`contiguous()`	否（如果不连续）	让数据在内存中变得连续
`clone()`	否	完完整整拷贝一份

为什么要关心这个？

因为：

视图操作几乎不花时间，也不占额外显存
拷贝操作既花时间又占显存
很多CUDA操作要求数据是连续的（contiguous），如果不连续，PyTorch会偷偷帮你拷贝------这可能就是性能瓶颈

🔍 调试技巧 ：怀疑哪里有不必要的数据拷贝？用x.data_ptr()看内存地址，如果地址变了，说明发生了拷贝。

三、autograd：自动微分的本质

autograd是PyTorch的核心------自动计算梯度。

3.1 基本用法

python 复制代码

x = torch.randn(4, 4, requires_grad=True)
y = (x * x).sum()
y.backward()
print(x.grad)  # dy/dx = 2x

原理很简单：

前向传播时，每个算子都记录自己的"梯度函数"
反向传播时（backward()），沿计算图从后往前，用链式法则把梯度传回去

3.2 Infra视角的关键认知

认知一：计算图保存中间激活 → 训练显存的大头

前向传播时，为了反向传播能算梯度，很多中间结果（激活）需要保存下来。这些激活占用的显存，往往比模型参数还大（尤其是大batch、长序列时）。

这就是为什么会有Activation Checkpointing------不存中间激活，用的时候再算一遍。牺牲计算，换取显存（又是不可能三角！）。

认知二：推理时必须关闭梯度

推理的时候不需要反向传播，所以：

用torch.no_grad()上下文管理器
或者用torch.inference_mode()（更彻底，更快）

python 复制代码

with torch.inference_mode():
    output = model(input)

不加的话，会白白保存计算图和中间激活------浪费显存、浪费时间。

四、Module与训练循环

4.1 Module树与Parameter收集

python 复制代码

model = MyModel().cuda()
opt = torch.optim.AdamW(model.parameters(), lr=1e-4)

for x, y in loader:
    opt.zero_grad(set_to_none=True)  # 清零梯度
    loss = model(x).loss(y)          # 前向
    loss.backward()                  # 反向
    opt.step()                       # 更新参数

几个Infra视角的细节：

model.parameters()：递归收集Module树上所有的Parameter（可训练参数）
zero_grad(set_to_none=True)：把梯度设为None而不是0，直接释放显存，还省了一次写操作
优化器状态：Adam有两个状态量（momentum和variance），和参数一样大------这就是为什么训练显存 ≈ 参数×16（混合精度+Adam）

4.2 模型保存与加载

python 复制代码

# 保存：推荐只存权重（state_dict）
torch.save(model.state_dict(), "model.pt")

# 加载：用map_location指定设备，避免先加载到CPU再搬到GPU
model.load_state_dict(torch.load("model.pt", map_location="cpu"))

为什么推荐存state_dict而不是整个model？

state_dict就是个字典：{name: tensor}，简洁清晰
整个model序列化容易出问题（类定义变了就加载不了）
跨设备加载一定要用map_location，不然会先加载到默认设备

分布式训练的checkpoint：

FSDP/ZeRO下，参数是分片的，保存时需要合并（full_state_dict）
或者保存分片格式（sharded），加载时再合并

五、Infra工程师的"听诊器"：调试工具

5.1 显存调试

OOM（Out Of Memory，显存不够）是AI Infra最常见的问题之一。怎么定位？

python 复制代码

# 打印详细的显存统计信息
print(torch.cuda.memory_summary())

# 获取峰值显存
peak_memory = torch.cuda.max_memory_allocated()
print(f"峰值显存: {peak_memory / 1024**3:.2f} GB")

# 重置峰值统计
torch.cuda.reset_peak_memory_stats()

用途：

定位OOM：看看显存都花在哪了
验证优化效果：量化/内存池/梯度检查点到底省了多少显存

5.2 Profiler：找出谁最慢

性能优化的第一步是什么？找瓶颈！

PyTorch内置的profiler就是干这个的：

python 复制代码

with torch.profiler.profile(
    activities=[
        torch.profiler.ProfilerActivity.CPU,
        torch.profiler.ProfilerActivity.CUDA
    ]
) as prof:
    model(x)

# 按CUDA时间排序，显示前10个最耗时的算子
print(prof.key_averages().table(sort_by="cuda_time_total", row_limit=10))

输出会告诉你：

每个算子花了多少时间（CPU时间、CUDA时间）
调用了多少次
占用了多少显存

🎯 性能优化铁律：先profile，再优化。不知道瓶颈在哪就瞎优化，90%都是无用功。

六、其他Infra常用特性

6.1 Activation Checkpointing

前面提到过，用计算换显存：

python 复制代码

from torch.utils.checkpoint import checkpoint

# 把某个模块用checkpoint包起来
output = checkpoint(model.block, x)

效果：这个模块的中间激活不存了，反向传播时重新算一遍。

牺牲：多花一次前向计算的时间
换取：节省大量激活显存

（又是不可能三角的应用！）

6.2 torch.compile

PyTorch 2.0引入的新特性，自动优化模型计算图：

python 复制代码

model = torch.compile(model)

原理：把PyTorch的动态图抓成静态图，然后做图优化（算子融合等）。

优点：不用改代码就能提速（通常10%-50%）
缺点：第一次运行有编译开销，有些算子不支持

💡 如果你做推理优化，torch.compile是一个"低投入、中产出"的优化手段------值得一试。

七、小白学习建议

不要只学"怎么用"，要想"为什么"：比如遇到OOM，不要只会"调小batch size"，要想想：为什么OOM了？显存都花在哪了？
多用工具 ：memory_summary和profiler是你的好朋友，遇到问题先让数据说话
读源码 ：从简单的开始读，比如torch.nn.Linear的实现，慢慢培养感觉

八、总结一下

今天这篇文章，我们从Infra工程师的视角重新认识了PyTorch：

Tensor三要素：dtype（精度与显存的起点）、device（数据搬运是性能杀手）、视图vs拷贝（避免不必要的复制）
autograd ：计算图保存中间激活 → 训练显存大头；推理时必须用inference_mode()
Module与训练循环 ：zero_grad(set_to_none=True)省显存、state_dict是模型保存的标准方式
调试工具 ：memory_summary查显存、profiler找瓶颈------Infra工程师的听诊器

下一篇预告：我们来深入GPU内部------SM是什么？Tensor Core又是什么？存储层次有几级？Memory Wall到底是什么？搞懂硬件，才能做好优化。

📚 本系列导航：

第1篇：一文搞懂什么是AI基础设施：五大方向与不可能三角
第2篇：大厂AI Infra组织架构大揭秘
第3篇：AI Infra岗位路线图：级别、薪资与成长路径
第4篇：AI Infra学习路线图：从入门到精通的完整指南
第5篇：AI Infra必备编程基础：Python、C++与Linux速通指南
第6篇：AI Infra数学基础速成：线代、概率与微积分（工程向）
第7篇：Transformer架构工程拆解：从Attention到MLA的优化密码
第8篇：PyTorch框架Infra实战指南：Tensor、autograd与性能调试（本文）
第9篇：GPU硬件概论：从SM到Memory Wall，搞懂优化的物理基础
更多内容持续更新中......