[JazyJiang]

Fixes #6480

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 Plugin smoke on NPU

Not pytest: the repo's own plugin tests cannot be collected in the Ascend image (see section 6). This is a standalone script that builds a tiny Llama (2 layers, hidden 128) and runs fwd/bwd/optimizer step twice per plugin, each plugin in its own process.

ASCEND_RT_VISIBLE_DEVICES=0 python -m torch.distributed.run \
  --nproc_per_node 1 --master_port 29531 plugin_smoke.py --plugin <name> --dtype bf16

Plugin	Before	After
TorchDDP	PASS (loss 6.923661, 6.976354)	PASS, unchanged
LowLevelZeRO-1	PASS (loss 6.923819, 6.976377)	PASS, unchanged
LowLevelZeRO-2	PASS (loss 6.923818, 6.976378)	PASS, unchanged
Gemini	FAIL (unrelated, see section 6)	FAIL, unchanged
TorchFSDP	FAIL at torch.cuda.current_device()	gets past device selection, then fails on the optimizer rebuild (section 6)

1 NPU, bf16, no FAST_TEST equivalent, no skips, no reruns. Script and raw JSON results can be attached on request.

5.3 GPU regression — NOT RUN

Not run on CUDA yet. Needed before merge:

pytest tests/test_booster/test_plugin/test_torch_fsdp_plugin.py
pytest tests/test_booster/test_plugin/test_low_level_zero_plugin.py
pytest tests/test_booster/test_plugin/test_torch_ddp_plugin.py

(batch 5; 2 GPUs). The hook itself needs no GPU, so only the one-line plugin change is at risk. Results will be posted here before this is marked ready.

The only evidence so far comes from a non-CUDA accelerator, so this PR should not be treated as verified until the above runs.

6. Known issues

6.1 FSDP still does not run end to end on a non-CUDA accelerator

• Test: TorchFSDP plugin smoke on Ascend 910B2

• Symptom: after this fix, configure() fails one step later at optimizer.__init__(fsdp_model.parameters(), **optimizer.defaults) with TypeError: AdamW.__init__() got an unexpected keyword argument 'decoupled_weight_decay'

• Introduced by this PR: no. It is a torch 2.13 change: AdamW.defaults carries decoupled_weight_decay, which AdamW.__init__ does not accept.

• Also affects CUDA: yes. Reproduced in the CI image 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 Repo plugin tests cannot be collected in the Ascend image

• Test: tests/test_booster/test_plugin/test_torch_fsdp_plugin.py
• Symptom: ModuleNotFoundError: No module named 'timm' at collection. Installing timm pulls torch 2.14.1 and breaks torch_npu; installing it with --no-deps then fails on torchvision, which has no matching build in this image.
• Introduced by this PR: no. The test dependencies assume the CUDA image.
• Handling: none here. It is a prerequisite for any future NPU CI, not for this change, which is verified by 5.1 and 5.3 on CUDA.

6.3 Gemini fails on NPU

Unrelated to this PR and not caused by it: colo_tensor.py:81 raises RuntimeError: The tensor has a non-zero number of elements, but its data is not allocated yet. during backward. Recorded here only so the smoke table above is not misread. Not yet confirmed whether it is accelerator-specific.

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.

[JazyJiang]

Superseded by #6482, which carries the same commit (3fa26fe) from a branch in this repository instead of a fork.

Reason: the GPU batches in Layer 2 require head.repo.full_name == github.repository, so on a fork PR they are skipped and this change — which touches tests/test_booster/test_plugin/test_torch_fsdp_plugin.py's subject — would have had no GPU coverage at all. Closing here to keep the discussion in one place.