aurora_configs: native FairChem XPU via collaborator fork; add esen + allscaip

docs/environments.md

@@ -120,7 +120,7 @@ Rebuilds (`rootstock install <name> --force`) install exactly the locked version
 # explicit = true
 ```
 
-Then `setup()` takes `device="xpu"`. See `aurora_configs/{mace,uma}.py` for complete examples (UMA also monkeypatches FairChem to accept `xpu` and forces FP64). PyTorch's XPU build ships its Intel runtime under the env's `lib/`, which the worker adds to `LD_LIBRARY_PATH` automatically.
+Then `setup()` takes `device="xpu"`. See `aurora_configs/{mace,uma,esen,allscaip}.py` for complete examples (the FairChem envs install fairchem-core from a fork with native XPU support — PyPI releases accept only `cpu`/`cuda` — and force FP64 inference). PyTorch's XPU build ships its Intel runtime under the env's `lib/`, which the worker adds to `LD_LIBRARY_PATH` automatically.
 
 ### `CHECKPOINTS` table
 

sample_model_configurations/aurora_configs/allscaip.py

@@ -0,0 +1,94 @@
+# /// script
+# requires-python = ">=3.11"
+# dependencies = [
+#     "fairchem-core",
+#     "ase>=3.26",
+#     # Intel XPU (Aurora PVC) torch build. >=2.13: older XPU wheels have far
+#     # slower FP64 kernels (see uma.py).
+#     "torch>=2.13",
+#     # torch's XPU wheels depend on this; it lives only on the XPU index, so it
+#     # must be a direct dep for [tool.uv.sources] to route it (the index is
+#     # explicit, so transitive-only deps aren't fetched from it).
+#     "triton-xpu",
+# ]
+#
+# [tool.uv.sources]
+# # Experimental: fairchem-core from a fork whose xpu-support branch adds
+# # native Intel-GPU device handling (device="xpu", torch.xpu seeding/cache
+# # management, XCCL collectives) -- PyPI fairchem-core accepts only
+# # "cpu"/"cuda" device strings. The branch tracks upstream main; switch back
+# # to a PyPI fairchem-core once XPU support merges upstream.
+# fairchem-core = { git = "https://github.com/abagusetty/fairchem.git", branch = "xpu-support", subdirectory = "packages/fairchem-core" }
+# torch = { index = "pytorch-xpu" }
+# triton-xpu = { index = "pytorch-xpu" }
+#
+# [[tool.uv.index]]
+# name = "pytorch-xpu"
+# url = "https://download.pytorch.org/whl/xpu"
+# explicit = true
+# ///
+"""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, and (3) InferenceSettings defaults to float32, so we set
+base_precision_dtype=float64 to match the FP64 reference. The all-to-all
+node attention path is untested on XPU hardware.
+
+Pin one PVC tile with ZE_AFFINITY_MASK in the job (the worker inherits it).
+OMol checkpoints expect `charge` and `spin` in `atoms.info`.
+"""
+
+CHECKPOINTS = {
+    "allscaip-md-conserving-all-omol": "allscaip-md-conserving-all-omol",
+    "allscaip-md-direct-all-omol": "allscaip-md-direct-all-omol",
+    # Your own fine-tuned weights: pair with weights= (loaded via setup_from_path).
+    "allscaip:custom": None,
+}
+
+
+def _fairchem_device(device: str) -> str:
+    """Translate an indexed device ("xpu:2") into what fairchem accepts.
+
+    MLIPPredictUnit normalizes the requested device to a bare type and resolves
+    the actual GPU itself from torch's current-device state
+    (torch.xpu.current_device()), so an index has to travel through
+    torch.xpu.set_device, not the argument -- the same constraint the CUDA
+    configs shim around for multi-GPU verifies.
+    """
+    if device.startswith("xpu:"):
+        import torch
+
+        torch.xpu.set_device(int(device.split(":", 1)[1]))
+        return "xpu"
+    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 setup(checkpoint: str, device: str = "xpu", **kwargs):
+    from fairchem.core import FAIRChemCalculator, pretrained_mlip
+
+    predictor = pretrained_mlip.get_predict_unit(
+        CHECKPOINTS[checkpoint],
+        device=_fairchem_device(device),
+        inference_settings=_fp64_settings(),
+    )
+    return FAIRChemCalculator(predictor, **kwargs)
+
+
+def setup_from_path(path: str, device: str = "xpu", **kwargs):
+    # Custom checkpoints (`:custom` ids with user weights): a weights *file* loads
+    # through load_predict_unit, not the registry-name lookup setup() uses.
+    from fairchem.core import FAIRChemCalculator
+    from fairchem.core.units.mlip_unit import load_predict_unit
+
+    predictor = load_predict_unit(
+        path, device=_fairchem_device(device), inference_settings=_fp64_settings()
+    )
+    return FAIRChemCalculator(predictor, **kwargs)

sample_model_configurations/aurora_configs/esen.py

@@ -0,0 +1,94 @@
+# /// script
+# requires-python = ">=3.11"
+# dependencies = [
+#     "fairchem-core",
+#     "ase>=3.26",
+#     # Intel XPU (Aurora PVC) torch build. >=2.13: older XPU wheels have far
+#     # slower FP64 kernels (see uma.py).
+#     "torch>=2.13",
+#     # torch's XPU wheels depend on this; it lives only on the XPU index, so it
+#     # must be a direct dep for [tool.uv.sources] to route it (the index is
+#     # explicit, so transitive-only deps aren't fetched from it).
+#     "triton-xpu",
+# ]
+#
+# [tool.uv.sources]
+# # Experimental: fairchem-core from a fork whose xpu-support branch adds
+# # native Intel-GPU device handling (device="xpu", torch.xpu seeding/cache
+# # management, XCCL collectives) -- PyPI fairchem-core accepts only
+# # "cpu"/"cuda" device strings. The branch tracks upstream main; switch back
+# # to a PyPI fairchem-core once XPU support merges upstream.
+# fairchem-core = { git = "https://github.com/abagusetty/fairchem.git", branch = "xpu-support", subdirectory = "packages/fairchem-core" }
+# torch = { index = "pytorch-xpu" }
+# triton-xpu = { index = "pytorch-xpu" }
+#
+# [[tool.uv.index]]
+# name = "pytorch-xpu"
+# url = "https://download.pytorch.org/whl/xpu"
+# explicit = true
+# ///
+"""eSEN env (Intel XPU / Aurora) - FAIRChem eSEN single-task checkpoints.
+
+Same as nvidia_configs/esen.py except (1) torch resolves from the Intel XPU
+wheel index, (2) fairchem-core installs from a fork with native XPU support,
+and (3) InferenceSettings defaults to float32, so we set
+base_precision_dtype=float64 to match the FP64 reference.
+
+Pin one PVC tile with ZE_AFFINITY_MASK in the job (the worker inherits it).
+OMol checkpoints expect `charge` and `spin` in `atoms.info`.
+"""
+
+CHECKPOINTS = {
+    "esen-md-direct-all-omol": "esen-md-direct-all-omol",
+    "esen-sm-conserving-all-omol": "esen-sm-conserving-all-omol",
+    "esen-sm-direct-all-omol": "esen-sm-direct-all-omol",
+    # Your own fine-tuned weights: pair with weights= (loaded via setup_from_path).
+    "esen:custom": None,
+}
+
+
+def _fairchem_device(device: str) -> str:
+    """Translate an indexed device ("xpu:2") into what fairchem accepts.
+
+    MLIPPredictUnit normalizes the requested device to a bare type and resolves
+    the actual GPU itself from torch's current-device state
+    (torch.xpu.current_device()), so an index has to travel through
+    torch.xpu.set_device, not the argument -- the same constraint the CUDA
+    configs shim around for multi-GPU verifies.
+    """
+    if device.startswith("xpu:"):
+        import torch
+
+        torch.xpu.set_device(int(device.split(":", 1)[1]))
+        return "xpu"
+    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 setup(checkpoint: str, device: str = "xpu", **kwargs):
+    from fairchem.core import FAIRChemCalculator, pretrained_mlip
+
+    predictor = pretrained_mlip.get_predict_unit(
+        CHECKPOINTS[checkpoint],
+        device=_fairchem_device(device),
+        inference_settings=_fp64_settings(),
+    )
+    return FAIRChemCalculator(predictor, **kwargs)
+
+
+def setup_from_path(path: str, device: str = "xpu", **kwargs):
+    # Custom checkpoints (`:custom` ids with user weights): a weights *file* loads
+    # through load_predict_unit, not the registry-name lookup setup() uses.
+    from fairchem.core import FAIRChemCalculator
+    from fairchem.core.units.mlip_unit import load_predict_unit
+
+    predictor = load_predict_unit(
+        path, device=_fairchem_device(device), inference_settings=_fp64_settings()
+    )
+    return FAIRChemCalculator(predictor, **kwargs)

sample_model_configurations/aurora_configs/uma.py

@@ -1,9 +1,8 @@
 # /// script
 # requires-python = ">=3.11"
 # dependencies = [
-#     # 2.22 is the first release with uma-s-1p2p1 in the registry.
-#     "fairchem-core>=2.22",
-#     "ase>=3.22",
+#     "fairchem-core",
+#     "ase>=3.26",
 #     # Intel XPU (Aurora PVC) torch build. >=2.13: older XPU wheels (e.g. 2.8)
 #     # have far slower FP64 kernels -- UMA's first forward took >40 min on
 #     # 2.8.0+xpu vs ~2 min on 2.13.0+xpu (PVC tile).
@@ -15,6 +14,12 @@
 # ]
 #
 # [tool.uv.sources]
+# # Experimental: fairchem-core from a fork whose xpu-support branch adds
+# # native Intel-GPU device handling (device="xpu", torch.xpu seeding/cache
+# # management, XCCL collectives) -- PyPI fairchem-core accepts only
+# # "cpu"/"cuda" device strings. The branch tracks upstream main; switch back
+# # to a PyPI fairchem-core once XPU support merges upstream.
+# fairchem-core = { git = "https://github.com/abagusetty/fairchem.git", branch = "xpu-support", subdirectory = "packages/fairchem-core" }
 # torch = { index = "pytorch-xpu" }
 # triton-xpu = { index = "pytorch-xpu" }
 #
@@ -26,12 +31,10 @@
 """UMA env (Intel XPU / Aurora) - Meta's UMA foundation model via FAIRChem.
 
 Same as nvidia_configs/uma.py except (1) torch resolves from the Intel XPU wheel
-index, and (2) two XPU-specific fixes in setup():
-
-  * FairChem's MLIPPredictUnit._setup_device asserts device in {cpu, cuda}; we
-    monkeypatch it to accept "xpu" (Intel GPU support is not yet upstream).
-  * InferenceSettings defaults to float32; we set base_precision_dtype=float64
-    so energies/forces match the FP64 reference (fp32 is wrong at ~1e-7).
+index, (2) fairchem-core installs from a fork with native XPU support (replacing
+the monkeypatch this env previously carried), and (3) InferenceSettings defaults
+to float32, so we set base_precision_dtype=float64 to match the FP64 reference
+(fp32 is wrong at ~1e-7).
 
 Pin one PVC tile with ZE_AFFINITY_MASK in the job (the worker inherits it).
 Requires HF_TOKEN for the gated facebook/UMA checkpoints (download on a login
@@ -76,27 +79,21 @@ def _require_task(task, checkpoint):
     return task
 
 
-def _enable_xpu() -> None:
-    """Teach FairChem's predict unit to accept device="xpu".
+def _fairchem_device(device: str) -> str:
+    """Translate an indexed device ("xpu:2") into what fairchem accepts.
 
-    Upstream MLIPPredictUnit._setup_device allows only "cpu"/"cuda". Idempotent:
-    only wraps the original once.
+    MLIPPredictUnit normalizes the requested device to a bare type and resolves
+    the actual GPU itself from torch's current-device state
+    (torch.xpu.current_device()), so an index has to travel through
+    torch.xpu.set_device, not the argument -- the same constraint the CUDA
+    configs shim around for multi-GPU verifies.
     """
-    import fairchem.core.units.mlip_unit.predict as _predict
-    import torch
-
-    if getattr(_predict.MLIPPredictUnit._setup_device, "_xpu_patched", False):
-        return
-    _orig = _predict.MLIPPredictUnit._setup_device
+    if device.startswith("xpu:"):
+        import torch
 
-    def _setup_device(self, device):
-        if str(device).startswith("xpu"):
-            self.device = torch.device(device)
-            return
-        return _orig(self, device)
-
-    _setup_device._xpu_patched = True
-    _predict.MLIPPredictUnit._setup_device = _setup_device
+        torch.xpu.set_device(int(device.split(":", 1)[1]))
+        return "xpu"
+    return device
 
 
 def _fp64_settings():
@@ -108,21 +105,23 @@ def _fp64_settings():
 
 def setup(checkpoint: str, device: str = "xpu", task: str | None = None, **kwargs):
     task = _require_task(task, checkpoint)
-    _enable_xpu()
     from fairchem.core import FAIRChemCalculator, pretrained_mlip
 
     predictor = pretrained_mlip.get_predict_unit(
-        CHECKPOINTS[checkpoint], device=device, inference_settings=_fp64_settings()
+        CHECKPOINTS[checkpoint],
+        device=_fairchem_device(device),
+        inference_settings=_fp64_settings(),
     )
     return FAIRChemCalculator(predictor, task_name=task, **kwargs)
 
 
 def setup_from_path(path: str, device: str = "xpu", task: str | None = None, **kwargs):
     # Custom checkpoints (`:custom` ids with user weights): a weights *file* loads
     # through load_predict_unit, not the registry-name lookup setup() uses.
-    _enable_xpu()
     from fairchem.core import FAIRChemCalculator
     from fairchem.core.units.mlip_unit import load_predict_unit
 
-    predictor = load_predict_unit(path, device=device, inference_settings=_fp64_settings())
+    predictor = load_predict_unit(
+        path, device=_fairchem_device(device), inference_settings=_fp64_settings()
+    )
     return FAIRChemCalculator(predictor, task_name=task, **kwargs)