[megatron, vllm, recipe] feat: add Qwen2.5-1.5B GDPO training on Ascend

examples/gdpo_trainer/README.md

@@ -6,13 +6,27 @@ GDPO is a multi-reward, rubric-style variant whose advantage estimator aggregate
 
 | Script                               | Infer | Train | Platform |
 |--------------------------------------|-------|-------|----------|
-| `run_qwen3_8b_fsdp.sh`          | vLLM  | FSDP  | NVIDIA   |
+| `run_qwen3_8b_fsdp.sh`               | vLLM        | FSDP     | NVIDIA    |
+| `run_qwen2_5_1_5b_megatron.sh`       | vLLM-Ascend | Megatron | Ascend NPU |
 
-Prepare a rubric-style dataset (e.g. `rlla_4k`) and point `DATA_DIR` to it.
+Prepare a rubric-style dataset (e.g. `rlla_4k`) and point `DATA_ROOT` to it.
 
 ## Key Flags
 
 - `algorithm.adv_estimator=gdpo`
 - `+algorithm.gdpo_reward_keys='["accuracy_reward", "format_reward"]'`
 - `reward.reward_manager.name=gdpo`
 - `reward.custom_reward_function.path=$REPO_ROOT/verl/utils/reward_score/rlla.py`
+
+## Megatron + vLLM-Ascend
+
+```bash
+MODEL_PATH=/path/to/Qwen2.5-1.5B-Instruct \
+DATA_ROOT=/path/to/data \
+NPUS_PER_NODE=4 \
+bash examples/gdpo_trainer/run_qwen2_5_1_5b_megatron.sh
+```
+
+Ensure the container provides enough `/dev/shm` capacity for the configured weight-transfer bucket.
+See [verl-ascend-recipe issue #18](https://github.com/verl-project/verl-ascend-recipe/issues/18)
+for the validated environment, training logs, and 100-step results.

examples/gdpo_trainer/run_qwen2_5_1_5b_megatron.sh

@@ -0,0 +1,214 @@
+#!/usr/bin/env bash
+# GDPO | Qwen2.5-1.5B-Instruct | Megatron training | vLLM-Ascend rollout | Ascend NPUs
+
+set -xeuo pipefail
+
+########################### environment ###########################
+
+export TOKENIZERS_PARALLELISM=${TOKENIZERS_PARALLELISM:-false}
+export HYDRA_FULL_ERROR=${HYDRA_FULL_ERROR:-1}
+export RAY_DEDUP_LOGS=${RAY_DEDUP_LOGS:-0}
+export VLLM_USE_V1=${VLLM_USE_V1:-1}
+export VLLM_ALLREDUCE_USE_SYMM_MEM=${VLLM_ALLREDUCE_USE_SYMM_MEM:-0}
+export VLLM_ASCEND_ENABLE_NZ=${VLLM_ASCEND_ENABLE_NZ:-0}
+export TASK_QUEUE_ENABLE=${TASK_QUEUE_ENABLE:-2}
+export CPU_AFFINITY_CONF=${CPU_AFFINITY_CONF:-1}
+export HCCL_OP_EXPANSION_MODE=${HCCL_OP_EXPANSION_MODE:-AIV}
+export HCCL_ASYNC_ERROR_HANDLING=${HCCL_ASYNC_ERROR_HANDLING:-0}
+export HCCL_EXEC_TIMEOUT=${HCCL_EXEC_TIMEOUT:-3600}
+export HCCL_CONNECT_TIMEOUT=${HCCL_CONNECT_TIMEOUT:-3600}
+
+########################### user-adjustable ###########################
+
+MODEL_PATH=${MODEL_PATH:-Qwen/Qwen2.5-1.5B-Instruct}
+DATA_ROOT=${DATA_ROOT:-${HOME}/data}
+NNODES=${NNODES:-1}
+NPUS_PER_NODE=${NPUS_PER_NODE:-4}
+
+train_batch_size=${TRAIN_BATCH_SIZE:-16}
+val_batch_size=${VAL_BATCH_SIZE:-16}
+ppo_mini_batch_size=${PPO_MINI_BATCH_SIZE:-8}
+micro_batch_size=${MICRO_BATCH_SIZE:-2}
+max_prompt_length=${MAX_PROMPT_LENGTH:-2048}
+max_response_length=${MAX_RESPONSE_LENGTH:-1024}
+max_token_length=${MAX_TOKEN_LENGTH:-$((max_prompt_length + max_response_length))}
+ppo_max_token_len_per_gpu=${PPO_MAX_TOKEN_LEN_PER_GPU:-$((max_token_length * micro_batch_size))}
+dataloader_num_workers=${DATALOADER_NUM_WORKERS:-0}
+seed=${SEED:-42}
+
+actor_lr=${ACTOR_LR:-1e-6}
+kl_loss_coef=${KL_LOSS_COEF:-0.001}
+
+actor_tp=${ACTOR_TP:-2}
+actor_pp=${ACTOR_PP:-1}
+rollout_tp=${ROLLOUT_TP:-2}
+rollout_n=${ROLLOUT_N:-4}
+rollout_gpu_memory_utilization=${ROLLOUT_GPU_MEMORY_UTILIZATION:-0.35}
+rollout_max_num_batched_tokens=${ROLLOUT_MAX_NUM_BATCHED_TOKENS:-8192}
+rollout_enforce_eager=${ROLLOUT_ENFORCE_EAGER:-True}
+weight_bucket_mb=${WEIGHT_BUCKET_MB:-512}
+
+rollout_world_size=$((NNODES * NPUS_PER_NODE))
+if (( rollout_tp <= 0 || rollout_world_size % rollout_tp != 0 )); then
+    echo "ROLLOUT_TP must be a positive divisor of NNODES * NPUS_PER_NODE." >&2
+    exit 2
+fi
+if (( rollout_n < 2 )); then
+    echo "GDPO requires ROLLOUT_N >= 2." >&2
+    exit 2
+fi
+rollout_replicas=$((rollout_world_size / rollout_tp))
+default_rollout_max_num_seqs=$(((train_batch_size * rollout_n + rollout_replicas - 1) / rollout_replicas))
+rollout_max_num_seqs=${ROLLOUT_MAX_NUM_SEQS:-${default_rollout_max_num_seqs}}
+
+offload=${OFFLOAD:-False}
+total_training_steps=${TOTAL_TRAINING_STEPS:-100}
+total_epochs=${TOTAL_EPOCHS:-1}
+save_freq=${SAVE_FREQ:--1}
+test_freq=${TEST_FREQ:--1}
+resume_mode=${RESUME_MODE:-auto}
+max_actor_ckpt_to_keep=${MAX_ACTOR_CKPT_TO_KEEP:-1}
+
+project_name=${PROJECT_NAME:-verl_gdpo_rlla4k}
+experiment_name=${EXPERIMENT_NAME:-qwen2_5_1_5b_gdpo_megatron_vllm_ascend}
+output_dir=${OUTPUT_DIR:-${PWD}/checkpoints/${experiment_name}}
+log_dir=${LOG_DIR:-${PWD}/logs}
+mkdir -p "${output_dir}" "${log_dir}"
+
+train_files="['${DATA_ROOT}/rlla_4k/train.parquet']"
+val_files="['${DATA_ROOT}/rlla_4k/test.parquet']"
+REPO_ROOT=$(cd -- "$(dirname "${BASH_SOURCE[0]}")/../.." >/dev/null 2>&1 && pwd)
+reward_function_path=${REWARD_FUNCTION_PATH:-${REPO_ROOT}/verl/utils/reward_score/rlla.py}
+
+########################### parameter arrays ###########################
+
+ALGORITHM=(
+    algorithm.adv_estimator=gdpo
+    '+algorithm.gdpo_reward_keys=["accuracy_reward","format_reward"]'
+    algorithm.use_kl_in_reward=False
+    algorithm.kl_ctrl.kl_coef=${kl_loss_coef}
+)
+
+DATA=(
+    "data.train_files=${train_files}"
+    "data.val_files=${val_files}"
+    data.train_batch_size=${train_batch_size}
+    data.val_batch_size=${val_batch_size}
+    data.max_prompt_length=${max_prompt_length}
+    data.max_response_length=${max_response_length}
+    data.dataloader_num_workers=${dataloader_num_workers}
+    data.seed=${seed}
+    data.filter_overlong_prompts=True
+    data.truncation=error
+)
+
+MODEL=(
+    "actor_rollout_ref.model.path=${MODEL_PATH}"
+    actor_rollout_ref.model.use_remove_padding=True
+    actor_rollout_ref.model.enable_gradient_checkpointing=True
+)
+
+ACTOR=(
+    actor_rollout_ref.actor.loss_agg_mode=token-mean
+    actor_rollout_ref.actor.optim.lr=${actor_lr}
+    actor_rollout_ref.actor.ppo_mini_batch_size=${ppo_mini_batch_size}
+    actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=${micro_batch_size}
+    actor_rollout_ref.actor.use_dynamic_bsz=True
+    actor_rollout_ref.actor.ppo_max_token_len_per_gpu=${ppo_max_token_len_per_gpu}
+    actor_rollout_ref.actor.use_kl_loss=False
+    actor_rollout_ref.actor.kl_loss_coef=${kl_loss_coef}
+    actor_rollout_ref.actor.kl_loss_type=low_var_kl
+    actor_rollout_ref.actor.entropy_coeff=0
+    actor_rollout_ref.actor.megatron.tensor_model_parallel_size=${actor_tp}
+    actor_rollout_ref.actor.megatron.pipeline_model_parallel_size=${actor_pp}
+    actor_rollout_ref.actor.megatron.param_offload=${offload}
+    actor_rollout_ref.actor.megatron.grad_offload=${offload}
+    actor_rollout_ref.actor.megatron.optimizer_offload=${offload}
+    actor_rollout_ref.actor.megatron.use_mbridge=True
+    actor_rollout_ref.actor.megatron.vanilla_mbridge=False
+    actor_rollout_ref.actor.megatron.dtype=bfloat16
+    +actor_rollout_ref.actor.megatron.override_transformer_config.apply_rope_fusion=True
+    +actor_rollout_ref.actor.megatron.override_transformer_config.position_embedding_type=rope
+    +actor_rollout_ref.actor.megatron.override_transformer_config.use_fused_rotary_pos_emb=True
+    +actor_rollout_ref.actor.megatron.override_transformer_config.normalization=RMSNorm
+    +actor_rollout_ref.actor.megatron.override_transformer_config.use_fused_rmsnorm=True
+    ++actor_rollout_ref.actor.megatron.override_transformer_config.attention_backend=flash
+    +actor_rollout_ref.actor.megatron.override_transformer_config.use_flash_attn=True
+)
+
+ROLLOUT=(
+    actor_rollout_ref.rollout.name=vllm
+    actor_rollout_ref.rollout.tensor_model_parallel_size=${rollout_tp}
+    actor_rollout_ref.rollout.gpu_memory_utilization=${rollout_gpu_memory_utilization}
+    actor_rollout_ref.rollout.max_model_len=${max_token_length}
+    actor_rollout_ref.rollout.max_num_seqs=${rollout_max_num_seqs}
+    actor_rollout_ref.rollout.max_num_batched_tokens=${rollout_max_num_batched_tokens}
+    actor_rollout_ref.rollout.checkpoint_engine.update_weights_bucket_megabytes=${weight_bucket_mb}
+    actor_rollout_ref.rollout.n=${rollout_n}
+    actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=${micro_batch_size}
+    actor_rollout_ref.rollout.log_prob_use_dynamic_bsz=True
+    actor_rollout_ref.rollout.log_prob_max_token_len_per_gpu=${ppo_max_token_len_per_gpu}
+    actor_rollout_ref.rollout.calculate_log_probs=False
+    actor_rollout_ref.rollout.enable_chunked_prefill=True
+    actor_rollout_ref.rollout.enable_prefix_caching=True
+    actor_rollout_ref.rollout.enforce_eager=${rollout_enforce_eager}
+    actor_rollout_ref.rollout.free_cache_engine=True
+    actor_rollout_ref.rollout.val_kwargs.n=1
+    actor_rollout_ref.rollout.val_kwargs.temperature=1.0
+    actor_rollout_ref.rollout.val_kwargs.top_p=0.7
+)
+
+REF=(
+    actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=${micro_batch_size}
+    actor_rollout_ref.ref.log_prob_use_dynamic_bsz=True
+    actor_rollout_ref.ref.log_prob_max_token_len_per_gpu=${ppo_max_token_len_per_gpu}
+    actor_rollout_ref.ref.megatron.tensor_model_parallel_size=${actor_tp}
+    actor_rollout_ref.ref.megatron.pipeline_model_parallel_size=${actor_pp}
+    actor_rollout_ref.ref.megatron.param_offload=${offload}
+    actor_rollout_ref.ref.megatron.use_mbridge=True
+    actor_rollout_ref.ref.megatron.vanilla_mbridge=False
+)
+
+REWARD=(
+    "reward.custom_reward_function.path=${reward_function_path}"
+    reward.custom_reward_function.name=compute_score
+    reward.reward_manager.name=gdpo
+)
+
+TRAINER=(
+    trainer.balance_batch=True
+    trainer.critic_warmup=0
+    'trainer.logger=["console"]'
+    trainer.project_name=${project_name}
+    trainer.experiment_name=${experiment_name}
+    trainer.n_gpus_per_node=${NPUS_PER_NODE}
+    trainer.nnodes=${NNODES}
+    trainer.device=npu
+    trainer.val_before_train=False
+    trainer.save_freq=${save_freq}
+    trainer.test_freq=${test_freq}
+    trainer.resume_mode=${resume_mode}
+    trainer.max_actor_ckpt_to_keep=${max_actor_ckpt_to_keep}
+    trainer.total_epochs=${total_epochs}
+    trainer.total_training_steps=${total_training_steps}
+    "trainer.default_local_dir=${output_dir}"
+)
+
+EXTRA=(
+    model_engine=megatron
+)
+
+########################### launch ###########################
+
+log_file="${log_dir}/${experiment_name}_$(date +%Y%m%d_%H%M%S).log"
+PYTHONUNBUFFERED=1 python3 -m verl.trainer.main_ppo \
+    "${DATA[@]}" \
+    "${ALGORITHM[@]}" \
+    "${MODEL[@]}" \
+    "${ROLLOUT[@]}" \
+    "${ACTOR[@]}" \
+    "${REF[@]}" \
+    "${REWARD[@]}" \
+    "${TRAINER[@]}" \
+    "${EXTRA[@]}" \
+    "$@" 2>&1 | tee "${log_file}"

tests/trainer/ppo/test_gdpo_on_cpu.py

@@ -0,0 +1,102 @@
+# Copyright 2024 Bytedance Ltd. and/or its affiliates
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""CPU coverage for GDPO advantage and reward components."""
+
+from __future__ import annotations
+
+import random
+
+import numpy as np
+import pytest
+import torch
+
+from verl.experimental.reward_loop.reward_manager import get_reward_manager_cls
+from verl.trainer.ppo.core_algos import (
+    AdvantageEstimator,
+    compute_gdpo_outcome_advantage,
+    get_adv_estimator_fn,
+)
+from verl.utils.reward_score import rlla
+
+
+def _gdpo_inputs():
+    token_rewards = torch.zeros((4, 2), dtype=torch.float32)
+    response_mask = torch.ones((4, 2), dtype=torch.float32)
+    index = np.asarray(["a", "a", "b", "b"], dtype=object)
+    non_tensor_batch = {
+        "accuracy_reward": np.asarray([1.0, 0.0, 0.0, 1.0]),
+        "format_reward": np.asarray([0.0, 1.0, 0.0, 1.0]),
+    }
+    batch = {
+        "prompts": torch.zeros((4, 2), dtype=torch.long),
+        "attention_mask": torch.ones((4, 4), dtype=torch.long),
+    }
+    config = {"gdpo_reward_keys": ["accuracy_reward", "format_reward"]}
+    return token_rewards, response_mask, index, non_tensor_batch, batch, config
+
+
+def test_gdpo_normalizes_each_reward_dimension_before_sum() -> None:
+    inputs = _gdpo_inputs()
+    advantages, returns = compute_gdpo_outcome_advantage(
+        token_level_rewards=inputs[0],
+        response_mask=inputs[1],
+        index=inputs[2],
+        non_tensor_batch=inputs[3],
+        batch=inputs[4],
+        config=inputs[5],
+    )
+
+    torch.testing.assert_close(advantages, returns)
+    torch.testing.assert_close(advantages[0], torch.zeros(2), atol=1e-6, rtol=0)
+    torch.testing.assert_close(advantages[1], torch.zeros(2), atol=1e-6, rtol=0)
+    assert torch.all(advantages[2] < 0)
+    assert torch.all(advantages[3] > 0)
+    assert advantages.mean().item() == pytest.approx(0.0, abs=1e-6)
+
+
+def test_gdpo_requires_declared_reward_components() -> None:
+    inputs = _gdpo_inputs()
+    with pytest.raises(AssertionError, match="format_reward"):
+        compute_gdpo_outcome_advantage(
+            token_level_rewards=inputs[0],
+            response_mask=inputs[1],
+            index=inputs[2],
+            non_tensor_batch={"accuracy_reward": inputs[3]["accuracy_reward"]},
+            batch=inputs[4],
+            config=inputs[5],
+        )
+
+
+def test_gdpo_registries() -> None:
+    assert get_adv_estimator_fn("gdpo") is compute_gdpo_outcome_advantage
+    assert get_adv_estimator_fn(AdvantageEstimator.GDPO) is compute_gdpo_outcome_advantage
+    assert get_reward_manager_cls("gdpo").__name__ == "GDPORewardManager"
+
+
+def test_rlla_reward_returns_total_and_decoupled_components(monkeypatch) -> None:
+    monkeypatch.setattr(random, "randint", lambda *_args: 2)
+    answer = '<think>Call the tool.</think>\n<tool_call>\n{"name":"lookup","parameters":{"id":1}}\n</tool_call>'
+    result = rlla.compute_score(
+        data_source="rlla",
+        solution_str=answer,
+        ground_truth=answer,
+        extra_info={"experiment_name": "qwen2_5_1_5b_gdpo"},
+    )
+
+    assert result == {
+        "score": 4.0,
+        "format_reward": 1.0,
+        "accuracy_reward": 3.0,
+    }