gad train

YTianZHU · YTianZHU · commit e39e564cff59 · 2025-11-14T22:43:35.000-08:00
diff --git a/verl/trainer/ppo/core_algos.py b/verl/trainer/ppo/core_algos.py
@@ -205,7 +205,6 @@ def compute_grpo_outcome_advantage(
     index: np.ndarray,
     epsilon: float = 1e-6,
     norm_adv_by_std_in_grpo: str = True,
-    compute_teacher: bool = False,
 ):
     """
     Compute advantage for GRPO, operating only on Outcome reward
@@ -230,7 +229,6 @@ def compute_grpo_outcome_advantage(
     scores = token_level_rewards.sum(dim=-1)
 
     id2score = defaultdict(list)
-    id2adv = defaultdict(list)
     id2mean = {}
     id2std = {}
 
@@ -252,12 +250,6 @@ def compute_grpo_outcome_advantage(
                 scores[i] = (scores[i] - id2mean[index[i]]) / (id2std[index[i]] + epsilon)
             else:
                 scores[i] = scores[i] - id2mean[index[i]]
-            id2adv[index[i]].append(scores[i])
-        
-        if compute_teacher:
-            for i in range(bsz):
-                scores[i] = max(id2adv[index[i]])
-        
         scores = scores.unsqueeze(-1) * response_mask
 
     return scores, scores
@@ -566,16 +558,6 @@ def agg_loss(loss_mat: torch.Tensor, loss_mask: torch.Tensor, loss_agg_mode: str
     return loss
 
 
-def compute_sft_loss(
-    log_prob,
-    response_mask,
-    loss_agg_mode: str = "token-mean",
-):
-    pg_loss = agg_loss(loss_mat=-log_prob, loss_mask=response_mask, loss_agg_mode=loss_agg_mode)
-
-    return pg_loss
-
-
 def compute_policy_loss(
     old_log_prob,
     log_prob,
diff --git a/verl/trainer/ppo/ray_trainer.py b/verl/trainer/ppo/ray_trainer.py
@@ -180,7 +180,7 @@ def apply_kl_penalty(data: DataProto, kl_ctrl: core_algos.AdaptiveKLController,
     return data, metrics
 
 
-def compute_response_mask(data: DataProto, compute_teacher=False):
+def compute_response_mask(data: DataProto):
     """Compute the attention mask for the response part of the sequence.
 
     This function extracts the portion of the attention mask that corresponds to the model's response,
@@ -192,16 +192,10 @@ def compute_response_mask(data: DataProto, compute_teacher=False):
     Returns:
         torch.Tensor: The attention mask for the response tokens.
     """
-    if compute_teacher:
-        responses = data.batch["teacher_response"]
-        response_length = responses.size(1)
-        attention_mask = data.batch["teacher_attention_mask"]
-        return attention_mask[:, -response_length:]
-    else:
-        responses = data.batch["responses"]
-        response_length = responses.size(1)
-        attention_mask = data.batch["attention_mask"]
-        return attention_mask[:, -response_length:]
+    responses = data.batch["responses"]
+    response_length = responses.size(1)
+    attention_mask = data.batch["attention_mask"]
+    return attention_mask[:, -response_length:]
 
 
 def compute_advantage(data: DataProto, adv_estimator, gamma=1.0, lam=1.0, num_repeat=1, multi_turn=False, norm_adv_by_std_in_grpo=True, config=None):
@@ -225,7 +219,7 @@ def compute_advantage(data: DataProto, adv_estimator, gamma=1.0, lam=1.0, num_re
     """
     # Back-compatible with trainers that do not compute response mask in fit
     if "response_mask" not in data.batch.keys():
-        data.batch["response_mask"] = compute_response_mask(data, compute_teacher=False)
+        data.batch["response_mask"] = compute_response_mask(data)
     # prepare response group
     if adv_estimator == AdvantageEstimator.GAE:
         # Compute advantages and returns using Generalized Advantage Estimation (GAE)
