diff --git a/docs/environments.md b/docs/environments.md index 010f6ae..d787def 100644 --- a/docs/environments.md +++ b/docs/environments.md @@ -120,7 +120,7 @@ Rebuilds (`rootstock install --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 diff --git a/sample_model_configurations/aurora_configs/allscaip.py b/sample_model_configurations/aurora_configs/allscaip.py new file mode 100644 index 0000000..09c6f39 --- /dev/null +++ b/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) diff --git a/sample_model_configurations/aurora_configs/esen.py b/sample_model_configurations/aurora_configs/esen.py new file mode 100644 index 0000000..37f5f60 --- /dev/null +++ b/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) diff --git a/sample_model_configurations/aurora_configs/uma.py b/sample_model_configurations/aurora_configs/uma.py index 8575ea6..41cc330 100644 --- a/sample_model_configurations/aurora_configs/uma.py +++ b/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,11 +105,12 @@ 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) @@ -120,9 +118,10 @@ def setup(checkpoint: str, device: str = "xpu", task: str | None = None, **kwarg 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)