diff --git a/verl/trainer/config/generation.yaml b/verl/trainer/config/generation.yaml index d3068e886c9..e940c423b49 100644 --- a/verl/trainer/config/generation.yaml +++ b/verl/trainer/config/generation.yaml @@ -40,6 +40,11 @@ rollout: disable_log_stats: True enable_chunked_prefill: True n: 1 + length_penalty: + enabled: False # Set to True to enable length penalty during generation + alpha: 0.0 # Alpha parameter: positive favors longer sequences, negative favors shorter + min_length: 0 # Minimum sequence length before applying penalty + max_length: null # Maximum sequence length (null means no maximum) actor: strategy: fsdp # This is for backward-compatibility ulysses_sequence_parallel_size: 1 # sp size diff --git a/verl/trainer/config/ppo_megatron_trainer.yaml b/verl/trainer/config/ppo_megatron_trainer.yaml index 22420d72062..f497d012153 100644 --- a/verl/trainer/config/ppo_megatron_trainer.yaml +++ b/verl/trainer/config/ppo_megatron_trainer.yaml @@ -223,6 +223,11 @@ reward_model: use_dynamic_bsz: ${critic.use_dynamic_bsz} max_length: null launch_reward_fn_async: False # custom reward function executed async on CPU, during log_prob + length_penalty: + enabled: False # Set to True to enable length penalty + alpha: 0.0 # Alpha parameter: positive favors longer sequences, negative favors shorter + min_length: 0 # Minimum sequence length before applying penalty + max_length: null # Maximum sequence length (null means no maximum) custom_reward_function: path: null @@ -252,6 +257,9 @@ trainer: nnodes: 1 n_gpus_per_node: 8 save_freq: -1 + track_token_lengths: true + positive_reward_threshold: 0.5 + token_length_tag_prefix: "token_length" # auto: find the last ckpt to resume. If can't find, start from scratch resume_mode: auto # or disable or resume_path if resume_from_path is set resume_from_path: null diff --git a/verl/trainer/config/ppo_trainer.yaml b/verl/trainer/config/ppo_trainer.yaml index ca14b378bf9..3723631419e 100644 --- a/verl/trainer/config/ppo_trainer.yaml +++ b/verl/trainer/config/ppo_trainer.yaml @@ -183,6 +183,11 @@ reward_model: forward_max_token_len_per_gpu: ${critic.forward_max_token_len_per_gpu} reward_manager: naive launch_reward_fn_async: False # custom reward function executed async on CPU, during log_prob + length_penalty: + enabled: False # Set to True to enable length penalty + alpha: 0.0 # Alpha parameter: positive favors longer sequences, negative favors shorter + min_length: 0 # Minimum sequence length before applying penalty + max_length: null # Maximum sequence length (null means no maximum) custom_reward_function: path: null @@ -214,6 +219,9 @@ trainer: nnodes: 1 n_gpus_per_node: 8 save_freq: -1 + track_token_lengths: true + positive_reward_threshold: 0.5 + token_length_tag_prefix: "token_length" # auto: find the last ckpt to resume. If can't find, start from scratch resume_mode: auto # or disable or resume_path if resume_from_path is set resume_from_path: null diff --git a/verl/trainer/ppo/ray_trainer.py b/verl/trainer/ppo/ray_trainer.py index 1637a40b8bf..ad796d4233f 100644 --- a/verl/trainer/ppo/ray_trainer.py +++ b/verl/trainer/ppo/ray_trainer.py @@ -825,6 +825,65 @@ def _balance_batch(self, batch: DataProto, metrics, logging_prefix="global_seqle global_balance_stats = log_seqlen_unbalance(seqlen_list=global_seqlen_lst, partitions=global_partition_lst, prefix=logging_prefix) metrics.update(global_balance_stats) + def log_token_lengths(self, data, step, metrics=None): + """Log average token lengths for positive and negative responses.""" + if not self.config.trainer.get('track_token_lengths', False): + return + + threshold = self.config.trainer.get('positive_reward_threshold', 0.5) + tag_prefix = self.config.trainer.get('token_length_tag_prefix', 'token_length') + + try: + if 'acc' in data.batch: + rewards = data.batch['acc'].cpu().tolist() + elif 'token_level_scores' in data.batch: + rewards = data.batch['token_level_scores'].sum(-1).cpu().tolist() + else: + print("Warning: Cannot track token lengths - no reward information found") + return + except Exception as e: + print(f"Warning: Error computing rewards for token length tracking: {e}") + return + + responses = data.batch['responses'] + attention_mask = data.batch['attention_mask'] + prompt_len = data.batch['prompts'].shape[-1] + valid_response_lengths = attention_mask[:, prompt_len:].sum(dim=-1) + + positive_tokens = [] + negative_tokens = [] + + for i, reward in enumerate(rewards): + length = valid_response_lengths[i].item() + response_tokens = responses[i][:length].tolist() + + if reward > threshold: + positive_tokens.append(response_tokens) + else: + negative_tokens.append(response_tokens) + + pos_avg_len = 0 + neg_avg_len = 0 + + if positive_tokens: + pos_avg_len = sum(len(tokens) for tokens in positive_tokens) / len(positive_tokens) + + if negative_tokens: + neg_avg_len = sum(len(tokens) for tokens in negative_tokens) / len(negative_tokens) + + log_dict = { + f"{tag_prefix}/positive_avg_token_length": pos_avg_len, + f"{tag_prefix}/negative_avg_token_length": neg_avg_len + } + + if metrics is not None: + metrics.update(log_dict) + + if hasattr(self, 'tracker') and self.tracker: + self.tracker.log(log_dict, step=step) + + return log_dict + def fit(self): """ The training loop of PPO. @@ -968,9 +1027,41 @@ def fit(self): reward_tensor, reward_extra_infos_dict = ray.get(future_reward) batch.batch["token_level_scores"] = reward_tensor + batch.batch["acc"] = batch.batch["token_level_scores"].sum(-1) + + token_metrics = self.log_token_lengths(batch, self.global_steps, metrics) + print(f"{list(reward_extra_infos_dict.keys())=}") + if reward_extra_infos_dict: - batch.non_tensor_batch.update({k: np.array(v) for k, v in reward_extra_infos_dict.items()}) + for key in ["length_penalty_applied", "length_penalty_alpha", + "original_rewards_mean", "penalized_rewards_mean", "penalty_ratio"]: + if key in reward_extra_infos_dict: + metrics[f"reward/length_penalty/{key}"] = reward_extra_infos_dict[key] + + for key, value in reward_extra_infos_dict.items(): + if isinstance(value, (int, float)): # Scalar values only + if key not in ["length_penalty_applied", "length_penalty_alpha", + "original_rewards_mean", "penalized_rewards_mean", "penalty_ratio"]: + if "score" in key: + metrics[f"reward/scores/{key}"] = value + elif "format" in key: + metrics[f"reward/format/{key}"] = value + elif "proof" in key: + metrics[f"reward/proof/{key}"] = value + else: + metrics[f"reward/other/{key}"] = value + + safe_dict = {} + for k, v in reward_extra_infos_dict.items(): + if isinstance(v, list) and len(v) == len(batch): + try: + safe_dict[k] = np.array(v) + except Exception as e: + print(f"Warning: Could not convert reward extra info '{k}' to numpy array: {e}") + + if safe_dict: + batch.non_tensor_batch.update(safe_dict) # compute rewards. apply_kl_penalty if available if self.config.algorithm.use_kl_in_reward: diff --git a/verl/trainer/ppo/reward.py b/verl/trainer/ppo/reward.py index 73d1c0a82f8..bd7a9b04b29 100644 --- a/verl/trainer/ppo/reward.py +++ b/verl/trainer/ppo/reward.py @@ -15,8 +15,10 @@ import os import ray +import torch from verl import DataProto +from verl.utils.length_penalty import apply_length_penalty def get_custom_reward_fn(config): @@ -75,6 +77,10 @@ def load_reward_manager(config, tokenizer, num_examine, **reward_kwargs): else: raise NotImplementedError + # Pass length penalty config to reward manager + length_penalty_config = dict(config.reward_model.get("length_penalty", {})) + reward_kwargs.update({"length_penalty_config": length_penalty_config}) + compute_score = get_custom_reward_fn(config) return reward_manager_cls( tokenizer=tokenizer, @@ -102,6 +108,52 @@ def compute_reward(data: DataProto, reward_fn): print(f"Error in reward_fn: {e}") reward_tensor = reward_fn(data) reward_extra_infos_dict = {} + + # Apply length penalty if configured and if sequence lengths are available + if hasattr(reward_fn, "length_penalty_config") and reward_fn.length_penalty_config.get("enabled", False): + sequence_lengths = None + if "response_lengths" in data.batch: + sequence_lengths = data.batch["response_lengths"] + elif "attention_mask" in data.batch and "prompts" in data.batch: + prompt_len = data.batch["prompts"].shape[-1] + attention_mask = data.batch["attention_mask"] + total_lengths = torch.sum(attention_mask, dim=1) + sequence_lengths = total_lengths - prompt_len + + if sequence_lengths is not None: + alpha = reward_fn.length_penalty_config.get("alpha", 0.0) + min_length = reward_fn.length_penalty_config.get("min_length", 0) + max_length = reward_fn.length_penalty_config.get("max_length", None) + + orig_rewards = reward_tensor.sum(-1) + + print(f"Applying length penalty with alpha={alpha}, min_length={min_length}, max_length={max_length}") + print(f"Mean sequence length: {sequence_lengths.float().mean().item():.1f}") + print(f"Original rewards mean: {orig_rewards.mean().item():.4f}") + + penalized_rewards = apply_length_penalty( + orig_rewards, + sequence_lengths, + alpha=alpha, + min_length=min_length, + max_length=max_length + ) + + print(f"Penalized rewards mean: {penalized_rewards.mean().item():.4f}") + print(f"Penalty ratio: {(penalized_rewards / (orig_rewards + 1e-8)).mean().item():.4f}") + + # Compute scaling factor for each sequence + scaling_factors = penalized_rewards / (orig_rewards + 1e-8) # Add small epsilon to avoid division by zero + + # Apply scaling factors to token-level rewards + reward_tensor = reward_tensor * scaling_factors.unsqueeze(-1) + + if reward_extra_infos_dict is not None: + reward_extra_infos_dict["length_penalty_applied"] = True + reward_extra_infos_dict["length_penalty_alpha"] = alpha + reward_extra_infos_dict["original_rewards_mean"] = orig_rewards.mean().item() + reward_extra_infos_dict["penalized_rewards_mean"] = penalized_rewards.mean().item() + reward_extra_infos_dict["penalty_ratio"] = (penalized_rewards / (orig_rewards + 1e-8)).mean().item() return reward_tensor, reward_extra_infos_dict @@ -113,4 +165,4 @@ def compute_reward_async(data: DataProto, config, tokenizer): This is meant to be run in a separate Ray worker. """ reward_fn = load_reward_manager(config, tokenizer, num_examine=0, **config.reward_model.get("reward_kwargs", {})) - return compute_reward(data, reward_fn) + return compute_reward(data, reward_fn) \ No newline at end of file diff --git a/verl/utils/length_penalty.py b/verl/utils/length_penalty.py new file mode 100644 index 00000000000..b2770a62297 --- /dev/null +++ b/verl/utils/length_penalty.py @@ -0,0 +1,114 @@ +# Copyright 2025 Individual Contributors +# +# 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. + +import torch +import numpy as np + +def compute_length_penalty(sequence_lengths, alpha=0.0, min_length=0, max_length=None): + """ + Compute length penalty for sequence generation. + + Args: + sequence_lengths: Tensor or numpy array of shape [batch_size] containing sequence lengths + alpha: Float controlling the strength and direction of the penalty: + - alpha > 0: Favor longer sequences + - alpha < 0: Favor shorter sequences + - alpha = 0: No length penalty + min_length: Minimum sequence length before applying penalty (no penalty below this) + max_length: Maximum sequence length (sequences longer than this get max penalty) + + Returns: + Tensor or numpy array of shape [batch_size] containing length penalties + """ + if alpha == 0.0: + if isinstance(sequence_lengths, torch.Tensor): + return torch.ones_like(sequence_lengths, dtype=torch.float32) + else: + return np.ones_like(sequence_lengths, dtype=np.float32) + + effective_lengths = sequence_lengths.copy() if isinstance(sequence_lengths, np.ndarray) else sequence_lengths.clone() + + if isinstance(effective_lengths, np.ndarray): + effective_lengths = effective_lengths.astype(np.float32) + else: + effective_lengths = effective_lengths.float() + + if min_length > 0: + if isinstance(effective_lengths, np.ndarray): + effective_lengths = np.maximum(effective_lengths - min_length, np.zeros_like(effective_lengths)) + else: + effective_lengths = torch.maximum(effective_lengths - min_length, + torch.zeros_like(effective_lengths)) + + if max_length is not None: + if isinstance(effective_lengths, np.ndarray): + effective_lengths = np.minimum(effective_lengths, np.ones_like(effective_lengths) * (max_length - min_length)) + else: + effective_lengths = torch.minimum(effective_lengths, + torch.ones_like(effective_lengths) * (max_length - min_length)) + + # Calculate penalty using the standard formula: ((5 + length)/6)^alpha + # This is similar to the formula used in Google Neural Machine Translation paper [https://arxiv.org/pdf/1609.08144] + if isinstance(effective_lengths, np.ndarray): + penalty = ((5.0 + effective_lengths) / 6.0) ** alpha + else: + penalty = ((5.0 + effective_lengths) / 6.0) ** alpha + + return penalty + +def apply_length_penalty(rewards, sequence_lengths, **kwargs): + """ + Apply length penalty to rewards. + + Args: + rewards: List, Tensor or numpy array of shape [batch_size] containing the original rewards + sequence_lengths: List, Tensor or numpy array of shape [batch_size] containing sequence lengths + **kwargs: Parameters for compute_length_penalty + + Returns: + List, Tensor or numpy array of shape [batch_size] with length-penalized rewards + """ + import numpy as np + import torch + + is_tensor = isinstance(rewards, torch.Tensor) + is_list = isinstance(rewards, list) + + if is_list: + rewards_np = np.array(rewards) + elif is_tensor: + rewards_np = rewards.cpu().numpy() + else: + rewards_np = rewards + + if isinstance(sequence_lengths, torch.Tensor): + sequence_lengths_np = sequence_lengths.cpu().numpy() + elif isinstance(sequence_lengths, list): + sequence_lengths_np = np.array(sequence_lengths) + else: + sequence_lengths_np = sequence_lengths + + length_penalties = compute_length_penalty(sequence_lengths_np, **kwargs) + + penalized_rewards_np = rewards_np * length_penalties + + if is_tensor: + device = rewards.device + penalized_rewards = torch.tensor(penalized_rewards_np, dtype=rewards.dtype, device=device) + elif is_list: + penalized_rewards = penalized_rewards_np.tolist() + else: + penalized_rewards = penalized_rewards_np + + return penalized_rewards \ No newline at end of file diff --git a/verl/workers/reward_manager/batch.py b/verl/workers/reward_manager/batch.py index 570fdd71dea..b5913f40d75 100644 --- a/verl/workers/reward_manager/batch.py +++ b/verl/workers/reward_manager/batch.py @@ -17,7 +17,7 @@ import torch from verl import DataProto - +from verl.utils.length_penalty import apply_length_penalty class BatchRewardManager: def __init__(self, tokenizer, num_examine, compute_score, reward_fn_key="data_source", **reward_kwargs): @@ -25,6 +25,8 @@ def __init__(self, tokenizer, num_examine, compute_score, reward_fn_key="data_so self.num_examine = num_examine self.compute_score = compute_score self.reward_fn_key = reward_fn_key + self.length_penalty_config = reward_kwargs.pop("length_penalty_config", {}) + self.use_length_penalty = self.length_penalty_config.get("enabled", False) self.reward_kwargs = reward_kwargs def verify(self, data): @@ -106,4 +108,4 @@ def __call__(self, data: DataProto, return_dict=False): if return_dict: return {"reward_tensor": reward_tensor, "reward_extra_info": reward_extra_info} else: - return reward_tensor + return reward_tensor \ No newline at end of file diff --git a/verl/workers/reward_model/base.py b/verl/workers/reward_model/base.py index cb719bd0f9b..38dee2a16ea 100644 --- a/verl/workers/reward_model/base.py +++ b/verl/workers/reward_model/base.py @@ -16,13 +16,21 @@ """ from abc import ABC, abstractmethod +import torch from verl import DataProto +from verl.utils.length_penalty import apply_length_penalty class BasePPORewardModel(ABC): def __init__(self, config): self.config = config + # Initialize length penalty parameters + self.length_penalty_config = config.reward_model.get("length_penalty", {}) + self.use_length_penalty = self.length_penalty_config.get("enabled", False) + self.length_penalty_alpha = self.length_penalty_config.get("alpha", 0.0) + self.length_penalty_min_length = self.length_penalty_config.get("min_length", 0) + self.length_penalty_max_length = self.length_penalty_config.get("max_length", None) @abstractmethod def compute_reward(self, data: DataProto) -> DataProto: @@ -42,3 +50,47 @@ def compute_reward(self, data: DataProto) -> DataProto: """ pass + + def apply_length_penalty(self, rewards, sequence_lengths): + """ + Apply length penalty to the rewards if enabled. + + Args: + rewards: Tensor of shape [batch_size] containing rewards + sequence_lengths: Tensor of shape [batch_size] containing sequence lengths + + Returns: + Tensor of shape [batch_size] with potentially modified rewards + """ + if not self.use_length_penalty: + return rewards + + return apply_length_penalty( + rewards, + sequence_lengths, + alpha=self.length_penalty_alpha, + min_length=self.length_penalty_min_length, + max_length=self.length_penalty_max_length + ) + + def _get_sequence_lengths(self, data): + """ + Extract response lengths from the data. + + Args: + data: DataProto object with input_ids and attention_mask + + Returns: + Tensor of shape [batch_size] containing sequence lengths + """ + if "response_lengths" in data: + return data["response_lengths"] + + # If response_lengths not provided, try to compute from attention_mask + if "attention_mask" in data and "prompt_lengths" in data: + # Calculate response length by subtracting prompt length from total length + total_lengths = torch.sum(data["attention_mask"], dim=1) + prompt_lengths = data["prompt_lengths"] + return total_lengths - prompt_lengths + + return None \ No newline at end of file