diff --git a/megatron/rl/agent/api.py b/megatron/rl/agent/api.py index dd52c7277bf..310ef106f1e 100644 --- a/megatron/rl/agent/api.py +++ b/megatron/rl/agent/api.py @@ -48,6 +48,10 @@ class Rollout(AgentBaseModel): reward: float = None env_id: str = '' problem_id: str | None = None + policy_staleness: list[list[int]] + kv_cache_staleness: list[list[int]] + completed_at_step: list[int] + num_evictions: list[int] class TokenRollout(AgentBaseModel): @@ -59,6 +63,10 @@ class TokenRollout(AgentBaseModel): logprobs: list[list[float]] | None = None env_id: str = '' problem_id: str | None = None + policy_staleness: list[list[int]] + kv_cache_staleness: list[list[int]] + completed_at_step: list[int] + num_evictions: list[int] class ContrastiveRollout(AgentBaseModel): diff --git a/megatron/rl/agent/reward_only_agent.py b/megatron/rl/agent/reward_only_agent.py index 4099406a98a..a945dd885e3 100644 --- a/megatron/rl/agent/reward_only_agent.py +++ b/megatron/rl/agent/reward_only_agent.py @@ -104,6 +104,10 @@ async def rollout_from_response( generation_mask=[generation_mask], env_id=self.env_id, problem_id=golden['problem_id'] if 'problem_id' in golden else None, + policy_staleness=[response.policy_staleness], + kv_cache_staleness=[response.kv_cache_staleness], + completed_at_step=[response.completed_at_step], + num_evictions=[response.num_evictions], ) else: rollout = Rollout( @@ -111,6 +115,10 @@ async def rollout_from_response( reward=await self.get_reward(response_text, golden), env_id=self.env_id, problem_id=golden['problem_id'] if 'problem_id' in golden else None, + policy_staleness=[response.policy_staleness], + kv_cache_staleness=[response.kv_cache_staleness], + completed_at_step=[response.completed_at_step], + num_evictions=[response.num_evictions], ) return rollout diff --git a/megatron/rl/inference/api.py b/megatron/rl/inference/api.py index 87f6b87f908..95e24a55e73 100644 --- a/megatron/rl/inference/api.py +++ b/megatron/rl/inference/api.py @@ -23,3 +23,7 @@ class InferenceResponse(BaseModel): token_ids: list[int] | None = None prompt_length: int | None = None logprobs: list[float] | None = None + policy_staleness: list[int] + kv_cache_staleness: list[int] + completed_at_step: int + num_evictions: int diff --git a/megatron/rl/inference/megatron.py b/megatron/rl/inference/megatron.py index ee31a16eeaa..b1a7dcd8b83 100644 --- a/megatron/rl/inference/megatron.py +++ b/megatron/rl/inference/megatron.py @@ -70,6 +70,10 @@ async def base_generate(self, request: InferenceRequest) -> InferenceResponse: token_ids=choice.prompt_token_ids + choice.generation_token_ids, logprobs=choice.generation_log_probs, prompt_length=len(choice.prompt_token_ids), + policy_staleness=choice.policy_staleness, + kv_cache_staleness=choice.kv_cache_staleness, + completed_at_step=args.curr_iteration, + num_evictions=getattr(choice, 'num_evictions', 0), ) @classmethod @@ -123,6 +127,10 @@ async def kill(self): await self._client.stop_engines() await self._inference_engine.stopped.wait() + def increment_staleness(self): + if dist.get_rank() == 0: + self._client.increment_staleness() + async def suspend(self): if dist.get_rank() == 0: self._client.suspend_engines() diff --git a/megatron/rl/rl_utils.py b/megatron/rl/rl_utils.py index 46436add1da..a7ca8d6c4d5 100644 --- a/megatron/rl/rl_utils.py +++ b/megatron/rl/rl_utils.py @@ -275,6 +275,10 @@ class RolloutStats: min_inf_prob: None | float max_inf_prob: None | float mean_inf_prob: None | float + policy_staleness: list[list[int]] + kv_cache_staleness: list[list[int]] + completed_at_steps: list[list[int]] + num_evictions: list[list[int]] # Runtime state container for RL-specific data that shouldn't be checkpointed @@ -520,6 +524,7 @@ def get_environment_rollouts( args.cuda_graph_impl, False, # offload optimizer during rollout collection is handled above training_model=model if has_separate_inference_model else None, + increment_staleness_on_suspend=True, ) as inference_interface: with nvtx_range("inference-setup"): @@ -733,11 +738,19 @@ def compute_group_stats( env_ids = [] group_reward_ids = [] num_turns = [] # num_turns per traj + all_policy_staleness = [] + all_kv_cache_staleness = [] + all_completed_at_steps = [] + all_num_evictions = [] for group in rollouts: group_rewards = [] group_traj_lengths = [] group_turn_lengths = [] group_num_turns = [] + group_policy_staleness = [] + group_kv_staleness = [] + group_completed_at_steps = [] + group_num_evictions = [] for rollout in group: if isinstance(rollout, TokenRollout): for turn_traj in rollout.trajectory: @@ -758,6 +771,14 @@ def compute_group_stats( roll_turn_lens = [len(t) for t in rollout.trajectory] group_turn_lengths.extend(roll_turn_lens) group_traj_lengths.append(sum(roll_turn_lens)) + group_policy_staleness.extend(s for turn in rollout.policy_staleness for s in turn) + group_kv_staleness.extend(s for turn in rollout.kv_cache_staleness for s in turn) + group_completed_at_steps.extend(rollout.completed_at_step) + group_num_evictions.append(sum(rollout.num_evictions)) + all_policy_staleness.append(group_policy_staleness) + all_kv_cache_staleness.append(group_kv_staleness) + all_completed_at_steps.append(group_completed_at_steps) + all_num_evictions.append(group_num_evictions) traj_lens.append(group_traj_lengths) turn_lens.append(group_turn_lengths) env_ids.append(group[0].env_id) # All rollouts in a group share the env_id by design. @@ -785,10 +806,44 @@ def compute_group_stats( min_inf_prob=None, max_inf_prob=None, mean_inf_prob=None, + policy_staleness=all_policy_staleness, + kv_cache_staleness=all_kv_cache_staleness, + completed_at_steps=all_completed_at_steps, + num_evictions=all_num_evictions, ) return stats +def compute_true_staleness( + per_token_staleness: list[list[int]], + completed_at_steps: list[list[int]], + turn_lens: list[list[int]], + current_iteration: int, +) -> list[int]: + """Compute true per-token staleness by adding the completion gap. + + Args: + per_token_staleness: Grouped flat list of per-token raw staleness values. + completed_at_steps: Grouped list of per-turn completion steps. + turn_lens: Grouped list of per-turn token counts. + current_iteration: Current training iteration. + + Returns: + Flat list of true staleness values (one per token across all groups). + """ + result = [] + for group_staleness, group_completed, group_turn_lens in zip( + per_token_staleness, completed_at_steps, turn_lens + ): + token_idx = 0 + for completed_at, num_tokens in zip(group_completed, group_turn_lens): + gap = current_iteration - completed_at + for _ in range(num_tokens): + result.append(group_staleness[token_idx] + gap) + token_idx += 1 + return result + + def prep_wandb_metrics( wandb_writer: wandb_run.Run, traj_lens: List[List[int]], @@ -796,8 +851,13 @@ def prep_wandb_metrics( rewards: List[List[float]], num_turns: List[List[int]], advantages: List[float], + policy_staleness: List[List[int]], + kv_cache_staleness: List[List[int]], + num_evictions: List[List[int]], + completed_at_steps: List[List[int]], + current_iteration: int, example_group: list[TokenRollout | Rollout] | None = None, - tokenizer: MegatronTokenizer | None = None + tokenizer: MegatronTokenizer | None = None, ): """Make a wandb-parseable dictionary of metrics for logging. @@ -809,8 +869,13 @@ 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_staleness: Grouped list of per-token policy staleness. + kv_cache_staleness: Grouped list of per-token KV cache staleness. + num_evictions: Grouped list of per-rollout number of evictions. + completed_at_steps: Grouped list of per-turn completed at steps. + current_iteration: Current training iteration. + example_group: A list of rollouts of one group to log examples of trajectories. tokenizer: Tokenizer to untokenize trajectories for logging. - example_groups: A list of rollouts of one group to log examples of trajectories. """ group_table = wandb_writer.Table( @@ -818,6 +883,11 @@ def prep_wandb_metrics( data=[[np.mean(g), np.std(g)] for g in rewards], ) + true_policy_staleness = compute_true_staleness( + policy_staleness, completed_at_steps, turn_lens, current_iteration) + true_kv_staleness = compute_true_staleness( + kv_cache_staleness, completed_at_steps, turn_lens, current_iteration) + metrics = { 'group_means_hist': wandb_writer.plot.histogram( group_table, 'group_means', 'Group Means' @@ -837,6 +907,14 @@ def prep_wandb_metrics( ), 'advantages', 'Advantages' ), + 'rollout_table': wandb_writer.Table( + columns=['reward', 'traj_length', 'num_evictions'], + 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], + )), + ), '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]), @@ -852,6 +930,15 @@ 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), + '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_at_steps for s in g]), } if example_group: if tokenizer is None: @@ -910,8 +997,15 @@ def maybe_log_training_metrics( rewards = group_stats.rewards num_turns = group_stats.num_turns advantages = group_stats.advantages - metrics = metrics | prep_wandb_metrics(wandb_writer=wandb_writer, - traj_lens=traj_lens, turn_lens=turn_lens, rewards=rewards, num_turns=num_turns, advantages=advantages) + policy_staleness = group_stats.policy_staleness + kv_cache_staleness = group_stats.kv_cache_staleness + num_evictions = group_stats.num_evictions + completed_at_steps = group_stats.completed_at_steps + + metrics = metrics | prep_wandb_metrics(wandb_writer=wandb_writer, + traj_lens=traj_lens, turn_lens=turn_lens, rewards=rewards, num_turns=num_turns, advantages=advantages, + policy_staleness=policy_staleness, kv_cache_staleness=kv_cache_staleness, num_evictions=num_evictions, + completed_at_steps=completed_at_steps, current_iteration=current_iteration) env_stats = lambda cont, idx: [cont[i] for i in idx] group_turn_counts = [sum(nt) for nt in num_turns] @@ -925,11 +1019,16 @@ def maybe_log_training_metrics( end = st + group_turn_counts[i] env_advantages.extend(advantages[st:end]) - env_metrics = prep_wandb_metrics(wandb_writer=wandb_writer, traj_lens=env_stats(traj_lens, env_idx), - turn_lens=env_stats(turn_lens, env_idx), + env_metrics = prep_wandb_metrics(wandb_writer=wandb_writer, traj_lens=env_stats(traj_lens, env_idx), + turn_lens=env_stats(turn_lens, env_idx), rewards=env_stats(rewards, env_idx), num_turns=env_stats(num_turns, env_idx), advantages=env_advantages, + policy_staleness=env_stats(policy_staleness, env_idx), + kv_cache_staleness=env_stats(kv_cache_staleness, env_idx), + num_evictions=env_stats(num_evictions, env_idx), + completed_at_steps=env_stats(completed_at_steps, env_idx), + current_iteration=current_iteration, example_group=example_groups[env_id], tokenizer=tokenizer, ) @@ -1153,7 +1252,7 @@ def prepare_data_for_update( tokenizer: MegatronTokenizer, sequence_packing: bool, is_correction: bool, -) -> RerunDataIterator: +) -> tuple[RerunDataIterator, RolloutStats, dict]: """Extract data for the update from raw rollouts. Args: @@ -1165,7 +1264,7 @@ def prepare_data_for_update( is_correction: Prepare data for IS correction if True. Returns: - Cycled iterator over dataset batches. In GRPO we might want to go over the same data multiple times. + Tuple of (cycled iterator over dataset batches, group stats, example groups per env). """ args = get_args() nvtx_range = get_nvtx_range() @@ -1418,15 +1517,7 @@ def prepare_data_for_update( loader = DataLoader(data, batch_size=args.micro_batch_size) - with nvtx_range("log-wandb-tb"): - maybe_log_training_metrics( - group_stats=group_stats, - current_iteration=args.curr_iteration, - tokenizer=tokenizer, - example_groups=example_groups, - ) - - return RerunDataIterator(itertools.cycle(loader)) + return RerunDataIterator(itertools.cycle(loader)), group_stats, example_groups def get_grpo_data_iterator( @@ -1466,27 +1557,38 @@ def get_grpo_data_iterator( RerunDataIterator for the current training step """ runtime_state = get_rl_runtime_state() + tokenizer = get_tokenizer() # We collect new rollouts when we've gone over the collected data 'grpo_iterations' times. - global_batches_per_collection = (grpo_prompts_per_step * grpo_group_size) // global_batch_size + global_batches_per_collection = (grpo_prompts_per_step * grpo_group_size) // global_batch_size if ( buffered_rollouts is None or - iteration == runtime_state.last_collection_iteration + + iteration == runtime_state.last_collection_iteration + (grpo_iterations * global_batches_per_collection) ): - buffered_rollouts = get_environment_rollouts( + rollouts = get_environment_rollouts( model, inference_model, optimizer, grpo_prompts_per_step, grpo_group_size ) - buffered_rollouts = prepare_data_for_update(model=model, - ref_state_dict=ref_state_dict, - rollouts=buffered_rollouts, - tokenizer=get_tokenizer(), + buffered_rollouts, group_stats, example_groups = prepare_data_for_update( + model=model, + ref_state_dict=ref_state_dict, + rollouts=rollouts, + tokenizer=tokenizer, sequence_packing=sequence_packing, is_correction=is_correction, - ) + ) + runtime_state.group_stats = group_stats + runtime_state.example_groups = example_groups runtime_state.reset_iteration_counters(iteration) + maybe_log_training_metrics( + group_stats=runtime_state.group_stats, + current_iteration=iteration, + tokenizer=tokenizer, + example_groups=runtime_state.example_groups, + ) + return buffered_rollouts @@ -1704,6 +1806,7 @@ def megatron_rl_inference_mode( cuda_graph_impl: str, offload_optimizer_during_inference: bool, training_model: Optional[list[LanguageModule]] = None, + increment_staleness_on_suspend: bool = False, ): """Manage the model inference context when collecting rollouts. @@ -1768,6 +1871,8 @@ def megatron_rl_inference_mode( with nvtx_range("suspend-engine"): loop.run_until_complete(inference_interface.suspend()) + if increment_staleness_on_suspend: + inference_interface.increment_staleness() if cuda_graph_impl != "none" and not args.rl_training_cuda_graphs: toggle_cuda_graphs(lang_module, 'none') diff --git a/tests/unit_tests/rl/test_rl_utils.py b/tests/unit_tests/rl/test_rl_utils.py index 6206fa93c61..8fc6ac8654d 100644 --- a/tests/unit_tests/rl/test_rl_utils.py +++ b/tests/unit_tests/rl/test_rl_utils.py @@ -316,6 +316,10 @@ def test_prepare_data_for_update(self, initialize_model_parallel): logprobs=[[0.1, 0.2, 0.3]], env_id='MEGAENV', problem_id="2", + policy_staleness=[[0, 0, 0]], + kv_cache_staleness=[[0, 0, 0]], + completed_at_step=[1], + num_evictions=[0], ) r2 = TokenRollout( trajectory=[[1, 2, 3, 4]], @@ -324,6 +328,10 @@ def test_prepare_data_for_update(self, initialize_model_parallel): logprobs=[[0.1, 0.2, 0.3, -1.2]], env_id='MEGAENV', problem_id="2", + policy_staleness=[[0, 0, 0, 0]], + kv_cache_staleness=[[0, 0, 0, 0]], + completed_at_step=[1], + num_evictions=[0], ) rollouts = [[r1, r2] for _ in range(dp)] @@ -342,6 +350,10 @@ def test_prepare_data_for_update(self, initialize_model_parallel): logprobs=torch.tensor([[-0.2, -0.3, -3.2]]).cuda(), env_id='MEGAENV', problem_id="2", + policy_staleness=[[0, 0, 0, 0]], + kv_cache_staleness=[[0, 0, 0, 0]], + completed_at_step=[1], + num_evictions=[0], ) r2 = TokenRollout( trajectory=torch.tensor([[1, 2, 234, tokenizer.eod]], dtype=torch.float).cuda(), @@ -350,9 +362,13 @@ def test_prepare_data_for_update(self, initialize_model_parallel): logprobs=torch.tensor([[-0.2, -0.3, -1.2]]), env_id='MEGAENV', problem_id="2", + policy_staleness=[[0, 0, 0, 0]], + kv_cache_staleness=[[0, 0, 0, 0]], + completed_at_step=[1], + num_evictions=[0], ) rollouts = [[r1, r2] for _ in range(dp)] - data_iter = rl_utils.prepare_data_for_update( + data_iter, _, _ = rl_utils.prepare_data_for_update( [model], {}, rollouts, tokenizer, sequence_packing=False, is_correction=False ) @@ -383,6 +399,10 @@ def test_prepare_trajectories(self, use_sequence_packing, num_turns): logprobs=[[0.1, 0.2, 0.3, 0.35]] * num_turns, env_id='MEGAENV', problem_id="1", + policy_staleness=[[0, 0, 0, 0]] * num_turns, + kv_cache_staleness=[[0, 0, 0, 0]] * num_turns, + completed_at_step=[0] * num_turns, + num_evictions=[0] * num_turns, ) r2 = TokenRollout( trajectory=[[4, 5, 6, 7, tokenizer.eod]] * num_turns, @@ -391,6 +411,10 @@ def test_prepare_trajectories(self, use_sequence_packing, num_turns): logprobs=[[0.4, 0.5, 0.6, 0.7, 0.75]] * num_turns, env_id='MEGAENV', problem_id="2", + policy_staleness=[[0, 0, 0, 0, 0]] * num_turns, + kv_cache_staleness=[[0, 0, 0, 0, 0]] * num_turns, + completed_at_step=[0] * num_turns, + num_evictions=[0] * num_turns, ) r3 = TokenRollout( trajectory=[[8, 9, tokenizer.eod]] * num_turns, @@ -399,6 +423,10 @@ def test_prepare_trajectories(self, use_sequence_packing, num_turns): logprobs=[[0.8, 0.9, 0.95]] * num_turns, env_id='MEGAENV', problem_id="3", + policy_staleness=[[0, 0, 0]] * num_turns, + kv_cache_staleness=[[0, 0, 0]] * num_turns, + completed_at_step=[0] * num_turns, + num_evictions=[0] * num_turns, ) rollouts = [r1, r2, r3] @@ -868,6 +896,30 @@ def test_get_logprobs_cuda_graphs(self, initialize_model_parallel): CudaGraphManager.fwd_mempools = None CudaGraphManager.bwd_mempools = None + def test_compute_true_staleness(self): + # Single group, 2 turns: turn 1 has 1 token, turn 2 has 2 tokens + result = rl_utils.compute_true_staleness( + per_token_staleness=[[1, 2, 3]], + completed_at_steps=[[8, 7]], + turn_lens=[[1, 2]], + current_iteration=10, + ) + # token 0: 1 + (10-8) = 3 + # token 1: 2 + (10-7) = 5 + # token 2: 3 + (10-7) = 6 + assert result == [3, 5, 6] + + # Multiple groups + result = rl_utils.compute_true_staleness( + per_token_staleness=[[0, 0], [1]], + completed_at_steps=[[5], [5]], + turn_lens=[[2], [1]], + current_iteration=6, + ) + # group 1: [0+1, 0+1] = [1, 1] + # group 2: [1+1] = [2] + assert result == [1, 1, 2] + @pytest.mark.parametrize( "initialize_model_parallel", [pytest.param((1, 1), id="tp1-pp1")], @@ -881,8 +933,25 @@ def test_prep_wandb_metrics(self, initialize_model_parallel): rewards = [[1, 1], [-1, 2]] num_turns = [[42, 2], [10, 8]] advantages = [0, 1] + # Per-token staleness (6 tokens in group 1, 3 in group 2; matching turn_lens) + policy_staleness = [[1, 2, 2, 0, 0, 3], [0, 3, 3]] + kv_cache_staleness = [[1, 1, 1, 0, 0, 2], [0, 2, 2]] + num_evictions = [[0, 1], [0, 0]] + # Per-turn completed_at_steps (5 turns in group 1, 2 in group 2; matching turn_lens) + completed_at_steps = [[5, 4, 5, 5, 3], [5, 3]] + current_iteration = 6 metrics = rl_utils.prep_wandb_metrics( - MagicMock(), traj_lens, turn_lens, rewards, num_turns, advantages + MagicMock(), + traj_lens, + turn_lens, + rewards, + num_turns, + advantages, + policy_staleness=policy_staleness, + kv_cache_staleness=kv_cache_staleness, + num_evictions=num_evictions, + completed_at_steps=completed_at_steps, + current_iteration=current_iteration, ) assert metrics["mean_reward"] == 0.75 assert metrics["mean_advantage"] == 0.5 @@ -898,3 +967,15 @@ def test_prep_wandb_metrics(self, initialize_model_parallel): assert metrics["mean_num_turns"] == 15.5 assert metrics["max_num_turns"] == 42 assert metrics["min_num_turns"] == 2 + # true_policy_staleness = [2, 4, 4, 1, 1, 6, 1, 6, 6] + assert metrics["mean_policy_staleness"] == np.mean([2, 4, 4, 1, 1, 6, 1, 6, 6]) + assert metrics["max_policy_staleness"] == 6 + assert metrics["min_policy_staleness"] == 1 + # true_kv_staleness = [2, 3, 3, 1, 1, 5, 1, 5, 5] + assert metrics["mean_kv_cache_staleness"] == np.mean([2, 3, 3, 1, 1, 5, 1, 5, 5]) + assert metrics["max_kv_cache_staleness"] == 5 + assert metrics["min_kv_cache_staleness"] == 1 + assert metrics["total_eviction_count"] == 1 + assert metrics["max_num_evictions"] == 1 + # completion_gaps per-turn = [1, 2, 1, 1, 3, 1, 3] + assert metrics["mean_completion_gap"] == np.mean([1, 2, 1, 1, 3, 1, 3])