> For clean Markdown of any page, append .md to the page URL. > For a complete documentation index, see https://docs.nvidia.com/sdgm/llms.txt. > For AI client integration (Claude Code, Cursor, etc.), connect to the MCP server at https://docs.nvidia.com/sdgm/_mcp/server. # Optimize PyG with torch.compile > A practical guide to optimizing PyG-based GNN training by eliminating re-compilations, graph breaks, and host-device synchronization **Source title:** Speeding Up Graph Learning Models with PyG and torch.compile\ **Authors:** Akihiro Nitta and Matthias Fey\ **Published:** April 2025 ## 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)](https://github.com/pyg-team/pytorch_geometric) for message passing, [PyTorch Frame](https://github.com/pyg-team/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. Graph learning workloads have irregular, variable-sized inputs that make `torch.compile` harder to apply than on standard vision or language models. The optimizations in this guide address that irregularity. ## 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. Think of the default mode as a conveyor belt sized exactly for the first box it sees, which must be rebuilt for every differently sized box. `dynamic=True` creates an adjustable belt that handles different box sizes from the start. ### 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. ```python # 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. Use `torch.compile(dynamic=True)` for variable input shapes and a tensor learning rate for scheduler compatibility. Both are one-line changes that remove common sources of re-compilation in GNN training. ## 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 ![Three steps to eliminate graph breaks: identify compiler breaks with TORCH\_LOGS, replace incompatible patterns such as tuple keys with compiler-friendly alternatives, then verify that the complete forward pass compiles as one graph for maximum kernel fusion.](https://fdr-prod-docs-files-public.s3.us-east-1.amazonaws.com/nvidia-sdgm.docs.buildwithfern.com/aa690e566b04ca842b2fd354aed816e05bf8fa989cf9f6999c16ad81c202718e/img/research/pyg-compile-graph-breaks.svg?X-Amz-Algorithm=AWS4-HMAC-SHA256&X-Amz-Content-Sha256=UNSIGNED-PAYLOAD&X-Amz-Credential=AKIA6KXJSKKNFOCF7G4B%2F20260920%2Fus-east-1%2Fs3%2Faws4_request&X-Amz-Date=20260920T164627Z&X-Amz-Expires=604800&X-Amz-Signature=07f60a07309844905e45596b15c4c019dbc8420e88c41daa5e963d2abcac651e&X-Amz-SignedHeaders=host&x-amz-checksum-mode=ENABLED&x-id=GetObject) 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: ```python # 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") ``` Update PyG and replace incompatible patterns such as tuple keys in `ModuleDict` to fix graph breaks. Disable CUDA Graphs for neighbor-sampled GNN training because variable mini-batch sizes make repeated kernel recording counterproductive. ## 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. ```python # 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. ```python # 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. Removing host-device synchronization is like letting a factory assembly line run continuously instead of stopping it to count finished products after every item. Keep metric accumulation on the GPU, provide `output_size` to `torch.repeat_interleave`, and update libraries that introduce hidden synchronization. ## The full optimization pipeline ![Five-stage torch.compile optimization pipeline: enable dynamic compilation for variable subgraphs, prevent scalar-triggered re-compilations, resolve graph breaks, disable CUDA Graphs for variable batches, and remove host-device synchronization.](https://fdr-prod-docs-files-public.s3.us-east-1.amazonaws.com/nvidia-sdgm.docs.buildwithfern.com/ad5b7166873ad922d3bd09720e016127174f517aaeae24ce8c4dfd5abbd4b729/img/research/pyg-compile-optimization-pipeline.svg?X-Amz-Algorithm=AWS4-HMAC-SHA256&X-Amz-Content-Sha256=UNSIGNED-PAYLOAD&X-Amz-Credential=AKIA6KXJSKKNFOCF7G4B%2F20260920%2Fus-east-1%2Fs3%2Faws4_request&X-Amz-Date=20260920T164627Z&X-Amz-Expires=604800&X-Amz-Signature=cc94e64c1a6e9905ffb4da53c41ee0beb4ee3830662996f3b587c1afe411b259&X-Amz-SignedHeaders=host&x-amz-checksum-mode=ENABLED&x-id=GetObject) 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. Re-compilation fixes address compiler input, graph-break fixes address compiler output, and synchronization removal addresses runtime behavior. Applying all three layers together produces the full 30-35% speedup. ## 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. | Task | Type | Description | | -------------------- | ------------------- | ----------------------------------------------------------------- | | `user-churn` | Node classification | Predict whether a customer will churn within one week. | | `item-sales` | Node regression | Estimate weekly article sales volume. | | `user-item-purchase` | Link prediction | Forecast which articles a customer will purchase over seven days. | ### 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. | Task | Mode | Speedup | Accuracy impact | | -------------------- | ---------------------------------- | -------------------- | --------------- | | `user-churn` | `torch.compile` with optimizations | Approximately 30-35% | None | | `item-sales` | `torch.compile` with optimizations | Approximately 30-35% | None | | `user-item-purchase` | `torch.compile` with optimizations | Approximately 30-35% | None | 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. The combined optimizations delivered 30-35% training speedups across real e-commerce classification, regression, and link-prediction tasks with no reported accuracy loss. ## Practical checklist ### 1. Enable compilation with dynamic shapes ```python 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 ```python 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](https://pytorch.org/tutorials/recipes/recipes/profiler_recipe.html) 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](https://pytorch.org/tutorials/recipes/recipes/tuning_guide.html) These are incremental, low-risk changes that do not require changing the model architecture. Start with `torch.compile(dynamic=True)`, then eliminate re-compilations, graph breaks, and synchronization while profiling each step.