diff --git a/.github/scripts/ci/check_no_hardcoded_device.py b/.github/scripts/ci/check_no_hardcoded_device.py
new file mode 100644
index 000000000000..ab81a8903862
--- /dev/null
+++ b/.github/scripts/ci/check_no_hardcoded_device.py
@@ -0,0 +1,90 @@
+#!/usr/bin/env python3
+"""Fail if guarded directories hardcode a specific accelerator or communication backend.
+
+Device-specific branches belong in the accelerator, communication backend and extension
+layers only. Elsewhere they fail silently: the code works on CUDA and breaks the first time
+it runs anywhere else. For example, TorchFSDPPlugin.configure() used to pass
+``torch.cuda.current_device()`` as the FSDP ``device_id``, so ``booster.boost()`` raised
+``AssertionError: Torch not compiled with CUDA enabled`` on an Ascend NPU.
+
+Directories are locked down one at a time: clean one up, add it to GUARDED, and this keeps
+it clean. Much of the tree still hardcodes CUDA (mostly under ``colossalai/legacy``) and is
+deliberately out of scope.
+
+Usage:
+ python .github/scripts/ci/check_no_hardcoded_device.py # check GUARDED
+ python .github/scripts/ci/check_no_hardcoded_device.py
... # check given dirs
+
+To keep a genuinely device-specific line, mark it with a reason:
+ some_call(torch.cuda.current_device()) # allow-hardcoded-device: CUDA-only FP8 path
+"""
+
+from __future__ import annotations
+
+import re
+import sys
+from pathlib import Path
+
+# Directories that have been cleaned up and must stay device-agnostic, relative to the repo root.
+GUARDED = [
+ "colossalai/booster/plugin",
+]
+
+# (pattern, what was found, what to use instead)
+PATTERNS = [
+ (re.compile(r"\btorch\.cuda\."), "torch.cuda.*", "the matching get_accelerator() method"),
+ (re.compile(r"\.cuda\(\)"), ".cuda()", ".to(get_accelerator().get_current_device())"),
+ (re.compile(r"""(? list[tuple[Path, int, str, str, str]]:
+ """Return (path, lineno, source line, what was found, suggested fix) for every violation."""
+ violations = []
+ for path in sorted(directory.rglob("*.py")):
+ try:
+ lines = path.read_text(encoding="utf-8").splitlines()
+ except (OSError, UnicodeDecodeError) as exc:
+ # Report instead of skipping silently: an unreadable file is not a passing file.
+ print(f"error: cannot read {path}: {exc}", file=sys.stderr)
+ raise SystemExit(2)
+ for lineno, line in enumerate(lines, 1):
+ if ALLOW.search(line):
+ continue
+ for pattern, found, fix in PATTERNS:
+ if pattern.search(line):
+ violations.append((path, lineno, line.strip(), found, fix))
+ return violations
+
+
+def main() -> int:
+ targets = sys.argv[1:] or GUARDED
+ missing = [t for t in targets if not (REPO_ROOT / t).is_dir()]
+ if missing:
+ print(f"error: no such directory: {', '.join(missing)} (repo root {REPO_ROOT})", file=sys.stderr)
+ return 2
+
+ violations = [v for target in targets for v in scan(REPO_ROOT / target)]
+ if not violations:
+ print(f"no hardcoded device found in {len(targets)} guarded director{'y' if len(targets) == 1 else 'ies'}")
+ return 0
+
+ print(f"Found {len(violations)} hardcoded device reference(s) in guarded directories:\n", file=sys.stderr)
+ for path, lineno, line, found, fix in violations:
+ print(f" {path.relative_to(REPO_ROOT)}:{lineno}: {found}", file=sys.stderr)
+ print(f" {line}", file=sys.stderr)
+ print(f" use {fix} instead\n", file=sys.stderr)
+ print(
+ "Device-specific branches belong in the accelerator, communication backend and\n"
+ "extension layers only. To keep a line that is genuinely device-specific, mark it\n"
+ "with '# allow-hardcoded-device: '.",
+ file=sys.stderr,
+ )
+ return 1
+
+
+if __name__ == "__main__":
+ sys.exit(main())
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
index 2c7073b47bf1..be2df68891ca 100644
--- a/.pre-commit-config.yaml
+++ b/.pre-commit-config.yaml
@@ -38,3 +38,12 @@ repos:
- id: end-of-file-fixer
- id: mixed-line-ending
args: ['--fix=lf']
+
+ - repo: local
+ hooks:
+ - id: no-hardcoded-device
+ name: no hardcoded device in guarded dirs
+ entry: python .github/scripts/ci/check_no_hardcoded_device.py
+ language: system
+ pass_filenames: false
+ always_run: true
diff --git a/colossalai/booster/plugin/torch_fsdp_plugin.py b/colossalai/booster/plugin/torch_fsdp_plugin.py
index 05cc432a673d..5c58e6fe9fca 100644
--- a/colossalai/booster/plugin/torch_fsdp_plugin.py
+++ b/colossalai/booster/plugin/torch_fsdp_plugin.py
@@ -25,6 +25,7 @@
from torch.optim.lr_scheduler import _LRScheduler as LRScheduler
from torch.utils.data import DataLoader
+from colossalai.accelerator import get_accelerator
from colossalai.checkpoint_io import CheckpointIndexFile, CheckpointIO, GeneralCheckpointIO, utils
from colossalai.checkpoint_io.utils import (
async_save_state_dict_shards,
@@ -537,7 +538,7 @@ def configure(
lr_scheduler: Optional[LRScheduler] = None,
) -> Tuple[nn.Module, OptimizerWrapper, Callable, DataLoader, LRScheduler]:
# wrap the model with PyTorch FSDP
- fsdp_model = TorchFSDPModel(model, device_id=torch.cuda.current_device(), **self.fsdp_kwargs)
+ fsdp_model = TorchFSDPModel(model, device_id=get_accelerator().get_current_device(), **self.fsdp_kwargs)
if self.fp8_communication:
from colossalai.quantization.utils import patch_fsdp_params_comm_hook