From edfe81a60527a35ca5a2bc327e995e2967b21745 Mon Sep 17 00:00:00 2001 From: Owen Price Skelly <21372141+OwenPriceSkelly@users.noreply.github.com> Date: Tue, 1 Sep 2026 14:07:51 -0500 Subject: [PATCH] =?UTF-8?q?aurora=5Fconfigs/allscaip:=20run=20float32=20?= =?UTF-8?q?=E2=80=94=20AllScAIP=20does=20not=20support=20FP64=20inference?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 --- .../aurora_configs/allscaip.py | 63 +++++-------------- 1 file changed, 16 insertions(+), 47 deletions(-) diff --git a/sample_model_configurations/aurora_configs/allscaip.py b/sample_model_configurations/aurora_configs/allscaip.py index f96a0af..f3e565c 100644 --- a/sample_model_configurations/aurora_configs/allscaip.py +++ b/sample_model_configurations/aurora_configs/allscaip.py @@ -30,11 +30,21 @@ """AllScAIP env (Intel XPU / Aurora) - FAIRChem scalable attention MLIP. Same as nvidia_configs/allscaip.py except (1) torch resolves from the Intel -XPU wheel index, (2) fairchem-core installs from a fork with native XPU -support, (3) InferenceSettings defaults to float32, so we set -base_precision_dtype=float64 to match the FP64 reference, and (4) FP64 -inference needs a dtype patch in AllScAIP's radius graph (see -_patch_image_id_dtype). +XPU wheel index and (2) fairchem-core installs from a fork with native XPU +support. + +Unlike uma.py/esen.py, this env does NOT force +InferenceSettings(base_precision_dtype=float64): AllScAIP does not support +FP64 inference. Its radius-graph construction creates tensors at the torch +default dtype that crash against a double batch (torch.mm and index_put +dtype mismatches), and even past those, the backbone hard-casts its node +representations to float32 before the output heads +(fairchem models/allscaip/AllScAIP.py), which then mismatches the doubled +head weights. So this env runs the fairchem default float32 -- the same +precision the NVIDIA deployments verify at. If XPU float32 kernels prove +numerically inadequate here (the reason uma.py forces FP64), that will +surface as a verification failure, and FP64 support has to land in +fairchem first. Pin one PVC tile with ZE_AFFINITY_MASK in the job (the worker inherits it). OMol checkpoints expect `charge` and `spin` in `atoms.info`. @@ -65,50 +75,12 @@ def _fairchem_device(device: str) -> str: return device -def _fp64_settings(): - import torch - from fairchem.core.units.mlip_unit.api.inference import InferenceSettings - - return InferenceSettings(base_precision_dtype=torch.float64, tf32=False) - - -def _patch_image_id_dtype(): - """Work around a dtype mismatch in AllScAIP's radius graph under FP64. - - biknn_radius_graph builds the PBC image-offset tensors (image_id) with - torch.get_default_dtype() -- float32 -- while base_precision_dtype=float64 - casts the batch, including cell, to double, so build_radius_graph's - `torch.mm(image_id, cell)` raises "expected mat1 and mat2 to have the - same dtype, but got: float != double" (fairchem - models/allscaip/utils/allscaip_radius_graph.py). Not device-specific: - any FP64 run hits it. Cast each image_id to its cell's dtype on the way - into batched_radius_graph; drop this once fairchem builds image_id with - the cell dtype. - """ - from fairchem.core.models.allscaip.utils import allscaip_radius_graph as graph_mod - - if getattr(graph_mod.batched_radius_graph, "_image_id_dtype_patched", False): - return - original = graph_mod.batched_radius_graph - - def patched(pos_list, cell_list, image_id_list, *args, **kwargs): - image_id_list = [ - image_id.to(cell.dtype) for image_id, cell in zip(image_id_list, cell_list) - ] - return original(pos_list, cell_list, image_id_list, *args, **kwargs) - - patched._image_id_dtype_patched = True - graph_mod.batched_radius_graph = patched - - def setup(checkpoint: str, device: str = "xpu", **kwargs): from fairchem.core import FAIRChemCalculator, pretrained_mlip - _patch_image_id_dtype() predictor = pretrained_mlip.get_predict_unit( CHECKPOINTS[checkpoint], device=_fairchem_device(device), - inference_settings=_fp64_settings(), ) return FAIRChemCalculator(predictor, **kwargs) @@ -119,8 +91,5 @@ def setup_from_path(path: str, device: str = "xpu", **kwargs): from fairchem.core import FAIRChemCalculator from fairchem.core.units.mlip_unit import load_predict_unit - _patch_image_id_dtype() - predictor = load_predict_unit( - path, device=_fairchem_device(device), inference_settings=_fp64_settings() - ) + predictor = load_predict_unit(path, device=_fairchem_device(device)) return FAIRChemCalculator(predictor, **kwargs)