feat: add core_algo for GDPO

Rhetee · Rhetee · commit d149dd9804c6 · 2026-02-27T17:14:02.000+08:00
diff --git a/examples/gdpo_trainer/run_gdpo.sh b/examples/gdpo_trainer/run_gdpo.sh
@@ -0,0 +1,50 @@
+export DATA_DIR="verl/dataset/rlla_4k"
+export BASE_MODEL="Qwen/Qwen2.5-1.5B-Instruct"
+export EXPERIMENT_NAME="qwen2.5-1.5B-GDPO"
+export CKPT_DIR="verl/results/gdpo"
+
+PROJECT_DIR="$(pwd)"
+
+trainer_n_gpus_per_node=8
+trainer_nnodes=1
+
+python3 -u -m verl.trainer.main_ppo \
+    algorithm.adv_estimator=gdpo \
+    data.train_files=$DATA_DIR/train.parquet \
+    data.val_files=$DATA_DIR/test.parquet \
+    data.train_batch_size=32 \
+    data.val_batch_size=16 \
+    data.max_prompt_length=2048 \
+    data.max_response_length=1024 \
+    actor_rollout_ref.model.path=$BASE_MODEL \
+    actor_rollout_ref.actor.optim.lr=1e-6 \
+    actor_rollout_ref.model.use_remove_padding=True \
+    actor_rollout_ref.actor.ppo_mini_batch_size=4 \
+    actor_rollout_ref.actor.use_dynamic_bsz=True \
+    actor_rollout_ref.actor.use_kl_loss=False \
+    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.grad_offload=False \
+    actor_rollout_ref.actor.fsdp_config.optimizer_offload=False \
+    actor_rollout_ref.rollout.tensor_model_parallel_size=1 \
+    actor_rollout_ref.rollout.name=vllm \
+    actor_rollout_ref.rollout.gpu_memory_utilization=0.6 \
+    actor_rollout_ref.rollout.n=4 \
+    actor_rollout_ref.ref.fsdp_config.param_offload=True \
+    algorithm.kl_ctrl.kl_coef=0.001 \
+    reward.custom_reward_function.path="$PROJECT_DIR/verl/experimental/reward_loop/reward_manager/gdpo.py" \
+    reward.custom_reward_function.name=compute_score \
+    trainer.critic_warmup=0 \
+    trainer.logger=['console'] \
+    trainer.project_name=Var_inspect \
+    trainer.n_gpus_per_node=$trainer_n_gpus_per_node \
+    trainer.experiment_name=$EXPERIMENT_NAME \
+    trainer.n_gpus_per_node=$trainer_nnodes \
+    trainer.nnodes=1 \
+    trainer.save_freq=5 \
+    trainer.test_freq=10 \
+    trainer.default_local_dir=$CKPT_DIR \
+    trainer.total_epochs=15 \
+    trainer.val_before_train=False 2>&1 | tee ${LOG_PATH}
diff --git a/verl/experimental/reward_loop/reward_manager/gdpo.py b/verl/experimental/reward_loop/reward_manager/gdpo.py
@@ -78,6 +78,14 @@ async def run_single(self, data: DataProto) -> dict:
                 ),
             )
 
+        # result = {
+        # "score": score,
+        # "score_list": [fomrat_score, correctness_score],
+        # }
+
+        # return = {"reward_score": reward, "reward_extra_info": reward_extra_info}
+        # reward_extra_info = {"score": score, "score_list": [fomrat_score, correctness_score]}
+
         reward_extra_info = {}
 
         score: float
diff --git a/verl/trainer/ppo/core_algos.py b/verl/trainer/ppo/core_algos.py
@@ -96,6 +96,7 @@ class AdvantageEstimator(str, Enum):
 
     GAE = "gae"
     GRPO = "grpo"
+    GDPO = "gdpo"
     REINFORCE_PLUS_PLUS = "reinforce_plus_plus"
     REINFORCE_PLUS_PLUS_BASELINE = "reinforce_plus_plus_baseline"
     REMAX = "remax"
