2026 年 1 月 20 日，2025 冬季 InfiniTensor 大模型与人工智能系统训练营 AI 编译器系列第三课《PyTorch 中图优化》开讲。

本节课主要讲解 PyTorch 核心设计理念与图优化技术体系演进。

PyTorch 基础介绍

动态计算图的优势与劣势

PyTorch 中的图优化

• 前端：在保持 Python 动态语义的前提下，对张量相关计算进行捕获与抽象，生成可优化的中间表示（IR）
• 后端：基于中间表示进行图级优化、调度与代码生成

1. TorchScript

• 介绍：通过torch.jit实现
• 两个接口：
  • torch.jit.trace
  • torch.jit.script

2. torch.fx

• 介绍：Python 层 IR 与图变换工具
• 技术亮点：

  • Symbolic Tracing：通过 Proxy 对象代替输入张量，记录所有张量运算的结构与路径。

  • fx IR：高度简洁，6 种核心指令，placeholder、get_attr、call_function、call_module、call_method、output。

• Python 源到源转换：
  • IR 反向映射回 Python 源码
  • 基于 FX Graph 修改模型结构，导出新的nn.Module代码

3. PyTorch 2.0 时代：torch.compile

• 介绍：统一编译接口，在尽量保留 Python 动态语义的前提下，通过图捕获与优化提升执行效率
• 组件与步骤：
  • TorchDynamo（前端）：解释器级 tracing 框架，在 Python 字节码层面捕获张量相关计算，生成受 guard 约束的 FX Graph
    
    • 工作原理：拦截 Python 函数执行，记录张量相关 IR
    • Guards 机制：监控变量类型与值，确保重用编译图
    • Graph Break：子图划分，处理无法转换为静态图的部分

• AOT Autograd：提前生成前向与反向联合计算图
  • 工作流程：Dispatcher 机制拦截与重定义 Tensor 操作，切分正反向计算图

• TorchInductor（后端）：后端优化，生成高效可执行代码
  • 工作流程：接受 FX Graph，进行算子融合、内存规划与调度优化，并生成高效的执行代码
  • 示例：CUDA 设备上的 Triton 后端代码生成过程

总结

本节课讲了 PyTorch 图优化技术体系的演进历程与核心特性。后续课程将讲解 AI 编译器中的后端优化。

💡 学习贴士

更多细节及完整课程内容，请观看直播或查看课程回放。

链接：https://www.infinitensor.com/camp/winter2025

📚 训练营报名持续进行中

扫码报名与加入群聊，或点击 「阅读原文」 直达官网报名