[Ascend NPU]add device abstraction and Ascend NPU support

[noemotiovon]

What does this PR do?

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

---

Key changes

• 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 (see docs/ascend/ascend_roll.md).

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 detailed accuracy and performance benchmarks (see below).

Accuracy comparison plan

• RL-based methods (Agentic, RLVR, etc.):
  Expect Mean Absolute Error (MAE) ≤ 4% between Ascend and A100 under identical configurations.

  $$ \text{Mean Error} = \frac{\sum_{i=1}^{N} \lvert reward_i^{npu} - reward_i^{gpu} \rvert}{N} \leq 0.04 $$

• Supervised methods (DPO, Distill, etc.):
  Expect Relative Error ≤ 2% between Ascend and A100 under identical configurations.

$$ \text{Mean Error} = \frac{\sum_{i=1}^{N} \lvert loss_i^{npu} - loss_i^{gpu} \rvert}{N} \leq 0.02 $$

Throughput comparison plan

• Compare Ascend NPU vs. A100 by averaging the throughput values (tpu) of the first 4 steps from logs.

• Define throughput ratio as:

$$ \text{tpu ratio} = \frac{\text{NPU average}}{\text{A100 average}} $$

---

Test Environment

• CANN: 8.1 RC1
• torch: 2.5.1
• torch-npu: 2.5.1
• transformers: 4.52.4
• vllm: 0.8.4
• vllm-ascend: 0.8.4rc2
• deepspeed: 0.16.4
• sglang: not support now
• megatron: not support now

Worth noting:

• All current validations on Ascend NPUs are based on separate mode.
• Colocate mode still has unresolved issues and is under active investigation.

---

Supported scripts

Script	Status
examples/qwen2.5-0.5B-agentic/run_agentic_pipeline_sokoban.sh	Support
examples/qwen2.5-0.5B-agentic/run_agentic_pipeline_frozen_lake.sh	Support
examples/start_agentic_rollout_pipeline.py	Support
examples/qwen2.5-3B-dpo_megatron/run_dpo_pipeline.sh (Use DeepSpeed)	Support
examples/qwen2.5-7B-rlvr_megatron/run_rlvr_pipeline.sh (Use DeepSpeed)	Support

[CLAassistant]

All committers have signed the CLA.

[noemotiovon]

Environment

• CANN: 8.1 RC1
• torch: 2.5.1
• torch-npu: 2.5.1
• vllm: 0.8.4
• vllm-ascend: 0.8.4rc2
• deepspeed: 0.16.4
• sglang: not support now
• megatron: not support now

[noemotiovon]

Test 1 Script

bash examples/agentic_demo/run_agentic_pipeline_frozen_lake_single_node_demo.sh

Yaml

defaults:
  - ../config/envs@_here_
  - ../config/deepspeed_zero@_here_
  - ../config/deepspeed_zero2@_here_
  - ../config/deepspeed_zero3@_here_
  - ../config/deepspeed_zero3_cpuoffload@_here_

hydra:
  run:
    dir: .
  output_subdir: null

exp_name: "agentic_pipeline"
seed: 42
logging_dir: ./output/logs
output_dir: ./output
render_save_dir: /home/lichenguang25/tmp/data/oss_bucket_0/yali/output/render
system_envs:
  USE_MODELSCOPE: '1'

#track_with: wandb
#tracker_kwargs:
#  api_key:
#  project: roll-agentic
#  name: ${exp_name}_frozen_lake
#  notes: "agentic_pipeline"
#  tags:
#    - agentic
#    - roll
#    - baseline


track_with: tensorboard
tracker_kwargs:
  log_dir: /home/lichenguang25/tmp/data/oss_bucket_0/yali/llm/tensorboard/roll_exp/agentic_sokoban

num_gpus_per_node: 1

max_steps: 100
save_steps: 10000
logging_steps: 1
eval_steps: 10
resume_from_checkpoint: false

rollout_batch_size: 16
val_batch_size: 16
sequence_length: 4096

reward_clip: 20
advantage_clip: 10.0
ppo_epochs: 1
adv_estimator: "reinforce"
#pg_clip: 0.1
#dual_clip_loss: True
init_kl_coef: 0.0
whiten_advantages: true
entropy_loss_coef: 0

pretrain: Qwen/Qwen2.5-0.5B-Instruct
reward_pretrain: Qwen/Qwen2.5-0.5B-Instruct

