move pg_loss into tis_function for icepop (#635)

Author: zhuzilin

examples/train_infer_mismatch_helper/mis.py

@@ -218,6 +218,7 @@ def compute_mis_weights(
 def compute_mis_weights_with_cp(
     args,
     *,
+    pg_loss: torch.Tensor,
     train_log_probs: list[torch.Tensor],
     rollout_log_probs: list[torch.Tensor],
     loss_masks: list[torch.Tensor],
@@ -274,7 +275,9 @@ def slice_cp_and_concat(
         values = slice_cp_and_concat(values, total_lengths, response_lengths)
         result_metrics[key_name] = values
 
-    return is_weights, result_metrics
+    pg_loss = pg_loss * is_weights
+
+    return pg_loss, result_metrics
 
 
 def add_ppl_metrics(

slime/backends/megatron_utils/loss.py

@@ -424,6 +424,7 @@ def policy_loss_function(
         def vanilla_tis_function(
             args,
             *,
+            pg_loss: torch.Tensor,
             train_log_probs: list[torch.Tensor],
             rollout_log_probs: list[torch.Tensor],
             **kwargs: Any,
@@ -439,13 +440,15 @@ def vanilla_tis_function(
                 "tis_clipfrac": tis_clipfrac.clone().detach(),
                 "tis_abs": tis_abs.clone().detach(),
             }
-            return tis_weights, metrics
+            pg_loss = pg_loss * tis_weights
+            return pg_loss, metrics
 
         assert "rollout_log_probs" in batch, "rollout_log_probs must be provided for TIS"
 
         ois = (-ppo_kl).exp()
         tis_kwargs = {
             "args": args,
+            "pg_loss": pg_loss,
             "train_log_probs": batch["log_probs"],
             "rollout_log_probs": batch["rollout_log_probs"],
             "loss_masks": batch["loss_masks"],
@@ -457,9 +460,7 @@ def vanilla_tis_function(
             tis_func = load_function(args.custom_tis_function_path)
         else:
             tis_func = vanilla_tis_function
-        tis_weights, tis_metrics = tis_func(**tis_kwargs)
-
-        pg_loss = pg_loss * tis_weights
+        pg_loss, tis_metrics = tis_func(**tis_kwargs)
 
     pg_loss = sum_of_sample_mean(pg_loss)
     pg_clipfrac = sum_of_sample_mean(pg_clipfrac)