[JazyJiang]

Fixes #6480

Related to #6455, which replaces hardcoded device selection with the accelerator API in colossalai/fx and colossalai/inference. No file overlap with this PR.

1. Background and problem

Current behaviour. TorchFSDPPlugin.configure() passes torch.cuda.current_device() as the FSDP device_id.

Problem. On any accelerator other than CUDA, booster.boost() raises AssertionError: Torch not compiled with CUDA enabled before FSDP is built, so the plugin cannot be used at all. This is the only torch.cuda call left under colossalai/booster/plugin; LowLevelZeroPlugin already goes through colossalai.accelerator.

Reproduce. Construct any model, wrap it with TorchFSDPPlugin and call booster.boost() on a non-CUDA build. Full snippet in the linked issue. Reproduced on an Ascend 910B2.

Expected behaviour. The device comes from the accelerator the framework has already detected.

Why now. The failure is invisible on CUDA, so it survives review and CI. The NPU port hit it on the first run of the plugin.

2. What changed

• colossalai/booster/plugin/torch_fsdp_plugin.py: use get_accelerator().get_current_device() for the FSDP device_id, and import get_accelerator.
• .github/scripts/ci/check_no_hardcoded_device.py (new): fails when a guarded directory references torch.cuda, .cuda() or a hardcoded nccl backend. Reports file, line, the offending source and the suggested replacement.
• .pre-commit-config.yaml: register it as the no-hardcoded-device hook, so it runs in the existing Layer 1 pre-commit job. No GPU and no new workflow.

Not included in this PR:

• Cleaning up the other hardcoded device references. Outside colossalai/legacy and colossalai/accelerator there are still ~130; only colossalai/booster/plugin is added to GUARDED. Locking everything now would turn main red. Directories go in one at a time as they are cleaned.
• The second FSDP blocker on non-CUDA accelerators (see section 6).
• Any NPU-specific code. Nothing in this PR is Ascend-specific.

3. User and compatibility impact

Area	Impact
Python API	None
Checkpoint format	None
Default behaviour	None on CUDA. CudaAccelerator.get_current_device() returns torch.device(f"cuda:{torch.cuda.current_device()}"), i.e. the same device, as a torch.device rather than an int. FSDP's device_id accepts both.
Model accuracy	None
Performance / memory	None
Python / PyTorch / CUDA versions	None
Custom C++/CUDA operators	None
Docker images	None
Existing training scripts	None

New contributor-facing behaviour: pre-commit now fails if someone adds a hardcoded device under colossalai/booster/plugin. The message names the file, line and replacement, and # allow-hardcoded-device: <reason> is the opt-out.

4. Test environment

Verified on Ascend NPU only so far. The template's NVIDIA fields do not apply; the Ascend equivalents are given instead.

OS	openEuler 22.03 LTS-SP2, Linux 5.10.0, aarch64, glibc 2.38
Python	3.12.13
Accelerator	8x Ascend 910B2 (64GB), 1-2 used
Ascend HDK (driver)	25.2.1, ascendhal 7.35.23
CANN	9.1.0 (V100R001C11SPC001B243)
PyTorch	2.13.0
torch_npu	2.13.0.rc1
Transformers	5.17.0
flash-attn	not installed
ColossalAI commit	cb3b223 (main) + this change
Container image	ascendai/torch-npu:2.13.0rc1-cann9.1.0-910b-openeuler24.03-py3.12

NVIDIA Driver / torch.version.cuda / nvcc: not applicable on this host.

5. Test commands and results

5.1 The hook

python .github/scripts/ci/check_no_hardcoded_device.py

• On main: exit code 1, reports colossalai/booster/plugin/torch_fsdp_plugin.py:540: torch.cuda.*
• With this change: exit code 0, no hardcoded device found in 1 guarded directory
• Re-introducing the old line makes it fail again (checked by reverting the line, running, and restoring it).

No GPU/NPU needed; runs on any machine.

5.2 Repository plugin tests on NPU

tests/test_booster/test_plugin/test_torch_fsdp_plugin.py passes on an Ascend 910B2 with this change:

FAST_TEST=1 pytest tests/test_booster/test_plugin/test_torch_fsdp_plugin.py
1 passed in 56.18s   (2 processes)

Reverting only the one line in this PR makes that same test fail at torch_fsdp_plugin.py, at device_id=torch.cuda.current_device(). That is the evidence for this change.

Three things outside this PR were needed to get the repository's tests running on that machine at all, and are not proposed here: installing torchvision/torchaudio from the CPU wheel index (the aarch64 wheels on PyPI link libcudart.so.13), replacing seven hardcoded .to("cuda") / torch.cuda.empty_cache() calls in the four plugin test files, and dropping a @torch.jit.script decorator in the ChatGLM modeling code that fails the TorchScript fuser on a non-CUDA device. The other three plugin tests also pass once those are in place (TorchDDP, LowLevelZero, Gemini).

