Optimize PyG with torch.compile

Why compilation matters for graph learning

PyTorch 2.0 introduced torch.compile, a JIT compiler that fuses GPU kernels, eliminates Python overhead, and can speed up standard deep learning workloads. Graph neural networks (GNNs) present unique challenges that make compilation less straightforward than calling torch.compile(model).

The core difficulty is irregular data. Every GNN mini-batch can have a different number of nodes, edges, and neighbors. Neighbor sampling produces subgraphs of varying shapes across iterations. That variability triggers re-compilations, graph breaks, and host-device synchronization that can erase the performance gains from compilation.

Kumo’s production stack for Relational Deep Learning (RDL) constructs a large heterogeneous graph from relational database tables. Each table becomes a node type, each row becomes a node, and primary key-foreign key pairs define edges. The system uses PyTorch Geometric (PyG) for message passing, PyTorch Frame for multi-modal feature encoding, and PyTorch Lightning for training orchestration.

This guide covers the optimizations that produced 30-35% training speedups on real workloads without sacrificing model accuracy. Each technique applies to PyG-based pipelines.

Avoid unnecessary re-compilations

The largest source of wasted time with torch.compile in graph learning is re-compilation. By default, PyTorch compiles an optimized kernel for the exact input shapes it sees. When the next mini-batch has different shapes, which happens every iteration in GNNs due to neighbor sampling, PyTorch re-compiles from scratch. That can make training slower than eager mode.

Dynamic input shapes

In a standard image pipeline, every batch has the same tensor dimensions. In graph learning, the number of sampled nodes and edges changes per batch. Without intervention, each new shape triggers a full re-compilation.

Use torch.compile(dynamic=True). This tells the compiler to generate kernels that handle variable shapes from the start, using symbolic shape analysis instead of hard-coding dimensions. The compiled kernel works for different input sizes without re-compiling.

Learning-rate schedulers

torch.compile treats floating-point scalar arguments as compile-time constants. When a learning-rate scheduler updates a floating-point learning rate at each step, the compiler can treat it as a new constant and re-compile. Wrap the learning rate in a tensor instead.

# Before: triggers re-compilation on every learning-rate update
optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
# After: the learning rate is a tensor, not a compile-time constant
optimizer = torch.optim.Adam(model.parameters(), lr=torch.tensor(0.01))

The tensor is treated as a runtime value, so scheduler updates no longer trigger re-compilation.

Graph breaks and CUDA Graphs

Even after re-compilations are removed, torch.compile can underperform when the computation graph is fragmented. The main considerations are graph breaks and whether to enable CUDA Graphs.

Graph breaks

A graph break occurs when torch.compile encounters an operation it cannot trace through. The compiler splits the computation into separately compiled subgraphs, preventing kernel fusion across the break boundary.

In PyG, tuple keys in ModuleDict lookups for heterogeneous graphs can cause graph breaks. Converting tuple keys to strings lets the compiler trace through the full forward pass. Staying on the latest PyG release also incorporates upstream torch.compile compatibility improvements.

CUDA Graphs: when not to use them

CUDA Graphs record a sequence of GPU kernels and replay them as one unit, removing per-kernel launch overhead. They are useful for fixed-shape workloads such as image classification or transformer inference with fixed sequence lengths.

For GNN training with neighbor sampling, CUDA Graphs are counterproductive. The node count in each mini-batch differs across iterations. CUDA Graphs must re-record the kernel sequence when shapes change, increasing memory usage and slowing training.

The Kumo pipeline disables CUDA Graphs:

# CUDA Graphs disabled: mini-batch shapes vary per iteration
model = torch.compile(model, dynamic=True)
# Do not use this mode: it enables CUDA Graphs internally.
# torch.compile(model, mode="reduce-overhead")

Eliminate host-device synchronization

