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>

# Conflicts:
#	roll/distributed/executor/cluster.py
#	roll/distributed/executor/worker.py
#	roll/distributed/scheduler/initialize.py

Author: FightingZhen

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
@@ -114,22 +115,32 @@ def _create_workers(self):
                 env_vars["MASTER_PORT"] = str(self.master_port)
             if deploy_pg["gpu_rank"] is not None:
                 RayUtils.update_env_vars_for_visible_devices(
-                    env_vars=env_vars, 
+                    env_vars=env_vars,
                     gpu_ranks=pg_zero_gpu_ranks)
             if "ROLL_LOG_DIR" in os.environ:
                 env_vars["ROLL_LOG_DIR"] = os.environ["ROLL_LOG_DIR"]
             env_vars.update(self.worker_config.system_envs)
 
             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)
+
+            worker_options = {
+                "scheduling_strategy": PlacementGroupSchedulingStrategy(placement_group=deploy_pg["placement_group"]),
+                "name": worker_name,
+                "namespace": RAY_NAMESPACE,
+                "runtime_env": runtime_env,
+                "num_cpus": 0.01,
+            }
+
+            if current_platform.ray_device_key == "GPU":
+                worker_options.update({"num_gpus": 0.01 if self.worker_config.device_mapping else 0})
+            elif current_platform.ray_device_key == "NPU":
+                worker_options.update({
+                    "num_cpus": 0,
+                    "resources": {current_platform.ray_device_key: 0.01 if self.worker_config.device_mapping else 0},
+                })
+
+            worker = self.worker_cls.options(**worker_options).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

@@ -18,6 +18,7 @@
 from roll.utils.offload_states import OffloadStateType
 from roll.platforms import current_platform
 from roll.utils.ray_utils import RayUtils
+from roll.platforms import current_platform
 
 
 @dataclass

roll/utils/gpu_utils.py

@@ -1,21 +1,51 @@
 import torch
 from enum import Enum
 
+
+def is_torch_npu_available() -> bool:
+    """Check the availability of NPU"""
+    try:
+        import torch_npu  # noqa: F401
+
+        return torch.npu.is_available()
+    except ImportError:
+        return False
+
+
+is_cuda_available = torch.cuda.is_available()
+is_npu_available = is_torch_npu_available()
+
+
+def get_device_name() -> str:
+    """Get device type name"""
+    if is_cuda_available:
+        return torch.cuda.get_device_name().upper()
+    elif is_npu_available:
+        return torch.npu.get_device_name().upper()
+    else:
+        raise RuntimeError("Can not find device type.")
+
+
 class DeviceType(Enum):
     NONE = "NONE"
     NVIDIA = "NVIDIA"
     AMD = "AMD"
+    ASCEND = "ASCEND"
     UNKNOWN = "UNKNOWN"
 
+
 class GPUUtils:
     @staticmethod
     def get_device_type() -> DeviceType:
-        if not torch.cuda.is_available():
+        if not is_cuda_available and not is_npu_available:
             return DeviceType.NONE
-        device_name = torch.cuda.get_device_name().upper()
+
+        device_name = get_device_name()
         if "NVIDIA" in device_name:
             return DeviceType.NVIDIA
-        if "AMD" in device_name:
+        elif "AMD" in device_name:
             return DeviceType.AMD
+        elif "ASCEND" in device_name:
+            return DeviceType.ASCEND
+
         return DeviceType.UNKNOWN
-    

roll/utils/ray_utils.py

@@ -41,6 +41,12 @@ def get_custom_env_env_vars(
                 "TORCHINDUCTOR_COMPILE_THREADS": "2",
                 "PYTORCH_HIP_ALLOC_CONF": "expandable_segments:True",
             }
+        elif DeviceType.ASCEND == device_type:
+            env_vars = {
+                "TORCHINDUCTOR_COMPILE_THREADS": "2",
+                "PYTORCH_NPU_ALLOC_CONF": "expandable_segments:True",
+                "PYTHONPATH": os.environ.get("PYTHONPATH", ""),
+            }
         elif DeviceType.UNKNOWN == device_type:
             env_vars = {
                 "TORCHINDUCTOR_COMPILE_THREADS": "2",
@@ -63,6 +69,11 @@ def update_env_vars_for_visible_devices(
                 "HIP_VISIBLE_DEVICES": ",".join(map(str, gpu_ranks)),
                 "RAY_EXPERIMENTAL_NOSET_HIP_VISIBLE_DEVICES": "1",
             }
+        elif DeviceType.ASCEND == device_type:
+            visible_devices_env_vars = {
+                "ASCEND_RT_VISIBLE_DEVICES": ",".join(map(str, gpu_ranks)),
+                "RAY_EXPERIMENTAL_NOSET_ASCEND_RT_VISIBLE_DEVICES": "1",
+            }
         else:
             visible_devices_env_vars = {
                 "CUDA_VISIBLE_DEVICES": ",".join(map(str, gpu_ranks)),
@@ -77,6 +88,8 @@ def get_visible_gpus(device_type: DeviceType | None = None) -> list:
             device_type = GPUUtils.get_device_type()
         if DeviceType.AMD == device_type:
             return os.environ.get("HIP_VISIBLE_DEVICES", "").split(",")
-        if DeviceType.NVIDIA == device_type or DeviceType.UNKNOWN == device_type:
+        elif DeviceType.ASCEND == device_type:
+            return os.environ.get("ASCEND_RT_VISIBLE_DEVICES", "").split(",")
+        elif DeviceType.NVIDIA == device_type or DeviceType.UNKNOWN == device_type:
             return os.environ.get("CUDA_VISIBLE_DEVICES", "").split(",")
         return []