Make training device-portable (XPU/CUDA/CPU)#4
Open
saforem2 wants to merge 1 commit into
Open
Conversation
Adds a small `utils.device` helper that auto-detects the active torch
accelerator (cuda / xpu / cpu) and exposes:
- `set_device(local_rank)` / `synchronize()` for cross-device parity
- `dist_backend()` returning nccl / ccl / gloo as appropriate
- no-op shims for cuda nvtx + profiler so the training loop runs
unchanged on non-cuda accelerators
`train.py` now uses these helpers in place of hard-coded `torch.cuda.*`
calls — CUDA hosts behave identically (nvtx ranges and profiler hooks
still real), while XPU/CPU hosts no-op them. AMP autocast / GradScaler
device strings and the DDP backend are picked from the active device.
Distributed init falls back to `ezpz.setup_torch()` when WORLD_SIZE
isn't already populated (mpiexec on PALS clusters like Aurora/Sunspot
don't set env vars the way SLURM/torchrun do). srun (Perlmutter) and
torchrun paths are unchanged.
Also fixes `utils.logging_utils.slurm_filter`: it crashed with KeyError
outside SLURM. Now falls back to RANK / PMI_RANK / PMIX_RANK, so the
"only rank 0 logs" behavior works under mpiexec too.
Tested on Sunspot (Intel XPU, PALS): 12 XPUs/node, BF16, pytorch
dataloader. Loss and throughput stable across epochs.
There was a problem hiding this comment.
Pull request overview
Note
Copilot was unable to run its full agentic suite in this review.
This PR improves portability across different accelerator types (CUDA/XPU/CPU) and launchers (SLURM/torchrun/MPI) by centralizing device selection/synchronization and making rank-aware logging/distributed init more robust.
Changes:
- Extend
slurm_filterto recognize multiple rank environment variables (SLURM/PMI/PMIx/torchrun). - Add
utils/device.pyto abstract device detection, synchronization, distributed backend choice, and profiling shims. - Update
train.pyto use the new device abstraction and to support non-env://(MPI/PALS) distributed setup viaezpz.
Reviewed changes
Copilot reviewed 3 out of 3 changed files in this pull request and generated 4 comments.
| File | Description |
|---|---|
| utils/logging_utils.py | Makes logging rank detection work across more launchers/environments. |
| utils/device.py | New device abstraction layer for CUDA/XPU/CPU + profiling/sync shims + backend selection. |
| train.py | Uses the new device helpers and updates distributed initialization logic for env- and MPI-based setups. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Comment on lines
+7
to
+9
| import contextlib | ||
| import os | ||
| import torch |
|
|
||
| if params.amp_dtype == torch.float16: | ||
| scaler = GradScaler("cuda") | ||
| scaler = GradScaler(_device.DEVICE_TYPE) |
Comment on lines
+356
to
+363
| if os.environ.get("WORLD_SIZE", "").isdigit(): | ||
| world_size = int(os.environ["WORLD_SIZE"]) | ||
| world_rank = int(os.environ.get("RANK", "0")) | ||
| local_rank = int(os.environ.get("LOCAL_RANK", "0")) | ||
| if world_size > 1: | ||
| torch.distributed.init_process_group( | ||
| backend=_device.dist_backend(), init_method="env://" | ||
| ) |
Comment on lines
+46
to
+51
| def dist_backend() -> str: | ||
| if DEVICE_TYPE == "cuda": | ||
| return "nccl" | ||
| if DEVICE_TYPE == "xpu": | ||
| return "ccl" | ||
| return "gloo" |
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
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.
Summary
Adds a small
utils.devicehelper that auto-detects the active torch accelerator (cuda / xpu / cpu) and exposes:set_device(local_rank)/synchronize()for cross-device paritydist_backend()returning nccl / ccl / gloo as appropriatetrain.pynow uses these helpers in place of hard-coded `torch.cuda.*` calls — CUDA hosts behave identically (nvtx ranges and profiler hooks still real), while XPU/CPU hosts no-op them. AMP autocast / GradScaler device strings and the DDP backend are picked from the active device.Distributed init falls back to `ezpz.setup_torch()` when WORLD_SIZE isn't already populated (mpiexec on PALS clusters like Aurora/Sunspot don't set env vars the way SLURM/torchrun do). srun (Perlmutter) and torchrun paths are unchanged — they hit the fast path.
Also fixes `utils.logging_utils.slurm_filter`: it crashed with KeyError outside SLURM. Now falls back to RANK / PMI_RANK / PMIX_RANK so the "only rank 0 logs" behavior works under mpiexec too.
Scope
train.pyand `utils/logging_utils.py` only. `train_mp.py` / `train_mp_graphs.py` and `networks/vit_te.py` (Transformer Engine) are unchanged — the model-parallel + TE paths still target CUDA exclusively.Test plan
Notes
The ezpz dependency is only imported lazily, on the non-SLURM/torchrun branch. CUDA + SLURM users won't see a new requirement.