actor_train:
  model_args:
    flash_attn: fa2
    disable_gradient_checkpointing: false
    dtype: fp16
    model_type: ~
  training_args:
    learning_rate: 1.0e-6
    weight_decay: 0
    per_device_train_batch_size: 1
    gradient_accumulation_steps: 16
    warmup_steps: 10
  data_args:
    template: qwen2_5
  strategy_args:
    strategy_name: deepspeed_train
    strategy_config: ${deepspeed_zero2}
    # strategy_name: megatron_train
    # strategy_config:
    #   tensor_model_parallel_size: 1
    #   pipeline_model_parallel_size: 1
    #   expert_model_parallel_size: 1
    #   use_distributed_optimizer: true
    #   recompute_granularity: full
  device_mapping: list(range(0,1))
  infer_batch_size: 1

actor_infer:
  model_args:
    flash_attn: fa2
    disable_gradient_checkpointing: true
    dtype: fp16
  generating_args:
    max_new_tokens: 32 # single-turn response length
    top_p: 0.99
    top_k: 100
    num_beams: 1
    temperature: 0.99
    num_return_sequences: 1
  data_args:
    template: qwen2_5
  strategy_args:
    strategy_name: vllm
    strategy_config:
      gpu_memory_utilization: 0.8
      block_size: 16
      load_format: auto
  device_mapping: list(range(0,1))
  infer_batch_size: 1

reference:
  model_args:
    flash_attn: fa2
    disable_gradient_checkpointing: true
    dtype: fp16
    model_type: ~
  data_args:
    template: qwen2_5
  strategy_args:
    strategy_name: hf_infer
    strategy_config: ~
  device_mapping: list(range(0,1))
  infer_batch_size: 1

enable_response_mask: True
action_sep: "||"
use_turn_scores: False # important to GAE when applying token-level rewards to token-level advantages. If False, will take the sum of scores as the reward for the last turn.
enable_think: False # False -> no think RL
max_actions_per_traj: 10
reward_normalization:
  grouping: tags # 可以tags(env_type)/traj_group_id(group)/batch(rollout_batch)... group_by计算reward/adv
  method: identity # asym_clip / identity / mean_std

custom_envs:
  SimpleSokoban:
    env_type: sokoban
    max_actions_per_traj:  ${max_actions_per_traj} # used in environment state manager to control the actual max actions executed per trajectory
    max_steps_per_traj: ${max_actions_per_traj}
    env_instruction: "You are solving the Sokoban puzzle. You are the player and you need to push all boxes to targets. When you are right next to a box, you can push it by moving in the same direction. You cannot push a box through a wall, and you cannot pull a box. The answer must be one of action in a turn, format is <answer>Right</answer>"
    max_tokens: 100 # used to curate llm prompt "max words", not used for rollout
    env_config: # keys should be a subset of SokobanConfig
      dim_x: 6
      dim_y: 6
      num_boxes: 1
      max_steps: ${max_actions_per_traj}
  LargerSokoban:
    env_type: sokoban
    max_actions_per_traj:  ${max_actions_per_traj}
    max_steps_per_traj: ${max_actions_per_traj}
    env_instruction: "You are solving the Sokoban puzzle. You are the player and you need to push all boxes to targets. When you are right next to a box, you can push it by moving in the same direction. You cannot push a box through a wall, and you cannot pull a box. The answer must be one of action in a turn, format is <answer>Right</answer>"
    max_tokens: 100
    env_config:
      dim_x: 8
      dim_y: 8
      num_boxes: 2
      max_steps: ${max_actions_per_traj}
      search_depth: 10
  SokobanDifferentGridVocab:
    env_type: sokoban
    max_actions_per_traj:  ${max_actions_per_traj}
    max_steps_per_traj: ${max_actions_per_traj}
    env_instruction: "You are solving the Sokoban puzzle. You are the player and you need to push all boxes to targets. When you are right next to a box, you can push it by moving in the same direction. You cannot push a box through a wall, and you cannot pull a box. The answer must be one of action in a turn, format is <answer>Right</answer>"
    max_tokens: 100
    env_config: # keys should be a subset of SokobanConfig
      search_depth: 30
      dim_x: 6
      dim_y: 6
      num_boxes: 1
      max_steps: ${max_actions_per_traj}
      grid_lookup: { 0: "W", 1: ".", 2: "G", 3: "C", 4: "B", 5: "A", 6: "@" }
      grid_vocab: { "W": "wall", ".": "empty", "G": "target", "C": "box on target", "B": "box", "A": "player", "@": "player on target" }
  FrozenLake:
    env_type: frozen_lake
    max_actions_per_traj:  ${max_actions_per_traj}
    max_steps_per_traj: ${max_actions_per_traj}
    env_instruction: "You are solving the FrozenLake puzzle. Forbid the whole and go to the target. You may move to the unintended direction due to the slippery ice. The answer must be one of action in a turn, format is <answer>Right</answer>"
    max_tokens: 100
    env_config:
      is_slippery: false