Removing unnecessary synchronization between the CPU (host) and GPU (device) is often the highest-impact optimization category. Each synchronization forces the GPU to finish queued work before the CPU can proceed, creating a pipeline bubble. These points can dominate training time more than computation itself.

Metric accumulation on the CPU

Accumulating metrics such as accuracy on the CPU forces synchronization each iteration.

# Before: forces GPU synchronization every iteration
metric += int((label == pred).sum())
# After: accumulate on the GPU and synchronize once at the end of the epoch
metric += (label == pred).sum()

The int() call transfers data from device to host and blocks until pending GPU kernels finish. Keeping the running total as a tensor transfers the final accumulated value to the CPU only once per epoch.

torch.repeat_interleave without output_size

GNN message passing often uses torch.repeat_interleave to expand node features along edges. Without output_size, PyTorch must synchronize to calculate the output tensor size before allocating memory and launching kernels.

# Before: synchronizes to compute output size
expanded = torch.repeat_interleave(x, repeats, dim=0)
# After: provide the known output size to avoid synchronization
expanded = torch.repeat_interleave(x, repeats, dim=0, output_size=known_size)

When edge counts are known from the graph structure, passing the size lets PyTorch allocate memory and launch kernels without waiting.

Third-party library synchronization

Even well-maintained libraries can introduce hidden synchronization. Older torchmetrics releases, for example, introduced device synchronization during metric computation. Profile before and after library upgrades, and pin versions that do not introduce synchronization.

The full optimization pipeline

These techniques are layered rather than a single switch. Each addresses a different layer of the stack, and the effects compound.

The RDL architecture has three stages that benefit from the optimizations:

1. Feature encoding with PyTorch Frame encodes multi-modal column data including numerical, categorical, text, images, and timestamps into dense embeddings.
2. Message passing with PyG performs heterogeneous message passing across the relational graph, propagating information between connected tables.
3. Training orchestration with PyTorch Lightning manages the training loop, distributed training, and checkpointing.

Feature encoding involves standard tensor operations that compile cleanly. Message passing has the variable-shape challenges described above. Training orchestration wraps the outer loop and benefits from synchronization removal.

Benchmark results

All benchmarks ran on a g6.4xlarge instance with PyTorch 2.7.0 and CUDA 12.8. The dataset was the Kaggle H&M recommendation challenge, a real-world e-commerce dataset with users, items, and transactions. The evaluation covered three prediction types.

Speedup results

Across all three tasks, torch.compile with the full optimization pipeline delivered 30% to 35% speedups over eager-mode execution. The source reports no reduction in predictive accuracy, with compiled models producing identical predictions to their eager-mode counterparts.

The speedups are consistent across classification, regression, and link prediction because the optimizations target infrastructure bottlenecks rather than a task-specific computation. A 35% training speedup turns a ten-hour training run into a 6.5-hour run, saving GPU time and speeding iteration across many experiments.

Practical checklist

1. Enable compilation with dynamic shapes

model = torch.compile(model, dynamic=True)

Do not use mode="reduce-overhead" for variable-shape workloads because it enables CUDA Graphs.

2. Prevent learning-rate re-compilation

optimizer = torch.optim.Adam(
    model.parameters(),
    lr=torch.tensor(0.01),  # tensor, not float
)

3. Eliminate graph breaks

• Update PyG to the latest version.
• Convert tuple keys to string keys in custom ModuleDict usage.
• Set TORCH_LOGS="graph_breaks" to find remaining breaks.

4. Remove host-device synchronization

• Replace int(tensor) and float(tensor) in training loops with tensor-native accumulation.
• Pass output_size to torch.repeat_interleave whenever the size is known.
• Upgrade torchmetrics and other metric libraries.

5. Profile and verify

• Use the PyTorch Profiler to confirm synchronization points are gone.
• Compare eager and compiled training-loss curves to verify numerical equivalence.
• Measure wall-clock time per epoch rather than per step because compilation overhead is amortized over an epoch.

Recommended resources

• PyTorch Performance Tuning Guide