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,184 @@
+# ROLL x Ascend
+
+Last updated: 08/15/2025.
+
+我们在 ROLL 上增加对华为昇腾设备的支持。
+
+## 硬件支持
+
+Atlas 200T A2 Box16
+
+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 |
+| tensordict                    | 0.8.3 (ARM)   |
+| transformer-engine[pytorch]   | not supported |
+
+1. 支持通过 transformers 使能 --flash_attention_2， transformers 需大于等于 4.52.0版本。
+2. 不支持通过 flash_attn 使能 flash attention 加速。
+3. 针对 ARM 服务器，tensordict 要求 0.8.3，可在依赖安装完成后再手动安装 tensordict。
+4. 暂不支持 transformer-engine[pytorch] 
+
+```
+pip install transformers==4.52.4
+pip install tensordict==0.8.3
+```
+
+## 快速开始，单节点部署指引
+
+正式使用前，建议您通过对单节点流水线的训练尝试以检验环境准备和安装的正确性。
+由于目前暂不支持 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后缀）。
+```
+
+## 支持现状
+
+
+**表1** NPU 已通过流水线验证
+
+| pipeline                                                      | hardware            |
+|---------------------------------------------------------------|---------------------|
+| examples/qwen2.5-0.5B-agentic/run_agentic_pipeline_sokoban.sh | Atlas 900 A2 PODc   |
+| examples/qwen2.5-0.5B-agentic/run_agentic_rollout_sokoban.sh  | Atlas 900 A2 PODc   |
+| examples/qwen2.5-1.5B-distill_ds/run_distill_pipeline.sh      | Atlas 900 A2 PODc   |
+| examples/qwen2.5-3B-dpo_megatron/run_dpo_pipeline.sh          | Atlas 900 A2 PODc   |
+| examples/qwen2.5-7B-rlvr_megatron/run_rlvr_pipeline.sh        | Atlas 900 A2 PODc   |
+
+**表2** NPU 待流水线验证
+
+| pipeline                                                      | hardware            |
+|---------------------------------------------------------------|---------------------|
+| examples/qwen2.5-vl-7B-distill/run_distill_vl_ds_pipeline.sh  | Atlas 900 A2 PODc   |
+| examples/qwen2.5-vl-7B-rlvr/run_rlvr_pipeline.sh              | Atlas 900 A2 PODc   |
+
+## 后续计划
+
+
+分别按照以下规则进行与 GPU 的精度与吞吐量的对比
+精度对比:
+根据经验，对于 Agentic 和 RLVR 等 RL 类算法，我们期望在相同配置下华为昇腾设备与 A100 的 rewards 平均绝对误差 <= 4%，计算方式参考下公式。
+```
+$ Mean Error = \frac{\sum_{i=1}^{N} |reward_i^{npu} - reward_{i}^{gpu}|}{N}  \leq 0.04 $
+```
+对于 DPO 和 Distill 等类算法，我们期望在相同配置下华为昇腾设备与 A100 的 loss 相对误差 <= 4%，计算方式参考下公式。
+```
+$ Mean Error = \frac{\sum_{i=1}^{N} |loss_i^{npu} - loss_{i}^{gpu}|}{N}  \leq 0.04 $
+```
+
+吞吐对比：Ascend npu 和 A100 分别取日志中前4个 step 的 throughput 的 tpu 值 做平均， tpu ratio = npu 平均值 / A100 平均值。
+
+
+## 声明
+-----------------------------------
+ROLL 中提供的 Ascend 支持代码皆为参考样例，商业使用请通过官方正式途径沟通，谢谢。

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