train_env_manager:
  format_penalty: -0.001
  env_groups: 1
  group_size: 1
  tags: [FrozenLake]
  n_groups: [1] # If not set, all env names divide nums equally. Under the same group, the env config and env seed (prompt) are equal in each generation

val_env_manager:
  env_groups: 2
  group_size: 1 # should be set to 1 because val temperature is set to 0 and same prompt leads to same output
  tags: [SimpleSokoban, FrozenLake]
  n_groups: [1, 1] # TODO: If not set, all env names divide nums equally. Under the same group, the env config and env seed (prompt) are equal in each generation

Result

[noemotiovon]

Test 2 Script

bash examples/qwen2.5-0.5B-agentic_ds/run_agentic_pipeline_sokoban.sh

Yaml

defaults:
  - ../config/traj_envs@_here_
  - ../config/deepspeed_zero@_here_
  - ../config/deepspeed_zero2@_here_
  - ../config/deepspeed_zero3@_here_
  - ../config/deepspeed_zero3_cpuoffload@_here_

hydra:
  run:
    dir: .
  output_subdir: null

exp_name: "agentic_pipeline"
seed: 42
logging_dir: ./output/logs
output_dir: ./output
render_save_dir: ./output/render
system_envs:
  USE_MODELSCOPE: '1'

#track_with: wandb
#tracker_kwargs:
#  api_key:
#  project: roll-agentic
#  name: ${exp_name}_sokoban
#  notes: "agentic_pipeline"
#  tags:
#    - agentic
#    - roll
#    - baseline

track_with: tensorboard
tracker_kwargs:
  log_dir: /home/lichenguang25/tmp/data/oss_bucket_0/yali/llm/tensorboard/roll_exp/agentic_sokoban


checkpoint_config:
  type: file_system
  output_dir: /home/lichenguang25/tmp/data/cpfs_0/rl_examples/models/${exp_name}

num_gpus_per_node: 4

max_steps: 1024
save_steps: 10000
logging_steps: 1
eval_steps: 10
resume_from_checkpoint: false

rollout_batch_size: 128
val_batch_size: 1024
sequence_length: 2048

advantage_clip: 0.2
ppo_epochs: 1
adv_estimator: "grpo"
#pg_clip: 0.1
#dual_clip_loss: True
init_kl_coef: 0.0
whiten_advantages: true
entropy_loss_coef: 0
max_grad_norm: 1.0

pretrain: Qwen/Qwen2.5-0.5B-Instruct
reward_pretrain: Qwen/Qwen2.5-0.5B-Instruct

actor_train:
  model_args:
    attn_implementation: fa2
    disable_gradient_checkpointing: false
    dtype: bf16
    model_type: ~
  training_args:
    learning_rate: 1.0e-6
    weight_decay: 0
    per_device_train_batch_size: 2
    gradient_accumulation_steps: 64
    warmup_steps: 10
    lr_scheduler_type: cosine
  data_args:
    template: qwen2_5
  strategy_args:
    strategy_name: deepspeed_train
    strategy_config: ${deepspeed_zero3}
    # strategy_name: megatron_train
    # strategy_config:
    #   tensor_model_parallel_size: 1
    #   pipeline_model_parallel_size: 1
    #   expert_model_parallel_size: 1
    #   use_distributed_optimizer: true
    #   recompute_granularity: full
  device_mapping: list(range(0,2))
  infer_batch_size: 1

actor_infer:
  model_args:
    disable_gradient_checkpointing: true
    dtype: bf16
  generating_args:
    max_new_tokens: 128 # single-turn response length
    top_p: 0.99
    top_k: 100
    num_beams: 1
    temperature: 0.99
    num_return_sequences: 1
  data_args:
    template: qwen2_5
  strategy_args:
    strategy_name: vllm
    strategy_config:
      gpu_memory_utilization: 0.8
      block_size: 16
      load_format: auto
  device_mapping: list(range(2,3))

reference:
  model_args:
    attn_implementation: fa2
    disable_gradient_checkpointing: true
    dtype: bf16
    model_type: ~
  data_args:
    template: qwen2_5
  strategy_args:
    strategy_name: hf_infer
    strategy_config: ~
  device_mapping: list(range(3,4))
  infer_batch_size: 1

