From 5fd5fcd5cb314f768e41f1db915f6b623adfbcd1 Mon Sep 17 00:00:00 2001 From: Teodor-Dumitru Ene Date: Wed, 1 Apr 2026 15:42:20 -0500 Subject: [PATCH 1/5] Add tables and histogram for staleness --- megatron/rl/rl_utils.py | 46 +++++++++++++++++++++++++++++------------ 1 file changed, 33 insertions(+), 13 deletions(-) diff --git a/megatron/rl/rl_utils.py b/megatron/rl/rl_utils.py index 30c941d61dc..ae86168a239 100644 --- a/megatron/rl/rl_utils.py +++ b/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) @@ -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,12 @@ 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 + rollout_policy_staleness = [current_iteration - min(r) for g in policy_epoch for r in g] + rollout_kv_staleness = [current_iteration - min(r) 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 +911,21 @@ 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' + ], 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, )), ), + '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,16 +941,24 @@ 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), '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]), } + metrics['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' + ) + metrics['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: raise ValueError("If you provide an example group to log, you need to provide a tokenizer too.") From df6630d8fb3b6d9ddb249a00a229d52ea1978f4e Mon Sep 17 00:00:00 2001 From: Teodor-Dumitru Ene Date: Thu, 2 Apr 2026 10:42:05 -0500 Subject: [PATCH 2/5] Fix unit tests --- tests/unit_tests/rl/test_rl_utils.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/tests/unit_tests/rl/test_rl_utils.py b/tests/unit_tests/rl/test_rl_utils.py index 7d05708e27c..8510c688cfd 100644 --- a/tests/unit_tests/rl/test_rl_utils.py +++ b/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], [2]], [[5], [0]]] + kv_cache_epoch = [[[4], [3]], [[5], [1]]] # Per-turn max epoch stamps (when each turn completed) completed_epochs = [[5, 3], [5, 1]] num_evictions = [[0, 1], [0, 0]] From d0cdcecacb29493f73ead9f2a406be5c4a5e325e Mon Sep 17 00:00:00 2001 From: Teodor-Dumitru Ene Date: Thu, 2 Apr 2026 14:21:16 -0500 Subject: [PATCH 3/5] Track staleness of both first and last tokens --- megatron/rl/rl_utils.py | 17 +++++++++++++++-- tests/unit_tests/rl/test_rl_utils.py | 12 ++++++++++-- 2 files changed, 25 insertions(+), 4 deletions(-) diff --git a/megatron/rl/rl_utils.py b/megatron/rl/rl_utils.py index ae86168a239..a2a86c5ebba 100644 --- a/megatron/rl/rl_utils.py +++ b/megatron/rl/rl_utils.py @@ -884,9 +884,12 @@ def prep_wandb_metrics( data=[[np.mean(g), np.std(g)] for g in rewards], ) - # Per-rollout staleness + # Per-rollout staleness (oldest token) rollout_policy_staleness = [current_iteration - min(r) for g in policy_epoch for r in g] rollout_kv_staleness = [current_iteration - min(r) for g in kv_cache_epoch for r in g] + # Per-rollout staleness (newest token) + rollout_policy_last_token_staleness = [current_iteration - max(r) for g in policy_epoch for r in g] + rollout_kv_last_token_staleness = [current_iteration - max(r) 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] @@ -912,7 +915,9 @@ def prep_wandb_metrics( ), 'rollout_table': wandb_writer.Table( columns=[ - 'reward', 'traj_length', 'num_evictions', 'policy_staleness', 'kv_staleness' + '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], @@ -920,6 +925,8 @@ def prep_wandb_metrics( [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, )), ), 'per_token_table': wandb_writer.Table( @@ -947,6 +954,12 @@ def prep_wandb_metrics( '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]), diff --git a/tests/unit_tests/rl/test_rl_utils.py b/tests/unit_tests/rl/test_rl_utils.py index 8510c688cfd..6bf6e994ffb 100644 --- a/tests/unit_tests/rl/test_rl_utils.py +++ b/tests/unit_tests/rl/test_rl_utils.py @@ -903,8 +903,8 @@ def test_prep_wandb_metrics(self, initialize_model_parallel): num_turns = [[42, 2], [10, 8]] advantages = [0, 1] # Per-token epoch stamps, grouped by group then rollout - policy_epoch = [[[4], [2]], [[5], [0]]] - kv_cache_epoch = [[[4], [3]], [[5], [1]]] + 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 From 50a47b1e75f5186c7f25b8db2eae5f4e2074c885 Mon Sep 17 00:00:00 2001 From: Teodor-Dumitru Ene Date: Fri, 3 Apr 2026 08:45:38 -0500 Subject: [PATCH 4/5] Self-review --- megatron/rl/rl_utils.py | 29 +++++++++++------------------ 1 file changed, 11 insertions(+), 18 deletions(-) diff --git a/megatron/rl/rl_utils.py b/megatron/rl/rl_utils.py index a2a86c5ebba..59c031112dc 100644 --- a/megatron/rl/rl_utils.py +++ b/megatron/rl/rl_utils.py @@ -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, @@ -929,6 +929,7 @@ def prep_wandb_metrics( 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)), @@ -963,15 +964,15 @@ def prep_wandb_metrics( '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' + ), } - metrics['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' - ) - metrics['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: raise ValueError("If you provide an example group to log, you need to provide a tokenizer too.") @@ -1549,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 From 79914a871b6d6d9206f212c6a9bee36e2031d116 Mon Sep 17 00:00:00 2001 From: Teodor-Dumitru Ene Date: Wed, 8 Apr 2026 08:05:30 -0500 Subject: [PATCH 5/5] Address reviewer comments --- megatron/rl/rl_utils.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/megatron/rl/rl_utils.py b/megatron/rl/rl_utils.py index 59c031112dc..8ae4b6c0469 100644 --- a/megatron/rl/rl_utils.py +++ b/megatron/rl/rl_utils.py @@ -885,11 +885,11 @@ def prep_wandb_metrics( ) # Per-rollout staleness (oldest token) - rollout_policy_staleness = [current_iteration - min(r) for g in policy_epoch for r in g] - rollout_kv_staleness = [current_iteration - min(r) for g in kv_cache_epoch for r in g] + 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 - max(r) for g in policy_epoch for r in g] - rollout_kv_last_token_staleness = [current_iteration - max(r) for g in kv_cache_epoch for r in g] + 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]