feat(torch.compile): PET FX compilation

[HaoZeke]

Summary

Full-graph FX compilation for PET training. Reworked from per-submodule compilation to the PhACE pattern (#921), and based off of the much cleaner to read UPET thing @sirmarcel has :)

make_fx(tracing_mode="symbolic") traces the entire PET forward pass (including force/stress via autograd.grad(create_graph=False)) into a single FX graph, then torch.compile(dynamic=True, fullgraph=True) compiles it. Single backward through the compiled graph, no double backward, SDPA (Scaled Dot-Product Attention) everywhere.

Key changes

• _forward_from_batch(): Pure-tensor input/output forward method on PET, bypassing metatensor TensorMap construction to make the model traceable.
• neighbor_atom_indices in systems_to_batch: Maps each edge slot to its neighbor atom index for force scattering in the compiled path.
• modules/compile.py: compile_pet_model() entry point. Wraps PET in _PETBatchForward, traces with make_fx, compiles with torch.compile(dynamic=True, fullgraph=True).
• _wrap_compiled_output(): Converts compiled function outputs back to Dict[str, TensorMap] for the loss function. Handles energy, forces, stress, and non-energy targets. Fixed device mismatch for GPU training.
• training.compile option in documentation/hypers (default: false).
• SiLU decomposition: Replaces nn.SiLU with x * sigmoid(x) to avoid silu_backward nodes that break inductor.

What gets compiled into one graph

• Edge distance computation (differentiable)
• All GNN (Graph Neural Network) layers (CartesianTransformer blocks with SDPA)
• Feature extraction and prediction heads
• Per-atom to per-structure energy aggregation (scatter_add)
• Force computation: autograd.grad + chain-rule scatter
• Stress computation: virial via einsum + scatter

Adaptive cutoff support

Works with num_neighbors_adaptive out of the box. Adaptive cutoffs cause max_edges_per_node (the second dimension of NEF (Node-Edge-Feature) tensors) to vary per batch, which typically breaks torch.compile. The NEF format reshapes edges into padded 2D arrays of shape [n_atoms, max_neighbors, ...], and the adaptive cutoff computation (data-dependent edge filtering) happens outside the FX graph in systems_to_batch(). The traced graph receives pre-filtered NEF tensors with varying but deterministic shapes. dynamic=True generates shape-generic kernels that handle the varying max_edges_per_node without recompilation.

Key workarounds that make this work:

• use_duck_shape = False during make_fx tracing prevents FX from aliasing n_atoms and max_edges_per_node symbolic dimensions when they happen to have the same concrete value in the tracing sample
• scatter_add everywhere (never index_add_) avoids an inductor miscompilation bug with dynamic indices that produces silent numerical corruption (documented in UPET fx_compile.py)
• NamedMemberAccessor instead of functional_call for parameter swapping during tracing

Benchmark Results

Benchmarks run via pet_bench/bench_compile.py in the pixi_envs repo, using the CUDA pixi environment:

# From pixi_envs/orgs/metatensor/metatrain/
pixi run -e cuda bench-compile            # both configs
pixi run -e cuda bench-compile-fixed      # fixed cutoff only
pixi run -e cuda bench-compile-adaptive   # adaptive cutoff only

Model: Default PET (2.9M params, 2 GNN layers, 8 heads, d_pet=128)
Dataset: qm9_reduced_100 (100 structures, batch_size=8, energy+forces+stress)

CUDA (RTX 4070 Ti SUPER, PyTorch 2.10.0+cu128)

Fixed Cutoff

Compilation Overhead

Phase	Time
FX trace (make_fx)	6.4 s
First call (Triton codegen)	3.4 s
Total	9.8 s

Step Timing (50 steps after warmup)

Metric	Eager	Compiled	Change
Median	24.9 ms	16.3 ms	-35%
Mean	25.1 ms	15.2 ms	-39%
Min / Max	24.4 / 26.9 ms	11.9 / 17.4 ms

Summary

Speedup	1.53x
Compile overhead	9.8 s (one-time)
Break-even	~1138 steps (~95 epochs)

Adaptive Cutoff (k=16)

With num_neighbors_adaptive=16, each atom keeps only its 16 nearest neighbors. The NEF second dimension (max_edges_per_node) varies per batch, exercising dynamic=True.

Compilation Overhead

Phase	Time
FX trace (make_fx)	5.2 s
First call (Triton codegen)	1.9 s
Total	7.1 s

Step Timing (50 steps after warmup)

Metric	Eager	Compiled	Change
Median	26.1 ms	16.4 ms	-37%
Mean	26.0 ms	16.4 ms	-37%
Min / Max	21.0 / 30.1 ms	16.2 / 17.5 ms

Summary

Speedup	1.59x
Compile overhead	7.1 s (one-time)
Break-even	~729 steps (~61 epochs)

Fixed vs Adaptive Comparison

Metric	Fixed	Adaptive (k=16)
Speedup	1.53x	1.59x
Compile overhead	9.8 s	7.1 s
Break-even (epochs)	95	61
Compiled step variance	2.0 ms std	0.3 ms std

Adaptive cutoff compiles faster and gives higher speedup with more stable step times.

CPU (PyTorch 2.9.1)

Compilation Overhead

Phase	Time
FX trace (make_fx)	5.5 s
First call (C++ codegen)	16.1 s
Total	21.6 s

Step Timing (50 steps after warmup)

Metric	Eager	Compiled	Change
Median	49.5 ms	40.0 ms	-19%
Mean	50.5 ms	42.1 ms	-17%
Min / Max	31.9 / 66.8 ms	25.4 / 62.5 ms

Summary

Speedup	1.24x
Compile overhead	21.6 s (one-time)
Break-even	~2274 steps (~190 epochs)

Contributor checklist

• Tests updated (for new features and bugfixes)?
• Documentation updated (for new features)?
• Issue referenced (for PRs that solve an issue)?

Maintainer/Reviewer checklist

• CHANGELOG updated with public API or any other important changes?
• GPU tests passed (maintainer comment: "cscs-ci run")?

[HaoZeke]

cscs-ci run

[HaoZeke]

cscs-ci run

[HaoZeke]

cscs-ci run

[HaoZeke]

cscs-ci run

[HaoZeke]

cscs-ci run

[HaoZeke]

cc @JonathanSchmidt1, any stress/beta testing would be amazing

[HaoZeke]

cscs-ci run

[HaoZeke]

cscs-ci run

[JonathanSchmidt1]

cc @JonathanSchmidt1, any stress/beta testing would be amazing

I have been using the branch and it seems fine so far (but only on a new ML task so no comparison to the previous version yet).

[JonathanSchmidt1]

@HaoZeke Have you experimented with making it work with DDP?
edit: I played a bit with compiled ddp in https://github.com/JonathanSchmidt1/metatrain/tree/max-atom-sampler
First tests seem to work though I had to set the triton/inductor cache to be in scratch instead of /tmp

[HaoZeke]

@HaoZeke Have you experimented with making it work with DDP? edit: I played a bit with compiled ddp in https://github.com/JonathanSchmidt1/metatrain/tree/max-atom-sampler First tests seem to work though I had to set the triton/inductor cache to be in scratch instead of /tmp

very cool, I haven't checked it out yet myself, but if you're building off of this maybe we should push to have this in main a bit faster

[HaoZeke]

cscs-ci run

[abmazitov]

cscs-ci run

[sirmarcel]

I've been looking into this a little bit with Claude, are you sure that the gradients are correctly flowing through the adaptive cutoff procedure?

[sirmarcel]

Claude's take:

What's happening. systems_to_batch computes cutoff_factors from edge_distances (which derives from positions), then the compiled FX graph receives cutoff_factors as a plain tensor input. The graph's force computation is (dE_dR,) = autograd.grad(energy.sum(), edge_vectors, create_graph=False), which differentiates only w.r.t. edge_vectors — cutoff_factors is held fixed. The hand-coded chain rule outside the graph (the center/neighbor scatter that maps dE/d(edge_vectors) to dE/d(positions)) likewise has no path through cutoff_factors. So the gradient contribution positions → edge_distances → cutoff_factors → energy is silently dropped.

Note that this is not a .detach(); nothing severs the chain. The chain simply never exists, because positions.requires_grad is False everywhere systems_to_batch is called from the compiled path (the requires_grad_(True) call at trainer.py:681 is on c_edge_vectors after systems_to_batch has already produced cutoff_factors). For comparison, the eager pet.forward(systems, ...) path works correctly because MetatomicCalculator sets positions.requires_grad_(True) before calling the model (see metatomic/torch/ase_calculator.py:520).

[JonathanSchmidt1]

Claude's take:

What's happening. systems_to_batch computes cutoff_factors from edge_distances (which derives from positions), then the compiled FX graph receives cutoff_factors as a plain tensor input. The graph's force computation is (dE_dR,) = autograd.grad(energy.sum(), edge_vectors, create_graph=False), which differentiates only w.r.t. edge_vectors — cutoff_factors is held fixed. The hand-coded chain rule outside the graph (the center/neighbor scatter that maps dE/d(edge_vectors) to dE/d(positions)) likewise has no path through cutoff_factors. So the gradient contribution positions → edge_distances → cutoff_factors → energy is silently dropped.

Note that this is not a .detach(); nothing severs the chain. The chain simply never exists, because positions.requires_grad is False everywhere systems_to_batch is called from the compiled path (the requires_grad_(True) call at trainer.py:681 is on c_edge_vectors after systems_to_batch has already produced cutoff_factors). For comparison, the eager pet.forward(systems, ...) path works correctly because MetatomicCalculator sets positions.requires_grad_(True) before calling the model (see metatomic/torch/ase_calculator.py:520).

Stupid question do we not have any tests for the forces where this should show up?

[sirmarcel]

@JonathanSchmidt1 Claude
Why the regression tests miss it: test_forward_from_batch only checks per-atom energy match, and test_compiled_vs_eager_backward_with_forces builds its "eager" reference by replicating compile.py's force formula via _forward_from_batch with create_graph=True. Both sides drop the same gradient term, so they agree perfectly. There is currently no test of the form "bare pet.forward(systems, outputs) with positions.requires_grad_(True) — i.e. the path every existing checkpoint and UPETCalculator evaluation goes through — agrees with the compiled path on position gradients.

[sirmarcel]

We caught this because I forced testing directly against the uPET calculator.

[JonathanSchmidt1]

Thank you for the reply @sirmarcel . We definitely need some regression tests against fixed numbers ...

[abmazitov]

I skimmed through the code quickly and there are a few things I don't understand as for now.

First, what's the logic behind having two separate forward functions? It looks to me like we have a plain forward which is unchanged, and a secondary _forward_from_batch which is supposed to be compiled, right? Shouldn't we call _forward_from_batch inside forward? Which one of those two should be exposed during model export then? If we export the non-scripted forward function, then does it mean we only torch.compile for training, not for inference?

When I thought about this problem, I had the following structure in my head:

class _PETBatchForward(torch.nn.Module):
    def __init__(self, hypers: ModelHypers, *args, **kwargs):
        ...
    def forward(self, *inputs: List[torch.Tensor]):
        # Implementation of all the GNN steps,
        # i.e. node features / edge_features calculation,
        # and finally - calculating the atomic_predictions
        results: Dict[str, Dict[str, torch.Tensor]] = {}
        # Filling in results
        return results

class PET:
    def __init__(self, hypers: ModelHypers, *args, **kwargs):
        self._core = _PETBatchForward(hypers, *args, **kwargs)
        ...
    def forward(self, systems: List[System], outputs: Dict[str, ModelOutput):
        inputs = systems_to_batch(systems, *args, **kwargs)
        results = self._core(inputs)
        return_dict = ... # convert results to TensorMaps
        return return_dict

So essentially we keep the same PET class as a part of the API, but internally run the torch-compiled core for pure-torch part of the evaluation. This looks quite natural and allows to, for example, get access to the pure-torch part of PET from the outside, which can be very useful for people not familiar with metatomic world and trying to first understand how the core works. Is there any particular reason behind a currently chosen design, and not what (as an example) I have suggested above?