fix: place collectives and new modules on the process accelerator

src/mcore_bridge/bridge/gpt_bridge.py

@@ -18,6 +18,7 @@
 from mcore_bridge.tuners import LoraParallelLinear
 from mcore_bridge.utils import (MxFp4Dequantizer, PackedDequantizer, SafetensorLazyLoader, StreamingSafetensorSaver,
                                 deep_getattr, gc_collect, get_logger, is_master, unwrap_model)
+from mcore_bridge.utils.accelerator import accelerator_device
 from mcore_bridge.utils.constants import EXPORT_CHUNK_BYTES
 
 logger = get_logger()
@@ -321,7 +322,7 @@ def _set_module(self, mg_module, hf_state_dict, hf_prefix: str, to_mcore: bool):
             if self.pp_size > 1:
                 src_rank = torch.tensor([0 if hf_state_dict is None else self.pp_rank],
                                         dtype=torch.int64,
-                                        device='cuda')
+                                        device=accelerator_device())
                 dist.all_reduce(src_rank, group=self.pp_group)
                 src_rank = dist.get_global_rank(self.pp_group, src_rank.item())
                 meta_data = [None] if hf_state_dict is None else [list(hf_state_dict.keys())]
@@ -426,16 +427,16 @@ def _chunked_broadcast_pp(self, tensor, shape, dtype, src_rank: int, pp_group):
             for start in range(0, shape[0], rows):
                 end = min(shape[0], start + rows)
                 send = tensor[start:end]
-                if not send.is_cuda or send.dtype != dtype or not send.is_contiguous():
-                    send = send.to(device='cuda', dtype=dtype).contiguous()
+                if send.device.type != accelerator_device().type or send.dtype != dtype or not send.is_contiguous():
+                    send = send.to(device=accelerator_device(), dtype=dtype).contiguous()
                 dist.broadcast(send, src=src_rank, group=pp_group)
             return tensor
         output = torch.empty(shape, dtype=dtype, device='cpu')
         buf = None
         for start in range(0, shape[0], rows):
             end = min(shape[0], start + rows)
             if buf is None or buf.shape[0] != end - start:
-                buf = torch.empty([end - start] + list(shape[1:]), device='cuda', dtype=dtype)
+                buf = torch.empty([end - start] + list(shape[1:]), device=accelerator_device(), dtype=dtype)
             dist.broadcast(buf, src=src_rank, group=pp_group)
             output[start:end] = buf.cpu()
         return output
@@ -446,10 +447,10 @@ def _broadcast_ep_pp(self, tensor, is_expert):
         pp_rank = self.ep_pp_rank if is_expert else self.pp_rank
         # pp/ep
         if pp_size > 1:
-            src_rank = torch.tensor([0 if tensor is None else pp_rank], dtype=torch.int64, device='cuda')
+            src_rank = torch.tensor([0 if tensor is None else pp_rank], dtype=torch.int64, device=accelerator_device())
             dist.all_reduce(src_rank, group=pp_group)
             src_rank = dist.get_global_rank(pp_group, src_rank.item())
-            meta_data = torch.zeros(10, dtype=torch.int64, device='cuda')
+            meta_data = torch.zeros(10, dtype=torch.int64, device=accelerator_device())
             dtype_mapping = [torch.float64, torch.float32, torch.float16, torch.bfloat16, torch.uint8, torch.int32]
             dtype_mapping_r = {v: k for k, v in enumerate(dtype_mapping)}
             if tensor is None:
@@ -460,11 +461,12 @@ def _broadcast_ep_pp(self, tensor, is_expert):
                 elem_size = torch.empty((), dtype=dtype).element_size()
                 if self._stream_to_cpu() and numel * elem_size > self.export_chunk_bytes and len(shape) > 0:
                     return self._chunked_broadcast_pp(None, shape, dtype, src_rank, pp_group)
-                tensor = torch.empty(shape, device='cuda', dtype=dtype)
+                tensor = torch.empty(shape, device=accelerator_device(), dtype=dtype)
                 dist.broadcast(tensor, src=src_rank, group=pp_group)
             else:
                 meta_data[0] = tensor.ndim
