diff --git a/nemo_rl/experience/interfaces.py b/nemo_rl/experience/interfaces.py index 7e87dc0f059..f9187af7977 100644 --- a/nemo_rl/experience/interfaces.py +++ b/nemo_rl/experience/interfaces.py @@ -58,3 +58,5 @@ class PromptGroupRecord: metadata: dict[str, Any] completions: list["Completion"] rollout_metrics: dict[str, Any] + # Prompt-level DatumSpec weight; defaults preserve manually-built records. + loss_multiplier: float = 1.0 diff --git a/nemo_rl/experience/payload.py b/nemo_rl/experience/payload.py index b0441937942..4c22faa20d4 100644 --- a/nemo_rl/experience/payload.py +++ b/nemo_rl/experience/payload.py @@ -105,9 +105,9 @@ def record_to_train_batch( Returns: BatchedDataDict with input_ids, input_lengths, generation_logprobs, - token_mask, an all-ones sample_mask, the raw mask_sample and truncated - flags, prompt_ids_for_adv, total_reward, violation counts, and optional - routed experts and message-violation masks. + token_mask, a prompt-loss-weighted sample_mask, the raw mask_sample and + truncated flags, prompt_ids_for_adv, total_reward, violation counts, + and optional routed experts and message-violation masks. """ # Lazy imports: grpo and llm_message_utils transitively pull # experience.rollouts, so importing at module top risks a cycle. @@ -154,7 +154,7 @@ def record_to_train_batch( ) mask_sample = _mask_sample_flags(c.env_extras for c in completions) truncated = torch.tensor([c.truncated for c in completions], dtype=torch.bool) - sample_mask = torch.ones(n, dtype=torch.float32) + sample_mask = torch.full((n,), float(record.loss_multiplier), dtype=torch.float32) train_data: dict[str, Any] = { "input_ids": flat["token_ids"], diff --git a/nemo_rl/experience/rollout_manager.py b/nemo_rl/experience/rollout_manager.py index 504f146c5f0..db9181b954c 100644 --- a/nemo_rl/experience/rollout_manager.py +++ b/nemo_rl/experience/rollout_manager.py @@ -435,6 +435,7 @@ async def run_rollout(self, input_sample: DatumSpec) -> PromptGroupRecord: prompt=input_sample["message_log"], extra_env_info=input_sample["extra_env_info"], metadata={"task_name": input_sample["task_name"]}, + loss_multiplier=float(input_sample["loss_multiplier"]), completions=completions, rollout_metrics=rollout_metrics, ) @@ -828,6 +829,7 @@ async def run_rollout(self, input_sample: DatumSpec) -> PromptGroupRecord: prompt=prompt_message_log, extra_env_info=input_sample["extra_env_info"], metadata={"task_name": "nemo_gym"}, + loss_multiplier=float(input_sample["loss_multiplier"]), completions=completions, rollout_metrics=rollout_metrics, ) diff --git a/tests/unit/experience/test_payload.py b/tests/unit/experience/test_payload.py index a30e8f86c9e..fe9ca610f71 100644 --- a/tests/unit/experience/test_payload.py +++ b/tests/unit/experience/test_payload.py @@ -82,7 +82,9 @@ def _completion( ) -def _record(completions: list[Completion]) -> PromptGroupRecord: +def _record( + completions: list[Completion], *, loss_multiplier: float = 1.0 +) -> PromptGroupRecord: return PromptGroupRecord( prompt_idx=0, prompt=[ @@ -96,6 +98,7 @@ def _record(completions: list[Completion]) -> PromptGroupRecord: metadata={"task_name": "test"}, completions=completions, rollout_metrics={}, + loss_multiplier=loss_multiplier, ) @@ -378,3 +381,33 @@ def test_pack_payload_stamps_violation_counts_on_tags() -> None: "num_assistant_messages": 0, }, ] + + +def test_record_to_train_batch_preserves_fractional_loss_multiplier() -> None: + record = _record( + [ + _completion(route_start=10, reward=1.0), + _completion(route_start=20, reward=0.0), + ], + loss_multiplier=0.25, + ) + + batch = record_to_train_batch( + record, + pad_value_dict={"token_ids": 0}, + include_message_violation_fields=False, + ) + + torch.testing.assert_close(batch["sample_mask"], torch.full((2,), 0.25)) + + +def test_record_to_train_batch_drops_zero_weight_prompt() -> None: + record = _record([_completion(route_start=10, reward=1.0)], loss_multiplier=0.0) + + batch = record_to_train_batch( + record, + pad_value_dict={"token_ids": 0}, + include_message_violation_fields=False, + ) + + torch.testing.assert_close(batch["sample_mask"], torch.zeros(1)) diff --git a/tests/unit/experience/test_rollout_manager.py b/tests/unit/experience/test_rollout_manager.py index 9553910460b..9eb97e65cc1 100644 --- a/tests/unit/experience/test_rollout_manager.py +++ b/tests/unit/experience/test_rollout_manager.py @@ -807,7 +807,8 @@ def test_async_rollout_manager( - completions hold independent (not aliased) message_log objects """ vllm_generation, tokenizer, task_to_env, _, _ = multi_step_setup_vllm_async - input_sample = single_multi_step_calculator_input_sample + input_sample = deepcopy(single_multi_step_calculator_input_sample) + input_sample["loss_multiplier"] = 0.25 num_generations = 2 max_seq_len = 1024 max_rollout_turns = input_sample["extra_env_info"]["max_steps"] + 1 @@ -827,6 +828,7 @@ def test_async_rollout_manager( vllm_generation.finish_generation() assert isinstance(record, PromptGroupRecord) + assert record.loss_multiplier == 0.25 assert len(record.completions) == num_generations, ( f"Expected {num_generations} completions, got {len(record.completions)}" ) @@ -905,7 +907,8 @@ def test_async_rollout_manager_matches_original( TODO: remove this test together with run_async_multi_turn_rollout when the legacy path is deleted. """ vllm_generation, tokenizer, task_to_env, _, _ = multi_step_setup_vllm_async - input_sample = single_multi_step_calculator_input_sample + input_sample = deepcopy(single_multi_step_calculator_input_sample) + input_sample["loss_multiplier"] = 0.25 num_generations = 2 max_seq_len = 1024 max_rollout_turns = input_sample["extra_env_info"]["max_steps"] + 1 @@ -951,6 +954,7 @@ def test_async_rollout_manager_matches_original( # Both should produce N results assert len(original_batch["message_log"]) == num_generations assert len(record.completions) == num_generations + assert record.loss_multiplier == 0.25 for i in range(num_generations): orig_msg_log = original_batch["message_log"][i] @@ -1066,7 +1070,7 @@ def test_async_nemo_gym_rollout_manager( "extra_env_info": input_batch["extra_env_info"][0], "task_name": "nemo_gym", "idx": 0, - "loss_multiplier": float(input_batch["loss_multiplier"][0]), + "loss_multiplier": 0.25, } num_generations = 2 @@ -1081,6 +1085,7 @@ def test_async_nemo_gym_rollout_manager( record = asyncio.run(manager.run_rollout(single_prompt)) assert isinstance(record, PromptGroupRecord) + assert record.loss_multiplier == 0.25 assert len(record.completions) == num_generations, ( f"Expected {num_generations} completions, got {len(record.completions)}" )