feat: add device abstraction and Ascend NPU support

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

Key changes include:

- Added a unified device abstraction interface for initialization, memory
  management, and synchronization, enabling extensibility for both CUDA
  and Ascend.
- Replaced direct usage of Ray CUDA resource APIs with the new abstraction
  layer to support heterogeneous multi-device environments.
- Integrated Ascend inference backend via vLLM + vLLM-Ascend.
- Added experimental training support with DeepSpeed on Ascend hardware.
- Added documentation for Ascend usage.

This enhancement lays the foundation for seamless switching between CUDA
and Ascend devices.

Future work:

- Add inference support for SGLang on Ascend NPUs.
- Add training support for Megatron on Ascend NPUs.
- Add training support for FSDP on Ascend NPUs.
- Add support for vLLM versions >= 0.10.
- Provide documentation with accuracy and performance benchmarks.

Co-authored-by: noemotiovon <757486878@qq.com>
Co-authored-by: FightingZhen <295632982@qq.com>
Co-authored-by: lowdy1 <xiahouweidong@gmail.com>

Author: 3 people

docs/ascend/ascend_roll.md

@@ -0,0 +1,151 @@
+# ROLL x Ascend
+
+Last updated: 09/04/2025.
+
+我们在 ROLL 上增加对华为昇腾设备的支持。
+
+## 硬件支持
+
+Atlas 900 A2 PODc
+
+
+## 安装
+
+
+### 基础环境准备
+
+| software  | version     |
+|-----------|-------------|
+| Python    |  3.10     |
+| CANN      |  8.1.RC1  |
+
+### 创建 conda 环境
+
+
+使用以下命令在 Miniconda 中创建新的 conda 环境：
+
+```
+conda create --name roll python=3.10
+conda activate roll
+```
+
+### 安装 torch & torch_npu:
+
+
+为了能在 ROLL 中正常使用 torch 和 torch_npu，需使用以下命令安装 torch 和 torch_npu。
+
+```
+# 安装 torch 的 CPU 版本
+pip install torch==2.5.1 torchvision==0.20.1 torchaudio==2.5.1 --index-url https://download.pytorch.org/whl/cpu
+
+# 安装 torch_npu
+pip install torch_npu==2.5.1
+```
+
+
+### 安装vllm & vllm-ascend:
+
+为了能够在 ROLL 中正常使用 vllm，需使用以下命令编译安装 vllm 和 vllm-ascend。
+
+```
+# vllm
+git clone -b v0.8.4 --depth 1 https://github.com/vllm-project/vllm.git
+cd vllm
+
+VLLM_TARGET_DEVICE=empty pip install -v -e .
+cd ..
+```
+
+``` 
+# vllm-ascend
+git clone -b v0.8.4rc2 --depth 1 https://github.com/vllm-project/vllm-ascend.git
+cd vllm-ascend
+
+export COMPILE_CUSTOM_KERNELS=1
+pip install -e .
+cd ..
+```
+
+如果在安装 vllm-ascend 时遇到类似以下问题：
+
+```
+RuntimeError: CMake configuration failed: Command '['/pathto/miniconda3/envs/roll/bin/python3.10', '-m', 'pybind11', '--cmake']' returned non-zero exit status 2.
+```
+
+可尝试在 vllm-ascend 目录下 setup.py 文件 151-158 行进行如下修改并重新进行编译：
+
+```
+try:
+    # if pybind11 is installed via pip
+    pybind11_cmake_path = (subprocess.check_output(
+        [python_executable, "-m", "pybind11",
+        "--cmakedir"]).decode().strip())
+except subprocess.CalledProcessError as e:
+    # else specify pybind11 path installed from source code on CI container
+    raise RuntimeError(f"CMake configuration failed: {e}")
+```
+
+### 安装 ROLL
+
+```
+git clone https://github.com/alibaba/ROLL.git
+cd ROLL
+pip install -r requirements_common.txt
+pip install deepspeed==0.16.0
+cd ..
+```
+
+### 其他三方库说明
+
+| software                       | description   |
+|-------------------------------|---------------|
+| transformers                  | v4.52.4       |
+| flash_attn                    | not supported |
+| transformer-engine[pytorch]   | not supported |
+
+1. 支持通过 transformers 使能 --flash_attention_2， transformers 需等于 4.52.4版本。
+2. 不支持通过 flash_attn 使能 flash attention 加速。
+3. 暂不支持 transformer-engine[pytorch] 
+
+```
+pip install transformers==4.52.4
+```
+
+## 快速开始，单节点部署指引
+
+正式使用前，建议您通过对单节点流水线的训练尝试以检验环境准备和安装的正确性。
+由于目前暂不支持 Megatron-LM 训练，请首先将对应文件中 
+strategy_args 参数修改为 deepspeed 选项。
+
+1. 使用 shell 执行单节点流水线
+
+```
+bash examples/agentic_demo/run_agentic_pipeline_frozen_lake_single_node_demo.sh  
+```
+
+2. 使用配置文件执行 agentic pipeline
+
+```
+# 确保当前位于ROLL项目目录的根目录下
+# export PYTHONPATH=$(pwd):$PYTHONPATH
+
+python examples/start_agentic_pipeline.py \
+        --config_path qwen2.5-0.5B-agentic \
+        --config_name agentic_val_sokoban
+
+- ``--config_path`` – 包含您的YAML配置文件的目录。
+- ``--config_name`` – 文件名（不含.yaml后缀）。
+```
+
+## 支持现状
+
+| Feature         | Example                                                      | Training Backend | Inference Backend | Hardware          |
+| --------------- | ------------------------------------------------------------ | ---------------- | ----------------- | ----------------- |
+| Agentic         | examples/qwen2.5-0.5B-agentic/run_agentic_pipeline_sokoban.sh | DeepSpeed        | vLLM              | Atlas 900 A2 PODc |
+| Agentic-Rollout | examples/qwen2.5-0.5B-agentic/run_agentic_rollout_sokoban.sh | DeepSpeed        | vLLM              | Atlas 900 A2 PODc |
+| DPO             | examples/qwen2.5-3B-dpo_megatron/run_dpo_pipeline.sh         | DeepSpeed        | vLLM              | Atlas 900 A2 PODc |
+| RLVR            | examples/qwen2.5-7B-rlvr_megatron/run_rlvr_pipeline.sh       | DeepSpeed        | vLLM              | Atlas 900 A2 PODc |
+
+
+## 声明
+ROLL 中提供的 Ascend 支持代码皆为参考样例，生产环境使用请通过官方正式途径沟通，谢谢。