@@ -555,18 +549,23 @@ def _create_dataloader(self, train_dataset, val_dataset, collate_fn, train_sampl
         except Exception as e:
             print(f"Warning: Could not set total_training_steps in config. Structure missing? Error: {e}")
 
-    def _dump_generations(self, sample_inputs, sample_outputs, teacher_outputs, dump_path):
+    def _dump_generations(self, inputs, outputs, scores, reward_extra_infos_dict, dump_path):
         """Dump rollout/validation samples as JSONL."""
         os.makedirs(dump_path, exist_ok=True)
-        filename = os.path.join(dump_path, f"generation_results.jsonl")
+        filename = os.path.join(dump_path, f"{self.global_steps}.jsonl")
 
-        n = len(sample_inputs)
+        n = len(inputs)
         base_data = {
-            "input": sample_inputs,
-            "output": sample_outputs,
-            "teacher_output": teacher_outputs,
+            "input": inputs,
+            "output": outputs,
+            "score": scores,
+            "step": [self.global_steps] * n,
         }
 
+        for k, v in reward_extra_infos_dict.items():
+            if len(v) == n:
+                base_data[k] = v
+
         lines = []
         for i in range(n):
             entry = {k: v[i] for k, v in base_data.items()}
@@ -700,9 +699,10 @@ def safe_rouge_score(ref, cand):
         val_data_dir = self.config.trainer.get("validation_data_dir", None)
         if val_data_dir:
             self._dump_generations(
-                sample_inputs=sample_inputs,
-                sample_outputs=sample_outputs,
-                teacher_outputs=teacher_outputs,
+                inputs=sample_inputs,
+                outputs=sample_outputs,
+                scores=sample_scores,
+                reward_extra_infos_dict=reward_extra_infos_dict,
                 dump_path=val_data_dir,
             )
 
@@ -886,11 +886,12 @@ def _load_checkpoint(self):
 
         actor_path = os.path.join(global_step_folder, "actor")
         critic_path = os.path.join(global_step_folder, "critic")
+        # NOTE: have directly loaded from actor_rollout_ref.model.path and critic.model.path
         # load actor
-        self.actor_rollout_wg.load_checkpoint(actor_path, del_local_after_load=self.config.trainer.del_local_ckpt_after_load)
+        # self.actor_rollout_wg.load_checkpoint(actor_path, del_local_after_load=self.config.trainer.del_local_ckpt_after_load)
         # load critic
-        if self.use_critic:
-            self.critic_wg.load_checkpoint(critic_path, del_local_after_load=self.config.trainer.del_local_ckpt_after_load)
+        # if self.use_critic and os.path.exists(critic_path):
+        #     self.critic_wg.load_checkpoint(critic_path, del_local_after_load=self.config.trainer.del_local_ckpt_after_load)
 
         # load dataloader,
         # TODO: from remote not implemented yet
@@ -1035,9 +1036,7 @@ def fit(self):
 
                     # recompute old_log_probs
                     with marked_timer("old_log_prob", timing_raw, color="blue"):
-                        batch.meta_info["compute_teacher"] = False
                         old_log_prob = self.actor_rollout_wg.compute_log_prob(batch)
-                        
                         entropys = old_log_prob.batch["entropys"]
                         response_masks = batch.batch["response_mask"]
                         loss_agg_mode = self.config.actor_rollout_ref.actor.loss_agg_mode
@@ -1084,18 +1083,18 @@ def fit(self):
                     with marked_timer("reward", timing_raw, color="yellow"):
                         future_reward = None
                         reward_extra_infos_dict = {}
-                        batch.meta_info["compute_teacher"] = False
                         values = self.critic_wg.compute_values(batch)
                         batch = batch.union(values)
                         reward_tensor = batch.batch["values"]
+                        # reward_tensor: (bsz, response_length)
 
                     with marked_timer("adv", timing_raw, color="brown"):
                         # we combine with rule-based rm
                         reward_extra_infos_dict: dict[str, list]
                         if self.config.reward_model.launch_reward_fn_async:
                             reward_tensor, reward_extra_infos_dict = ray.get(future_reward)
                         batch.batch["token_level_scores"] = reward_tensor
-                        
+
                         if reward_extra_infos_dict:
                             batch.non_tensor_batch.update({k: np.array(v) for k, v in reward_extra_infos_dict.items()})
 
@@ -1107,6 +1106,7 @@ def fit(self):
                             batch.batch["token_level_rewards"] = batch.batch["token_level_scores"]
 
                         # compute advantages, executed on the driver process
+
                         norm_adv_by_std_in_grpo = self.config.algorithm.get("norm_adv_by_std_in_grpo", True)  # GRPO adv normalization factor
 
                         batch = compute_advantage(
diff --git a/verl/utils/dataset/rl_dataset.py b/verl/utils/dataset/rl_dataset.py
@@ -267,7 +267,7 @@ def __getitem__(self, item):
         
         if teacher_response is not None:
             teacher_response = self.tokenizer(teacher_response, return_tensors="pt", add_special_tokens=False)
-        
+
         if not self.processor_type == "MiniCPMVImageProcessor":
             input_ids, attention_mask = verl_F.postprocess_data(
                 input_ids=input_ids,
diff --git a/verl/workers/actor/dp_actor.py b/verl/workers/actor/dp_actor.py
@@ -28,7 +28,7 @@
 
 import verl.utils.torch_functional as verl_F
 from verl import DataProto
-from verl.trainer.ppo.core_algos import agg_loss, compute_policy_loss, compute_sft_loss, get_policy_loss_fn, kl_penalty
+from verl.trainer.ppo.core_algos import agg_loss, compute_policy_loss, get_policy_loss_fn, kl_penalty
 from verl.utils.debug import GPUMemoryLogger
 from verl.utils.device import get_device_id, get_device_name, is_cuda_available, is_npu_available
 from verl.utils.fsdp_utils import FSDPModule, fsdp2_clip_grad_norm_
@@ -79,16 +79,13 @@ def __init__(self, config, actor_module: nn.Module, actor_optimizer: torch.optim
         )
         self.device_name = get_device_name()
 
-    def _forward_micro_batch(self, micro_batch, temperature, compute_teacher, calculate_entropy=False) -> Tuple[torch.Tensor, torch.Tensor]:
+    def _forward_micro_batch(self, micro_batch, temperature, calculate_entropy=False) -> Tuple[torch.Tensor, torch.Tensor]:
         """
         Returns:
             entropy: # (bs, response_len)
             log_probs: # (bs, response_len)
         """
-        if compute_teacher:
-            response_length = micro_batch["teacher_response"].size(-1)
-        else:
-            response_length = micro_batch["responses"].size(-1)
+        response_length = micro_batch["responses"].size(-1)
         multi_modal_inputs = {}
         if "multi_modal_inputs" in micro_batch.keys():
             for key in micro_batch["multi_modal_inputs"][0].keys():
@@ -101,16 +98,10 @@ def _forward_micro_batch(self, micro_batch, temperature, compute_teacher, calcul
                     multi_modal_inputs[key] = torch.cat([inputs[key] for inputs in micro_batch["multi_modal_inputs"]], dim=0)
 
         with torch.autocast(device_type=self.device_name, dtype=torch.bfloat16):
-            if compute_teacher:
-                input_ids = micro_batch["teacher_input_ids"]
-                batch_size, seqlen = input_ids.shape
-                attention_mask = micro_batch["teacher_attention_mask"]
-                position_ids = micro_batch["teacher_position_ids"]
-            else:
-                input_ids = micro_batch["input_ids"]
-                batch_size, seqlen = input_ids.shape
-                attention_mask = micro_batch["attention_mask"]
-                position_ids = micro_batch["position_ids"]
+            input_ids = micro_batch["input_ids"]
+            batch_size, seqlen = input_ids.shape
+            attention_mask = micro_batch["attention_mask"]
+            position_ids = micro_batch["position_ids"]
             entropy = None
             if position_ids.dim() == 3:  # qwen2vl mrope
                 position_ids = position_ids.transpose(0, 1)  # (bsz, 3, seqlen) -> (3, bsz, seqlen)
@@ -315,7 +306,6 @@ def compute_log_prob(self, data: DataProto, calculate_entropy=False) -> torch.Te
         Returns:
             torch.Tensor: the log_prob tensor
         """
-        compute_teacher = data.meta_info["compute_teacher"]
         # set to eval
         self.actor_module.eval()
 
@@ -324,10 +314,7 @@ def compute_log_prob(self, data: DataProto, calculate_entropy=False) -> torch.Te
         use_dynamic_bsz = data.meta_info["use_dynamic_bsz"]
 
         def _get_micro_batches(data: DataProto) -> Tuple[list, list | None]:
-            if compute_teacher:
-                select_keys = ["teacher_response", "teacher_input_ids", "teacher_attention_mask", "teacher_position_ids"]
-            else:
-                select_keys = ["responses", "input_ids", "attention_mask", "position_ids"] 
+            select_keys = ["responses", "input_ids", "attention_mask", "position_ids"]
             batch = data.select(batch_keys=select_keys).batch
             has_multi_modal_inputs = "multi_modal_inputs" in data.non_tensor_batch
 
@@ -352,7 +339,7 @@ def _get_micro_batches(data: DataProto) -> Tuple[list, list | None]:
                     return micro_batches_dp, None
             elif use_dynamic_bsz:
                 max_token_len = data.meta_info["max_token_len"] * self.ulysses_sequence_parallel_size
-                micro_batches, indices = rearrange_micro_batches(batch=batch, max_token_len=max_token_len, compute_teacher=compute_teacher)
+                micro_batches, indices = rearrange_micro_batches(batch=batch, max_token_len=max_token_len)
                 return micro_batches, indices
             else:
                 micro_batches = batch.split(micro_batch_size)
@@ -366,7 +353,7 @@ def _get_micro_batches(data: DataProto) -> Tuple[list, list | None]:
             if isinstance(micro_batch, DataProto):
                 micro_batch = {**micro_batch.batch, **micro_batch.non_tensor_batch}
             with torch.no_grad():
-                entropy, log_probs = self._forward_micro_batch(micro_batch, compute_teacher=compute_teacher, temperature=temperature, calculate_entropy=calculate_entropy)
+                entropy, log_probs = self._forward_micro_batch(micro_batch, temperature=temperature, calculate_entropy=calculate_entropy)
             log_probs_lst.append(log_probs)
             if calculate_entropy:
                 entropy_lst.append(entropy)
@@ -387,17 +374,13 @@ def _get_micro_batches(data: DataProto) -> Tuple[list, list | None]:
 
     @GPUMemoryLogger(role="dp actor", logger=logger)
     def update_policy(self, data: DataProto):
-
         # make sure we are in training mode
         self.actor_module.train()
 
         temperature = data.meta_info["temperature"]  # temperature must be in the data.meta_info to avoid silent error
         multi_turn = data.meta_info.get("multi_turn", False)
 
-        select_keys = [
-            "responses", "input_ids", "attention_mask", "position_ids", "old_log_probs", "advantages",
-            "teacher_response", "teacher_input_ids", "teacher_attention_mask", "teacher_position_ids"
-        ]
+        select_keys = ["responses", "input_ids", "attention_mask", "position_ids", "old_log_probs", "advantages"]
         if multi_turn:
             select_keys.append("loss_mask")
         if self.config.use_kl_loss:
@@ -439,7 +422,7 @@ def update_policy(self, data: DataProto):
                         micro_batches = data.select(select_keys, non_tensor_select_keys).chunk(num_micro_batches)
                 elif self.config.use_dynamic_bsz:
                     max_token_len = self.config.ppo_max_token_len_per_gpu * self.ulysses_sequence_parallel_size
-                    micro_batches, _ = rearrange_micro_batches(batch=mini_batch, max_token_len=max_token_len, compute_teacher=False)
+                    micro_batches, _ = rearrange_micro_batches(batch=mini_batch, max_token_len=max_token_len)
                 else:
                     self.gradient_accumulation = self.config.ppo_mini_batch_size // self.config.ppo_micro_batch_size_per_gpu
                     # split batch into micro_batches
@@ -462,16 +445,12 @@ def update_policy(self, data: DataProto):
                     else:
                         data = data.to(get_device_id())  # actor device is cpu when using offload
                     responses = data["responses"]
-                    teacher_response = data["teacher_response"]
                     response_length = responses.size(1)
-                    teacher_response_length = teacher_response.size(1)
                     attention_mask = data["attention_mask"]
-                    teacher_attention_mask = data["teacher_attention_mask"]
                     if multi_turn:
                         response_mask = data["loss_mask"][:, -response_length:]
                     else:
                         response_mask = attention_mask[:, -response_length:]
-                        teacher_response_mask = teacher_attention_mask[:, -teacher_response_length:]
 
                     old_log_prob = data["old_log_probs"]
                     advantages = data["advantages"]
@@ -487,17 +466,22 @@ def update_policy(self, data: DataProto):
                     calculate_entropy = False
                     if entropy_coeff != 0:
                         calculate_entropy = True
-                    teacher_entropy, teacher_log_prob = self._forward_micro_batch(micro_batch=data, compute_teacher=True, temperature=temperature, calculate_entropy=calculate_entropy)
+                    entropy, log_prob = self._forward_micro_batch(micro_batch=data, temperature=temperature, calculate_entropy=calculate_entropy)
 
                     loss_mode = self.config.policy_loss.get("loss_mode", "vanilla")
 
                     if self.config.policy_loss.loss_mode == "vanilla":
-                        teacher_pg_loss = compute_sft_loss(
-                            log_prob=teacher_log_prob,
-                            response_mask=teacher_response_mask,
+                        pg_loss, pg_clipfrac, ppo_kl, pg_clipfrac_lower = compute_policy_loss(
+                            old_log_prob=old_log_prob,
+                            log_prob=log_prob,
+                            advantages=advantages,
+                            response_mask=response_mask,
+                            cliprange=clip_ratio,
+                            cliprange_low=clip_ratio_low,
+                            cliprange_high=clip_ratio_high,
+                            clip_ratio_c=clip_ratio_c,
                             loss_agg_mode=loss_agg_mode,
                         )
-                        pg_loss = teacher_pg_loss
                     else:
                         policy_loss_fn = get_policy_loss_fn(loss_mode)
                         pg_loss, pg_clipfrac, ppo_kl, pg_clipfrac_lower = policy_loss_fn(old_log_prob, log_prob, advantages, response_mask, loss_agg_mode, self.config)
@@ -510,6 +494,16 @@ def update_policy(self, data: DataProto):
                     else:
                         policy_loss = pg_loss
 
+                    if self.config.use_kl_loss:
+                        ref_log_prob = data["ref_log_prob"]
+                        # compute kl loss
+                        kld = kl_penalty(logprob=log_prob, ref_logprob=ref_log_prob, kl_penalty=self.config.kl_loss_type)
+                        kl_loss = agg_loss(loss_mat=kld, loss_mask=response_mask, loss_agg_mode=loss_agg_mode)
+
+                        policy_loss = policy_loss + kl_loss * self.config.kl_loss_coef
+                        metrics["actor/kl_loss"] = kl_loss.detach().item()
+                        metrics["actor/kl_coef"] = self.config.kl_loss_coef
+
                     if self.config.use_dynamic_bsz:
                         # relative to the dynamic bsz
                         loss = policy_loss * (len(data) / self.config.ppo_mini_batch_size)
@@ -519,7 +513,9 @@ def update_policy(self, data: DataProto):
 
                     data = {
                         "actor/pg_loss": pg_loss.detach().item(),
-                        "actor/teacher_pg_loss": teacher_pg_loss.detach().item(),
+                        "actor/pg_clipfrac": pg_clipfrac.detach().item(),
+                        "actor/ppo_kl": ppo_kl.detach().item(),
+                        "actor/pg_clipfrac_lower": pg_clipfrac_lower.detach().item(),
                     }
                     append_to_dict(metrics, data)
 
diff --git a/verl/workers/critic/dp_critic.py b/verl/workers/critic/dp_critic.py
diff --git a/verl/workers/fsdp_workers.py b/verl/workers/fsdp_workers.py
