aurora_configs/allscaip: run float32 — AllScAIP does not support FP64 inference - #224
Merged
Merged
Conversation
… 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
force-pushed
the
aurora-allscaip-fp64-default-dtype
branch
from
September 1, 2026 19:30
ea6278d to
edfe81a
Compare
OwenPriceSkelly
changed the title
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.
Supersedes the approach in #221 (merged) after CPU probing against fairchem main showed AllScAIP cannot run at
base_precision_dtype=float64at all — the Aurora verify crash was only the first of three layers:image_idbuilt at the torch default dtype →torch.mm(image_id, cell)crash (what the Aurora sync hit; aurora_configs/allscaip: patch radius-graph image_id dtype for FP64 inference #221 wrapped this one).index_putdtype crash (found testing the upstream fix; a wrapper can't reach between these internal sites).torch.float32before the output heads (AllScAIP.py,neighbor_reps[:, 0].to(torch.float32)), which then mismatch the doubled head weights. Verified by a full tiny-model FP64 forward on CPU: crashes identically withmodel.double()alone and withtorch.set_default_dtype(float64)— no env-side workaround can reach an explicit cast.Layer 3 makes FP64 AllScAIP an upstream design question (should the head seam honor
base_precision_dtype?), so this PR stops forcing FP64 and runs the fairchem default float32 — exactly whatnvidia_configs/allscaip.pyruns and verifies green at._fairchem_device(the xpu index shim) stays.Known risk, explicitly accepted: uma.py forces FP64 because XPU float32 gave wrong energies for UMA (#217). Whether AllScAIP's float32 is numerically adequate on XPU is exactly what the next sync's verification adjudicates — if it fails tolerance, allscaip-on-Aurora is blocked on upstream FP64 support and we drop the env until then.
Validation:
ruff format --checkclean, AST parses, deps unchanged. Runtime proof is the next Aurora sync (both allscaip verifies retry automatically).🤖 Generated with Claude Code