reward_normalization:
  grouping: traj_group_id # 可以tags(env_type)/traj_group_id(group)/batch(rollout_batch)... group_by计算reward/adv
  method: mean_std # asym_clip / identity / mean_std

train_env_manager:
  format_penalty: -0.15 # sokoban env penalty_for_step=-0.1
  max_env_num_per_worker: 16
  num_env_groups: 128
  # under the same group, the env config and env seed are ensured to be equal
  group_size: 8
  tags: [SimpleSokoban]
  num_groups_partition: [128] # If not set, all env names divide nums equally. Under the same group, the env config and env seed (prompt) are equal in each generation

val_env_manager:
  max_env_num_per_worker: 32
  num_env_groups: 1024
  group_size: 1 # should be set to 1 because val temperature is set to 0 and same prompt leads to same output
  tags: [SimpleSokoban, LargerSokoban, SokobanDifferentGridVocab, FrozenLake]
  num_groups_partition: [256, 256, 256, 256] # TODO: If not set, all env names divide nums equally. Under the same group, the env config and env seed (prompt) are equal in each generation


# Here, you can override variables defined in the imported envs. max_tokens_per_step: 128 in custom_env.SimpleSokoban, here replaced by 64
max_tokens_per_step: 64

custom_envs:
  SimpleSokoban:
    ${custom_env.SimpleSokoban}
  LargerSokoban:
    ${custom_env.LargerSokoban}
  SokobanDifferentGridVocab:
    ${custom_env.SokobanDifferentGridVocab}
  FrozenLake:
    ${custom_env.FrozenLake}
  FrozenLakeThink:
    ${custom_env.FrozenLakeThink}

Result

[HuangJoJo]

Thanks for your Contribution to ROLL! This Design helps a lot for other Hardware support work.

[noemotiovon]

Test 3 Script

bash examples/qwen2.5-0.5B-agentic/run_agentic_pipeline_frozen_lake.sh

Yaml

defaults:
  - ../config/traj_envs@_here_
  - ../config/deepspeed_zero@_here_
  - ../config/deepspeed_zero2@_here_
  - ../config/deepspeed_zero3@_here_
  - ../config/deepspeed_zero3_cpuoffload@_here_

hydra:
  run:
    dir: .
  output_subdir: null

exp_name: "agentic_pipeline"
seed: 42
logging_dir: ./output/logs
output_dir: ./output
render_save_dir: ./output/render
system_envs:
  USE_MODELSCOPE: '1'

#track_with: wandb
#tracker_kwargs:
#  api_key:
#  project: roll-agentic
#  name: ${exp_name}_sokoban
#  notes: "agentic_pipeline"
#  tags:
#    - agentic
#    - roll
#    - baseline

track_with: tensorboard
tracker_kwargs:
  log_dir: /home/lichenguang25/tmp/data/oss_bucket_0/yali/llm/tensorboard/roll_exp/agentic_frozen_lake


checkpoint_config:
  type: file_system
  output_dir: /home/lichenguang25/tmp/data/cpfs_0/rl_examples/models/${exp_name}

num_gpus_per_node: 4

max_steps: 1024
save_steps: 10000
logging_steps: 1
eval_steps: 10
resume_from_checkpoint: false

rollout_batch_size: 128
val_batch_size: 1024
sequence_length: 2048

advantage_clip: 0.2
ppo_epochs: 1
adv_estimator: "grpo"
#pg_clip: 0.1
#dual_clip_loss: True
init_kl_coef: 0.0
whiten_advantages: true
entropy_loss_coef: 0
max_grad_norm: 1.0

pretrain: Qwen/Qwen2.5-0.5B-Instruct
reward_pretrain: Qwen/Qwen2.5-0.5B-Instruct

actor_train:
  model_args:
    attn_implementation: fa2
    disable_gradient_checkpointing: false
    dtype: bf16
    model_type: ~
  training_args:
    learning_rate: 1.0e-6
    weight_decay: 0
    per_device_train_batch_size: 2
    gradient_accumulation_steps: 64
    warmup_steps: 10
    lr_scheduler_type: cosine
  data_args:
    template: qwen2_5
  strategy_args:
    strategy_name: deepspeed_train
    strategy_config: ${deepspeed_zero3}
    # strategy_name: megatron_train
    # strategy_config:
    #   tensor_model_parallel_size: 1
    #   pipeline_model_parallel_size: 1
    #   expert_model_parallel_size: 1
    #   use_distributed_optimizer: true
    #   recompute_granularity: full
  device_mapping: list(range(0,2))
  infer_batch_size: 1

