Add tables and histogram for RL staleness

megatron/rl/rl_utils.py

@@ -801,8 +801,8 @@ def compute_group_stats(
             group_traj_lengths.append(sum(roll_turn_lens))
             assert rollout.policy_epoch, "Rollout has no policy_epoch data"
             assert rollout.kv_cache_epoch, "Rollout has no kv_cache_epoch data"
-            group_policy_epoch.append(min(turn[0][1] for turn in rollout.policy_epoch))
-            group_kv_epoch.append(min(turn[0][1] for turn in rollout.kv_cache_epoch))
+            group_policy_epoch.append([epoch for turn in rollout.policy_epoch for _, epoch in turn])
+            group_kv_epoch.append([epoch for turn in rollout.kv_cache_epoch for _, epoch in turn])
             group_completed_epochs.extend(turn[-1][1] for turn in rollout.policy_epoch)
             group_num_evictions.append(sum(rollout.num_evictions))
         all_policy_epoch.append(group_policy_epoch)
@@ -852,8 +852,8 @@ def prep_wandb_metrics(
         rewards: List[List[float]],
         num_turns: List[List[int]],
         advantages: List[float],
-        policy_epoch: List[List[int]],
-        kv_cache_epoch: List[List[int]],
+        policy_epoch: List[List[List[int]]],
+        kv_cache_epoch: List[List[List[int]]],
         completed_epochs: List[List[int]],
         num_evictions: List[List[int]],
         current_iteration: int,
@@ -870,8 +870,8 @@ def prep_wandb_metrics(
         rewards: Grouped list of rewards.
         num_turns: Grouped list of number of turns in the trajectories.
         advantages: Flattened list of advantages.
-        policy_epoch: Grouped list of per-rollout min policy epoch stamps.
-        kv_cache_epoch: Grouped list of per-rollout min KV cache epoch stamps.
+        policy_epoch: Grouped list of per-token policy epoch stamps.
+        kv_cache_epoch: Grouped list of per-token KV cache epoch stamps.
         completed_epochs: Grouped list of per-turn max policy epoch stamps.
         num_evictions: Grouped list of per-rollout number of evictions.
         current_iteration: Current training iteration.
@@ -884,8 +884,15 @@ def prep_wandb_metrics(
         data=[[np.mean(g), np.std(g)] for g in rewards],
     )
 
-    true_policy_staleness = [current_iteration - s for g in policy_epoch for s in g]
-    true_kv_staleness = [current_iteration - s for g in kv_cache_epoch for s in g]
+    # Per-rollout staleness (oldest token)
+    rollout_policy_staleness = [current_iteration - r[0] for g in policy_epoch for r in g]
+    rollout_kv_staleness = [current_iteration - r[0] for g in kv_cache_epoch for r in g]
+    # Per-rollout staleness (newest token)
+    rollout_policy_last_token_staleness = [current_iteration - r[-1] for g in policy_epoch for r in g]
+    rollout_kv_last_token_staleness = [current_iteration - r[-1] for g in kv_cache_epoch for r in g]
+    # Per-token staleness
+    per_token_policy_staleness = [current_iteration - e for g in policy_epoch for r in g for e in r]
+    per_token_kv_staleness = [current_iteration - e for g in kv_cache_epoch for r in g for e in r]
 
     metrics = {
             'group_means_hist': wandb_writer.plot.histogram(
@@ -907,13 +914,26 @@ def prep_wandb_metrics(
                 'advantages', 'Advantages'
             ),
             'rollout_table': wandb_writer.Table(
-                columns=['reward', 'traj_length', 'num_evictions'],
+                columns=[
+                    'reward', 'traj_length', 'num_evictions',
+                    'policy_staleness', 'kv_staleness',
+                    'policy_last_token_staleness', 'kv_last_token_staleness',
+                ],
                 data=list(zip(
                     [r for g in rewards for r in g],
                     [l for g in traj_lens for l in g],
                     [e for g in num_evictions for e in g],
+                    rollout_policy_staleness,
+                    rollout_kv_staleness,
+                    rollout_policy_last_token_staleness,
+                    rollout_kv_last_token_staleness,
                 )),
             ),
+            # NOTE: This table can get very large (one row per token across all rollouts).
+            'per_token_table': wandb_writer.Table(
+                columns=['policy_staleness', 'kv_staleness'],
+                data=list(zip(per_token_policy_staleness, per_token_kv_staleness)),
+            ),
             'mean_turn_length': np.mean([np.mean(g) for g in turn_lens]),
             'mean_turn_length_std': np.mean([np.std(g) for g in turn_lens]),
             'max_turn_length': max([max(g) for g in turn_lens]),
@@ -929,15 +949,29 @@ def prep_wandb_metrics(
             'mean_advantage': np.mean(advantages),
             'nonzero_groups_ratio': np.count_nonzero(advantages)
             / len(advantages),
-            'mean_policy_staleness': np.mean(true_policy_staleness),
-            'max_policy_staleness': max(true_policy_staleness),
-            'min_policy_staleness': min(true_policy_staleness),
-            'mean_kv_cache_staleness': np.mean(true_kv_staleness),
-            'max_kv_cache_staleness': max(true_kv_staleness),
-            'min_kv_cache_staleness': min(true_kv_staleness),
+            'mean_policy_staleness': np.mean(rollout_policy_staleness),
+            'max_policy_staleness': max(rollout_policy_staleness),
+            'min_policy_staleness': min(rollout_policy_staleness),
+            'mean_kv_cache_staleness': np.mean(rollout_kv_staleness),
+            'max_kv_cache_staleness': max(rollout_kv_staleness),
+            'min_kv_cache_staleness': min(rollout_kv_staleness),
+            'mean_policy_last_token_staleness': np.mean(rollout_policy_last_token_staleness),
+            'max_policy_last_token_staleness': max(rollout_policy_last_token_staleness),
+            'min_policy_last_token_staleness': min(rollout_policy_last_token_staleness),
+            'mean_kv_cache_last_token_staleness': np.mean(rollout_kv_last_token_staleness),
+            'max_kv_cache_last_token_staleness': max(rollout_kv_last_token_staleness),
+            'min_kv_cache_last_token_staleness': min(rollout_kv_last_token_staleness),
             'total_eviction_count': sum([sum(g) for g in num_evictions]),
             'max_num_evictions': max([max(g) for g in num_evictions]),
             'mean_completion_gap': np.mean([current_iteration - s for g in completed_epochs for s in g]),
+            'per_token_policy_staleness_hist': wandb_writer.plot.histogram(
+                wandb_writer.Table(columns=['staleness'], data=[[s] for s in per_token_policy_staleness]),
+                'staleness', 'Per-Token Policy Staleness'
+            ),
+            'per_token_kv_staleness_hist': wandb_writer.plot.histogram(
+                wandb_writer.Table(columns=['staleness'], data=[[s] for s in per_token_kv_staleness]),
+                'staleness', 'Per-Token KV Cache Staleness'
+            ),
     }
     if example_group:
         if tokenizer is None:
@@ -1516,14 +1550,6 @@ def prepare_data_for_update(
                 data = TensorDataset(*dataset_tensors)
                 loader = DataLoader(data, batch_size=args.micro_batch_size)
 
-        with nvtx_range("rl/log-wandb-tb", time=True):
-            maybe_log_training_metrics(
-                group_stats=group_stats,
-                current_iteration=args.curr_iteration,
-                tokenizer=tokenizer,
-                example_groups=example_groups,
-            )
-
     return RerunDataIterator(itertools.cycle(loader)), group_stats, example_groups
 
 

tests/unit_tests/rl/test_rl_utils.py

@@ -902,9 +902,9 @@ def test_prep_wandb_metrics(self, initialize_model_parallel):
         rewards = [[1, 1], [-1, 2]]
         num_turns = [[42, 2], [10, 8]]
         advantages = [0, 1]
-        # Per-rollout min iteration stamps (2 rollouts in group 1, 2 in group 2)
-        policy_epoch = [[4, 2], [5, 0]]
-        kv_cache_epoch = [[4, 3], [5, 1]]
+        # Per-token epoch stamps, grouped by group then rollout
+        policy_epoch = [[[4, 5], [2, 3]], [[5], [0, 1]]]
+        kv_cache_epoch = [[[4, 5], [3, 4]], [[5], [1, 2]]]
         # Per-turn max epoch stamps (when each turn completed)
         completed_epochs = [[5, 3], [5, 1]]
         num_evictions = [[0, 1], [0, 0]]
@@ -944,6 +944,14 @@ def test_prep_wandb_metrics(self, initialize_model_parallel):
         assert metrics["mean_kv_cache_staleness"] == np.mean([2, 3, 1, 5])
         assert metrics["max_kv_cache_staleness"] == 5
         assert metrics["min_kv_cache_staleness"] == 1
+        # last_token (max epoch per rollout): policy=[5, 3, 5, 1] -> staleness=[1, 3, 1, 5]
+        assert metrics["mean_policy_last_token_staleness"] == np.mean([1, 3, 1, 5])
+        assert metrics["max_policy_last_token_staleness"] == 5
+        assert metrics["min_policy_last_token_staleness"] == 1
+        # last_token (max epoch per rollout): kv=[5, 4, 5, 2] -> staleness=[1, 2, 1, 4]
+        assert metrics["mean_kv_cache_last_token_staleness"] == np.mean([1, 2, 1, 4])
+        assert metrics["max_kv_cache_last_token_staleness"] == 4
+        assert metrics["min_kv_cache_last_token_staleness"] == 1
         assert metrics["total_eviction_count"] == 1
         assert metrics["max_num_evictions"] == 1
         # mean_completion_gap = mean([6-5, 6-3, 6-5, 6-1]) = mean([1, 3, 1, 5]) = 2.5