@@ -2370,3 +2371,70 @@ def compute_policy_loss_bypass_mode(
     pg_metrics.update(rollout_metrics)
 
     return pg_loss, pg_metrics
+
+
+@register_adv_est(AdvantageEstimator.GDPO)  # or simply: @register_adv_est("gdpo")
+def compute_gdpo_outcome_advantage(
+    token_level_rewards: torch.Tensor,
+    response_mask: torch.Tensor,
+    index: np.ndarray,
+    epsilon: float = 1e-6,
+    norm_adv_by_std_in_grpo: bool = True,
+    config: Optional[AlgoConfig] = None,
+    score_list: Optional[list[torch.Tensor]] = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """
+    Compute advantage for GDPO, operating only on Outcome reward
+    (with only one scalar reward for each response).
+
+    Args:
+        token_level_rewards: `(torch.Tensor)`
+            shape is (bs, response_length)
+        response_mask: `(torch.Tensor)`
+            shape is (bs, response_length)
+        index: `(np.ndarray)`
+            index array for grouping
+        epsilon: `(float)`
+            small value to avoid division by zero
+        norm_adv_by_std_in_grpo: `(bool)`
+            whether to scale the GRPO advantage
+        config: `(Optional[AlgoConfig])`
+            algorithm configuration object
+        score_list: `(Optional[list[torch.Tensor]])`
+            multi scores for GDPO
+
+    Note:
+        Ref GDPO (https://arxiv.org/abs/2601.05242).
+
+    Returns:
+        advantages: `(torch.Tensor)`
+            shape is (bs, response_length)
+        Returns: `(torch.Tensor)`
+            shape is (bs, response_length)
+    """
+    if score_list is None:
+        score_list = [token_level_rewards]
+        # for debug
+        print("------no multi-score-find---------")
+    num_scores = len(score_list)
+    new_advantage = None
+    for i in range(num_scores):
+        token_level_scores = score_list[i]
+
+        normalized_score, _ = compute_grpo_outcome_advantage(
+            token_level_rewards=token_level_scores,
+            response_mask=response_mask,
+            index=index,
+            epsilon=epsilon,
+            norm_adv_by_std_in_grpo=norm_adv_by_std_in_grpo,
+            config=config,
+        )
+
+        if new_advantage is None:
+            new_advantage = normalized_score
+        else:
+            new_advantage += normalized_score
+
+    advantages = verl_F.masked_whiten(new_advantage, response_mask) * response_mask
+
+    return advantages, advantages
diff --git a/verl/trainer/ppo/ray_trainer.py b/verl/trainer/ppo/ray_trainer.py
@@ -211,6 +211,58 @@ def compute_advantage(
             rollout_is_weights = data.batch.get("rollout_is_weights", None)
             adv_kwargs["rollout_is_weights"] = rollout_is_weights
 
+        if adv_estimator == AdvantageEstimator.GDPO:
+            assert "score_list" in data.batch, (
+                "GDPO need multi-scores. "
+                "Please change the config: reward.custom_reward_function.path to point gdpo.py or "
+                "change the reward function to compute multi-scores."
+            )
+
+            # prompt_length = prompt_ids.size(1)
+            # response_length = attention_mask[:, prompt_length:].sum(dim=1) - 1
+            # rm_scores = torch.zeros_like(response_mask, dtype=torch.float32)
+            # rm_scores[torch.arange(response_mask.size(0)), response_length] =
+            # torch.tensor(scores, dtype=torch.float32)
+            # batch["rm_scores"] = rm_scores
+
+            #     batch = TensorDict(
+            #     {
+            #         "prompts": prompt_ids,  # [bsz, prompt_length]
+            #         "responses": response_ids,  # [bsz, response_length]
+            #         "response_mask": response_mask,  # [bsz, response_length]
+            #         "input_ids": input_ids,  # [bsz, prompt_length + response_length]
+            #         "attention_mask": attention_mask,  # [bsz, prompt_length + response_length]
+            #         # position_ids: [bsz, 3, prompt_length + response_length]
+            # or [bsz, prompt_length + response_length]
+            #         "position_ids": position_ids,
+            #         **optional_outputs,
+            #     },
+            #     batch_size=len(inputs),
+            # )
+            score_list = []
+            multi_score_tensor = torch.tensor(
+                data.non_tensor_batch["score_list"], dtype=torch.float32
+            )  # # [bsz, score_num, 1]
+            print(f"----------multi_score_tensor:{multi_score_tensor.shape}")
+
+            for i in range(multi_score_tensor.shape[1]):
+                rm_score = multi_score_tensor[:, i]
+                prompt_length = data.batch["prompts"].size(1)
+                response_length = data.batch["attention_mask"][:, prompt_length:].sum(dim=1) - 1
+                rm_scores = torch.zeros_like(data.batch["response_mask"], dtype=torch.float32)
+                rm_scores[torch.arange(data.batch["response_mask"].size(0)), response_length] = torch.tensor(
+                    rm_score, dtype=torch.float32
+                )
+                score_list.append(rm_scores)
+
+            # sum_score_tensor = data.batch["token_level_rewards"]
+
+            # rm_scores[torch.arange(rm_scores.size(0)), valid_response_length - 1] = torch.tensor(
+            #     scores, dtype=torch.float32
+            # )
+            adv_kwargs["score_list"] = score_list
+
+            # np.array([[format_score,correct_score] for info in reward_extra_infos])
         # calculate advantage estimator
         advantages, returns = adv_estimator_fn(**adv_kwargs)
         data.batch["advantages"] = advantages
diff --git a/verl/utils/reward_score/rlla.py b/verl/utils/reward_score/rlla.py
@@ -306,9 +306,6 @@ def compute_score(solution_str, ground_truth, step=0):
     format_max_possible = 1.0
     format_min_possible = 0.0
 
-    length_max_possible = 1.0
-    length_min_possible = 0.0
-
     completions = [[{"role": "assistant", "content": predict_str}]]
     answer = [ground_truth]
 
@@ -317,14 +314,11 @@ def compute_score(solution_str, ground_truth, step=0):
         completions, answer, step, tool_max_possible, tool_min_possible
     )[0]
 
-    if str(os.getenv("WITHLENGTH", 0)) == "1":
-        print("WITHLENGTH is set to 1, so length score is set!")
-        length_score = customize_length_reward_func(
-            completions, answer, step, length_max_possible, length_min_possible
-        )[0]
-    else:
-        length_score = 0
+    score = fomrat_score + correctness_score
 
-    score = fomrat_score + correctness_score + length_score
+    result = {
+        "score": score,
+        "score_list": [fomrat_score, correctness_score],
+    }
 
-    return score, fomrat_score, correctness_score, length_score
+    return result