5.3 Standalone plugin smoke on NPU

A small script building a tiny Llama (2 layers, hidden 128) and running fwd/bwd/step twice per plugin, one plugin per process:

Plugin	Before	After
TorchDDP	PASS	PASS, unchanged
LowLevelZeRO-1 / -2	PASS	PASS, unchanged
TorchFSDP	FAIL at torch.cuda.current_device()	gets past device selection

1 and 2 NPUs, bf16. A separate check confirmed 2-card DDP gradients match a single process on the same global batch (relative 5.2e-7 over 590k elements), so the collective path is exercised.

5.4 GPU regression

Run on CUDA through the repository's layered CI (Layer 2) on gpu-h20-5, inside the qualified dependency image sha256:ae8ff34bf340a692803edca6e0f38a576fc1a2b206c01b4ef6f7c2a40dfd8c10, torch 2.13.0+cu130 / CUDA 13.0, on NVIDIA H20 (sm 9.0):

Batch	Scope	Result
2 (02-single-gpu-core, 1 GPU)	Accelerator, native FP16, LazyTensor, inference data structures	23 passed, 2 skipped (262s)
4a (1 GPU)	Container smoke	22 passed (3.8s)
5 (05-booster-lora-2gpu, 2 GPUs)	test_dp_plugin_base.py, test_torch_ddp_plugin.py, test_torch_fsdp_plugin.py, test_low_level_zero_plugin.py, test_lora.py	5 passed (554s)

Batch 5 is the batch that covers this change: test_torch_fsdp_plugin.py::test_torch_fsdp_plugin passed in 118.40s on 2 GPUs.

Run: https://github.com/hpcaitech/ColossalAI/actions/runs/40308859635

That is attempt 3. The two earlier attempts executed no tests at all: they stopped at the runner's idle-GPU check with Need N idle GPU(s), found 0, which reports host capacity rather than anything about the code.

What CI tested is GitHub's merge commit 7e9551dec40fa57003be73de761950fc736aad9c, that is this PR's head 3fa26fe5 merged into base cb3b2234 (the base recorded in section 4). main has moved on since, and this has not been re-run against the newer base.

Still unverified on CUDA: the FSDP + AdamW path described in section 6.1, since batch 5's FSDP test constructs SGD; gradient checkpointing, lazy init and checkpoint save/load, which live in other batches; and fp16.

6. Known issues

6.1 TorchFSDP + AdamW is broken on torch 2.13, on every device

• Symptom: configure() rebuilds the optimizer with optimizer.__init__(fsdp_model.parameters(), **optimizer.defaults). Under torch 2.13 AdamW.defaults carries decoupled_weight_decay, which AdamW.__init__ does not accept, so this raises TypeError.

• Scope: specific to optimizers whose defaults do not match their __init__ signature. test_torch_fsdp_plugin.py uses SGD and is unaffected, which is why it passes above.

• Introduced by this PR: no, and not accelerator-specific.

• Also affects CUDA: yes. In wangzhijian/colossalai-ci:torch2.13-cu130-tf5.17.0-exp-20261019:

  torch 2.13.0+cu130
  AdamW.__init__ accepts decoupled_weight_decay: False
  defaults contains it: True
  rebuild: FAIL -> AdamW.__init__() got an unexpected keyword argument 'decoupled_weight_decay'
  

• Handling: out of scope here, to keep this PR to one change.

• Issue: to be filed

• Owner / timing: to be assigned

6.2 The repository's tests assume a CUDA environment

Getting them to run on a non-CUDA accelerator needed the three changes listed in 5.2. Two of those are device assumptions in the repository itself: seven hardcoded CUDA calls in the four plugin test files, and a @torch.jit.script decorator whose fuser rejects the device. tests/ still contains roughly 900 hardcoded CUDA references across 206 files.

• Introduced by this PR: no.
• Handling: none here. This PR guards colossalai/booster/plugin only; tests/ is a much larger cleanup that deserves its own discussion rather than riding along.

7. Risk and rollback

Risk. Low. One line, and the replacement resolves to the same CUDA device. The hook could annoy contributors if a legitimate device-specific line lands in a guarded directory; the # allow-hardcoded-device: opt-out covers that.

The hook runs on whole directories rather than changed files, so an unrelated PR that adds a hardcoded device under colossalai/booster/plugin will fail pre-commit rather than the author's own files. That is the intent, but it does mean the failure can appear in a PR that did not expect it.

Rollback. Revert the commit. The two parts are independent: dropping the hook (the .pre-commit-config.yaml entry plus the script) leaves the one-line fix in place, and vice versa.

8. Documentation and release impact

README	No
Installation docs	No
API docs	No
Docker docs	No
Changelog	No
version.txt	No
PyPI	No
DockerHub image	No

The hook's own usage and the opt-out marker are documented in the script docstring.