feat: add initial support for Ascend devices with unified device abstraction

This commit introduces native support for Ascend NPUs in the ROLL project,
while preserving compatibility with existing CUDA-based infrastructure.

Key changes include:
- Introduced a unified device abstraction interface to encapsulate device
  initialization, memory management, and synchronization, enabling extensibility
  for both CUDA and Ascend devices.
- Replaced direct usage of  and Ray CUDA resource APIs with the new
  abstraction layer to support multi-device environments.
- Integrated Ascend inference backend via vLLM + vLLM-ascend.
- Added experimental support for training with MindSpeed on Ascend hardware.

This enhancement lays the groundwork for seamless switching across CUDA and
Ascend devices.

Signed-off-by: noemotiovon <757486878@qq.com>

Author: noemotiovon

mcore_adapter/src/mcore_adapter/initialize.py

@@ -8,6 +8,8 @@
 from .training_args import TrainingArguments
 from .utils import get_logger
 
+from roll.platforms import current_platform
+
 
 logger = get_logger(__name__)
 
@@ -28,7 +30,7 @@ def _set_random_seed(seed_):
         random.seed(seed)
         np.random.seed(seed)
         torch.manual_seed(seed)
-        if torch.cuda.device_count() > 0:
+        if current_platform.device_count() > 0:
             tensor_parallel.model_parallel_cuda_manual_seed(seed)
     else:
         raise ValueError("Seed ({}) should be a positive integer.".format(seed))
@@ -45,10 +47,10 @@ def _initialize_distributed(args: "TrainingArguments"):
     logger.info(f"Initializing mpu on device {args.device}")
     if not torch.distributed.is_initialized():
         # Manually set the device ids.
