[NPU] feat: Support FSDP worker and vLLM Ascend

.github/workflows/e2e_ascend.yml

@@ -49,6 +49,25 @@ jobs:
           npu-smi info
       - name: Checkout volcengine/verl repo
         uses: actions/checkout@v4
-      - name: Run test
+      - name: Install the current repository
         run: |
-          lscpu
+          pip3 install hf_transfer peft
+          pip3 install -e .[test]
+      - name: Prepare gsm8k dataset
+        run: |
+          ray stop --force
+          python3 examples/data_preprocess/gsm8k.py
+      - name: Running deepseek generation tests on ASCEND NPU
+        run: |
+          ray stop --force
+          bash tests/generation/run_gen_deepseek_v2_lite_math_npu.sh
+      - name: Running gsm8k e2e training tests with LoRA on ASCEND NPU
+        run: |
+          ray stop --force
+          bash tests/sft/run_sft_qwen05_peft.sh 8 $HOME/ckpts
+          rm -rf $HOME/ckpts
+      - name: Running gsm8k e2e training tests with GRPO on ASCEND NPU
+        run: |
+          ray stop --force
+          bash tests/e2e/run_qwen05_gsm8k_grpo_npu.sh
+          rm -rf $HOME/ckpts

docs/ascend/ascend_vllm0.7.3.rst

