aurora_configs/allscaip: patch radius-graph image_id dtype for FP64 inference - #221
Merged
Merged
Conversation
…nference
The first Aurora sync (2026-08-28) failed both allscaip checkpoints at
verify with:
RuntimeError: expected mat1 and mat2 to have the same dtype,
but got: float != double
in build_radius_graph's torch.mm(image_id, cell). AllScAIP builds its PBC
image-offset tensors with torch.get_default_dtype() (float32), while this
env's InferenceSettings(base_precision_dtype=float64) casts the batch --
including cell -- to double. Not XPU-specific: any FP64 run of AllScAIP
hits it; the nvidia config runs fp32 and never did.
Cast image_id to the cell dtype via a batched_radius_graph wrapper in
setup()/setup_from_path(), to be dropped once fairchem builds image_id
with the cell dtype (reported upstream to the fork maintainer).
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
OwenPriceSkelly
added a commit
that referenced
this pull request
Sep 1, 2026
… inference Reverts the FP64 forcing (and #221's image_id patch rationale) after CPU probing against fairchem main showed AllScAIP cannot run at base_precision_dtype=float64 at all, in three layers: 1. the radius graph builds image_id at the torch default dtype -> torch.mm(image_id, cell) dtype crash (the Aurora verify failure); 2. past that, the padded disp/envelope buffers are also default-dtype -> index_put dtype crash; 3. past both (fixed upstream in facebookresearch/fairchem#2184), the backbone hard-casts its node representations to float32 before the output heads (AllScAIP.py), which then mismatch the doubled head weights -- an explicit cast no default-dtype workaround can reach. So FP64 AllScAIP needs upstream design work, not an env shim. Run the fairchem default float32 instead -- the precision the NVIDIA deployments verify at. If XPU float32 proves numerically inadequate (the reason uma.py forces FP64), verification will say so and allscaip-on-Aurora is blocked on upstream FP64 support. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
OwenPriceSkelly
added a commit
that referenced
this pull request
Sep 1, 2026
… inference (#224) Reverts the FP64 forcing (and #221's image_id patch rationale) after CPU probing against fairchem main showed AllScAIP cannot run at base_precision_dtype=float64 at all, in three layers: 1. the radius graph builds image_id at the torch default dtype -> torch.mm(image_id, cell) dtype crash (the Aurora verify failure); 2. past that, the padded disp/envelope buffers are also default-dtype -> index_put dtype crash; 3. past both (fixed upstream in facebookresearch/fairchem#2184), the backbone hard-casts its node representations to float32 before the output heads (AllScAIP.py), which then mismatch the doubled head weights -- an explicit cast no default-dtype workaround can reach. So FP64 AllScAIP needs upstream design work, not an env shim. Run the fairchem default float32 instead -- the precision the NVIDIA deployments verify at. If XPU float32 proves numerically inadequate (the reason uma.py forces FP64), verification will say so and allscaip-on-Aurora is blocked on upstream FP64 support. Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
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.
Fixes the two
allscaipverify failures from the first Aurora sync (manifest pushed 2026-08-28T22:47Z).Failure
Both
allscaip-md-{conserving,direct}-all-omolverifies died on xpu with:Root cause
biknn_radius_graphbuilds its PBC image-offset tensors withtorch.get_default_dtype()(float32) — both the PBCtorch.cartesian_prod(torch.arange(...))path and the non-PBC identity image — while this env'sInferenceSettings(base_precision_dtype=float64)casts the batch, includingcell, to double. Not XPU-specific: any FP64 run of AllScAIP hits it. The nvidia config runs default fp32 and never trips it; eSEN's graph path doesn't share the pattern, which is whyesenverified 3/3 on the same sync.Fix
A guarded wrapper around
batched_radius_graph(applied insetup()/setup_from_path(), same precedent as the_setup_devicemonkeypatch this env carried before #220) that casts eachimage_idto itscell's dtype. Droppable once fairchem buildsimage_idwith the cell dtype — being reported to the fork maintainer (abagusetty) as the proper upstream fix.Validation
ruff format --checkclean; AST parses; deps unchanged (no rebuild needed — sync re-registers the source and re-verifies).batched_radius_graph(pos_list, cell_list, image_id_list, ...), module-global lookup so the wrapper takes effect), but the proof is the next sync — the two failed verifies retry automatically.🤖 Generated with Claude Code