企业官网建设流程全解析

【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：精度与显存的起点

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

2.2 device：数据在哪

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

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

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

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

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

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

为什么要关心这个？

因为：

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

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

三、autograd：自动微分的本质

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

3.1 基本用法

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

原理很简单：

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

3.2 Infra视角的关键认知

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

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

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

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

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

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

withtorch.inference_mode():output=model(input)

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

四、Module与训练循环

4.1 Module树与Parameter收集

model=MyModel().cuda()opt=torch.optim.AdamW(model.parameters(),lr=1e-4)forx,yinloader: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 模型保存与加载

# 保存：推荐只存权重（state_dict）torch.save(model.state_dict(),"model.pt")# 加载：用map_location指定设备，避免先加载到CPU再搬到GPUmodel.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最常见的问题之一。怎么定位？

# 打印详细的显存统计信息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就是干这个的：

withtorch.profiler.profile(activities=[torch.profiler.ProfilerActivity.CPU,torch.profiler.ProfilerActivity.CUDA])asprof: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

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

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

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

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

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

6.2 torch.compile

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

model=torch.compile(model)

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

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

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

七、小白学习建议

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

八、总结一下

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

1. Tensor三要素：dtype（精度与显存的起点）、device（数据搬运是性能杀手）、视图vs拷贝（避免不必要的复制）
2. autograd：计算图保存中间激活 → 训练显存大头；推理时必须用inference_mode()
3. Module与训练循环：zero_grad(set_to_none=True)省显存、state_dict是模型保存的标准方式
4. 调试工具：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，搞懂优化的物理基础
• 更多内容持续更新中……