@@ -0,0 +1,100 @@
+========
+verl x Ascend
+========
+
+我们在 verl 上增加对华为昇腾设备的支持。
+
+=======
+硬件支持
+=======
+
+* Atlas 800T A2
+* Atlas 200T A2 Box16
+
+=======
+安装
+=======
+
+------
+环境准备
+------
+
+| 软件      | 版本         |
+| --------- | ----------- |
+| Python    | == 3.10     |
+| torch     | == 2.5.1    |
+| torch_npu | == 2.5.1rc1 |
+| CANN      | == 8.1.RC1  |
+
+1. 使用 vLLM，需遵循 vllm-ascend 的安装教程 <https://vllm-ascend.readthedocs.io/en/v0.7.3/installation.html>。
+2. 为了能够在 ASCEND NPU 上正常使能 flash_attention_2， transformers 版本需要大于等于 4.51.4(not released)，因此暂时使用源码安装。
+
+.. code-block:: bash
+    git clone --depth 1 https://github.com/huggingface/transformers.git
+    cd transformers
+    pip install -e .
+
+3. 目前支持 LLM 模型的 GRPO 训练，VLM模型的 GRPO 训练因为 vllm-ascend 的问题将会在后续支持，涉及到的issue为：
+
+https://github.com/vllm-project/vllm-ascend/issues/809
+https://github.com/vllm-project/vllm-ascend/issues/825
+
+
+------
+源码安装
+------
+
+.. code-block:: bash
+    git clone https://github.com/volcengine/verl.git
+    cd verl
+    pip install -r requirements-npu.txt
+    pip install -e .
+
+------
+vLLM
+------
+
+为了保证能够在 verl 上正常使用 vLLM，需要安装 vLLM Ascend 插件（`vllm-ascend`）。关于在华为昇腾上支持的 vLLM 版本以及和 vLLM Ascend 的配套关系请参考`安装教程 <https://vllm-ascend.readthedocs.io/en/v0.7.1rc1/installation.html>`_。
+
+------
+Ray
+------
+
+------
+其他第三方库说明
+------
+
++--------------+--------+
+| 软件          | 说明   |
++==============+========+
+| flash_attn   | 不支持  |
++--------------+--------+
+| liger-kernel | 不支持  |
++--------------+--------+
+
+------
+精度对比
+------
+
+根据经验，对于SFT等微调算法，我们期望在相同配置下，在华为昇腾设备上的 Loss 与英伟达 GPU 的 Loss 平均误差小于 2%，具体计算方式如下：
+
+.. image:: https://github.com/eric-haibin-lin/verl-community/tree/main/docs/loss_comparison.png
+   :alt: Alt text
+
+其中，N 表示训练的步数。更多信息请参考[精度计算说明](https://www.hiascend.com/document/detail/zh/Pytorch/600/ptmoddevg/trainingmigrguide/LMaccuracy_0001.html)。
+
+根据经验，对于GRPO等强化学习算法，我们期望在相同配置下，在华为昇腾设备上的 reward 与英伟达 GPU 的 reward 平均绝对误差小于 4%，具体计算参考 Loss 计算。
+
+------
+进展
+------
+
++--------+--------+
+| 算法    | 进展   |
++========+========+
+| SFT    | 已支持  |
++--------+--------+
+| PPO    | 已支持  |
++--------+--------+
+| GRPO   | 已支持  |
++--------+--------+

requirements-npu.txt

@@ -0,0 +1,19 @@
+# requirements.txt records the full set of dependencies for development
+accelerate
+codetiming
+datasets
+dill
+hydra-core
+numpy
+pandas
+peft
+pyarrow>=15.0.0
+pybind11
+pylatexenc
+ray
+tensordict<0.6
+wandb
+mathruler
+torchdata
+einops
+qwen_vl_utils

tests/e2e/run_qwen05_gsm8k_grpo_npu.sh

@@ -0,0 +1,44 @@
+# Tested with 1 & 8 NPUs
+
+set -x
+
+export VLLM_ATTENTION_BACKEND=XFORMERS
+
+python3 -m verl.trainer.main_ppo \
+    algorithm.adv_estimator=grpo \
+    data.train_files=$HOME/data/gsm8k/train.parquet \
+    data.val_files=$HOME/data/gsm8k/test.parquet \
+    data.train_batch_size=128 \
+    data.max_prompt_length=512 \
+    data.max_response_length=128 \
+    data.filter_overlong_prompts=True \
+    data.truncation='error' \
+    actor_rollout_ref.model.path=Qwen/Qwen2.5-0.5B-Instruct \
+    actor_rollout_ref.actor.optim.lr=1e-6 \
+    actor_rollout_ref.model.use_remove_padding=False \
+    actor_rollout_ref.actor.ppo_mini_batch_size=64 \
+    actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=20 \
+    actor_rollout_ref.actor.use_kl_loss=True \
+    actor_rollout_ref.actor.kl_loss_coef=0.001 \
+    actor_rollout_ref.actor.kl_loss_type=low_var_kl \
+    actor_rollout_ref.model.enable_gradient_checkpointing=True \
+    actor_rollout_ref.actor.fsdp_config.param_offload=False \
+    actor_rollout_ref.actor.fsdp_config.optimizer_offload=False \
+    actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=40 \
+    actor_rollout_ref.rollout.enable_chunked_prefill=False \
+    actor_rollout_ref.rollout.tensor_model_parallel_size=2 \
+    actor_rollout_ref.rollout.name=vllm \
+    actor_rollout_ref.rollout.gpu_memory_utilization=0.6 \
+    actor_rollout_ref.rollout.n=5 \
+    actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=40 \
+    actor_rollout_ref.ref.fsdp_config.param_offload=True \
+    algorithm.kl_ctrl.kl_coef=0.001 \
+    trainer.critic_warmup=0 \
+    trainer.logger=['console'] \
+    trainer.project_name='verl_grpo_example_gsm8k' \
+    trainer.experiment_name='qwen2_7b_function_rm' \
+    trainer.n_gpus_per_node=8 \
+    trainer.nnodes=1 \
+    trainer.save_freq=-1 \
+    trainer.test_freq=5 \
+    trainer.total_epochs=1 $@

tests/generation/run_gen_deepseek_v2_lite_math_npu.sh

@@ -0,0 +1,30 @@
+# Tested with 1 & 8 NPUs
+set -x
+
+if [ "$#" -lt 2 ]; then
+    echo "Usage: run_gen_deepseek_v2_lite_math_npu.sh <nproc_per_node> <save_path> [other_configs...]"
+    exit 1
+fi
+
+nproc_per_node=$1
+save_path=$2
+
+# Shift the arguments so $@ refers to the rest
+shift 2
+
+python3 -m verl.trainer.main_generation \
+    trainer.nnodes=1 \
+    trainer.n_gpus_per_node=8 \
+    data.path=~/data/rlhf/gsm8k/test.parquet \
+    data.prompt_key=prompt \
+    data.n_samples=1 \
+    data.output_path=~/data/rlhf/math/deepseek_v2_lite_gen_test.parquet \
+    model.path=deepseek-ai/deepseek-llm-7b-chat \
+    +model.trust_remote_code=True \
+    rollout.temperature=1.0 \
+    rollout.top_k=50 \
+    rollout.top_p=0.7 \
+    rollout.prompt_length=2048 \
+    rollout.response_length=1024 \
+    rollout.tensor_model_parallel_size=2 \
+    rollout.gpu_memory_utilization=0.8

verl/__init__.py

@@ -14,9 +14,13 @@
 
 import logging
 import os
+import pkg_resources
 
+from pkg_resources import DistributionNotFound
+from packaging.version import parse as parse_version
 from .protocol import DataProto
 from .utils.logging_utils import set_basic_config
+from verl.utils.device import is_npu_available
 
 version_folder = os.path.dirname(os.path.join(os.path.abspath(__file__)))
 
@@ -38,3 +42,19 @@
     from modelscope.utils.hf_util import patch_hub
 
     patch_hub()
+
+if is_npu_available:
+    from .utils import npu_patch
+
+    package_name = 'transformers'
+    required_version_spec = '4.51.0'
+    try:
+        installed_version = pkg_resources.get_distribution(package_name).version
+        installed = parse_version(installed_version)
+        required = parse_version(required_version_spec)
+
+        if not installed >= required:
+            raise ValueError(f"{package_name} version required >= {required_version_spec} on ASCEND NPU, current version is {installed}.")
+    except DistributionNotFound as e:
+        raise ImportError(
+            f"{package_name} not installed。please run pip install {package_name}=={required_version_spec}") from e

verl/models/transformers/qwen2_vl.py

@@ -28,7 +28,7 @@
 )
 
 try:
