Torch 内部机制（第一部分）—— FX 图

这是一个新的博客系列，我将尝试深入解析 PyTorch 和 PyTorch Compile 的生态系统。我一直想找个理由深入研究，弄清楚 PyTorch 为什么这么做、怎么做，以及它做得有多出色。希望到这个系列结束时，你能更深入地理解 PyTorch 2.0 生态系统的运作方式，并能更从容地调试程序。

到底什么是 FX 图？

什么是 FX 图？我为什么要关心它？我想读 torch.compile，跳过这个无聊的图。坏消息是：理解 FX 图是理解编译生态系统中所有其他组件的基础。从 Dynamo 到 Inductor，每个组件都以某种方式使用了 FX 图。好消息是：我将在一篇博客里全部分析清楚！

FX 图是 PyTorch 模块提供的一种中间表示（IR），它将 Python 代码转换成一个有向无环图（DAG）。

经典的 DAG

DAG 代表有向无环图，基本由节点和有向边组成。这是一个酷炫的 DAG——

三个核心对象

FX 图有三个你需要理解的概念——Graph、Node 和 GraphModule。我来展示一下我的意思。下面是一个用 torch.fx.symbolic_trace 追踪的简单模型：

import torch
import torch.fx

class MyModel(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.linear = torch.nn.Linear(4, 8)

    def forward(self, x):
        return torch.relu(self.linear(x))

gm = torch.fx.symbolic_trace(MyModel())
gm.graph.print_tabular()

打印结果：

opcode         name    target           args        kwargs
-------------  ------  ---------------  ----------  --------
placeholder    x       x                ()          {}
call_module    linear  linear           (x,)        {}
call_function  relu    torch.relu       (linear,)   {}
output         output  output           (relu,)     {}

FX 图将 Python 代码展平为一系列节点。乍看可能觉得信息很多，但别担心，我们会逐一解析每个字段。

Graph

torch.fx.Graph 是一个计算图，内部保存着一系列节点。你可以打印它来查看完整结构：

graph():
    %x : [num_users=1] = placeholder[target=x]
    %linear : [num_users=1] = call_module[target=linear](args = (%x,), kwargs = {})
    %relu : [num_users=1] = call_function[target=torch.relu](args = (%linear,), kwargs = {})
    return relu

这是 FX 图的文本表示形式。下一节会有更多介绍。

Node

节点是表示输入、函数调用或输出的基本单元。每个节点主要有 6 个你会用到的字段。

node.op

node.op 告诉你节点的大类。有且仅有六个可能的值：

• placeholder - 图输入
• get_attr - 例如读取参数
• call_function - 调用 torch 运算或 Python 函数
• call_method - 对张量调用方法（如 .view()）
• call_module - 调用子模块（如 self.linear）
• output - 图输出

node.name

node.name 是打印图时看到的以 % 为前缀的临时变量名。例如：

%add = call_function[target=operator.add](args = (%x, %y), kwargs = {})
# node.name == "add"

node.args

node.args 是一个元组，包含该节点依赖的节点。这就是你追踪数据流的方式：

%add = call_function[target=operator.add](args = (%x, %y), kwargs = {})   # add 依赖于 x 和 y
%relu = call_function[target=torch.relu](args = (%add,), kwargs = {})     # relu 依赖于 add

node.kwargs

与 args 类似，但用于关键字参数。

node.users

node.users 是该节点被其他节点使用的次数。

# 输入 x 被 2 个节点使用
%x : [num_users=2] = placeholder[target=x]
# ...
%add = call_function[target=operator.add](args = (%x, %y))
%neg = call_function[target=torch.neg](args = (%x,))

# node.users[x] 会映射到 {add: 1, neg: 1}

这对图变换至关重要。删除节点前，必须检查 len(node.users) == 0，或者先将它的用户重定向到另一个节点。这也是检测死代码的方式。

node.target

这是最重要的字段——也是最容易混淆的，因为它的含义取决于 node.op。你不能孤立地读取 target。比如 target="linear"。如果 op 是 call_module，意思是“调用 self.linear”。如果 op 是 get_attr，意思是“读取 self.linear”（这会出错，因为 linear 是模块，不是张量）。所以一定要先检查 op 再检查 target。

node.meta

node.meta 保存节点的元数据信息。

FX 图是如何构建的

到目前为止，我们看到 torch.fx 神奇地将 Python 代码变成了这个图，但它是怎么做到的？它通过一种叫符号追踪的概念来实现。

符号追踪

首先，为什么我们需要符号追踪？当普通的 Python 代码运行时，它立即执行代码，执行完后你得到了结果，但不知道你是如何得到这个结果的——基本上就是计算图。因此我们转向符号追踪。

通过符号追踪，torch.fx 拦截这个调用，说“嘿，我看到了一个 relu 操作”，然后创建一个 Node。没有实际计算发生。它只模拟计算。这种拦截是通过 Proxy 对象完成的。

Proxy 包装每个输入值，并重写所有可能在前向传播中使用的 Python 特殊方法：

• __add__、__mul__、__sub__ 用于算术运算
• __getattr__ 用于属性访问，如 .shape、.view()
• __getitem__ 用于索引，如 x[0]
• __torch_function__ 用于 torch 运算，如 torch.relu(x)
• __call__ 用于子模块调用，如 self.linear(x)

每当这些钩子触发时，Proxy 就会创建一个新的 Node 并记录操作。Node 存储操作码、目标函数和输入（args/kwargs）。当 forward 返回时，我们就得到了一个完整的图。

import torch
import torch.fx

class MyModel(torch.nn.Module):
    def forward(self, x):
        z = x + 1  # Proxy.__add__ 触发 → 创建 Node(op=call_function, target=operator.add)
        return torch.relu(z)  # __torch_function__ 触发 → 创建 Node(op=call_function, target=torch.relu)

这太棒了！但这种 Proxy 方法也有一些局限性。

符号追踪的局限性

1. 动态控制流——如果条件依赖于运行时变化（基于变量、数据或用户输入）的程序执行路径，则无法追踪，因为 Proxy 没有实际值来评估条件：

def forward(self, x):
    if x.sum() > 0:  # proxy 无法评估这个！
        return torch.relu(x)
    else:
        return torch.neg(x)

2. 非 torch 的 Python 函数——普通的 Python 函数不会被拦截：

from math import sqrt
def normalize(x):
    return x / sqrt(len(x))

traced = torch.fx.symbolic_trace(normalize)  # 失败！

3. 静态控制流没有问题——如果条件在初始化时已知，追踪可以完美工作：

class MyModule(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.do_activation = False
        self.linear = torch.nn.Linear(512, 512)

    def forward(self, x):
        x = self.linear(x)
        if self.do_activation:  # 静态的——不依赖数据
            x = torch.relu(x)
        return x
# 这个追踪正确

总结

理解 FX 图是理解整个 PyTorch 2.0 编译栈的第一步。所以，恭喜你读到了这里。在下一篇文章中，我们将讨论 TorchDynamo，一种字节码级别的 JIT，以及它如何处理这些不同之处。