-        torch.cuda.set_device(args.device)
+        current_platform.set_device(args.device)
         # Call the init process
         torch.distributed.init_process_group(
-            backend=args.ddp_backend or "nccl",
+            backend=args.ddp_backend or current_platform.communication_backend,
             rank=int(os.getenv("RANK", "0")),
             world_size=int(os.getenv("WORLD_SIZE", "1")),
             timeout=args.ddp_timeout_delta,

mcore_adapter/src/mcore_adapter/models/converter/convert_utils.py

@@ -7,6 +7,7 @@
 import torch.distributed as dist
 from megatron.core import mpu
 from packaging.version import Version as PkgVersion
+from roll.platforms import current_platform
 
 
 if TYPE_CHECKING:
@@ -232,7 +233,7 @@ class StackedTensors:
 
 
 class TensorBucket:
-    def __init__(self, bucket_size, device="cuda"):
+    def __init__(self, bucket_size, device=current_platform.device_type):
         self.buffer = torch.empty(bucket_size, dtype=torch.int8, device=device)
         self.device = device
         self.bucket_size = bucket_size

mcore_adapter/src/mcore_adapter/models/model_factory.py

@@ -25,6 +25,7 @@
 from .converter.model_converter import ModelConverter
 from .model_config import McaModelConfig
 from .model_utils import ModuleUtilsMixin, RMSNorm, exists_hf_config, exists_mca_config
+from roll.platforms import current_platform
 
 
 if TYPE_CHECKING:
@@ -279,7 +280,7 @@ def __init__(self, config: "McaModelConfig", **kwargs):
         for param in self.parameters():
             tensor_parallel.set_defaults_if_not_set_tensor_model_parallel_attributes(param)
         if not config.use_cpu_initialization:
-            self.cuda(torch.cuda.current_device())
+            self.cuda(current_platform.current_device())
 
     def _get_transformer_layer_spec(self, config: Optional["McaModelConfig"]=None):
         config = config or self.config

mcore_adapter/src/mcore_adapter/models/model_utils.py

@@ -7,7 +7,7 @@
 
 from ..constants import MCA_CONFIG_NAME
 from ..utils import get_logger
-
+from roll.platforms import current_platform
 
 if TYPE_CHECKING:
     from megatron.core.transformer import TransformerConfig
@@ -91,7 +91,7 @@ def floating_point_ops(
 class RMSNorm(nn.Module):
     def __init__(self, config: "TransformerConfig", hidden_size, eps=1e-6, **kwargs):
         super().__init__()
-        device = torch.cuda.current_device() if not config.use_cpu_initialization else None
+        device = current_platform.current_device() if not config.use_cpu_initialization else None
         self.weight = torch.nn.Parameter(torch.ones(hidden_size, dtype=config.params_dtype, device=device))
         self.variance_epsilon = eps
 

mcore_adapter/src/mcore_adapter/models/qwen2_5_vl/modeling_qwen2_5_vl.py

@@ -4,6 +4,7 @@
 from megatron.core import mpu
 from megatron.core.transformer.attention import SelfAttention
 from torch import nn
+from roll.platforms import current_platform
 
 from ..auto.modeling_auto import register_model
 from ..model_factory import McaGPTModel
@@ -58,7 +59,7 @@ def __init__(
         if rotary_percent < 1.0:
             dim = int(dim * rotary_percent)
 
-        device = "cpu" if use_cpu_initialization else torch.cuda.current_device()
+        device = "cpu" if use_cpu_initialization else current_platform.current_device()
         self.inv_freq = 1.0 / (rotary_base ** (torch.arange(0, dim, 2, dtype=torch.float32, device=device) / dim))
 
     @torch.no_grad()
@@ -199,7 +200,7 @@ def __init__(self, config: "Qwen2_5_VLConfig", **kwargs):
                 Qwen2_5_VLVisionConfig(**config.vision_config),
                 attn_implementation="flash_attention_2",
                 torch_dtype=self.config.params_dtype,
-            ).to(torch.cuda.current_device())
+            ).to(current_platform.current_device())
             for param in self.vision_model.parameters():
                 setattr(param, "sequence_parallel", config.sequence_parallel)
 

mcore_adapter/src/mcore_adapter/models/qwen2_vl/modeling_qwen2_vl.py

@@ -4,6 +4,7 @@
 from megatron.core import mpu
 from megatron.core.transformer.attention import SelfAttention
 from torch import nn
+from roll.platforms import current_platform
 
 from ..auto.modeling_auto import register_model
 from ..model_factory import McaGPTModel
@@ -61,7 +62,7 @@ def __init__(
         self.rotary_interleaved = rotary_interleaved
 
         self.seq_len_interpolation_factor = seq_len_interpolation_factor
-        device = "cpu" if use_cpu_initialization else torch.cuda.current_device()
+        device = "cpu" if use_cpu_initialization else current_platform.current_device()
         self.inv_freq = 1.0 / (rotary_base ** (torch.arange(0, dim, 2, dtype=torch.float32, device=device) / dim))
 
     @torch.no_grad()
@@ -204,7 +205,7 @@ def __init__(self, config: "Qwen2VLConfig", **kwargs):
                 Qwen2VLVisionConfig(**config.vision_config),
                 attn_implementation="sdpa",
                 torch_dtype=self.config.params_dtype,
-            ).to(torch.cuda.current_device())
+            ).to(current_platform.current_device())
             for param in self.vision_model.parameters():
                 setattr(param, "sequence_parallel", config.sequence_parallel)
 

mcore_adapter/src/mcore_adapter/trainer/trainer.py

@@ -36,6 +36,7 @@
     seed_worker,
     speed_metrics,
 )
+from roll.platforms import current_platform
 
 from ..checkpointing import get_checkpoint_dir, load_state_dict_from_checkpoint
 from ..constants import DIST_OPTIMIZER_DIR, IGNORE_INDEX
@@ -501,7 +502,7 @@ def _save_rng_state(self, output_dir):
             "random_rng_state": random.getstate(),
             "np_rng_state": np.random.get_state(),
             "torch_rng_state": torch.get_rng_state(),
-            "cuda_rng_state": torch.cuda.get_rng_state(),
+            "cuda_rng_state": current_platform.get_rng_state(),
             "rng_tracker_states": tensor_parallel.get_cuda_rng_tracker().get_states(),
         }
         if self.args.world_size <= 1:
@@ -537,7 +538,7 @@ def _load_rng_state(self, checkpoint):
         random.setstate(checkpoint_rng_state["random_rng_state"])
         np.random.set_state(checkpoint_rng_state["np_rng_state"])
         torch.set_rng_state(checkpoint_rng_state["torch_rng_state"])
-        torch.cuda.set_rng_state(checkpoint_rng_state["cuda_rng_state"])
+        current_platform.set_rng_state(checkpoint_rng_state["cuda_rng_state"])
         # Check for empty states array
         if not checkpoint_rng_state["rng_tracker_states"]:
             raise KeyError

roll/distributed/executor/cluster.py

@@ -17,6 +17,7 @@
     collect_all_to_all,
     dispatch_one_to_all,
 )
+from roll.platforms import current_platform
 from roll.utils.constants import RAY_NAMESPACE
 from roll.distributed.scheduler.resource_manager import ResourceManager
 from roll.utils.import_utils import safe_import_class
@@ -122,14 +123,25 @@ def _create_workers(self):
 
             runtime_env = RuntimeEnv(env_vars=env_vars)
             self.worker_config.resource_placement_groups = pgs
-            worker = self.worker_cls.options(
-                scheduling_strategy=PlacementGroupSchedulingStrategy(placement_group=deploy_pg["placement_group"]),
-                name=worker_name,
-                namespace=RAY_NAMESPACE,
-                runtime_env=runtime_env,
-                num_cpus=0.01,
-                num_gpus=0.01 if self.worker_config.device_mapping else 0,
-            ).remote(worker_config=self.worker_config)
+            if current_platform.ray_device_key == "GPU":
+                worker = self.worker_cls.options(
+                    scheduling_strategy=PlacementGroupSchedulingStrategy(placement_group=deploy_pg["placement_group"]),
+                    name=worker_name,
+                    namespace=RAY_NAMESPACE,
+                    runtime_env=runtime_env,
+                    num_cpus=0.01,
+                    num_gpus=0.01 if self.worker_config.device_mapping else 0,
+                ).remote(worker_config=self.worker_config)
+            else:
+                worker = self.worker_cls.options(
+                    scheduling_strategy=PlacementGroupSchedulingStrategy(placement_group=deploy_pg["placement_group"]),
+                    name=worker_name,
+                    namespace=RAY_NAMESPACE,
+                    runtime_env=runtime_env,
+                    num_cpus=0.01,
+                    num_gpus=0,
+                    resources={current_platform.ray_device_key: 0.01 if self.worker_config.device_mapping else 0},
+                ).remote(worker_config=self.worker_config)
             self.workers.append(worker)
             if rank == 0:
                 self.master_addr, self.master_port = ray.get(worker.get_master_addr_and_port.remote())

roll/distributed/executor/worker.py

@@ -16,6 +16,7 @@
 from roll.utils.context_managers import state_offload_manger
 from roll.utils.logging import get_logger
 from roll.utils.offload_states import OffloadStateType
+from roll.platforms import current_platform
 from roll.utils.ray_utils import RayUtils
 
 

roll/distributed/scheduler/decorator.py

@@ -16,6 +16,7 @@
 
 from roll.distributed.scheduler.protocol import DataProto, ObjectRefWrap
 from roll.utils.logging import get_logger
+from roll.platforms import current_platform
 
 logger = get_logger()
 
@@ -274,9 +275,9 @@ async def inner_async(*args, **kwargs):
                     result = await func(*args, **kwargs)
                     if clear_cache:
                         try:
-                            torch._C._cuda_clearCublasWorkspaces()
+                            current_platform._cuda_clearCublasWorkspaces()
                             gc.collect()
-                            torch.cuda.empty_cache()
+                            current_platform.empty_cache()
                         except Exception as oe:
                             pass
 
@@ -295,9 +296,9 @@ def inner(*args, **kwargs):
                     result = func(*args, **kwargs)
                     if clear_cache:
                         try:
-                            torch._C._cuda_clearCublasWorkspaces()
+                            current_platform._cuda_clearCublasWorkspaces()
                             gc.collect()
-                            torch.cuda.empty_cache()
+                            current_platform.empty_cache()
                         except Exception as oe:
                             pass