-    from flash_attn import flash_attn_func, flash_attn_varlen_func
+    from transformers.modeling_flash_attention_utils import flash_attn_func, flash_attn_varlen_func
 
     _flash_supports_window_size = "window_size" in list(inspect.signature(flash_attn_func).parameters)
 except ImportError:

verl/protocol.py

@@ -36,6 +36,7 @@
 
 from verl.utils.py_functional import union_two_dict
 from verl.utils.torch_functional import allgather_dict_tensors
+from verl.utils.device import get_torch_device
 
 __all__ = ["DataProto", "union_tensor_dict"]
 
@@ -272,7 +273,7 @@ def __setstate__(self, data):
 
         batch_deserialized_bytes, non_tensor_batch, meta_info = data
         batch_deserialized = io.BytesIO(initial_bytes=batch_deserialized_bytes)
-        batch = torch.load(batch_deserialized, weights_only=False, map_location="cpu" if not torch.cuda.is_available() else None)
+        batch = torch.load(batch_deserialized, weights_only=False, map_location="cpu" if not get_torch_device().is_available() else None)
         self.batch = batch
         self.non_tensor_batch = non_tensor_batch
         self.meta_info = meta_info
@@ -802,7 +803,7 @@ def all_gather_data_proto(data: DataProto, process_group):
     group_size = torch.distributed.get_world_size(group=process_group)
     assert isinstance(data, DataProto)
     prev_device = data.batch.device
-    data.batch = data.batch.cuda(device=torch.cuda.current_device())
+    data.batch = data.batch.to(get_torch_device().current_device())
     data.batch = allgather_dict_tensors(data.batch.contiguous(), size=group_size, group=process_group, dim=0)
     data.batch = data.batch.to(prev_device)
     # all gather non_tensor_batch

verl/single_controller/base/worker.py

@@ -23,6 +23,7 @@
 import ray
 
 from .decorator import Dispatch, Execute, register
+from verl.utils.device import get_torch_device
 
 
 @dataclass
@@ -127,10 +128,9 @@ def __init__(self, cuda_visible_devices=None) -> None:
 
         import torch
         from packaging import version
-
         ###
         # [SUPPORT AMD: torch]
-        if torch.cuda.is_available() and "AMD" in torch.cuda.get_device_name() and version.parse(ray.__version__) < version.parse("2.45.0"):
+        if torch.cuda.is_available() and "AMD" in get_torch_device().get_device_name() and version.parse(ray.__version__) < version.parse("2.45.0"):
             os.environ["CUDA_VISIBLE_DEVICES"] = os.environ.get("ROCR_VISIBLE_DEVICES")
             os.environ["LOCAL_RANK"] = os.environ.get("RAY_LOCAL_RANK")
         ###
@@ -148,7 +148,7 @@ def __init__(self, cuda_visible_devices=None) -> None:
 
         ###
         # [SUPPORT AMD: torch]
-        if torch.cuda.is_available() and "AMD" in torch.cuda.get_device_name() and version.parse(ray.__version__) < version.parse("2.45.0"):
+        if torch.cuda.is_available() and "AMD" in get_torch_device().get_device_name() and version.parse(ray.__version__) < version.parse("2.45.0"):
             self.local_rank = int(os.environ["LOCAL_RANK"])
             cuda_visible_devices = str(local_rank)
         ###
@@ -168,8 +168,8 @@ def __init__(self, cuda_visible_devices=None) -> None:
 
         ###
         # [SUPPORT AMD: torch]
-        if torch.cuda.is_available() and "AMD" in torch.cuda.get_device_name() and version.parse(ray.__version__) < version.parse("2.45.0"):
-            torch.cuda.set_device(int(cuda_visible_devices))
+        if torch.cuda.is_available() and "AMD" in get_torch_device().get_device_name() and version.parse(ray.__version__) < version.parse("2.45.0"):
+            get_torch_device().set_device(int(cuda_visible_devices))
         ###
 
         self.fused_worker_dict = {}