From 3fa26fe518a5de5a54fa5a0342c87a676e5feec6 Mon Sep 17 00:00:00 2001 From: JazyJiang Date: Mon, 26 Oct 2026 16:34:11 +0800 Subject: [PATCH] [booster] use the accelerator API for the TorchFSDP device id TorchFSDPPlugin.configure() passed torch.cuda.current_device() as the FSDP device_id, so on a non-CUDA accelerator booster.boost() raised "AssertionError: Torch not compiled with CUDA enabled" before FSDP was even built. Use get_accelerator().get_current_device() instead, which is what LowLevelZeroPlugin already does. This was the only torch.cuda call under colossalai/booster/plugin. Add a pre-commit hook, no-hardcoded-device, that fails when a guarded directory references torch.cuda, .cuda() or a hardcoded nccl backend, so the same mistake does not come back. Only colossalai/booster/plugin is guarded: the rest of the tree still has around 130 such references outside colossalai/legacy and colossalai/accelerator, and directories can be added to GUARDED as they are cleaned up. A line that is genuinely device-specific can opt out with '# allow-hardcoded-device: '. Checked on an Ascend 910B2: the hook reports the call on main and passes with this change, and TorchFSDPPlugin now gets past device selection. --- .../scripts/ci/check_no_hardcoded_device.py | 90 +++++++++++++++++++ .pre-commit-config.yaml | 9 ++ .../booster/plugin/torch_fsdp_plugin.py | 3 +- 3 files changed, 101 insertions(+), 1 deletion(-) create mode 100644 .github/scripts/ci/check_no_hardcoded_device.py 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