actor_infer:
  model_args:
    disable_gradient_checkpointing: true
    dtype: bf16
  generating_args:
    max_new_tokens: 128 # single-turn response length
    top_p: 0.99
    top_k: 100
    num_beams: 1
    temperature: 0.99
    num_return_sequences: 1
  data_args:
    template: qwen2_5
  strategy_args:
    strategy_name: vllm
    strategy_config:
      gpu_memory_utilization: 0.8
      block_size: 16
      load_format: auto
  device_mapping: list(range(2,3))

reference:
  model_args:
    attn_implementation: fa2
    disable_gradient_checkpointing: true
    dtype: bf16
    model_type: ~
  data_args:
    template: qwen2_5
  strategy_args:
    strategy_name: hf_infer
    strategy_config: ~
  device_mapping: list(range(3,4))
  infer_batch_size: 1

reward_normalization:
  grouping: traj_group_id # 可以tags(env_type)/traj_group_id(group)/batch(rollout_batch)... group_by计算reward/adv
  method: mean_std # asym_clip / identity / mean_std

train_env_manager:
  format_penalty: -0.15 # sokoban env penalty_for_step=-0.1
  max_env_num_per_worker: 16
  num_env_groups: 128
  # under the same group, the env config and env seed are ensured to be equal
  group_size: 8
  tags: [FrozenLake]
  num_groups_partition: [128] # If not set, all env names divide nums equally. Under the same group, the env config and env seed (prompt) are equal in each generation

val_env_manager:
  max_env_num_per_worker: 32
  num_env_groups: 1024
  group_size: 1 # should be set to 1 because val temperature is set to 0 and same prompt leads to same output
  tags: [SimpleSokoban, LargerSokoban, SokobanDifferentGridVocab, FrozenLake]
  num_groups_partition: [256, 256, 256, 256] # TODO: If not set, all env names divide nums equally. Under the same group, the env config and env seed (prompt) are equal in each generation


# Here, you can override variables defined in the imported envs. max_tokens_per_step: 128 in custom_env.SimpleSokoban, here replaced by 64
max_tokens_per_step: 64

custom_envs:
  SimpleSokoban:
    ${custom_env.SimpleSokoban}
  LargerSokoban:
    ${custom_env.LargerSokoban}
  SokobanDifferentGridVocab:
    ${custom_env.SokobanDifferentGridVocab}
  FrozenLake:
    ${custom_env.FrozenLake}
  FrozenLakeThink:
    ${custom_env.FrozenLakeThink}
  FrozenLakeLocallyDefineExamples:  # Can import from unified envs config or define dict locally
    env_type: frozen_lake
    max_tokens_per_step: ${max_tokens_per_step}
    user_prompt_format: ${user_prompt_think_format}
    env_manager_cls: ${env_manager_cls}
    use_thread_lock: true
    env_config:
      env_instruction: "You are solving the FrozenLake puzzle. Forbid the whole and go to the target. You may move to the unintended direction due to the slippery ice. The answer must be one of action in a turn, format is <answer>Right</answer>"
      action_pattern: ${think_action_pattern}
      max_steps: ${max_actions_per_traj}
      is_slippery: false

Result

[lowdy1]

Test 4 Script

python examples/start_agentic_rollout_pipeline.py --config_path qwen2.5-0.5B-agentic  --config_name agentic_rollout_sokoban

yaml

defaults:
  - ../config/traj_envs@_here_
  - ../config/deepspeed_zero@_here_
  - ../config/deepspeed_zero2@_here_
  - ../config/deepspeed_zero3@_here_
  - ../config/deepspeed_zero3_cpuoffload@_here_

hydra:
  run:
    dir: .
  output_subdir: null

exp_name: "agentic_pipeline"
seed: 42
render_save_dir: ./oss_bucket_0/yali/llm/output/render

#track_with: wandb
#tracker_kwargs:
#  api_key:
#  project: roll-agentic
#  name: ${exp_name}_sokoban
#  notes: "agentic_pipeline"
#  tags:
#    - agentic
#    - roll
#    - baseline

track_with: tensorboard
tracker_kwargs:
  log_dir: ./oss_bucket_0/yali/llm/tensorboard/roll_exp/agentic_sokoban

num_gpus_per_node: 4

max_steps: 128
save_steps: 10000

rollout_batch_size: 16
sequence_length: 1024

pretrain: Qwen/Qwen2.5-0.5B-Instruct

