Add support to Intel GPUs on Aurora - #2182
Open
abagusetty wants to merge 11 commits into
Open
abagusetty wants to merge 11 commits into
abagusetty wants to merge 11 commits into
Conversation
fairchem hard-codes CUDA as *the* accelerator, so it cannot run on Intel
Data Center GPU Max hardware (Aurora, Sunspot). The block was never deep --
the tree has no .cu/.cpp/.cuh sources, it is pure Python over torch -- but
it was absolute: MLIPPredictUnit._setup_device asserted
assert device in ["cpu", "cuda"]
and, worse, resolved anything that was not "cuda" to CPU. Bypassing the
assert alone would therefore have produced a silently slow CPU run rather
than an error.
This adds a central device abstraction and routes every CUDA-specific call
site through it, rather than sprinkling `or xpu` across the codebase.
New: fairchem/core/common/device_utils.py
Built on torch's own generic APIs (torch.get_device_module,
torch.accelerator, torch.amp.autocast), so a further backend is a table
entry rather than new call sites. Provides device detection/validation,
per-backend collective selection, memory/seed/stream helpers, and the
Triton opt-in gate.
Ported:
* predict.py -- assert widened; device normalised once in __init__ so
"auto"/"cuda:1" resolve before use; seeding, cache release, Ray GPU
reservation and the collective backend all device-agnostic.
* distutils.py -- assign/get_device_for_local_rank drive the detected
accelerator; visible-device logging names the right vendor env var
(ZE_AFFINITY_MASK / ONEAPI_DEVICE_SELECTOR, not CUDA_VISIBLE_DEVICES).
* rotation_cuda_graph.py -- graph capture via torch.xpu.make_graphed_callables
as well as torch.cuda's; both expose the same Stream surface.
* layer_norm / normalizer / element_references -- autocast-disable
decorators gained their xpu sibling, which the cuda+cpu pair missed.
* graph_parallel_a2a.py -- the native-all_to_all fast path is chosen by
backend *capability*, so oneCCL keeps it instead of being demoted to
pairwise send/recv.
* launchers -- DeviceType.XPU; backend follows the device type.
* public device Literal hints, benchmark/entry-point defaults.
Collectives use "xccl", which is oneCCL upstreamed into PyTorch as a native
backend -- libtorch_xpu.so links libccl.so directly. The legacy out-of-tree
name ("ccl", via oneccl_bindings_for_pytorch) is kept as a fallback for
builds that still ship it. NCCL is never selected for XPU.
UMA-S's fused Triton path stays opt-in off CUDA
(FAIRCHEM_ENABLE_TRITON_XPU=1). Those kernels are autotuned for NVIDIA
occupancy and may use backend-specific intrinsics, so compiling elsewhere
does not imply correct or fast.
Behaviour on NVIDIA is unchanged: get_available_accelerator() returns
"cuda" there, and every default resolves as before.
Verified on Sunspot (Intel Data Center GPU Max 1550, torch
2.13.0a0+gitcf30153, 12 devices): eSCN-MD forward and backward, conservative
forces from on-device autograd, all 79 parameter gradients finite and
non-zero, and energies matching CPU to 1e-3.
The graph-parallel GPU tests could not run on Intel GPUs at all: 15 of them
failed before reaching any assertion. Three separate causes, none of them in
the model code.
1. Hard-coded collective backend. The tests build
PGConfig(backend="nccl", ...), and NCCL is absent from an XPU PyTorch
build, so init_process_group died immediately. Now the backend follows the
detected device -- nccl on NVIDIA, oneCCL ("xccl") on Intel. Tests also
skip cleanly when no accelerator exists rather than failing.
2. Hard-coded tensor placement. _to_cuda() built torch.device(f"cuda:{rank}"),
and the shared test bodies built their partition tensors with
torch.arange(natoms) / partition_atoms_index_split(..., torch.device("cpu"))
while operating on device tensors -- so once NCCL stopped being the
blocker, they failed with "indices should be either on cpu or on the same
device as the indexed tensor". These now follow the data's device, which is
a no-op on CPU and correct on any accelerator.
(test_a2a_correctness.py already did this right via pos.device; only
test_graph_parallel.py mixed devices.)
3. Results could not cross the process boundary. spawn_multi_process returns
each rank's result through a multiprocessing.Manager dict, which pickles
it. CUDA tensors survive that via CUDA IPC; XPU has no equivalent and
raises "_share_fd_: only available on CPU". Results are now detached to
host memory before the handoff -- correct on every backend, free at test
sizes.
Also extends tests/core/models/uma/test_xpu_accelerator.py to cover the full
MLIPPredictUnit path, which is the code that used to hard-stop. It asserts the
unit lands on the accelerator *and* drives a real predict() to finite energy
and forces -- because the pre-port failure mode was not only the assert but
the silent resolve-to-CPU behind it, which an assert-only test would miss.
Device placement is lazy (_lazy_init runs on first predict), so inspecting
parameters before a prediction would report CPU on a correctly configured
unit.
Sunspot, 12 PVC tiles: tests/core/common/parallelism now 63 passed,
1 skipped, 0 failed. Was 15 failed.
Adds fairchem/core/common/device_utils.py and routes CUDA-specific call sites through it, so device selection, collective backend, memory and seeding helpers all follow the detected accelerator. CUDA behavior is unchanged.
Graph-parallel GPU tests now take the collective backend and tensor placement from the detected accelerator, and move results to host memory before crossing the process boundary since XPU has no CUDA-IPC equivalent.
assign_device_for_local_rank() autodetected the accelerator while the collective backend was chosen from the caller's requested device type. On a node whose hardware differs from the request the two diverged: the model bound to one accelerator while the backend expected another, surfacing later as "No backend type associated with device type ..." from inside DDP. Takes an optional device_type (defaulting to autodetection, so existing callers are unchanged) and threads the requested value through the paths that know it. An explicit request for absent hardware now raises instead of silently binding something else.
Three related corrections to the device layer.
memory_allocated(), max_memory_allocated() and reset_peak_memory_stats()
accept a device argument but were called with none, so an indexed spec was
silently answered for the current device instead: memory_allocated("xpu:1")
reported 0 MB while 64 MB was allocated there. They now forward the index via
the new device_index_of(). A bare device type still means "current device".
The hasattr() guards around five torch device calls are replaced by
_device_api(), which raises naming the backend and the missing call. Every
guarded attribute exists on both torch.cuda and torch.xpu, so the guards never
fired; on a future backend they would have reported 0 bytes of memory or
skipped seeding, both of which read as success.
umas_fast_gpu's Triton kernels no longer require FAIRCHEM_ENABLE_TRITON_XPU.
Triton supports Intel GPUs and the kernels are numerically correct there --
the fused-edgewise and execution-backend suites, including the
matches_pytorch reference comparisons and gradcheck, pass on XPU. The opt-in
guarded against a portability problem that does not exist, so XPU now takes
the same path as CUDA with no vendor branch.
The gpu and compile_gpu markers gated on torch.cuda.is_available(), and the test bodies hard-coded device="cuda" regardless, so the whole GPU suite was unrunnable on any other accelerator. The markers now admit any accelerator, and each affected module resolves an ACCELERATOR constant from the hardware present. Beyond the string literals, this also covers the idioms a literal search misses: .cuda() method calls, torch.cuda.* guards that silently skipped seeding or cache release, torch.load of CUDA-saved tensors (now map_location), and Ray reservations that requested zero GPUs off NVIDIA. Adds a cuda_only marker for tests that genuinely depend on NVIDIA, so those skip rather than fail. It is deliberately unused so far: the Triton kernels were the expected case and they turned out to be correct on XPU. On Intel Data Center GPU Max the gpu-marked suite goes from 60 failed to 74 passed, 0 failed. Notably test_full_train_eval_from_cli_aselmdb_gpu matches its expected_loss reference (13.66177, rel_tol 1e-4) computed on NVIDIA, so training converges identically across the two backends. test_rotation_cuda_graph.py keeps its .cuda() calls: that module is skipped unconditionally today, so converting it would be unverifiable.
|
Hi @abagusetty! Thank you for your pull request and welcome to our community. Action RequiredIn order to merge any pull request (code, docs, etc.), we require contributors to sign our Contributor License Agreement, and we don't seem to have one on file for you. ProcessIn order for us to review and merge your suggested changes, please sign at https://code.facebook.com/cla. If you are contributing on behalf of someone else (eg your employer), the individual CLA may not be sufficient and your employer may need to sign the corporate CLA. Once the CLA is signed, our tooling will perform checks and validations. Afterwards, the pull request will be tagged with If you have received this in error or have any questions, please contact us at cla@meta.com. Thanks! |
|
Thank you for signing our Contributor License Agreement. We can now accept your code for this (and any) Meta Open Source project. Thanks! |
zulissimeta
added
enhancement
…kernels Resolves the merge in favour of the accelerator-agnostic device layer, then carries that treatment to the Triton work upstream added in facebookresearch#2154: packed_gate required CUDA tensors outright, and the new tests hard-coded CUDA devices. Both now follow whatever accelerator is present. Upstream's execution_backends.py is kept verbatim apart from dropping the XPU Triton opt-in; an earlier resolution had discarded its _dense_l2_wigner and pack_blocks helpers, which broke the compact Wigner layout its kernels expect. Verified on Intel Data Center GPU Max: the gpu-marked suite is 93 passed, 0 failed, including the packed-gate and compact-Wigner tests that compare against the materialized reference implementations and the torch.compile dynamic-shape paths.
Takes upstream's graph-parallel test fixes (facebookresearch#2165) and generalizes its _requires_gpus helper from CUDA device counts to whichever accelerator is present, so the rank-per-device gating works on Intel GPUs too.
abagusetty
force-pushed
the
xpu-support
branch
from
September 16, 2026 17:30
b89f2a2 to
52d93db
Compare
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
2 participants
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.
cudatogeneric-deviceshelps to systematically support CUDA and XPU leaving room for future target or rocm devicesDisclaimer: LLMs were used in preparing the PR