fx tracer / rope init hard-code the device, so NPU runs fall back to CPU (or crash)
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()return torch.device("cuda:0") if torch.cuda.is_available() else torch.device("cpu")colossalai/fx/passes/meta_info_prop.py::metainfo_trace()— same expressioncolossalai/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:
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 enabledget_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 enabledExpected 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.
Source: hpcaitech/ColossalAI