🎓 A1Learner
🔍

PyTorch An FX Graph Transform

A graph transform is a rewrite of a traced graph that keeps the module's answers the same. The closing transform fuses a Linear layer and its ReLU into one node.

Find the pair to fuse

The graph lists its nodes by name and kind. The fc1 and relu nodes sit side by side.

import torch


class TinyNet(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = torch.nn.Linear(4, 8)
        self.fc2 = torch.nn.Linear(8, 2)

    def forward(self, x):
        x = self.fc1(x)
        x = torch.relu(x)
        return self.fc2(x)


torch.manual_seed(0)
net = TinyNet()
traced = torch.fx.symbolic_trace(net)
graph = traced.graph

for node in graph.nodes:
    print(node.name, node.op)

Both nodes are present

Name checks confirm the pair exists before any edit. A missing name would fail loudly here.

names = [node.name for node in graph.nodes]
print("Linear node found:", "fc1" in names)
print("Relu node found:", "relu" in names)

Who consumes the relu output

A node's users are the nodes that read its value. Only fc2 reads the relu result, so only fc2 needs rewiring.

print("Users of the relu node:", [u.name for u in list(graph.nodes)[2].users])

Example

import torch


class TinyNet(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = torch.nn.Linear(4, 8)
        self.fc2 = torch.nn.Linear(8, 2)

    def forward(self, x):
        x = self.fc1(x)
        x = torch.relu(x)
        return self.fc2(x)


def make_fused(linear):
    def fused(x):
        return torch.relu(linear(x))

    return fused


torch.manual_seed(0)
net = TinyNet()
traced = torch.fx.symbolic_trace(net)
graph = traced.graph
print("Nodes before:", len(list(graph.nodes)))

nodes = {node.name: node for node in graph.nodes}
relu_node = nodes["relu"]

with graph.inserting_after(relu_node):
    fused_node = graph.call_function(make_fused(net.fc1), args=(nodes["x"],))
relu_node.replace_all_uses_with(fused_node)
graph.eliminate_dead_code()
graph.lint()

fused_model = torch.fx.GraphModule(traced, graph)
print("Nodes after:", len(list(graph.nodes)))

x = torch.randn(3, 4)
with torch.no_grad():
    plain = net(x)
    fused_out = fused_model(x)
print("Outputs equal:", torch.equal(plain, fused_out))

Key points

  • Use replace_all_uses_with to point every consumer at a new node
  • eliminate_dead_code() removes nodes that nothing uses any more
  • Use graph.lint() to check a rewritten graph is well formed
  • A GraphModule built on the rewritten graph runs the new steps
  • The fused node runs the Linear layer and the ReLU in one call
  • Node counts before and after prove the rewrite changed the graph
  • In short, a transform edits the step list and then proves the answers survived
  • The fused callable here is a plain Python function held by the node

In this example, the traced graph of the same two-layer module is rewritten. A fused callable is built from the fc1 layer, and it runs the layer and then ReLU in one call. A new call_function node is inserted with the input as its argument. Every use of the old relu node is pointed at the fused node. Dead code elimination then removes the relu and fc1 nodes, and lint checks the result. The count falls from five nodes to four. The rebuilt module runs and matches the original outputs exactly.

Note

Serious fusion work targets compiled backends, and this lesson stays with the graph edit itself.

Try it Yourself

Edit the code, then press Run.

Running is simulated · Python Playground

Exercise

What proves a graph transform is safe to keep?