From 6bfbf0b283a08e4bd89a57b6f65a4968e38dce12 Mon Sep 17 00:00:00 2001 From: Zeyu Zhou Date: Wed, 24 Jun 2026 10:54:02 -0700 Subject: [PATCH 01/12] feat: support R3 router replay for async Gym rollouts Signed-off-by: Zeyu Zhou --- nemo_rl/algorithms/grpo.py | 1 + nemo_rl/environments/nemo_gym.py | 73 +++++++++++---- .../generation/vllm/vllm_worker_async.py | 80 ++++++++++++++++- .../test_nemo_gym_router_replay.py | 90 +++++++++++++++++++ 4 files changed, 226 insertions(+), 18 deletions(-) create mode 100644 tests/unit/environments/test_nemo_gym_router_replay.py diff --git a/nemo_rl/algorithms/grpo.py b/nemo_rl/algorithms/grpo.py index 5435df757a9..77a6d55effa 100644 --- a/nemo_rl/algorithms/grpo.py +++ b/nemo_rl/algorithms/grpo.py @@ -496,6 +496,7 @@ def _spinup_nemo_gym(base_urls, model_name): base_urls=base_urls, invalid_tool_call_patterns=invalid_tool_call_patterns, thinking_tags=thinking_tags, + require_routed_experts=router_replay_enabled(policy_config), initial_global_config_dict=nemo_gym_dict, ) nemo_gym_opts = {} diff --git a/nemo_rl/environments/nemo_gym.py b/nemo_rl/environments/nemo_gym.py index 611751af362..b3f8dcbfbd1 100644 --- a/nemo_rl/environments/nemo_gym.py +++ b/nemo_rl/environments/nemo_gym.py @@ -77,6 +77,9 @@ class NemoGymConfig(TypedDict): thinking_tags: NotRequired[ List[str] | None ] # Thinking tags to check for malformed usage + require_routed_experts: NotRequired[ + bool + ] # Require Gym output items to carry R3 routed_experts def _detect_invalid_tool_call_and_malformed_thinking( @@ -343,15 +346,48 @@ def _postprocess_nemo_gym_to_nemo_rl_result( prompt_token_ids = output_item_dict.pop("prompt_token_ids") generation_token_ids = output_item_dict.pop("generation_token_ids") generation_log_probs = output_item_dict.pop("generation_log_probs") + routed_experts_raw = output_item_dict.pop("routed_experts", None) new_prompt_token_ids = prompt_token_ids[len(seen_token_ids) :] - nemo_rl_message_log.append( - { - "role": "user", - "content": "", - "token_ids": torch.tensor(new_prompt_token_ids), - } - ) + routed_experts = None + if routed_experts_raw is not None: + routed_experts = torch.as_tensor(routed_experts_raw, dtype=torch.int32) + if routed_experts.dim() != 3: + raise ValueError( + "NeMo Gym returned routed_experts with invalid shape. " + "Expected [tokens, num_moe_layers, topk], got " + f"{tuple(routed_experts.shape)}." + ) + expected_tokens = len(prompt_token_ids) + len(generation_token_ids) + if routed_experts.shape[0] < expected_tokens: + raise ValueError( + "NeMo Gym returned too few routed_experts rows for a " + "trainable output item: " + f"routes={routed_experts.shape[0]}, expected_at_least=" + f"{expected_tokens}." + ) + elif self.cfg.get("require_routed_experts", False): + raise ValueError( + "policy.router_replay.enabled=true requires NeMo Gym output " + "items to include routed_experts, but the field was missing. " + "Make sure the Gym repo includes routed_experts propagation " + "and the NeMo-RL vLLM OpenAI-compatible server is configured " + "with enable_return_routed_experts." + ) + + prompt_start = len(seen_token_ids) + prompt_end = len(prompt_token_ids) + generation_start = prompt_end + generation_end = prompt_end + len(generation_token_ids) + + user_message = { + "role": "user", + "content": "", + "token_ids": torch.tensor(new_prompt_token_ids), + } + if routed_experts is not None: + user_message["routed_experts"] = routed_experts[prompt_start:prompt_end] + nemo_rl_message_log.append(user_message) # Valid tool calls go through the structured API (tool_calls field) and get # executed by NeMo-Gym. If tool call patterns appear in the text content instead, # the call was invalid and never executed — flag it so training can penalize it. @@ -365,16 +401,19 @@ def _postprocess_nemo_gym_to_nemo_rl_result( ) ) - nemo_rl_message_log.append( - { - "role": "assistant", - "content": "", - "token_ids": torch.tensor(generation_token_ids), - "generation_logprobs": torch.tensor(generation_log_probs), - "is_invalid_tool_call": is_invalid_tool_call, - "has_malformed_thinking": has_malformed_thinking, - } - ) + assistant_message = { + "role": "assistant", + "content": "", + "token_ids": torch.tensor(generation_token_ids), + "generation_logprobs": torch.tensor(generation_log_probs), + "is_invalid_tool_call": is_invalid_tool_call, + "has_malformed_thinking": has_malformed_thinking, + } + if routed_experts is not None: + assistant_message["routed_experts"] = routed_experts[ + generation_start:generation_end + ] + nemo_rl_message_log.append(assistant_message) seen_token_ids.extend(new_prompt_token_ids) seen_token_ids.extend(generation_token_ids) diff --git a/nemo_rl/models/generation/vllm/vllm_worker_async.py b/nemo_rl/models/generation/vllm/vllm_worker_async.py index 591b929fbf4..e849df0a273 100644 --- a/nemo_rl/models/generation/vllm/vllm_worker_async.py +++ b/nemo_rl/models/generation/vllm/vllm_worker_async.py @@ -671,8 +671,86 @@ class NeMoRLChatCompletionRequest( # vLLM 0.20 routes both /v1/chat/completions and /tokenize through # OpenAIServingRender.preprocess_chat, so the prefix-token override # belongs on the render subclass. + return_routed_experts = bool( + self.cfg.get("vllm_kwargs", {}).get("enable_return_routed_experts", False) + ) + class NeMoRLOpenAIServingChat(OpenAIServingChat): - pass + async def chat_completion_full_generator( + self, + request, + result_generator, + *args, + **kwargs, + ): + final_res = None + + async def capture_result_generator(): + nonlocal final_res + async for res in result_generator: + final_res = res + yield res + + response = await super().chat_completion_full_generator( + request, + capture_result_generator(), + *args, + **kwargs, + ) + if ( + not return_routed_experts + or not isinstance(response, ChatCompletionResponse) + or final_res is None + ): + return response + + outputs_by_index = { + output.index: output for output in getattr(final_res, "outputs", []) + } + prompt_token_count = len( + getattr(final_res, "prompt_token_ids", []) or [] + ) + + for choice in response.choices: + generation_details = outputs_by_index.get(choice.index) + if generation_details is None: + continue + + generation_token_count = len( + getattr(generation_details, "token_ids", []) or [] + ) + routed_experts, r3_stats = pad_and_align_routed_expert_indices( + final_res, + generation_details, + valid_length=prompt_token_count + generation_token_count, + padded_length=prompt_token_count + generation_token_count, + device=torch.device("cpu"), + require_complete_routed_experts=True, + return_stats=True, + ) + if routed_experts is None: + raise RuntimeError( + "vLLM was asked to return routed experts for the " + "OpenAI-compatible chat endpoint but the generation " + "output did not include routed_experts." + ) + if r3_stats["missing_routes"] > 0: + LOGGER.warning( + "R3 router replay fallback: vLLM returned incomplete " + "routed_experts for chat choice_idx=%d, " + "missing_token_routes=%d, actual_routes=%d, " + "expected_routes=%d. Megatron will use its own router " + "for those missing token routes.", + choice.index, + r3_stats["missing_routes"], + r3_stats["actual_routes"], + r3_stats["expected_routes"], + ) + choice.message.routed_experts = routed_experts.to( + dtype=torch.int32 + ).tolist() + + return response class NeMoRLOpenAIServingRender(NeMoRLOpenAIServingMixin, OpenAIServingRender): pass diff --git a/tests/unit/environments/test_nemo_gym_router_replay.py b/tests/unit/environments/test_nemo_gym_router_replay.py new file mode 100644 index 00000000000..fdc7a021f7e --- /dev/null +++ b/tests/unit/environments/test_nemo_gym_router_replay.py @@ -0,0 +1,90 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import pytest + +from nemo_rl.environments.nemo_gym import NemoGym + + +class _Tokenizer: + def batch_decode(self, batch): + return [" ".join(map(str, token_ids)) for token_ids in batch] + + +def _routes(num_tokens: int) -> list[list[list[int]]]: + return [[[token_idx, token_idx + 100]] for token_idx in range(num_tokens)] + + +def test_nemo_gym_postprocess_slices_routed_experts(): + nemo_gym_result = { + "response": { + "output": [ + { + "prompt_token_ids": [1, 2], + "generation_token_ids": [3], + "generation_log_probs": [-0.1], + "routed_experts": _routes(3), + }, + { + "prompt_token_ids": [1, 2, 3, 4, 5], + "generation_token_ids": [6, 7], + "generation_log_probs": [-0.2, -0.3], + "routed_experts": _routes(7), + }, + ] + }, + "responses_create_params": {"input": []}, + } + + class _MockSelf: + cfg = {"require_routed_experts": True} + + result = ( + NemoGym.__ray_metadata__.modified_class._postprocess_nemo_gym_to_nemo_rl_result( + _MockSelf(), nemo_gym_result, _Tokenizer() + ) + ) + + message_log = result["message_log"] + assert message_log[0]["token_ids"].tolist() == [1, 2] + assert message_log[0]["routed_experts"].tolist() == _routes(2) + assert message_log[1]["token_ids"].tolist() == [3] + assert message_log[1]["routed_experts"].tolist() == _routes(3)[2:3] + assert message_log[2]["token_ids"].tolist() == [4, 5] + assert message_log[2]["routed_experts"].tolist() == _routes(7)[3:5] + assert message_log[3]["token_ids"].tolist() == [6, 7] + assert message_log[3]["routed_experts"].tolist() == _routes(7)[5:7] + + +def test_nemo_gym_postprocess_requires_routed_experts_when_configured(): + nemo_gym_result = { + "response": { + "output": [ + { + "prompt_token_ids": [1, 2], + "generation_token_ids": [3], + "generation_log_probs": [-0.1], + }, + ] + }, + "responses_create_params": {"input": []}, + } + + class _MockSelf: + cfg = {"require_routed_experts": True} + + with pytest.raises(ValueError, match="requires NeMo Gym output items"): + NemoGym.__ray_metadata__.modified_class._postprocess_nemo_gym_to_nemo_rl_result( + _MockSelf(), nemo_gym_result, _Tokenizer() + ) From 87b3c9ed8841467012f926834906b1a66f08fcc2 Mon Sep 17 00:00:00 2001 From: Zeyu Zhou Date: Wed, 24 Jun 2026 14:34:32 -0700 Subject: [PATCH 02/12] test: cover async vLLM routed expert chat attachment Signed-off-by: Zeyu Zhou --- nemo_rl/models/generation/vllm/utils.py | 59 +++++++++++++++++ .../generation/vllm/vllm_worker_async.py | 52 ++------------- .../unit/models/generation/test_vllm_utils.py | 64 +++++++++++++++++++ 3 files changed, 129 insertions(+), 46 deletions(-) diff --git a/nemo_rl/models/generation/vllm/utils.py b/nemo_rl/models/generation/vllm/utils.py index 349d36fabfb..daa7d09eb39 100644 --- a/nemo_rl/models/generation/vllm/utils.py +++ b/nemo_rl/models/generation/vllm/utils.py @@ -192,6 +192,65 @@ def pad_and_align_routed_expert_indices( return (full, stats) if return_stats else full +def attach_routed_experts_to_chat_response_choices( + response: Any, + final_request_output: Any, + *, + device: torch.device, + logger: Any = None, +) -> Any: + """Attach aligned routed experts to OpenAI chat response choices.""" + outputs_by_index = { + output.index: output for output in getattr(final_request_output, "outputs", []) + } + prompt_token_count = len( + getattr(final_request_output, "prompt_token_ids", []) or [] + ) + + for choice in getattr(response, "choices", []): + generation_details = outputs_by_index.get(choice.index) + if generation_details is None: + continue + + generation_token_count = len(getattr(generation_details, "token_ids", []) or []) + routed_result = pad_and_align_routed_expert_indices( + final_request_output, + generation_details, + valid_length=prompt_token_count + generation_token_count, + padded_length=prompt_token_count + generation_token_count, + device=device, + require_complete_routed_experts=True, + return_stats=True, + ) + if not isinstance(routed_result, tuple): + raise RuntimeError( + "Expected routed_experts alignment to return stats for the " + "OpenAI-compatible chat endpoint." + ) + routed_experts, r3_stats = routed_result + if routed_experts is None: + raise RuntimeError( + "vLLM was asked to return routed experts for the " + "OpenAI-compatible chat endpoint but the generation " + "output did not include routed_experts." + ) + if r3_stats["missing_routes"] > 0 and logger is not None: + logger.warning( + "R3 router replay fallback: vLLM returned incomplete " + "routed_experts for chat choice_idx=%d, " + "missing_token_routes=%d, actual_routes=%d, " + "expected_routes=%d. Megatron will use its own router " + "for those missing token routes.", + choice.index, + r3_stats["missing_routes"], + r3_stats["actual_routes"], + r3_stats["expected_routes"], + ) + choice.message.routed_experts = routed_experts.to(dtype=torch.int32).tolist() + + return response + + def aggregate_spec_decode_counters( worker_metrics: list[dict[str, float | list[float]]], ) -> dict[str | tuple[str, int], float]: diff --git a/nemo_rl/models/generation/vllm/vllm_worker_async.py b/nemo_rl/models/generation/vllm/vllm_worker_async.py index e849df0a273..9d7da35aab1 100644 --- a/nemo_rl/models/generation/vllm/vllm_worker_async.py +++ b/nemo_rl/models/generation/vllm/vllm_worker_async.py @@ -41,6 +41,7 @@ verify_right_padding, ) from nemo_rl.models.generation.vllm.utils import ( + attach_routed_experts_to_chat_response_choices, format_prompt_for_vllm_generation, pad_and_align_routed_expert_indices, ) @@ -704,54 +705,13 @@ async def capture_result_generator(): ): return response - outputs_by_index = { - output.index: output for output in getattr(final_res, "outputs", []) - } - prompt_token_count = len( - getattr(final_res, "prompt_token_ids", []) or [] + return attach_routed_experts_to_chat_response_choices( + response, + final_res, + device=torch.device("cpu"), + logger=LOGGER, ) - for choice in response.choices: - generation_details = outputs_by_index.get(choice.index) - if generation_details is None: - continue - - generation_token_count = len( - getattr(generation_details, "token_ids", []) or [] - ) - routed_experts, r3_stats = pad_and_align_routed_expert_indices( - final_res, - generation_details, - valid_length=prompt_token_count + generation_token_count, - padded_length=prompt_token_count + generation_token_count, - device=torch.device("cpu"), - require_complete_routed_experts=True, - return_stats=True, - ) - if routed_experts is None: - raise RuntimeError( - "vLLM was asked to return routed experts for the " - "OpenAI-compatible chat endpoint but the generation " - "output did not include routed_experts." - ) - if r3_stats["missing_routes"] > 0: - LOGGER.warning( - "R3 router replay fallback: vLLM returned incomplete " - "routed_experts for chat choice_idx=%d, " - "missing_token_routes=%d, actual_routes=%d, " - "expected_routes=%d. Megatron will use its own router " - "for those missing token routes.", - choice.index, - r3_stats["missing_routes"], - r3_stats["actual_routes"], - r3_stats["expected_routes"], - ) - choice.message.routed_experts = routed_experts.to( - dtype=torch.int32 - ).tolist() - - return response - class NeMoRLOpenAIServingRender(NeMoRLOpenAIServingMixin, OpenAIServingRender): pass diff --git a/tests/unit/models/generation/test_vllm_utils.py b/tests/unit/models/generation/test_vllm_utils.py index 8e21abbf46e..a9683adda00 100644 --- a/tests/unit/models/generation/test_vllm_utils.py +++ b/tests/unit/models/generation/test_vllm_utils.py @@ -13,6 +13,7 @@ # limitations under the License. import math +from types import SimpleNamespace import pytest import torch @@ -21,6 +22,7 @@ from nemo_rl.models.generation.vllm.utils import ( R3_MISSING_ROUTE_SENTINEL, aggregate_spec_decode_counters, + attach_routed_experts_to_chat_response_choices, compute_spec_decode_metrics, format_prompt_for_vllm_generation, pad_and_align_routed_expert_indices, @@ -298,6 +300,68 @@ class Output: ) +def test_attach_routed_experts_to_chat_response_choices_reassociates_by_choice_index(): + final_res = SimpleNamespace( + prompt_token_ids=[101, 102, 103], + prompt_routed_experts=torch.tensor([[[10]], [[11]]], dtype=torch.int32), + outputs=[ + SimpleNamespace( + index=1, + token_ids=[201, 202], + routed_experts=torch.tensor([[[31]], [[32]]], dtype=torch.int32), + ), + SimpleNamespace( + index=0, + token_ids=[200], + routed_experts=torch.tensor([[[30]]], dtype=torch.int32), + ), + ], + ) + response = SimpleNamespace( + choices=[ + SimpleNamespace(index=0, message=SimpleNamespace()), + SimpleNamespace(index=1, message=SimpleNamespace()), + ] + ) + + attach_routed_experts_to_chat_response_choices( + response, + final_res, + device=torch.device("cpu"), + ) + + assert response.choices[0].message.routed_experts == [ + [[10]], + [[11]], + [[30]], + [[0]], + ] + assert response.choices[1].message.routed_experts == [ + [[10]], + [[11]], + [[31]], + [[32]], + [[0]], + ] + + +def test_attach_routed_experts_to_chat_response_choices_requires_routed_experts(): + final_res = SimpleNamespace( + prompt_token_ids=[101, 102], + outputs=[SimpleNamespace(index=0, token_ids=[200])], + ) + response = SimpleNamespace( + choices=[SimpleNamespace(index=0, message=SimpleNamespace())] + ) + + with pytest.raises(RuntimeError, match="did not include routed_experts"): + attach_routed_experts_to_chat_response_choices( + response, + final_res, + device=torch.device("cpu"), + ) + + @pytest.mark.vllm def test_vllm_speculative_decoding_patch_removed(): # The speculative decoding patch was fixed upstream in vLLM >= 0.14.0: From 1e39dd63268b0ad7ed57763fb9ee2736e84d683a Mon Sep 17 00:00:00 2001 From: Zeyu Zhou Date: Thu, 25 Jun 2026 10:21:12 -0700 Subject: [PATCH 03/12] fix: preserve R3 routes in async vLLM chat JSON Signed-off-by: Zeyu Zhou --- nemo_rl/models/generation/vllm/utils.py | 32 ++++++++++++++++++- .../generation/vllm/vllm_worker_async.py | 25 +++++++++------ .../unit/models/generation/test_vllm_utils.py | 28 ++++++++++++++++ 3 files changed, 74 insertions(+), 11 deletions(-) diff --git a/nemo_rl/models/generation/vllm/utils.py b/nemo_rl/models/generation/vllm/utils.py index daa7d09eb39..9635fa389de 100644 --- a/nemo_rl/models/generation/vllm/utils.py +++ b/nemo_rl/models/generation/vllm/utils.py @@ -207,10 +207,13 @@ def attach_routed_experts_to_chat_response_choices( getattr(final_request_output, "prompt_token_ids", []) or [] ) - for choice in getattr(response, "choices", []): + choices = list(getattr(response, "choices", [])) + attached_choice_indices = set() + for choice in choices: generation_details = outputs_by_index.get(choice.index) if generation_details is None: continue + attached_choice_indices.add(choice.index) generation_token_count = len(getattr(generation_details, "token_ids", []) or []) routed_result = pad_and_align_routed_expert_indices( @@ -248,9 +251,36 @@ def attach_routed_experts_to_chat_response_choices( ) choice.message.routed_experts = routed_experts.to(dtype=torch.int32).tolist() + if len(attached_choice_indices) != len(choices): + missing_choice_indices = sorted( + choice.index + for choice in choices + if choice.index not in attached_choice_indices + ) + raise RuntimeError( + "vLLM was asked to return routed experts for the " + "OpenAI-compatible chat endpoint but response choices could not be " + "matched to generation outputs: " + f"missing_choice_indices={missing_choice_indices}." + ) + return response +def model_dump_chat_response_with_routed_experts(response: Any) -> dict[str, Any]: + """Dump a vLLM OpenAI chat response while preserving dynamic R3 fields.""" + response_dict = response.model_dump() + for choice, choice_dict in zip( + getattr(response, "choices", []), response_dict.get("choices", []) + ): + routed_experts = getattr( + getattr(choice, "message", None), "routed_experts", None + ) + if routed_experts is not None: + choice_dict.setdefault("message", {})["routed_experts"] = routed_experts + return response_dict + + def aggregate_spec_decode_counters( worker_metrics: list[dict[str, float | list[float]]], ) -> dict[str | tuple[str, int], float]: diff --git a/nemo_rl/models/generation/vllm/vllm_worker_async.py b/nemo_rl/models/generation/vllm/vllm_worker_async.py index 9d7da35aab1..36f49d3ec2c 100644 --- a/nemo_rl/models/generation/vllm/vllm_worker_async.py +++ b/nemo_rl/models/generation/vllm/vllm_worker_async.py @@ -43,6 +43,7 @@ from nemo_rl.models.generation.vllm.utils import ( attach_routed_experts_to_chat_response_choices, format_prompt_for_vllm_generation, + model_dump_chat_response_with_routed_experts, pad_and_align_routed_expert_indices, ) from nemo_rl.models.generation.vllm.vllm_worker import BaseVllmGenerationWorker @@ -217,6 +218,14 @@ def __init__( self.llm = None self.vllm_device_ids = None + def _return_routed_experts_enabled(self) -> bool: + engine_args = getattr(self, "llm_async_engine_args", None) + if bool(getattr(engine_args, "enable_return_routed_experts", False)): + return True + return bool( + self.cfg.get("vllm_kwargs", {}).get("enable_return_routed_experts", False) + ) + def _reserve_port(self) -> None: """Bind and listen on a TCP socket to reserve a free port from the OS. @@ -672,9 +681,7 @@ class NeMoRLChatCompletionRequest( # vLLM 0.20 routes both /v1/chat/completions and /tokenize through # OpenAIServingRender.preprocess_chat, so the prefix-token override # belongs on the render subclass. - return_routed_experts = bool( - self.cfg.get("vllm_kwargs", {}).get("enable_return_routed_experts", False) - ) + worker_self = self class NeMoRLOpenAIServingChat(OpenAIServingChat): async def chat_completion_full_generator( @@ -699,7 +706,7 @@ async def capture_result_generator(): **kwargs, ) if ( - not return_routed_experts + not worker_self._return_routed_experts_enabled() or not isinstance(response, ChatCompletionResponse) or final_res is None ): @@ -791,7 +798,9 @@ async def create_chat_completion( ) elif isinstance(generator, ChatCompletionResponse): - return JSONResponse(content=generator.model_dump()) + return JSONResponse( + content=model_dump_chat_response_with_routed_experts(generator) + ) return StreamingResponse(content=generator, media_type="text/event-stream") @@ -1109,11 +1118,7 @@ async def process_single_sample(sample_idx): generation_details = final_request_output.outputs[0] generated_token_ids = list(generation_details.token_ids) num_generated_tokens = len(generated_token_ids) - return_routed_experts = bool( - self.cfg.get("vllm_kwargs", {}).get( - "enable_return_routed_experts", False - ) - ) + return_routed_experts = self._return_routed_experts_enabled() original_input_ids_single_row = input_ids_batch[sample_idx] final_output_tensor_len = current_input_actual_length + num_generated_tokens diff --git a/tests/unit/models/generation/test_vllm_utils.py b/tests/unit/models/generation/test_vllm_utils.py index a9683adda00..674a23b5fb7 100644 --- a/tests/unit/models/generation/test_vllm_utils.py +++ b/tests/unit/models/generation/test_vllm_utils.py @@ -25,6 +25,7 @@ attach_routed_experts_to_chat_response_choices, compute_spec_decode_metrics, format_prompt_for_vllm_generation, + model_dump_chat_response_with_routed_experts, pad_and_align_routed_expert_indices, ) @@ -362,6 +363,33 @@ def test_attach_routed_experts_to_chat_response_choices_requires_routed_experts( ) +def test_model_dump_chat_response_with_routed_experts_preserves_dynamic_field(): + routed_experts = [[[1]], [[2]]] + + class Response: + choices = [ + SimpleNamespace( + message=SimpleNamespace(routed_experts=routed_experts), + ) + ] + + def model_dump(self): + return { + "choices": [ + { + "message": { + "role": "assistant", + "content": "hello", + } + } + ] + } + + response_dict = model_dump_chat_response_with_routed_experts(Response()) + + assert response_dict["choices"][0]["message"]["routed_experts"] == routed_experts + + @pytest.mark.vllm def test_vllm_speculative_decoding_patch_removed(): # The speculative decoding patch was fixed upstream in vLLM >= 0.14.0: From c25a0bf186fa9a2f77c3eb2cab3547bf709bc443 Mon Sep 17 00:00:00 2001 From: Zeyu Zhou Date: Thu, 25 Jun 2026 12:27:10 -0700 Subject: [PATCH 04/12] refactor: move async chat R3 hook into mixin Signed-off-by: Zeyu Zhou --- nemo_rl/models/generation/vllm/vllm_worker_async.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/nemo_rl/models/generation/vllm/vllm_worker_async.py b/nemo_rl/models/generation/vllm/vllm_worker_async.py index 36f49d3ec2c..fab2e1330e8 100644 --- a/nemo_rl/models/generation/vllm/vllm_worker_async.py +++ b/nemo_rl/models/generation/vllm/vllm_worker_async.py @@ -683,7 +683,7 @@ class NeMoRLChatCompletionRequest( # belongs on the render subclass. worker_self = self - class NeMoRLOpenAIServingChat(OpenAIServingChat): + class NeMoRLOpenAIServingChatMixin: async def chat_completion_full_generator( self, request, @@ -719,6 +719,9 @@ async def capture_result_generator(): logger=LOGGER, ) + class NeMoRLOpenAIServingChat(NeMoRLOpenAIServingChatMixin, OpenAIServingChat): + pass + class NeMoRLOpenAIServingRender(NeMoRLOpenAIServingMixin, OpenAIServingRender): pass From aefdfaf9042dd1d42659beccd6d2fef5a863c418 Mon Sep 17 00:00:00 2001 From: Zeyu Zhou Date: Tue, 30 Jun 2026 09:43:52 -0700 Subject: [PATCH 05/12] test: cover async chat R3 route attachment guards Signed-off-by: Zeyu Zhou --- .../unit/models/generation/test_vllm_utils.py | 67 +++++++++++++++++++ 1 file changed, 67 insertions(+) diff --git a/tests/unit/models/generation/test_vllm_utils.py b/tests/unit/models/generation/test_vllm_utils.py index 674a23b5fb7..38790c689d6 100644 --- a/tests/unit/models/generation/test_vllm_utils.py +++ b/tests/unit/models/generation/test_vllm_utils.py @@ -14,6 +14,7 @@ import math from types import SimpleNamespace +from unittest.mock import MagicMock import pytest import torch @@ -363,6 +364,72 @@ def test_attach_routed_experts_to_chat_response_choices_requires_routed_experts( ) +def test_attach_routed_experts_to_chat_response_choices_warns_on_missing_routes(): + final_res = SimpleNamespace( + prompt_token_ids=[101, 102, 103], + outputs=[ + SimpleNamespace( + index=0, + token_ids=[200, 201], + routed_experts=torch.tensor([[[10]], [[11]]], dtype=torch.int32), + ) + ], + ) + response = SimpleNamespace( + choices=[SimpleNamespace(index=0, message=SimpleNamespace())] + ) + logger = MagicMock() + + attach_routed_experts_to_chat_response_choices( + response, + final_res, + device=torch.device("cpu"), + logger=logger, + ) + + logger.warning.assert_called_once_with( + "R3 router replay fallback: vLLM returned incomplete " + "routed_experts for chat choice_idx=%d, " + "missing_token_routes=%d, actual_routes=%d, " + "expected_routes=%d. Megatron will use its own router " + "for those missing token routes.", + 0, + 2, + 2, + 4, + ) + assert response.choices[0].message.routed_experts == [ + [[10]], + [[11]], + [[R3_MISSING_ROUTE_SENTINEL]], + [[R3_MISSING_ROUTE_SENTINEL]], + [[0]], + ] + + +def test_attach_routed_experts_to_chat_response_choices_raises_for_unmatched_choice(): + final_res = SimpleNamespace( + prompt_token_ids=[101, 102], + outputs=[ + SimpleNamespace( + index=1, + token_ids=[200], + routed_experts=torch.tensor([[[10]], [[11]]], dtype=torch.int32), + ) + ], + ) + response = SimpleNamespace( + choices=[SimpleNamespace(index=0, message=SimpleNamespace())] + ) + + with pytest.raises(RuntimeError, match=r"missing_choice_indices=\[0\]"): + attach_routed_experts_to_chat_response_choices( + response, + final_res, + device=torch.device("cpu"), + ) + + def test_model_dump_chat_response_with_routed_experts_preserves_dynamic_field(): routed_experts = [[[1]], [[2]]] From 8aa2f2ff246eefd8eb74cbe415a01de23a8a338e Mon Sep 17 00:00:00 2001 From: Zeyu Zhou Date: Tue, 30 Jun 2026 10:30:26 -0700 Subject: [PATCH 06/12] test: add R3 nightly coverage configs Signed-off-by: Zeyu Zhou --- ...ba3b-10n8g-megatron-cp2-r3-async-notq.yaml | 29 +++++++++ ...30ba3b-8n8g-megatron-cp2-r3-tq_simple.yaml | 11 ++++ ...rpo-qwen3-30ba3b-8n8g-megatron-cp2-r3.yaml | 2 +- ...-16n8g-megatron-cp2-r3-async-gym-notq.yaml | 35 +++++++++++ ...30ba3b-10n8g-megatron-cp2-r3-async-notq.sh | 41 ++++++++++++ ...3-30ba3b-8n8g-megatron-cp2-r3-tq_simple.sh | 41 ++++++++++++ ...e1-16n8g-megatron-cp2-r3-async-gym-notq.sh | 62 +++++++++++++++++++ 7 files changed, 220 insertions(+), 1 deletion(-) create mode 100644 examples/configs/recipes/llm/grpo-qwen3-30ba3b-10n8g-megatron-cp2-r3-async-notq.yaml create mode 100644 examples/configs/recipes/llm/grpo-qwen3-30ba3b-8n8g-megatron-cp2-r3-tq_simple.yaml create mode 100644 examples/configs/recipes/llm/grpo-qwen3-30ba3b-thinking-swe1-16n8g-megatron-cp2-r3-async-gym-notq.yaml create mode 100755 tests/test_suites/llm/grpo-qwen3-30ba3b-10n8g-megatron-cp2-r3-async-notq.sh create mode 100755 tests/test_suites/llm/grpo-qwen3-30ba3b-8n8g-megatron-cp2-r3-tq_simple.sh create mode 100755 tests/test_suites/llm/grpo-qwen3-30ba3b-thinking-swe1-16n8g-megatron-cp2-r3-async-gym-notq.sh diff --git a/examples/configs/recipes/llm/grpo-qwen3-30ba3b-10n8g-megatron-cp2-r3-async-notq.yaml b/examples/configs/recipes/llm/grpo-qwen3-30ba3b-10n8g-megatron-cp2-r3-async-notq.yaml new file mode 100644 index 00000000000..721528d92e8 --- /dev/null +++ b/examples/configs/recipes/llm/grpo-qwen3-30ba3b-10n8g-megatron-cp2-r3-async-notq.yaml @@ -0,0 +1,29 @@ +defaults: ./grpo-qwen3-30ba3b-8n8g-megatron-cp2-r3.yaml + +grpo: + async_grpo: + enabled: true + max_trajectory_age_steps: 1 + in_flight_weight_updates: true +loss_fn: + use_importance_sampling_correction: true +checkpointing: + checkpoint_dir: results/grpo-qwen3-30ba3b-10n8g-megatron-cp2-r3-async-notq +data_plane: + enabled: false +policy: + generation: + colocated: + enabled: false + resources: + num_nodes: 2 + gpus_per_node: 8 + vllm_cfg: + async_engine: true +logger: + log_dir: logs/grpo-qwen3-30ba3b-10n8g-megatron-cp2-r3-async-notq + wandb: + name: grpo-qwen3-30ba3b-10n8g-megatron-cp2-r3-async-notq +cluster: + num_nodes: 10 + gpus_per_node: 8 diff --git a/examples/configs/recipes/llm/grpo-qwen3-30ba3b-8n8g-megatron-cp2-r3-tq_simple.yaml b/examples/configs/recipes/llm/grpo-qwen3-30ba3b-8n8g-megatron-cp2-r3-tq_simple.yaml new file mode 100644 index 00000000000..37206e24718 --- /dev/null +++ b/examples/configs/recipes/llm/grpo-qwen3-30ba3b-8n8g-megatron-cp2-r3-tq_simple.yaml @@ -0,0 +1,11 @@ +defaults: ./grpo-qwen3-30ba3b-8n8g-megatron-cp2-r3.yaml + +checkpointing: + checkpoint_dir: results/grpo-qwen3-30ba3b-8n8g-megatron-cp2-r3-tq_simple +data_plane: + enabled: true + backend: simple +logger: + log_dir: logs/grpo-qwen3-30ba3b-8n8g-megatron-cp2-r3-tq_simple + wandb: + name: grpo-qwen3-30ba3b-8n8g-megatron-cp2-r3-tq_simple diff --git a/examples/configs/recipes/llm/grpo-qwen3-30ba3b-8n8g-megatron-cp2-r3.yaml b/examples/configs/recipes/llm/grpo-qwen3-30ba3b-8n8g-megatron-cp2-r3.yaml index 6ae8abd1188..ec04203ebb8 100644 --- a/examples/configs/recipes/llm/grpo-qwen3-30ba3b-8n8g-megatron-cp2-r3.yaml +++ b/examples/configs/recipes/llm/grpo-qwen3-30ba3b-8n8g-megatron-cp2-r3.yaml @@ -30,7 +30,7 @@ policy: data: max_input_seq_length: 2048 data_plane: - enabled: true + enabled: false logger: log_dir: logs/grpo-qwen3-30ba3b-8n8g-megatron-cp2-r3 wandb: diff --git a/examples/configs/recipes/llm/grpo-qwen3-30ba3b-thinking-swe1-16n8g-megatron-cp2-r3-async-gym-notq.yaml b/examples/configs/recipes/llm/grpo-qwen3-30ba3b-thinking-swe1-16n8g-megatron-cp2-r3-async-gym-notq.yaml new file mode 100644 index 00000000000..879c1ab17e4 --- /dev/null +++ b/examples/configs/recipes/llm/grpo-qwen3-30ba3b-thinking-swe1-16n8g-megatron-cp2-r3-async-gym-notq.yaml @@ -0,0 +1,35 @@ +defaults: ../../../nemo_gym/grpo_qwen3_30ba3b_thinking_swe1.yaml + +grpo: + num_prompts_per_step: 16 + num_generations_per_prompt: 8 + max_val_samples: 32 + val_batch_size: 32 + skip_reference_policy_logprobs_calculation: true + async_grpo: + enabled: true + max_trajectory_age_steps: 1 + in_flight_weight_updates: true +loss_fn: + use_importance_sampling_correction: true +checkpointing: + checkpoint_dir: results/grpo-qwen3-30ba3b-thinking-swe1-16n8g-megatron-cp2-r3-async-gym-notq + enabled: false +data_plane: null +policy: + train_global_batch_size: 128 + logprob_batch_size: 1 + router_replay: + enabled: true + generation: + vllm_cfg: + enable_prefix_caching: false + vllm_kwargs: + enable_chunked_prefill: false +logger: + log_dir: logs/grpo-qwen3-30ba3b-thinking-swe1-16n8g-megatron-cp2-r3-async-gym-notq + wandb: + name: grpo-qwen3-30ba3b-thinking-swe1-16n8g-megatron-cp2-r3-async-gym-notq +cluster: + num_nodes: 16 + gpus_per_node: 8 diff --git a/tests/test_suites/llm/grpo-qwen3-30ba3b-10n8g-megatron-cp2-r3-async-notq.sh b/tests/test_suites/llm/grpo-qwen3-30ba3b-10n8g-megatron-cp2-r3-async-notq.sh new file mode 100755 index 00000000000..69374ca0d87 --- /dev/null +++ b/tests/test_suites/llm/grpo-qwen3-30ba3b-10n8g-megatron-cp2-r3-async-notq.sh @@ -0,0 +1,41 @@ +#!/bin/bash +SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd) +source $SCRIPT_DIR/common.env + +export NRL_IGNORE_TP_ACCURACY_CHECK=1 +export NRL_ROUTER_REPLAY_VALIDATE=1 + +# ===== BEGIN CONFIG ===== +NUM_NODES=10 +GPUS_PER_NODE=8 +STEPS_PER_RUN=10 +MAX_STEPS=10 +NUM_RUNS=$(( (MAX_STEPS + STEPS_PER_RUN - 1) / STEPS_PER_RUN )) # Round up +NUM_MINUTES=60 +# ===== END CONFIG ===== + +exit_if_max_steps_reached + +cd $PROJECT_ROOT +uv run examples/run_grpo.py \ + --config $CONFIG_PATH \ + grpo.max_num_steps=$MAX_STEPS \ + logger.log_dir=$LOG_DIR \ + logger.wandb_enabled=True \ + logger.wandb.project=nemo-rl \ + logger.wandb.name=$EXP_NAME \ + logger.monitor_gpus=True \ + logger.tensorboard_enabled=True \ + checkpointing.enabled=True \ + checkpointing.checkpoint_dir=$CKPT_DIR \ + $@ \ + 2>&1 | tee $RUN_LOG + +uv run tests/json_dump_tb_logs.py $LOG_DIR --output_path $JSON_METRICS + +if [[ $(jq 'to_entries | .[] | select(.key == "train/loss") | .value | keys | map(tonumber) | max' $JSON_METRICS) -ge $MAX_STEPS ]]; then + uv run tests/check_metrics.py $JSON_METRICS \ + 'median(data["train/token_mult_prob_error"]) < 1.02' + + rm -rf "$CKPT_DIR" +fi diff --git a/tests/test_suites/llm/grpo-qwen3-30ba3b-8n8g-megatron-cp2-r3-tq_simple.sh b/tests/test_suites/llm/grpo-qwen3-30ba3b-8n8g-megatron-cp2-r3-tq_simple.sh new file mode 100755 index 00000000000..490f00fb665 --- /dev/null +++ b/tests/test_suites/llm/grpo-qwen3-30ba3b-8n8g-megatron-cp2-r3-tq_simple.sh @@ -0,0 +1,41 @@ +#!/bin/bash +SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd) +source $SCRIPT_DIR/common.env + +export NRL_IGNORE_TP_ACCURACY_CHECK=1 +export NRL_ROUTER_REPLAY_VALIDATE=1 + +# ===== BEGIN CONFIG ===== +NUM_NODES=8 +GPUS_PER_NODE=8 +STEPS_PER_RUN=10 +MAX_STEPS=10 +NUM_RUNS=$(( (MAX_STEPS + STEPS_PER_RUN - 1) / STEPS_PER_RUN )) # Round up +NUM_MINUTES=60 +# ===== END CONFIG ===== + +exit_if_max_steps_reached + +cd $PROJECT_ROOT +uv run examples/run_grpo.py \ + --config $CONFIG_PATH \ + grpo.max_num_steps=$MAX_STEPS \ + logger.log_dir=$LOG_DIR \ + logger.wandb_enabled=True \ + logger.wandb.project=nemo-rl \ + logger.wandb.name=$EXP_NAME \ + logger.monitor_gpus=True \ + logger.tensorboard_enabled=True \ + checkpointing.enabled=True \ + checkpointing.checkpoint_dir=$CKPT_DIR \ + $@ \ + 2>&1 | tee $RUN_LOG + +uv run tests/json_dump_tb_logs.py $LOG_DIR --output_path $JSON_METRICS + +if [[ $(jq 'to_entries | .[] | select(.key == "train/loss") | .value | keys | map(tonumber) | max' $JSON_METRICS) -ge $MAX_STEPS ]]; then + uv run tests/check_metrics.py $JSON_METRICS \ + 'median(data["train/token_mult_prob_error"]) < 1.02' + + rm -rf "$CKPT_DIR" +fi diff --git a/tests/test_suites/llm/grpo-qwen3-30ba3b-thinking-swe1-16n8g-megatron-cp2-r3-async-gym-notq.sh b/tests/test_suites/llm/grpo-qwen3-30ba3b-thinking-swe1-16n8g-megatron-cp2-r3-async-gym-notq.sh new file mode 100755 index 00000000000..f39a9b2f069 --- /dev/null +++ b/tests/test_suites/llm/grpo-qwen3-30ba3b-thinking-swe1-16n8g-megatron-cp2-r3-async-gym-notq.sh @@ -0,0 +1,62 @@ +#!/bin/bash +SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd) +source $SCRIPT_DIR/common.env + +export NRL_ROUTER_REPLAY_VALIDATE=1 + +# ===== BEGIN CONFIG ===== +NUM_NODES=16 +GPUS_PER_NODE=8 +STEPS_PER_RUN=10 +MAX_STEPS=10 +NUM_RUNS=$(( (MAX_STEPS + STEPS_PER_RUN - 1) / STEPS_PER_RUN )) # Round up +NUM_MINUTES=240 +# ===== END CONFIG ===== + +exit_if_max_steps_reached + +cd $PROJECT_ROOT + +DATA_DIR=$EXP_DIR/data +RAW_DATA_DIR=$DATA_DIR/raw +TRAIN_PATH=$DATA_DIR/swe1_train.jsonl +VALIDATION_PATH=$DATA_DIR/swe1_validation.jsonl +mkdir -p $RAW_DATA_DIR + +if [[ ! -f $RAW_DATA_DIR/swe1.jsonl ]]; then + uv run hf download nvidia/Nemotron-RL-Super-Training-Blends swe1.jsonl \ + --repo-type dataset \ + --local-dir $RAW_DATA_DIR +fi + +if [[ ! -f $TRAIN_PATH ]]; then + head -n 512 $RAW_DATA_DIR/swe1.jsonl > $TRAIN_PATH +fi +if [[ ! -f $VALIDATION_PATH ]]; then + tail -n 32 $RAW_DATA_DIR/swe1.jsonl > $VALIDATION_PATH +fi + +uv run examples/nemo_gym/run_grpo_nemo_gym.py \ + --config $CONFIG_PATH \ + grpo.max_num_steps=$MAX_STEPS \ + logger.log_dir=$LOG_DIR \ + logger.wandb_enabled=True \ + logger.wandb.project=nemo-rl \ + logger.wandb.name=$EXP_NAME \ + logger.monitor_gpus=True \ + logger.tensorboard_enabled=True \ + checkpointing.enabled=True \ + checkpointing.checkpoint_dir=$CKPT_DIR \ + data.train.data_path=$TRAIN_PATH \ + data.validation.data_path=$VALIDATION_PATH \ + $@ \ + 2>&1 | tee $RUN_LOG + +uv run tests/json_dump_tb_logs.py $LOG_DIR --output_path $JSON_METRICS + +if [[ $(jq 'to_entries | .[] | select(.key == "train/loss") | .value | keys | map(tonumber) | max' $JSON_METRICS) -ge $MAX_STEPS ]]; then + uv run tests/check_metrics.py $JSON_METRICS \ + 'median(data["train/token_mult_prob_error"]) < 1.02' + + rm -rf "$CKPT_DIR" +fi From 0a2718dc5f0cafbb233d96cbc64892f255d1f075 Mon Sep 17 00:00:00 2001 From: Zeyu Zhou Date: Tue, 30 Jun 2026 10:35:44 -0700 Subject: [PATCH 07/12] test: simplify R3 async nightly names Signed-off-by: Zeyu Zhou --- ...l => grpo-qwen3-30ba3b-10n8g-megatron-cp2-r3-async.yaml} | 6 +++--- ...ba3b-thinking-swe1-16n8g-megatron-cp2-r3-async-gym.yaml} | 6 +++--- ....sh => grpo-qwen3-30ba3b-10n8g-megatron-cp2-r3-async.sh} | 0 ...30ba3b-thinking-swe1-16n8g-megatron-cp2-r3-async-gym.sh} | 0 4 files changed, 6 insertions(+), 6 deletions(-) rename examples/configs/recipes/llm/{grpo-qwen3-30ba3b-10n8g-megatron-cp2-r3-async-notq.yaml => grpo-qwen3-30ba3b-10n8g-megatron-cp2-r3-async.yaml} (79%) rename examples/configs/recipes/llm/{grpo-qwen3-30ba3b-thinking-swe1-16n8g-megatron-cp2-r3-async-gym-notq.yaml => grpo-qwen3-30ba3b-thinking-swe1-16n8g-megatron-cp2-r3-async-gym.yaml} (93%) rename tests/test_suites/llm/{grpo-qwen3-30ba3b-10n8g-megatron-cp2-r3-async-notq.sh => grpo-qwen3-30ba3b-10n8g-megatron-cp2-r3-async.sh} (100%) rename tests/test_suites/llm/{grpo-qwen3-30ba3b-thinking-swe1-16n8g-megatron-cp2-r3-async-gym-notq.sh => grpo-qwen3-30ba3b-thinking-swe1-16n8g-megatron-cp2-r3-async-gym.sh} (100%) diff --git a/examples/configs/recipes/llm/grpo-qwen3-30ba3b-10n8g-megatron-cp2-r3-async-notq.yaml b/examples/configs/recipes/llm/grpo-qwen3-30ba3b-10n8g-megatron-cp2-r3-async.yaml similarity index 79% rename from examples/configs/recipes/llm/grpo-qwen3-30ba3b-10n8g-megatron-cp2-r3-async-notq.yaml rename to examples/configs/recipes/llm/grpo-qwen3-30ba3b-10n8g-megatron-cp2-r3-async.yaml index 721528d92e8..4efd10672cf 100644 --- a/examples/configs/recipes/llm/grpo-qwen3-30ba3b-10n8g-megatron-cp2-r3-async-notq.yaml +++ b/examples/configs/recipes/llm/grpo-qwen3-30ba3b-10n8g-megatron-cp2-r3-async.yaml @@ -8,7 +8,7 @@ grpo: loss_fn: use_importance_sampling_correction: true checkpointing: - checkpoint_dir: results/grpo-qwen3-30ba3b-10n8g-megatron-cp2-r3-async-notq + checkpoint_dir: results/grpo-qwen3-30ba3b-10n8g-megatron-cp2-r3-async data_plane: enabled: false policy: @@ -21,9 +21,9 @@ policy: vllm_cfg: async_engine: true logger: - log_dir: logs/grpo-qwen3-30ba3b-10n8g-megatron-cp2-r3-async-notq + log_dir: logs/grpo-qwen3-30ba3b-10n8g-megatron-cp2-r3-async wandb: - name: grpo-qwen3-30ba3b-10n8g-megatron-cp2-r3-async-notq + name: grpo-qwen3-30ba3b-10n8g-megatron-cp2-r3-async cluster: num_nodes: 10 gpus_per_node: 8 diff --git a/examples/configs/recipes/llm/grpo-qwen3-30ba3b-thinking-swe1-16n8g-megatron-cp2-r3-async-gym-notq.yaml b/examples/configs/recipes/llm/grpo-qwen3-30ba3b-thinking-swe1-16n8g-megatron-cp2-r3-async-gym.yaml similarity index 93% rename from examples/configs/recipes/llm/grpo-qwen3-30ba3b-thinking-swe1-16n8g-megatron-cp2-r3-async-gym-notq.yaml rename to examples/configs/recipes/llm/grpo-qwen3-30ba3b-thinking-swe1-16n8g-megatron-cp2-r3-async-gym.yaml index 879c1ab17e4..877172a2e28 100644 --- a/examples/configs/recipes/llm/grpo-qwen3-30ba3b-thinking-swe1-16n8g-megatron-cp2-r3-async-gym-notq.yaml +++ b/examples/configs/recipes/llm/grpo-qwen3-30ba3b-thinking-swe1-16n8g-megatron-cp2-r3-async-gym.yaml @@ -13,7 +13,7 @@ grpo: loss_fn: use_importance_sampling_correction: true checkpointing: - checkpoint_dir: results/grpo-qwen3-30ba3b-thinking-swe1-16n8g-megatron-cp2-r3-async-gym-notq + checkpoint_dir: results/grpo-qwen3-30ba3b-thinking-swe1-16n8g-megatron-cp2-r3-async-gym enabled: false data_plane: null policy: @@ -27,9 +27,9 @@ policy: vllm_kwargs: enable_chunked_prefill: false logger: - log_dir: logs/grpo-qwen3-30ba3b-thinking-swe1-16n8g-megatron-cp2-r3-async-gym-notq + log_dir: logs/grpo-qwen3-30ba3b-thinking-swe1-16n8g-megatron-cp2-r3-async-gym wandb: - name: grpo-qwen3-30ba3b-thinking-swe1-16n8g-megatron-cp2-r3-async-gym-notq + name: grpo-qwen3-30ba3b-thinking-swe1-16n8g-megatron-cp2-r3-async-gym cluster: num_nodes: 16 gpus_per_node: 8 diff --git a/tests/test_suites/llm/grpo-qwen3-30ba3b-10n8g-megatron-cp2-r3-async-notq.sh b/tests/test_suites/llm/grpo-qwen3-30ba3b-10n8g-megatron-cp2-r3-async.sh similarity index 100% rename from tests/test_suites/llm/grpo-qwen3-30ba3b-10n8g-megatron-cp2-r3-async-notq.sh rename to tests/test_suites/llm/grpo-qwen3-30ba3b-10n8g-megatron-cp2-r3-async.sh diff --git a/tests/test_suites/llm/grpo-qwen3-30ba3b-thinking-swe1-16n8g-megatron-cp2-r3-async-gym-notq.sh b/tests/test_suites/llm/grpo-qwen3-30ba3b-thinking-swe1-16n8g-megatron-cp2-r3-async-gym.sh similarity index 100% rename from tests/test_suites/llm/grpo-qwen3-30ba3b-thinking-swe1-16n8g-megatron-cp2-r3-async-gym-notq.sh rename to tests/test_suites/llm/grpo-qwen3-30ba3b-thinking-swe1-16n8g-megatron-cp2-r3-async-gym.sh From ca5ea9bf8c6427e87ec18c9c1b4d9fc64e0afd24 Mon Sep 17 00:00:00 2001 From: Zeyu Zhou Date: Tue, 30 Jun 2026 10:40:41 -0700 Subject: [PATCH 08/12] test: shorten R3 Gym nightly Signed-off-by: Zeyu Zhou --- ...n3-30ba3b-thinking-swe1-16n8g-megatron-cp2-r3-async-gym.sh | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/test_suites/llm/grpo-qwen3-30ba3b-thinking-swe1-16n8g-megatron-cp2-r3-async-gym.sh b/tests/test_suites/llm/grpo-qwen3-30ba3b-thinking-swe1-16n8g-megatron-cp2-r3-async-gym.sh index f39a9b2f069..3bc2d5f4700 100755 --- a/tests/test_suites/llm/grpo-qwen3-30ba3b-thinking-swe1-16n8g-megatron-cp2-r3-async-gym.sh +++ b/tests/test_suites/llm/grpo-qwen3-30ba3b-thinking-swe1-16n8g-megatron-cp2-r3-async-gym.sh @@ -7,8 +7,8 @@ export NRL_ROUTER_REPLAY_VALIDATE=1 # ===== BEGIN CONFIG ===== NUM_NODES=16 GPUS_PER_NODE=8 -STEPS_PER_RUN=10 -MAX_STEPS=10 +STEPS_PER_RUN=5 +MAX_STEPS=5 NUM_RUNS=$(( (MAX_STEPS + STEPS_PER_RUN - 1) / STEPS_PER_RUN )) # Round up NUM_MINUTES=240 # ===== END CONFIG ===== From 25c2f00ddde05b25834f80f9fd24ef855bff586a Mon Sep 17 00:00:00 2001 From: Zeyu Zhou Date: Tue, 30 Jun 2026 10:52:03 -0700 Subject: [PATCH 09/12] test: add R3 recipes to nightly suite Signed-off-by: Zeyu Zhou --- ...-30ba3b-thinking-swe1-16n8g-megatron-cp2-r3-async-gym.sh | 6 +++--- tests/test_suites/nightly.txt | 5 ++++- tests/unit/test_recipes_and_test_suites.py | 6 +++--- 3 files changed, 10 insertions(+), 7 deletions(-) diff --git a/tests/test_suites/llm/grpo-qwen3-30ba3b-thinking-swe1-16n8g-megatron-cp2-r3-async-gym.sh b/tests/test_suites/llm/grpo-qwen3-30ba3b-thinking-swe1-16n8g-megatron-cp2-r3-async-gym.sh index 3bc2d5f4700..43130d82f89 100755 --- a/tests/test_suites/llm/grpo-qwen3-30ba3b-thinking-swe1-16n8g-megatron-cp2-r3-async-gym.sh +++ b/tests/test_suites/llm/grpo-qwen3-30ba3b-thinking-swe1-16n8g-megatron-cp2-r3-async-gym.sh @@ -7,10 +7,10 @@ export NRL_ROUTER_REPLAY_VALIDATE=1 # ===== BEGIN CONFIG ===== NUM_NODES=16 GPUS_PER_NODE=8 -STEPS_PER_RUN=5 -MAX_STEPS=5 +STEPS_PER_RUN=3 +MAX_STEPS=3 NUM_RUNS=$(( (MAX_STEPS + STEPS_PER_RUN - 1) / STEPS_PER_RUN )) # Round up -NUM_MINUTES=240 +NUM_MINUTES=90 # ===== END CONFIG ===== exit_if_max_steps_reached diff --git a/tests/test_suites/nightly.txt b/tests/test_suites/nightly.txt index 2c2fa417dd4..1a6373a9d49 100644 --- a/tests/test_suites/nightly.txt +++ b/tests/test_suites/nightly.txt @@ -73,8 +73,11 @@ tests/test_suites/llm/grpo-qwen3-30ba3b-4n4g-megatron-qa-nvfp4.sh # CISPO async lag-1 high-off-policy run (Qwen3-30B-A3B, Megatron + non-colocated vLLM) tests/test_suites/llm/grpo-cispo-mm1-async-lag1-highoffpolicy-qwen3-30ba3b-3n8g-megatron-cispo.sh -# R3 router replay regression guard (Qwen3-30B-A3B, Megatron CP2 + EP4) — short 10-step smoke +# R3 router replay regression guards (Qwen3-30B-A3B, Megatron CP2 + EP4) tests/test_suites/llm/grpo-qwen3-30ba3b-8n8g-megatron-cp2-r3.sh +tests/test_suites/llm/grpo-qwen3-30ba3b-8n8g-megatron-cp2-r3-tq_simple.sh +tests/test_suites/llm/grpo-qwen3-30ba3b-10n8g-megatron-cp2-r3-async.sh +tests/test_suites/llm/grpo-qwen3-30ba3b-thinking-swe1-16n8g-megatron-cp2-r3-async-gym.sh # FP8 tests/test_suites/llm/grpo-llama3.1-8b-instruct-1n8g-megatron-fp8-rollouts.v3.sh diff --git a/tests/unit/test_recipes_and_test_suites.py b/tests/unit/test_recipes_and_test_suites.py index 3583b7b19be..48907bd900d 100644 --- a/tests/unit/test_recipes_and_test_suites.py +++ b/tests/unit/test_recipes_and_test_suites.py @@ -235,7 +235,7 @@ def test_all_recipe_yamls_accounted_for_in_test_suites( ) -def test_nightly_compute_stays_below_2300_hours(nightly_test_suite, tracker): +def test_nightly_compute_stays_below_2650_hours(nightly_test_suite, tracker): command = f"DRYRUN=1 HF_HOME=... HF_DATASETS_CACHE=... CONTAINER= ACCOUNT= PARTITION= ./tools/launch {' '.join(nightly_test_suite)}" print(f"Running command: {command}") @@ -267,8 +267,8 @@ def test_nightly_compute_stays_below_2300_hours(nightly_test_suite, tracker): f"Last line of output was not as expected: '{last_line}'" ) total_gpu_hours = float(last_line.split(":")[-1].strip()) - assert total_gpu_hours <= 2300, ( - f"Total GPU hours exceeded 2300: {last_line}. We should revisit the test suites to reduce the total GPU hours." + assert total_gpu_hours <= 2650, ( + f"Total GPU hours exceeded 2650: {last_line}. We should revisit the test suites to reduce the total GPU hours." ) tracker.track("total_nightly_gpu_hours", total_gpu_hours) From 6b8df9808bdee6bfa5cf09b8d17b2fdd58674090 Mon Sep 17 00:00:00 2001 From: Zeyu Zhou Date: Tue, 30 Jun 2026 11:16:23 -0700 Subject: [PATCH 10/12] test: use updated nightly GPU-hour cap Signed-off-by: Zeyu Zhou --- tests/unit/test_recipes_and_test_suites.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/tests/unit/test_recipes_and_test_suites.py b/tests/unit/test_recipes_and_test_suites.py index 48907bd900d..93af891db64 100644 --- a/tests/unit/test_recipes_and_test_suites.py +++ b/tests/unit/test_recipes_and_test_suites.py @@ -235,7 +235,7 @@ def test_all_recipe_yamls_accounted_for_in_test_suites( ) -def test_nightly_compute_stays_below_2650_hours(nightly_test_suite, tracker): +def test_nightly_compute_stays_below_3000_hours(nightly_test_suite, tracker): command = f"DRYRUN=1 HF_HOME=... HF_DATASETS_CACHE=... CONTAINER= ACCOUNT= PARTITION= ./tools/launch {' '.join(nightly_test_suite)}" print(f"Running command: {command}") @@ -267,8 +267,8 @@ def test_nightly_compute_stays_below_2650_hours(nightly_test_suite, tracker): f"Last line of output was not as expected: '{last_line}'" ) total_gpu_hours = float(last_line.split(":")[-1].strip()) - assert total_gpu_hours <= 2650, ( - f"Total GPU hours exceeded 2650: {last_line}. We should revisit the test suites to reduce the total GPU hours." + assert total_gpu_hours <= 3000, ( + f"Total GPU hours exceeded 3000: {last_line}. We should revisit the test suites to reduce the total GPU hours." ) tracker.track("total_nightly_gpu_hours", total_gpu_hours) From 43d34b5255b50e56a32fe0ed7156735031f3f255 Mon Sep 17 00:00:00 2001 From: Zeyu Zhou Date: Tue, 30 Jun 2026 11:31:47 -0700 Subject: [PATCH 11/12] test: minimize R3 nightly recipes Signed-off-by: Zeyu Zhou --- ...po-qwen3-30ba3b-10n8g-megatron-cp2-r3-async.yaml | 5 ----- ...qwen3-30ba3b-8n8g-megatron-cp2-r3-tq_simple.yaml | 2 -- .../llm/grpo-qwen3-30ba3b-8n8g-megatron-cp2-r3.yaml | 2 -- ...inking-swe1-16n8g-megatron-cp2-r3-async-gym.yaml | 13 ------------- 4 files changed, 22 deletions(-) diff --git a/examples/configs/recipes/llm/grpo-qwen3-30ba3b-10n8g-megatron-cp2-r3-async.yaml b/examples/configs/recipes/llm/grpo-qwen3-30ba3b-10n8g-megatron-cp2-r3-async.yaml index 4efd10672cf..e4c2b9e8094 100644 --- a/examples/configs/recipes/llm/grpo-qwen3-30ba3b-10n8g-megatron-cp2-r3-async.yaml +++ b/examples/configs/recipes/llm/grpo-qwen3-30ba3b-10n8g-megatron-cp2-r3-async.yaml @@ -1,16 +1,12 @@ defaults: ./grpo-qwen3-30ba3b-8n8g-megatron-cp2-r3.yaml - grpo: async_grpo: enabled: true - max_trajectory_age_steps: 1 in_flight_weight_updates: true loss_fn: use_importance_sampling_correction: true checkpointing: checkpoint_dir: results/grpo-qwen3-30ba3b-10n8g-megatron-cp2-r3-async -data_plane: - enabled: false policy: generation: colocated: @@ -26,4 +22,3 @@ logger: name: grpo-qwen3-30ba3b-10n8g-megatron-cp2-r3-async cluster: num_nodes: 10 - gpus_per_node: 8 diff --git a/examples/configs/recipes/llm/grpo-qwen3-30ba3b-8n8g-megatron-cp2-r3-tq_simple.yaml b/examples/configs/recipes/llm/grpo-qwen3-30ba3b-8n8g-megatron-cp2-r3-tq_simple.yaml index 37206e24718..f48450c8ef2 100644 --- a/examples/configs/recipes/llm/grpo-qwen3-30ba3b-8n8g-megatron-cp2-r3-tq_simple.yaml +++ b/examples/configs/recipes/llm/grpo-qwen3-30ba3b-8n8g-megatron-cp2-r3-tq_simple.yaml @@ -1,10 +1,8 @@ defaults: ./grpo-qwen3-30ba3b-8n8g-megatron-cp2-r3.yaml - checkpointing: checkpoint_dir: results/grpo-qwen3-30ba3b-8n8g-megatron-cp2-r3-tq_simple data_plane: enabled: true - backend: simple logger: log_dir: logs/grpo-qwen3-30ba3b-8n8g-megatron-cp2-r3-tq_simple wandb: diff --git a/examples/configs/recipes/llm/grpo-qwen3-30ba3b-8n8g-megatron-cp2-r3.yaml b/examples/configs/recipes/llm/grpo-qwen3-30ba3b-8n8g-megatron-cp2-r3.yaml index ec04203ebb8..e6a87889631 100644 --- a/examples/configs/recipes/llm/grpo-qwen3-30ba3b-8n8g-megatron-cp2-r3.yaml +++ b/examples/configs/recipes/llm/grpo-qwen3-30ba3b-8n8g-megatron-cp2-r3.yaml @@ -29,8 +29,6 @@ policy: enable_chunked_prefill: false data: max_input_seq_length: 2048 -data_plane: - enabled: false logger: log_dir: logs/grpo-qwen3-30ba3b-8n8g-megatron-cp2-r3 wandb: diff --git a/examples/configs/recipes/llm/grpo-qwen3-30ba3b-thinking-swe1-16n8g-megatron-cp2-r3-async-gym.yaml b/examples/configs/recipes/llm/grpo-qwen3-30ba3b-thinking-swe1-16n8g-megatron-cp2-r3-async-gym.yaml index 877172a2e28..514873446a7 100644 --- a/examples/configs/recipes/llm/grpo-qwen3-30ba3b-thinking-swe1-16n8g-megatron-cp2-r3-async-gym.yaml +++ b/examples/configs/recipes/llm/grpo-qwen3-30ba3b-thinking-swe1-16n8g-megatron-cp2-r3-async-gym.yaml @@ -1,24 +1,14 @@ defaults: ../../../nemo_gym/grpo_qwen3_30ba3b_thinking_swe1.yaml - grpo: num_prompts_per_step: 16 - num_generations_per_prompt: 8 max_val_samples: 32 val_batch_size: 32 - skip_reference_policy_logprobs_calculation: true - async_grpo: - enabled: true - max_trajectory_age_steps: 1 - in_flight_weight_updates: true -loss_fn: - use_importance_sampling_correction: true checkpointing: checkpoint_dir: results/grpo-qwen3-30ba3b-thinking-swe1-16n8g-megatron-cp2-r3-async-gym enabled: false data_plane: null policy: train_global_batch_size: 128 - logprob_batch_size: 1 router_replay: enabled: true generation: @@ -30,6 +20,3 @@ logger: log_dir: logs/grpo-qwen3-30ba3b-thinking-swe1-16n8g-megatron-cp2-r3-async-gym wandb: name: grpo-qwen3-30ba3b-thinking-swe1-16n8g-megatron-cp2-r3-async-gym -cluster: - num_nodes: 16 - gpus_per_node: 8 From 58b6af1eba7cf0d934f95ffe7af4f675eb86ad1b Mon Sep 17 00:00:00 2001 From: Yuki Huang Date: Tue, 30 Jun 2026 22:29:26 -0700 Subject: [PATCH 12/12] increase nightly 3000 -> 3270 (current 3232) Signed-off-by: Yuki Huang --- tests/unit/test_recipes_and_test_suites.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/tests/unit/test_recipes_and_test_suites.py b/tests/unit/test_recipes_and_test_suites.py index 93af891db64..174bbc8e02e 100644 --- a/tests/unit/test_recipes_and_test_suites.py +++ b/tests/unit/test_recipes_and_test_suites.py @@ -235,7 +235,7 @@ def test_all_recipe_yamls_accounted_for_in_test_suites( ) -def test_nightly_compute_stays_below_3000_hours(nightly_test_suite, tracker): +def test_nightly_compute_stays_below_3270_hours(nightly_test_suite, tracker): command = f"DRYRUN=1 HF_HOME=... HF_DATASETS_CACHE=... CONTAINER= ACCOUNT= PARTITION= ./tools/launch {' '.join(nightly_test_suite)}" print(f"Running command: {command}") @@ -267,8 +267,8 @@ def test_nightly_compute_stays_below_3000_hours(nightly_test_suite, tracker): f"Last line of output was not as expected: '{last_line}'" ) total_gpu_hours = float(last_line.split(":")[-1].strip()) - assert total_gpu_hours <= 3000, ( - f"Total GPU hours exceeded 3000: {last_line}. We should revisit the test suites to reduce the total GPU hours." + assert total_gpu_hours <= 3270, ( + f"Total GPU hours exceeded 3270: {last_line}. We should revisit the test suites to reduce the total GPU hours." ) tracker.track("total_nightly_gpu_hours", total_gpu_hours)