actor_infer:
  model_args:
    disable_gradient_checkpointing: true
    dtype: bf16
  generating_args:
    max_new_tokens: 128 # single-turn response length
    top_p: 0.99
    top_k: 100
    num_beams: 1
    temperature: 0.99
    num_return_sequences: 1
  data_args:
    template: qwen2_5
  strategy_args:
    strategy_name: vllm
    strategy_config:
      gpu_memory_utilization: 0.8
      block_size: 16
      load_format: auto # should set 'auto' here, because default load_format is 'dummy'
  device_mapping: list(range(0,4))

train_env_manager:
  format_penalty: -0.15 # sokoban env penalty_for_step=-0.1
  max_env_num_per_worker: 1
  num_env_groups: 1
  # under the same group, the env config and env seed are ensured to be equal
  group_size: 1
  tags: [SimpleSokoban]
  num_groups_partition: [1] # If not set, all env names divide nums equally. Under the same group, the env config and env seed (prompt) are equal in each generation


# Here, you can override variables defined in the imported envs. max_tokens_per_step: 128 in custom_env.SimpleSokoban, here replaced by 64
max_tokens_per_step: 64

custom_envs:
  SimpleSokoban:
    ${custom_env.SimpleSokoban}
  LargerSokoban:
    ${custom_env.LargerSokoban}
  SokobanDifferentGridVocab:
    ${custom_env.SokobanDifferentGridVocab}
  FrozenLake:
    ${custom_env.FrozenLake}
  FrozenLakeThink:
    ${custom_env.FrozenLakeThink}

[noemotiovon]

Test 5 Script

bash examples/qwen2.5-3B-dpo_megatron/run_dpo_pipeline.sh 

Yaml

defaults:
  - ../config/deepspeed_zero@_here_
  - ../config/deepspeed_zero2@_here_
  - ../config/deepspeed_zero3@_here_
  - ../config/deepspeed_zero3_cpuoffload@_here_

hydra:
  run:
    dir: .
  output_subdir: null

num_gpus_per_node: 4

exp_name: "distill_zero3"
seed: 42
logging_dir: ./output/logs
output_dir: ./output

checkpoint_config:
  type: file_system
  output_dir: /home/lichenguang25/tmp/data/oss_bucket_0/chuye/roll/distill/models/zero3

save_steps: 100
logging_steps: 1
resume_from_checkpoint: false

student_pretrain: Qwen/Qwen2.5-0.5B-Instruct
teacher_pretrain: Qwen/Qwen2.5-1.5B-Instruct

# distill config
distill_loss_weight: 0.85
kd_objective: forward_kl
distill_on_prompt: True

sequence_length: 1024
max_grad_norm: 1.0

question_key: question_zh
answer_key: answer_zh


student:
  model_args:
    attn_implementation: fa2
    disable_gradient_checkpointing: false
    dtype: bf16
    model_type: ~
  training_args:
    learning_rate: 2.0e-5
    weight_decay: 1.0e-2
    lr_scheduler_type: constant
    per_device_train_batch_size: 1
    gradient_accumulation_steps: 1
    warmup_steps: 0
    num_train_epochs: 1
  data_args:
    template: qwen2_5
    file_name:
      - data/GSM8K_zh/GSM8K_zh.json        #https://huggingface.co/datasets/meta-math/GSM8K_zh
    preprocessing_num_workers: 4

  strategy_args:
    strategy_name: deepspeed_train
    strategy_config: ${deepspeed_zero3}
  device_mapping: list(range(0,2))

teacher:
  model_args:
    attn_implementation: fa2
    disable_gradient_checkpointing: true
    dtype: bf16
  data_args:
    template: qwen2_5
  strategy_args:
    strategy_name: deepspeed_infer
    strategy_config: ${deepspeed_zero3}
  device_mapping: list(range(2,4))

system_envs:
  RAY_PROFILING: "0"

Result

[noemotiovon]

Test 6 Script

bash examples/qwen2.5-7B-rlvr_megatron/run_rlvr_pipeline.sh 

Yaml

defaults:
  - ../config/traj_envs@_here_
  - ../config/deepspeed_zero@_here_
  - ../config/deepspeed_zero2@_here_
  - ../config/deepspeed_zero3@_here_
  - ../config/deepspeed_zero3_cpuoffload@_here_

hydra:
  run:
    dir: .
  output_subdir: null

exp_name: "qwen2.5-7B-rlvr-config"
seed: 42
logging_dir: ./output/logs
output_dir: ./output
system_envs:
  USE_MODELSCOPE: '1'