docs/qa.md

@@ -63,7 +63,7 @@ profiler_output_dir: /data/oss_bucket_0/yali/llm/profile/${exp_name}
 
 0. **如何debug代码**
 
-在RayUtils的env中设置 "RAY_DEBUG": "legacy" ， 就可以采用pdb进行单步调试
+在对应后端的Platform的env中设置 "RAY_DEBUG": "legacy" ， 就可以采用pdb进行单步调试
 
 
 0. **如果出现这种错误: self.node2pg[node_rank] KeyError: 1**

docs_roll/docs/English/QuickStart/qa_issues.md

@@ -92,7 +92,7 @@ Then use the [Perfetto UI](https://ui.perfetto.dev/) tool for analysis.
 
 ### How to debug code?
 
-Set `"RAY_DEBUG": "legacy"` in RayUtils' env, and then you can use pdb for step-by-step debugging.
+Set `"RAY_DEBUG": "legacy"` in Platform' env, and then you can use pdb for step-by-step debugging.
 
 ## Common Errors and Solutions
 

docs_roll/docs/简体中文/快速开始/qa_issues.md

@@ -92,7 +92,7 @@ profiler_output_dir: /data/oss_bucket_0/yali/llm/profile/${exp_name}
 
 ### 如何 debug 代码？
 
-在 RayUtils 的 env 中设置 `"RAY_DEBUG": "legacy"`，就可以采用 pdb 进行单步调试。
+在 Platform 的 env 中设置 `"RAY_DEBUG": "legacy"`，就可以采用 pdb 进行单步调试。
 
 ## 常见错误及解决方案
 

mcore_adapter/src/mcore_adapter/initialize.py

@@ -8,6 +8,8 @@
 from .training_args import TrainingArguments
 from .utils import get_logger
 
+from .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 ...platforms import current_platform
 
 
 if TYPE_CHECKING:
@@ -286,7 +287,7 @@ def is_full(self):
 class SendBucketManager:
     def __init__(self, bucket_size):
         self.bucket_size = bucket_size
-        self.bucket = TensorBucket(bucket_size)
+        self.bucket = TensorBucket(bucket_size, current_platform.device_type)
 
     def push_tensor(self, tensor: "torch.Tensor", name: str):
         tensor_start = 0

mcore_adapter/src/mcore_adapter/models/model_factory.py

@@ -20,6 +20,7 @@
 from .converter.model_converter import ModelConverter
 from .model_config import McaModelConfig
 from .model_utils import ModuleUtilsMixin, RMSNorm, exists_hf_config, exists_mca_config, get_thd_data_on_this_cp_rank
+from ..platforms import current_platform
 
 
 if TYPE_CHECKING:
@@ -233,7 +234,8 @@ def get_batch_on_this_cp_rank(self, batch: Dict[str, "torch.Tensor"], dim3_keys:
                         val.shape[seq_dim] // (2 * cp_size),
                         *val.shape[(seq_dim + 1) :],
                     )
-                    index = torch.tensor([cp_rank, (2 * cp_size - cp_rank - 1)], device="cpu", pin_memory=True).cuda(
+                    index = torch.tensor([cp_rank, (2 * cp_size - cp_rank - 1)], device="cpu", pin_memory=True).to(
+                        current_platform.device_type,
                         non_blocking=True
                     )
                     val = val.index_select(seq_dim, index)
@@ -271,7 +273,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.to(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

@@ -9,7 +9,7 @@
 
 from ..constants import MCA_CONFIG_NAME
 from ..utils import get_logger
-
+from ..platforms import current_platform
 
 if TYPE_CHECKING:
     from megatron.core.transformer import TransformerConfig
@@ -93,7 +93,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

@@ -3,6 +3,7 @@
 import torch
 from megatron.core import mpu
 
+from ...platforms import current_platform
 from ..auto.modeling_auto import register_model
 from ..model_factory import McaGPTModel
 from ..model_utils import ModuleUtilsMixin
@@ -24,7 +25,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

@@ -3,6 +3,7 @@
 import torch
 from megatron.core import mpu
 
+from ...platforms import current_platform
 from ..auto.modeling_auto import register_model
 from ..model_factory import McaGPTModel
 from ..model_utils import ModuleUtilsMixin
@@ -24,7 +25,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)