#6454·ColossalAI

fx tracer / rope init hard-code the device, so NPU runs fall back to CPU (or crash)

Author: li-lizheCreated Sep 26, 2026Updated Sep 26, 2026

Describe the bug

ColossalAI already ships an accelerator abstraction (colossalai.accelerator.get_accelerator(), including an NpuAccelerator) and most of the code base uses it, but a few spots still pick the device by hand with a CUDA-or-CPU expression:

  • colossalai/fx/tracer/experimental.py::default_device()
    python
    return torch.device("cuda:0") if torch.cuda.is_available() else torch.device("cpu")
  • colossalai/fx/passes/meta_info_prop.py::metainfo_trace() — same expression
  • colossalai/inference/utils.py::init_to_get_rotary() — unconditional .cuda() on the rope cache

On an Ascend NPU host torch.cuda.is_available() is False, so default_device() returns cpu: the FX tracer silently puts the graph module and the sample inputs on the CPU instead of the NPU, i.e. it never uses the accelerator even though get_accelerator() correctly reports npu:0. init_to_get_rotary() is worse — it raises instead of degrading.

Reproduction

Ascend 910B4, torch 2.15.0.dev20260917+cpu + torch_npu, CANN, single process:

python
import torch, torch_npu
from colossalai.accelerator import get_accelerator
print(get_accelerator().name, get_accelerator().get_current_device())   # npu npu:0
from colossalai.fx.tracer.experimental import default_device
print(default_device())                                                 # cpu  <-- should be npu:0
from colossalai.inference.utils import init_to_get_rotary
class Cfg: hidden_size = 64; num_attention_heads = 8; max_position_embeddings = 128
class M: dtype = torch.float16
m = M(); m.config = Cfg()
init_to_get_rotary(m)   # AssertionError: Torch not compiled with CUDA enabled
get_accelerator().name = npu
get_accelerator().get_current_device() = npu:0
default_device() = cpu
  File ".../colossalai/inference/utils.py", line 63, in init_to_get_rotary
    self._cos_cached = torch.cos(freqs).to(self.dtype).cuda()
AssertionError: Torch not compiled with CUDA enabled

Expected behavior

The device should come from the accelerator abstraction: get_accelerator().get_current_device() returns npu:0 on Ascend, cuda:<current device> on CUDA and cpu on a CPU-only box, so CUDA behaviour is unchanged while NPU/XPU stop falling back to the CPU.

A PR that does exactly that for the three call sites above is on its way.