checkpoint_config:
  type: file_system
  output_dir: /home/lichenguang25/tmp/data/cpfs_0/rl_examples/models/${exp_name}

#track_with: wandb
#tracker_kwargs:
#  api_key:
#  project: roll_examples
#  notes: roll_examples
#  tags:
#    - rlvr
#    - baseline

track_with: tensorboard
tracker_kwargs:
  log_dir: /home/lichenguang25/tmp/data/oss_bucket_0/rl_examples/llm/tensorboard/roll_exp/rlvr

num_gpus_per_node: 4

max_steps: 50
save_steps: 10
logging_steps: 1
eval_steps: 10
resume_from_checkpoint: false


rollout_batch_size: 64  # prompt
prompt_length: 2048
response_length: 4096

num_return_sequences_in_group: 4
ppo_epochs: 1
adv_estimator: "reinforce"

# clip
value_clip: 0.5
reward_clip: 10
advantage_clip: 2.0
dual_clip_loss: true

# normalize
reward_norm: null
reward_shift: false
reward_scale: false

# data mask
max_len_mask: true
difficulty_mask: true
difficulty_low_threshold: 0.1
difficulty_high_threshold: 0.95
error_max_len_clip: false

# data weight
difficulty_loss_weight: false
length_loss_weight: false

# reward
add_token_level_kl: false

# advantage
whiten_advantages: true

# dynamic sampling scheduler
# use_additional_prompts: true
# max_running_requests: 256
# is_num_return_sequences_expand: false

pretrain: Qwen/Qwen2.5-0.5B-Instruct
reward_pretrain: Qwen/Qwen2.5-0.5B-Instruct

validation:
  data_args:
    template: qwen2_5
    file_name:
      - data/math_benchmarks.jsonl
  generating_args:
    max_new_tokens: ${response_length}
    top_p: 0.6
    top_k: 50
    num_beams: 1
    temperature: 0.6
    num_return_sequences: 1


actor_train:
  model_args:
    disable_gradient_checkpointing: false
    dtype: bf16
    model_type: ~
  training_args:
    learning_rate: 1.0e-6
    weight_decay: 0
    per_device_train_batch_size: 1
    gradient_accumulation_steps: 32
    warmup_steps: 20
    num_train_epochs: 50
  data_args:
    template: qwen2_5
    file_name:
      - data/code_KodCode_data.jsonl
      - data/llm_judge_Multi-subject-RLVR_deal_new.jsonl
      - data/math_deepmath_deal.jsonl
      - data/general_ifeval_train_deal.jsonl
      - data/general_CrossThink-QA_deal.jsonl
    domain_interleave_probs:
      math_rule: 0.4
      code_sandbox: 0.3
      llm_judge: 0.1
      crossthinkqa: 0.1
      ifeval: 0.1
    dataset_dir: data
    messages: messages
    interleave_probs: "1.0"
    preprocessing_num_workers: 1
  strategy_args:
    strategy_name: deepspeed_train
    strategy_config: ${deepspeed_zero3}
    # strategy_name: megatron_train
    # strategy_config:
    #   tensor_model_parallel_size: 1
    #   pipeline_model_parallel_size: 1
    #   expert_model_parallel_size: 1
    #   use_distributed_optimizer: true
    #   recompute_granularity: full
  device_mapping: list(range(0,1))
  infer_batch_size: 1

actor_infer:
  model_args:
    disable_gradient_checkpointing: true
    dtype: bf16
  generating_args:
    max_new_tokens: ${response_length}
    top_p: 0.99
    top_k: 100
    num_beams: 1
    temperature: 0.99
    num_return_sequences: ${num_return_sequences_in_group}
  data_args:
    template: qwen2_5
  strategy_args:
    strategy_name: vllm
    strategy_config:
      gpu_memory_utilization: 0.8
      block_size: 16
      max_model_len: 8000
  device_mapping: list(range(1,2))
  infer_batch_size: 1

reference:
  model_args:
    disable_gradient_checkpointing: true
    dtype: bf16
    model_type: ~
  data_args:
    template: qwen2_5
  strategy_args:
    strategy_name: deepspeed_infer
    strategy_config: ${deepspeed_zero3}
  device_mapping: list(range(2,3))
  infer_batch_size: 1

