Fix AOTI specialisation bug - #168
Merged
Daniel Zügner (danielzuegner) merged 18 commits intoJul 27, 2026
Merged
Conversation
…_body_indices on AOTI compile)
three_body_indices has shape [T, 2], where T is the number of angle
terms. T changes with the topology on every MD step, so it must stay
symbolic in the exported graph.
Dim.AUTO does not guarantee that: it silently *accepts* a specialisation
guard, so any op that materialises T as a Python int (e.g. via
aten.size.default, which guard_int()s the SymInt) bakes the compile-time
value into the .pt2 artifact with no warning. At runtime that indexes out
of bounds when T is smaller than at compile time (CUDA device-side
assert) and silently truncates three-body terms when T is larger,
returning wrong energies and forces.
Replace Dim.AUTO with a named Dim("num_triples"), which turns a
specialisation into a ConstraintViolationError at export time, and add
assert_threebody_dim_is_dynamic() as a second line of defence after
torch.export.export(): it inspects the three_body_indices placeholder and
refuses to package an artifact whose leading dim came back as a concrete
int. The error message points at the TORCH_LOGS invocation that locates
the offending op.
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
M3GNetForAOTI.forward precomputed bond_index_bias and
three_body_edge_map and passed them in input_dict, but M3Gnet.forward
reads neither -- it derives everything it needs from three_body_indices,
num_bonds and num_triple_ij. The two entries were pure overhead.
They were not free, either: both are torch.repeat_interleave calls with
a *tensor* repeats argument, whose output size is data-dependent. Each
one introduced an unbacked symbol into the exported program, which
showed up as u0/u1 in range_constraints and forced the tracer to guess
at several size comparisons ("could not evaluate Eq(u0, 0) due to data
dependency"). Removing them leaves range_constraints with only the four
genuine input dims.
No functional change: the values were computed and then ignored.
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
The cache key covered the model version, device, output set and torch version, but not the mattersim version. Any change to the traced graph that comes from mattersim itself rather than torch therefore reused the previously compiled .pt2 silently. That is exactly the situation created by the two preceding commits: an existing artifact compiled before the three-body fix has a frozen T and returns wrong energies and forces, but its cache key is unchanged, so a user upgrading into the fix keeps loading the broken artifact until they manually delete ~/.cache/mattersim/aoti/*.pt2. Fold __version__ into the key so upgrades invalidate stale artifacts on their own. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Rely on the named export dimension to reject specialization, and keep graph metadata inspection in the regression test rather than production code. Fingerprint model configuration and weights behind an explicit cache ABI so stale or custom-model artifacts cannot be reused. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 3724d1c0-700a-4e62-956e-5249677cc756
Contributor
|
Pushed maintainer refinement f8a3118:
The full AOTI regression file passes locally on CUDA (8 tests, including eager/AOTI parity above and below the compile-time triple count). Ubuntu, CodeQL, and docs checks pass. The macOS job has the same unrelated MPS BTE/relax numerical failures as the previous revision. |
Daniel Zügner (danielzuegner)
self-requested a review
July 27, 2026 12:20
Daniel Zügner (danielzuegner)
approved these changes
Jul 27, 2026
added 11 commits
July 27, 2026 12:25
Use the plain helper for AOTI tracing on CPU and CUDA, but preserve the original scripted execution path on MPS where eager evaluation changes MatterSim predictions. Exercise eager-versus-scripted layer parity on every available device. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 3724d1c0-700a-4e62-956e-5249677cc756
Temporarily compare eager and scripted spherical-harmonics outputs on the macOS CI runner and report the active Torch version and initial MatterSim prediction. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 3724d1c0-700a-4e62-956e-5249677cc756
Temporarily resolve the CI matrix to the versions available during the last green macOS run and remove the completed eager-versus-scripted diagnostic. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 3724d1c0-700a-4e62-956e-5249677cc756
Allow the resolver to choose published Torchaudio and Torchvision versions while holding Torch at 2.12 for the MPS regression boundary test. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 3724d1c0-700a-4e62-956e-5249677cc756
Restore unconstrained Torch resolution and temporarily pin the macOS job to the previously green macOS 15 ARM64 image, isolating the numerical regression from package-version changes. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 3724d1c0-700a-4e62-956e-5249677cc756
Temporarily compare identical CPU and MPS graph inputs at each M3GNet stage on the macOS 26 runner. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 3724d1c0-700a-4e62-956e-5249677cc756
Feed the exact CPU three-body interaction tensor and indices through scatter_sum on CPU and MPS to test the first functional reduction in graph_conv.0. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 3724d1c0-700a-4e62-956e-5249677cc756
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 3724d1c0-700a-4e62-956e-5249677cc756
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 3724d1c0-700a-4e62-956e-5249677cc756
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 3724d1c0-700a-4e62-956e-5249677cc756
Remove the temporary macOS diagnostics after isolating incorrect results to the MPS functional linear path on macOS 26. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 3724d1c0-700a-4e62-956e-5249677cc756
Contributor
|
Lower-level MPS investigation is complete. The failure is a macOS 26/PyTorch MPS regression, not an AOTI or spherical-harmonics regression:
I removed all temporary diagnostics and retained a |
Daniel Zügner (danielzuegner)
merged commit Jul 27, 2026
fd70170
into
microsoft:main
7 checks passed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Problem
AOTI compilation freezes in / specialises the
three_body_indicescount. This can result in crashes or silent incorrect energy and force predictions because the model will not usually have the right three-body terms, especially in MD runs.As an example:
Fix
Remove the
@torch.jit.scriptdecorator from_spherical_harmonicsand declare the triple count with a namedDim, so export fails rather than silently specialising it. Also: