Keep UMA atom and edge dimensions dynamic - #2127
Draft
mlazos wants to merge 6 commits into
Draft
Conversation
Custom operator boundaries caused Dynamo to specialize the first several edge counts despite dynamic compilation. Mark UMA atom and edge dimensions explicitly and scope inference-only Dynamo configuration to the inference call so one compiled predictor can serve changing system sizes. Test Plan: ``` TORCHINDUCTOR_CACHE_DIR=$PWD/../dynamic_branch_cache PYTHONPATH=$PWD/src:$PWD/..:$PYTHONPATH python <two-size uma-s-1p2 no-recompile check> ruff check src/fairchem/core/units/mlip_unit/predict.py ``` The real checkpoint retained 9 unique graphs when changing from 2 to 4 atoms. Authored with assistance from Codex. [ghstack-poisoned]
This was referenced Jul 31, 2026
Target configuration: uma-s-1p2 with `torch.compile(dynamic=True)`, `external_graph_gen=True`, `merge_mole=True`, and `umas_fast_gpu`, reusing one predictor across the 2-to-2048-atom size sweep. This issue is active in that target sweep, but causes guard failures and recompilation rather than a Dynamo graph break in a single fixed-size invocation. Custom-operator boundaries do not expose enough shape relationships for Dynamo to infer that every atom and edge axis should remain dynamic. For example, if the first `edge_index` has shape `[2, 156]`, Dynamo can guard the edge count as 156; a later `[2, 2496]` input invalidates that guard and recompiles. Explicitly mark the atom axes and edge axes of graph tensors before the first compiled call. `edge_index` dimension 1 and `cell_offsets` dimension 0 then remain symbolic, as do dimension 0 of atom-indexed tensors. Scope `recompile_limit` and float specialization with a Dynamo config decorator on inference instead of leaking process-global assignments from lazy setup. Test Plan: ``` TORCHINDUCTOR_CACHE_DIR=$PWD/../dynamic_branch_cache PYTHONPATH=$PWD/src:$PWD/..:$PYTHONPATH python <two-size uma-s-1p2 no-recompile check> ruff check src/fairchem/core/units/mlip_unit/predict.py ``` The real checkpoint retained nine unique graphs when changing from two to four atoms instead of compiling shape-specialized replacements. Authored with assistance from Codex. [ghstack-poisoned]
Target configuration: uma-s-1p2 with `torch.compile(dynamic=True)`,
`external_graph_gen=True`, `merge_mole=True`, and `umas_fast_gpu`, reusing one
predictor across the 2-to-2048-atom size sweep. This issue is active in that
target sweep, but causes guard failures and recompilation rather than a Dynamo
graph break in a single fixed-size invocation.
Custom-operator boundaries do not expose enough shape relationships for Dynamo
to infer that every atom and edge axis should remain dynamic. For example, if
the first `edge_index` has shape `[2, 156]`, Dynamo can guard the edge count as
156; a later `[2, 2496]` input invalidates that guard and recompiles.
Explicitly mark the atom axes and edge axes of graph tensors before the first
compiled call. `edge_index` dimension 1 and `cell_offsets` dimension 0 then
remain symbolic, as do dimension 0 of atom-indexed tensors. Scope
`recompile_limit` and float specialization with a Dynamo config decorator on
inference instead of leaking process-global assignments from lazy setup.
**Activation**
Dynamic atom and edge dimensions are enabled automatically for a compiled backend that advertises the fused edgewise path, including `umas_fast_gpu`:
```python
settings = InferenceSettings(
compile=True,
execution_mode="umas_fast_gpu",
)
```
The predictor calls `torch.compile(dynamic=True)` unless fixed-shape `full_cudagraph=True` is selected. With `full_cudagraph=True`, set `full_cudagraph_dynamic=True` to retain dynamic compilation and use the per-shape CUDA graph cache.
Test Plan:
```
TORCHINDUCTOR_CACHE_DIR=$PWD/../dynamic_branch_cache PYTHONPATH=$PWD/src:$PWD/..:$PYTHONPATH python <two-size uma-s-1p2 no-recompile check>
ruff check src/fairchem/core/units/mlip_unit/predict.py
```
The real checkpoint retained nine unique graphs when changing from two to four
atoms instead of compiling shape-specialized replacements.
Authored with assistance from Codex.
[ghstack-poisoned]
Target configuration: uma-s-1p2 with `torch.compile(dynamic=True)`,
`external_graph_gen=True`, `merge_mole=True`, and `umas_fast_gpu`, reusing one
predictor across the 2-to-2048-atom size sweep. Changing atom or edge counts in that target sweep causes guard failures
and recompilation rather than a Dynamo
graph break in a single fixed-size invocation.
Custom-operator boundaries do not expose enough shape relationships for Dynamo
to infer that every atom and edge axis should remain dynamic. For example, if
the first `edge_index` has shape `[2, 156]`, Dynamo can guard the edge count as
156; a later `[2, 2496]` input invalidates that guard and recompiles.
Explicitly mark the atom axes and edge axes of graph tensors before the first
compiled call. `edge_index` dimension 1 and `cell_offsets` dimension 0 then
remain symbolic, as do dimension 0 of atom-indexed tensors. Scope
`recompile_limit` and float specialization with a Dynamo config decorator on
inference instead of leaking process-global assignments from lazy setup.
**Activation**
Dynamic atom and edge dimensions are enabled automatically for a compiled backend that advertises the fused edgewise path, including `umas_fast_gpu`:
```python
settings = InferenceSettings(
compile=True,
execution_mode="umas_fast_gpu",
)
```
The predictor calls `torch.compile(dynamic=True)` unless fixed-shape `full_cudagraph=True` is selected. With `full_cudagraph=True`, set `full_cudagraph_dynamic=True` to retain dynamic compilation and use the per-shape CUDA graph cache.
Test Plan:
```
TORCHINDUCTOR_CACHE_DIR=$PWD/../dynamic_branch_cache PYTHONPATH=$PWD/src:$PWD/..:$PYTHONPATH python <two-size uma-s-1p2 no-recompile check>
ruff check src/fairchem/core/units/mlip_unit/predict.py
```
The real checkpoint retained nine unique graphs when changing from two to four
atoms instead of compiling shape-specialized replacements.
Authored with assistance from Codex.
[ghstack-poisoned]
Target configuration: uma-s-1p2 with `torch.compile(dynamic=True)`,
`external_graph_gen=True`, `merge_mole=True`, and `umas_fast_gpu`, reusing one
predictor across the 2-to-2048-atom size sweep. Changing atom or edge counts in that target sweep causes guard failures
and recompilation rather than a Dynamo
graph break in a single fixed-size invocation.
Custom-operator boundaries do not expose enough shape relationships for Dynamo
to infer that every atom and edge axis should remain dynamic. For example, if
the first `edge_index` has shape `[2, 156]`, Dynamo can guard the edge count as
156; a later `[2, 2496]` input invalidates that guard and recompiles.
Explicitly mark the atom axes and edge axes of graph tensors before the first
compiled call. `edge_index` dimension 1 and `cell_offsets` dimension 0 then
remain symbolic, as do dimension 0 of atom-indexed tensors. Scope
`recompile_limit` and float specialization with a Dynamo config decorator on
inference instead of leaking process-global assignments from lazy setup.
**Activation**
Dynamic atom and edge dimensions are enabled automatically for a compiled backend that advertises the fused edgewise path, including `umas_fast_gpu`:
```python
settings = InferenceSettings(
compile=True,
execution_mode="umas_fast_gpu",
)
```
The predictor calls `torch.compile(dynamic=True)` unless fixed-shape `full_cudagraph=True` is selected. With `full_cudagraph=True`, set `full_cudagraph_dynamic=True` to retain dynamic compilation and use the per-shape CUDA graph cache.
Test Plan:
```
TORCHINDUCTOR_CACHE_DIR=$PWD/../dynamic_branch_cache PYTHONPATH=$PWD/src:$PWD/..:$PYTHONPATH python <two-size uma-s-1p2 no-recompile check>
ruff check src/fairchem/core/units/mlip_unit/predict.py
```
The real checkpoint retained nine unique graphs when changing from two to four
atoms instead of compiling shape-specialized replacements.
Authored with assistance from Codex.
[ghstack-poisoned]
Target configuration: uma-s-1p2 with `torch.compile(dynamic=True)`,
`external_graph_gen=True`, `merge_mole=True`, and `umas_fast_gpu`, reusing one
predictor across the 2-to-2048-atom size sweep. Changing atom or edge counts in that target sweep causes guard failures
and recompilation rather than a Dynamo
graph break in a single fixed-size invocation.
Custom-operator boundaries do not expose enough shape relationships for Dynamo
to infer that every atom and edge axis should remain dynamic. For example, if
the first `edge_index` has shape `[2, 156]`, Dynamo can guard the edge count as
156; a later `[2, 2496]` input invalidates that guard and recompiles.
Explicitly mark the atom axes and edge axes of graph tensors before the first
compiled call. `edge_index` dimension 1 and `cell_offsets` dimension 0 then
remain symbolic, as do dimension 0 of atom-indexed tensors. Scope
`recompile_limit` and float specialization with a Dynamo config decorator on
inference instead of leaking process-global assignments from lazy setup.
**Activation**
Dynamic atom and edge dimensions are enabled automatically for a compiled backend that advertises the fused edgewise path, including `umas_fast_gpu`:
```python
settings = InferenceSettings(
compile=True,
execution_mode="umas_fast_gpu",
)
```
The predictor calls `torch.compile(dynamic=True)` for this configuration.
Test Plan:
```
PYTHONPATH=$PWD/src:$PYTHONPATH pytest -q tests/core/units/mlip_unit/test_predict.py -k test_mark_dynamic_input_dimensions_requires_fast_backend
```
The real checkpoint retained nine unique graphs when changing from two to four
atoms instead of compiling shape-specialized replacements.
Authored with assistance from Codex.
[ghstack-poisoned]
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
1 participant
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.
Stack from ghstack (oldest at bottom):
Custom-operator boundaries do not expose enough shape relationships for Dynamo
to infer that every atom and edge axis should remain dynamic. For example, if
the first
edge_indexhas shape[2, 156], Dynamo can guard the edge count as156; a later
[2, 2496]input invalidates that guard and recompiles.Explicitly mark the atom axes and edge axes of graph tensors before the first
compiled call.
edge_indexdimension 1 andcell_offsetsdimension 0 thenremain symbolic, as do dimension 0 of atom-indexed tensors. Scope
recompile_limitand float specialization with a Dynamo config decorator oninference instead of leaking process-global assignments from lazy setup.
Activation
Dynamic atom and edge dimensions are enabled automatically for a compiled backend that advertises the fused edgewise path, including
umas_fast_gpu:The predictor calls
torch.compile(dynamic=True)for this configuration.Test Plan:
The real checkpoint retained nine unique graphs when changing from two to four
atoms instead of compiling shape-specialized replacements.
Authored with assistance from Codex.