rewards:
  crossthinkqa:
    worker_cls: roll.pipeline.rlvr.rewards.crossthinkqa_rule_reward_worker.CrossThinkQARuleRewardWorker
    reward_type: soft
    response_length_penalty_coef: 0.0
    model_args:
      model_name_or_path: ${reward_pretrain}
    data_args:
      template: qwen2_5
    tag_included: [crossthinkqa]
    world_size: 1
    infer_batch_size: 1
  ifeval:
    worker_cls: roll.pipeline.rlvr.rewards.ifeval_rule_reward_worker.GeneralRuleRewardWorker
    reward_type: soft
    model_args:
      model_name_or_path: ${reward_pretrain}
    data_args:
      template: qwen2_5
    tag_included: [ifeval]
    world_size: 1
    infer_batch_size: 1
  math_rule:
    worker_cls: roll.pipeline.rlvr.rewards.math_rule_reward_worker.MathRuleRewardWorker
    model_args:
      model_name_or_path: ${reward_pretrain}
    data_args:
      template: qwen2_5
    tag_included: [deepmath_103k, aime]
    world_size: 1
    infer_batch_size: 1
# dynamic filter config
#    query_filter_config:
#      type: mean_filter
#      filter_args:
#        threshold_up: 0.9
#        threshold_down: 0.1
  code_sandbox:
    use_local: true
    worker_cls: roll.pipeline.rlvr.rewards.code_sandbox_reward_worker.CodeSandboxRewardWorker
    tag_included: [KodCode]
    model_args:
      model_name_or_path: ${reward_pretrain}
    data_args:
      template: qwen2_5
    world_size: 1
    infer_batch_size: 1
#    query_filter_config:
#      type: std_filter
#      filter_args:
#        std_threshold: 0
  llm_judge:
    # NOTE: llm as judge 也需要gpu, 不能和actor infer共享gpu
    worker_cls: roll.pipeline.rlvr.rewards.llm_judge_reward_worker.LLMJudgeRewardWorker
    judge_prompt: Qwen2.5-7B-Instruct-RLVR-prompt
    judge_model_type: inference
    tag_included: [RLVR]
    model_args:
      model_name_or_path: virtuoussy/Qwen2.5-7B-Instruct-RLVR
      attn_implementation: fa2
      disable_gradient_checkpointing: true
      dtype: bf16
      model_type: trl
    generating_args:
      max_new_tokens: 100
      top_p: 0.8
      top_k: 50
      num_beams: 1
      temperature: 0.8
      num_return_sequences: 1
    data_args:
      template: qwen2_5
    strategy_args:
      strategy_name: hf_infer
      strategy_config: null
    device_mapping: list(range(3,4))
    infer_batch_size: 1

Result

[noemotiovon]

Hi @HuangJoJo ,
The device-agnostic construction-related code has been completed, without any destructive changes to the GPU devices. Functional tests of ROLL on NPU have been conducted as described above. There are still a few tasks that need to be supported:

1. Support for shared-GPU mode on NPU
2. SGLang inference support
3. Support for vLLM + vLLM-Ascend version ≥ 0.10
4. FSDP and Megatron training support
5. Accuracy and performance testing

We have updated these tasks in the ISSUE to facilitate their completion and tracking.

We would like to thank @FightingZhen and @lowdy1 for their assistance, as well as the ROLL community ❤️ for providing resources and solutions!

[noemotiovon]

@FightingZhen, @lowdy1, Could you also help review the code, please?

[noemotiovon]

@FightingZhen , Thank you for your review, your suggestions were very helpful to me!

[noemotiovon]

Thank you for your suggestion. I have revised it based on the review comments.

[FightingZhen]

Looks good to go, let's start official review 👍

[noemotiovon]

Test on Nvidia A100:

examples/qwen2.5-0.5B-agentic/agent_val_frozen_lake.yaml

examples/qwen3_agentic_gem/gem_math_hotpotqa.yaml

examples/qwen2.5-7B-sft_megatron/sft_config.yaml

[tt0718]

@noemotiovon 👋hello from ROLL！感谢您对我们的关注和贡献
🤝ROLL期待和社区开发者建立更多的联系，共建交流，一起进步！后续项目也会组织定期的活动、给活跃用户发放礼品等
📧欢迎添加我的微信进一步沟通：tt19960718tt
如果您没有微信，也可以通过邮箱联系：tangtang.tt@alibaba-inc.com

👋 Hello from ROLL! Thank you for your interest and contributions to our project!
🤝 We at ROLL are eager to build stronger connections with our community developers. Let's foster collaboration and grow together! Looking ahead, the ROLL project will also organize regular events and offer gifts to active contributors.
📧 For further discussion:

• Add me on WeChat: tt19960718tt
• Or reach me via email: tangtang.tt@alibaba-inc.com