-                meta_data[1:1 + tensor.ndim] = torch.tensor(tensor.shape, dtype=torch.int64, device='cuda')
+                meta_data[1:1 + tensor.ndim] = torch.tensor(
+                    tensor.shape, dtype=torch.int64, device=accelerator_device())
                 meta_data[-1] = dtype_mapping_r[tensor.dtype]
                 dist.broadcast(meta_data, src=src_rank, group=pp_group)
                 if self._stream_to_cpu() and tensor.numel() * tensor.element_size() > (
@@ -559,7 +561,7 @@ def _set_state_dict(self,
         is_lora = isinstance(sub_module, LoraParallelLinear)
         is_modules_to_save = isinstance(sub_module, ModulesToSaveWrapper)
         if not to_mcore:
-            state = torch.tensor([is_lora, is_modules_to_save], dtype=torch.bool, device='cuda')
+            state = torch.tensor([is_lora, is_modules_to_save], dtype=torch.bool, device=accelerator_device())
             if is_expert and self.ep_pp_size > 1:
                 dist.all_reduce(state, group=self.ep_pp_group)
             elif not is_expert and self.pp_size > 1:
@@ -635,7 +637,7 @@ def _filter_prefix(state_dict, prefix: str):
     def _reduce_tensor_pp_group(self, tensor, to_mcore, dtype=torch.bool, op=dist.ReduceOp.MAX):
         if to_mcore:
             return tensor
-        tensor = torch.tensor([tensor], dtype=dtype, device='cuda')
+        tensor = torch.tensor([tensor], dtype=dtype, device=accelerator_device())
         if self.pp_size > 1:
             dist.all_reduce(tensor, group=self.pp_group, op=op)
         tensor = tensor.item()
@@ -695,7 +697,7 @@ def _set_qkv(self, mg_attn, hf_state_dict, to_mcore: bool, **kwargs):
             kv_block = kv_dim // self.fp8_block_size
             is_lora = False if mg_attn is None else isinstance(mg_attn.linear_qkv,
                                                                LoraParallelLinear) and self._peft_format
-            is_lora = torch.tensor([is_lora], dtype=torch.bool, device='cuda')
+            is_lora = torch.tensor([is_lora], dtype=torch.bool, device=accelerator_device())
             if self.pp_size > 1:
                 dist.all_reduce(is_lora, group=self.pp_group)
             if is_lora:
@@ -1064,7 +1066,7 @@ def _set_mlp_state(
         else:
             is_lora = False if mg_mlp is None else isinstance(mg_mlp.linear_fc1,
                                                               LoraParallelLinear) and self._peft_format
-            is_lora = torch.tensor([is_lora], dtype=torch.bool, device='cuda')
+            is_lora = torch.tensor([is_lora], dtype=torch.bool, device=accelerator_device())
             if is_expert and self.ep_pp_size > 1:
                 dist.all_reduce(is_lora, group=self.ep_pp_group)
             elif not is_expert and self.pp_size > 1:
@@ -1302,7 +1304,7 @@ def _set_mlp_state(
             else:
                 is_lora = False if mg_mlp is None else isinstance(mg_mlp.linear_fc2,
                                                                   LoraParallelLinear) and self._peft_format
-                is_lora = torch.tensor([is_lora], dtype=torch.bool, device='cuda')
+                is_lora = torch.tensor([is_lora], dtype=torch.bool, device=accelerator_device())
                 if is_expert and self.ep_pp_size > 1:
                     dist.all_reduce(is_lora, group=self.ep_pp_group)
                 elif not is_expert and self.pp_size > 1:
@@ -1451,7 +1453,7 @@ def _set_linear_decoupled_in_proj(self, mg_attn, hf_state_dict, to_mcore: bool):
             qkv_dim = key_dim * 2 + value_dim
             is_lora = False if mg_attn is None else isinstance(mg_attn.in_proj_qkvz,
                                                                LoraParallelLinear) and self._peft_format
-            is_lora = torch.tensor([is_lora], dtype=torch.bool, device='cuda')
+            is_lora = torch.tensor([is_lora], dtype=torch.bool, device=accelerator_device())
             if self.pp_size > 1:
                 dist.all_reduce(is_lora, group=self.pp_group)
             if is_lora:
@@ -1518,7 +1520,7 @@ def _set_linear_decoupled_in_proj(self, mg_attn, hf_state_dict, to_mcore: bool):
             a_dim = config.linear_num_value_heads // num_key_heads
             is_lora = False if mg_attn is None else isinstance(mg_attn.in_proj_ba,
                                                                LoraParallelLinear) and self._peft_format
-            is_lora = torch.tensor([is_lora], dtype=torch.bool, device='cuda')
+            is_lora = torch.tensor([is_lora], dtype=torch.bool, device=accelerator_device())
             if self.pp_size > 1:
                 dist.all_reduce(is_lora, group=self.pp_group)
             if is_lora:
@@ -1586,7 +1588,7 @@ def _set_linear_in_proj(self, mg_attn, hf_state_dict, to_mcore: bool):
             a_dim = config.linear_num_value_heads // num_key_heads
             is_lora = False if mg_attn is None else isinstance(mg_attn.in_proj,
                                                                LoraParallelLinear) and self._peft_format
-            is_lora = torch.tensor([is_lora], dtype=torch.bool, device='cuda')
+            is_lora = torch.tensor([is_lora], dtype=torch.bool, device=accelerator_device())
             if self.pp_size > 1:
                 dist.all_reduce(is_lora, group=self.pp_group)
             if is_lora:
@@ -1732,7 +1734,7 @@ def _set_layer_mlp(self, mg_layer, hf_state_dict, layer_idx: int, to_mcore: bool
         mg_mlp = None if mg_layer is None else mg_layer.mlp
         is_moe = True if hasattr(mg_mlp, 'experts') else False
         if not to_mcore:
-            is_moe = torch.tensor([is_moe], dtype=torch.bool, device='cuda')
+            is_moe = torch.tensor([is_moe], dtype=torch.bool, device=accelerator_device())
             if self.pp_size > 1:
                 dist.all_reduce(is_moe, group=self.pp_group)
         if is_moe:
@@ -1903,7 +1905,7 @@ def _convert(self, mg_models, hf_state_dict, hf_prefix: str, to_mcore: bool, tqd
                 else:
                     mg_layer = None
             if not to_mcore and self.pp_size > 1:
-                has_model = torch.tensor([mg_layer is not None], dtype=torch.bool, device='cuda')
+                has_model = torch.tensor([mg_layer is not None], dtype=torch.bool, device=accelerator_device())
                 dist.all_reduce(has_model, group=self.pp_group)
                 if not has_model:
                     mg_model = next(mg_models)  # compat vpp

src/mcore_bridge/model/gpts/bailing_moe.py

@@ -4,6 +4,7 @@
 
 from mcore_bridge.bridge import GPTBridge
 from mcore_bridge.tuners import LoraParallelLinear
+from mcore_bridge.utils.accelerator import accelerator_device
 
 from ..constant import ModelType
 from ..register import ModelMeta, register_model
@@ -70,7 +71,7 @@ def mg_to_hf(w, per_head_rows, last_dim):
         else:
             is_lora = False if mg_attn is None else (isinstance(mg_attn.linear_qkv, LoraParallelLinear)
                                                      and self._peft_format)
-            is_lora = torch.tensor([is_lora], dtype=torch.bool, device='cuda')
+            is_lora = torch.tensor([is_lora], dtype=torch.bool, device=accelerator_device())
             if self.pp_size > 1:
                 dist.all_reduce(is_lora, group=self.pp_group)
             if is_lora:

src/mcore_bridge/model/gpts/qwen3_next_gdn.py

@@ -10,6 +10,7 @@
 from mcore_bridge.bridge import GPTBridge
 from mcore_bridge.tuners import LoraParallelLinear
 from mcore_bridge.utils import get_env_args
+from mcore_bridge.utils.accelerator import accelerator_device
 
 from ..constant import ModelType
 from ..modules import GatedDeltaNet, GatedSelfAttention
@@ -75,7 +76,7 @@ def _set_linear_in_proj(self, mg_attn, hf_state_dict, to_mcore: bool):
             qkvz_dim = key_dim * 2 + value_dim * 2
             is_lora = False if mg_attn is None else isinstance(mg_attn.in_proj,
                                                                LoraParallelLinear) and self._peft_format
-            is_lora = torch.tensor([is_lora], dtype=torch.bool, device='cuda')
+            is_lora = torch.tensor([is_lora], dtype=torch.bool, device=accelerator_device())
             if self.pp_size > 1:
                 dist.all_reduce(is_lora, group=self.pp_group)
             if is_lora:

src/mcore_bridge/model/gpts/qwen4_exp.py

@@ -24,6 +24,7 @@
 from typing import List, Optional
 
 from mcore_bridge.utils import get_env_args, get_local_layer_specs, get_logger
+from mcore_bridge.utils.accelerator import accelerator_device
 from mcore_bridge.utils.megatron_utils import reconstruct_tensor_cp
 
 from ..modules import (QSA_SPARSE_KERNEL_ENV, GatedDeltaNet, MultiTokenPredictionLayer, QSAIndexer,
@@ -419,7 +420,7 @@ def _set_layer_mlp(self, mg_layer, hf_state_dict, layer_idx: int, to_mcore: bool
         mg_mlp = None if mg_layer is None else mg_layer.mlp
         is_moe = mg_mlp is not None and hasattr(mg_mlp, 'experts')
         if not to_mcore:
-            is_moe = torch.tensor([is_moe], dtype=torch.bool, device='cuda')
+            is_moe = torch.tensor([is_moe], dtype=torch.bool, device=accelerator_device())
             if self.pp_size > 1:
                 dist.all_reduce(is_moe, group=self.pp_group)
         if is_moe:
@@ -450,7 +451,7 @@ def _get_tp_split_dim(self, mg_key):
 
     def _get_pp_src_rank(self, has_module: bool) -> int:
         """Global rank of the PP stage holding the module (all-reduce MAX)."""
-        holder = torch.tensor([dist.get_rank() if has_module else -1], dtype=torch.long, device='cuda')
+        holder = torch.tensor([dist.get_rank() if has_module else -1], dtype=torch.long, device=accelerator_device())
         if self.pp_size > 1:
             dist.all_reduce(holder, op=dist.ReduceOp.MAX, group=self.pp_group)
         return int(holder.item())

src/mcore_bridge/model/mm_gpts/deepseek_v41.py

@@ -61,6 +61,7 @@ def mark_keep_in_fp32(tensor):
 from mcore_bridge.model.modules.engram import (DeepseekV41Engram, DeepseekV41TransformerLayer,
                                                build_deepseek_v41_engram_config, has_native_engram)
 from mcore_bridge.utils import is_master
+from mcore_bridge.utils.accelerator import accelerator_device
 
 from ..constant import ModelType
 from ..gpts.deepseek_v4 import DeepseekV4Bridge, DeepseekV4Loader, DSv4HybridSelfAttention, _apply_mla_rope
@@ -131,7 +132,7 @@ def __init__(self, config, *args, **kwargs):
         world_size = parallel_state.get_tensor_model_parallel_world_size()
         if config.num_attention_heads % world_size:
             raise ValueError('DSpark attention heads must be divisible by tensor parallel size.')
-        device = 'cpu' if config.use_cpu_initialization else torch.cuda.current_device()
+        device = 'cpu' if config.use_cpu_initialization else accelerator_device()
         self.attn_sink = mark_keep_in_fp32(
             nn.Parameter(torch.zeros(config.num_attention_heads // world_size, dtype=torch.float32, device=device)))
 
@@ -520,7 +521,7 @@ def __init__(self, config):
         self.image_start = nn.Parameter(torch.empty(config.hidden_size))
         self.image_end = nn.Parameter(torch.empty(config.hidden_size))
         self.image_newline = nn.Parameter(torch.empty(config.hidden_size))
-        target_device = torch.cuda.current_device() if torch.cuda.is_available() else None
+        target_device = accelerator_device()
         self.to(device=target_device, dtype=config.params_dtype)
         # Official RMSNorm scales remain fp32 even when the remaining ViT is bf16.
         for module in self.modules():
@@ -1541,7 +1542,7 @@ def _convert(self, mg_models, hf_state_dict, hf_prefix: str, to_mcore: bool, tqd
                 else:
                     mg_layer = None
             if not to_mcore and self.pp_size > 1:
-                has_model = torch.tensor([mg_layer is not None], dtype=torch.bool, device='cuda')
+                has_model = torch.tensor([mg_layer is not None], dtype=torch.bool, device=accelerator_device())
                 dist.all_reduce(has_model, group=self.pp_group)
                 if not has_model:
                     mg_model = next(mg_models)  # compat vpp

src/mcore_bridge/model/mm_gpts/utils.py

@@ -7,6 +7,7 @@
 
 from mcore_bridge.config import ModelConfig
 from mcore_bridge.utils import safe_ddp_context
+from mcore_bridge.utils.accelerator import accelerator_device
 
 
 @contextmanager
@@ -57,7 +58,7 @@ def __init__(self, config: ModelConfig):
             else:
                 self.prepare_model(hf_config)
 
-        self.to(device='cuda')
+        self.to(device=accelerator_device())
 
     @abstractmethod
     def prepare_model(self, hf_config: PretrainedConfig):

src/mcore_bridge/model/modules/dspark.py

@@ -16,6 +16,8 @@
 from torch import nn
 from typing import Callable, Optional, Sequence
 
+from mcore_bridge.utils.accelerator import accelerator_device
+
 
 class DeepseekV41DSparkRMSNorm(nn.Module):
     """RMSNorm matching the fp32 accumulation used by the reference model."""
@@ -66,7 +68,7 @@ def __init__(self, config):
             input_is_parallel=False,
             skip_bias_add=False,
         )
-        device = None if config.use_cpu_initialization else torch.cuda.current_device()
+        device = None if config.use_cpu_initialization else accelerator_device()
         self.main_norm = DeepseekV41DSparkRMSNorm(
             config.hidden_size,
             config.layernorm_epsilon,
@@ -160,7 +162,7 @@ class DeepseekV41DSparkConfidenceHead(nn.Module):
     def __init__(self, config):
         super().__init__()
         input_size = config.hidden_size + config.dspark_markov_rank
-        device = None if config.use_cpu_initialization else torch.cuda.current_device()
+        device = None if config.use_cpu_initialization else accelerator_device()
         self.proj = nn.Linear(
             input_size,
             1,
@@ -264,7 +266,7 @@ class DeepseekV41DSparkOutput(nn.Module):
 
     def __init__(self, config):
         super().__init__()
-        device = None if config.use_cpu_initialization else torch.cuda.current_device()
+        device = None if config.use_cpu_initialization else accelerator_device()
         self.block_size = config.dspark_block_size
         self.norm = DeepseekV41DSparkRMSNorm(
             config.hidden_size,

src/mcore_bridge/model/modules/ple.py

@@ -14,6 +14,7 @@
 from typing import List, Optional
 
 from ...utils import get_env_args, get_logger
+from ...utils.accelerator import accelerator_device
 from ...utils.constants import EXPORT_CHUNK_BYTES
 from ...utils.megatron_utils import get_num_samples, reconstruct_tensor_cp, split_cp_inputs
 from .hyper_connection_gated import Qwen4ExpTextGroupedRMSNorm
@@ -255,7 +256,7 @@ def iter_export_table_to_hf(self, prefix=''):
             if self.cpu_offload:
                 table = self.host_table
                 tp_start, tp_end = self.vocab_start, self.vocab_end
-                device = torch.cuda.current_device()
+                device = accelerator_device()
             else:
                 table = self.ngram_embedding.weight
                 tp_start = tp_rank * self.ngram_embedding.num_embeddings_per_partition

src/mcore_bridge/utils/accelerator.py

@@ -0,0 +1,21 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+"""Process accelerator, without assuming CUDA."""
+import torch
+
+
+def accelerator_device() -> torch.device:
+    """Device used for collectives and newly built modules.
+
+    ``torch.device('cuda')`` and ``torch.cuda.current_device()`` fail, or target
+    the wrong chip, when the process is on Ascend NPU or CPU. CUDA keeps the
+    current CUDA device, which is what ``device='cuda'`` already meant.
+    """
+    npu = getattr(torch, 'npu', None)
+    if npu is not None and npu.is_available():
+        return torch.device('npu', npu.current_device())
+    if torch.cuda.is_available():
+        return torch.device('cuda', torch.cuda.current_device())
+    mps = getattr(torch.backends, 'mps', None)
+    if mps is not None and mps.is_available():
+        return torch.device('mps')
+    return torch.device('cpu')