[FIX] Drop only the sample truncated instead of the entire group (#744)

Co-authored-by: Felipe Mello <felipemello@fb.com>

Author: felipemello1

apps/grpo/main.py

@@ -232,37 +232,19 @@ async def continuous_rollouts():
                     "episode/min_response_tokens", response_tokens, Reduce.MIN
                 )
 
-            # drop episodes if
-            # 1> reward std-dev is very small (including all 0s and all 1s)
-            # 2> any response was truncated (didn't end with EOS)
-            # TODO: change it to filter only truncated episodes instead of dropping entire group
+            # Drop entire group if reward std-dev is very small (no learning signal)
             rewards = [e.reward for e in episodes]
             rewards_std = torch.std(torch.tensor(rewards))
             is_low_variance = rewards_std < 1e-3
-            num_truncated = sum(
-                1 for e in episodes if e.completion.stop_reason == "length"
-            )
-            is_truncated = num_truncated > 0
-            drop = is_low_variance or is_truncated
 
             n = len(episodes)
             record_metric(
                 "main/continuous_rollouts/episodes_dropped/low_variance",
                 n if is_low_variance else 0,
                 Reduce.SUM,
             )
-            record_metric(
-                "main/continuous_rollouts/episodes_dropped/truncated",
-                num_truncated,
-                Reduce.SUM,
-            )
-            record_metric(
-                "main/continuous_rollouts/episodes_dropped/total",
-                n if drop else 0,
-                Reduce.SUM,
-            )
+            if is_low_variance:
 
-            if drop:
                 del input_ids, episodes
                 continue
 
@@ -283,8 +265,16 @@ async def continuous_rollouts():
             del input_ids
 
             advantages = await compute_advantages.compute.call_one(episodes)
+            num_truncated = 0
             for episode, advantage in zip(episodes, advantages):
                 episode.advantage = advantage
+
+                # Zero out loss_mask for truncated episodes so they don't contribute to gradient
+                # TODO: evaluate if we should drop truncated episodes instead
+                if episode.completion.stop_reason == "length":
+                    episode.loss_mask = torch.zeros_like(episode.loss_mask)
+                    num_truncated += 1
+
                 await replay_buffer.add.call_one(episode)
 
                 sample = episode.to_dict(
@@ -301,6 +291,11 @@ async def continuous_rollouts():
                     sample,
                     Reduce.SAMPLE,
                 )
+            record_metric(
+                "main/continuous_rollouts/episodes_dropped/truncated",
+                num_truncated,
+                Reduce.SUM,
+            )
 
             rollout_count += 1
             record_metric(