From 80a84eb99d31c73aa659973ba06814a2d73ef8ce Mon Sep 17 00:00:00 2001 From: Teodor-Dumitru Ene Date: Wed, 5 Aug 2026 17:43:36 -0500 Subject: [PATCH 01/10] perf(megatron): skip redundant weight load (#3397) Signed-off-by: Teodor-Dumitru Ene --- nemo_rl/algorithms/grpo.py | 1 + .../megatron/megatron_generation.py | 6 +++ nemo_rl/models/megatron/setup.py | 5 +- nemo_rl/models/policy/lm_policy.py | 3 ++ .../policy/workers/megatron_policy_worker.py | 6 +++ .../generation/test_megatron_generation.py | 51 +++++++++++++++++-- 6 files changed, 67 insertions(+), 5 deletions(-) diff --git a/nemo_rl/algorithms/grpo.py b/nemo_rl/algorithms/grpo.py index 4159e5d5b60..a646d7357b7 100644 --- a/nemo_rl/algorithms/grpo.py +++ b/nemo_rl/algorithms/grpo.py @@ -1047,6 +1047,7 @@ def init_megatron_generation(policy=None): tokenizer=tokenizer, processor=processor, weights_path=weights_path, + skip_weight_load=True, ) return mg, time.perf_counter() - t0 diff --git a/nemo_rl/models/generation/megatron/megatron_generation.py b/nemo_rl/models/generation/megatron/megatron_generation.py index 5565a6b3461..707f2230c6c 100644 --- a/nemo_rl/models/generation/megatron/megatron_generation.py +++ b/nemo_rl/models/generation/megatron/megatron_generation.py @@ -67,6 +67,7 @@ def __init__( name_prefix: str = "megatron_generation", processor: Optional[AutoProcessor] = None, weights_path: Optional[str] = None, + skip_weight_load: bool = False, ): """Initialize a MegatronGeneration instance. @@ -80,6 +81,7 @@ def __init__( name_prefix: Prefix for naming the worker group (non-colocated only). processor: Optional processor for VLMs (non-colocated only). weights_path: Optional path to model weights (non-colocated only). + skip_weight_load: Do not load the weights from the checkpoint; refit will do it. """ # Import here to avoid circular imports from nemo_rl.models.policy.lm_policy import Policy @@ -87,6 +89,9 @@ def __init__( assert (cluster is None) != (policy is None), ( "Provide exactly one of `cluster` or `policy`." ) + assert not (skip_weight_load and policy is not None), ( + "skip_weight_load only applies to the dedicated inference policy." + ) # `self.cfg` exposes the `generation` that matches the `GenerationInterface` contract. # `self._policy_config` keeps a reference to the full PolicyConfig. @@ -120,6 +125,7 @@ def __init__( init_optimizer=False, init_reference_model=False, weights_path=weights_path, + skip_weight_load=skip_weight_load, ) # Start the persistent inference engine + HTTP server during construction. diff --git a/nemo_rl/models/megatron/setup.py b/nemo_rl/models/megatron/setup.py index 4b888cd9ea5..2776acffbef 100644 --- a/nemo_rl/models/megatron/setup.py +++ b/nemo_rl/models/megatron/setup.py @@ -1314,6 +1314,7 @@ def setup_model_and_optimizer( get_position_embedding_ranks=None, pre_load_checkpoint_hook: Optional[Callable] = None, additional_pre_wrap_hooks: Optional[list[Callable]] = None, + load_weights: bool = True, ): state = GlobalState() _patch_bridge_signal_handler_for_worker_threads() @@ -1497,7 +1498,9 @@ def composed_peft_hook(model: list[MegatronModule]) -> list[MegatronModule]: print("Model, optimizer, and learning rate scheduler built") torch.distributed.barrier() - if megatron_cfg.peft is not None: + if not load_weights: + should_load_checkpoint = False + elif megatron_cfg.peft is not None: should_load_checkpoint = resume_checkpoint_exists if should_load_checkpoint: # The finetune toggle is explicitly set to True in order to avoid loading optimizer and RNG states diff --git a/nemo_rl/models/policy/lm_policy.py b/nemo_rl/models/policy/lm_policy.py index cdfa144b23c..6a184ded3e9 100644 --- a/nemo_rl/models/policy/lm_policy.py +++ b/nemo_rl/models/policy/lm_policy.py @@ -97,6 +97,7 @@ def __init__( init_reference_model: bool = True, processor: Optional[AutoProcessor] = None, worker_extension_cls_fqn: Optional[str] = None, + skip_weight_load: bool = False, ): if weights_path: weights_path = os.path.abspath(weights_path) @@ -258,6 +259,8 @@ def __init__( worker_sharding_annotations=self.sharding_annotations, pre_init_communication_queue=pre_init_queue, ) + if skip_weight_load: + worker_kwargs["skip_weight_load"] = True if use_v2: # DTensor v2 workers reconstruct tokenizer/processor locally to avoid diff --git a/nemo_rl/models/policy/workers/megatron_policy_worker.py b/nemo_rl/models/policy/workers/megatron_policy_worker.py index 52fefad7cdf..6a01a4cdbe1 100644 --- a/nemo_rl/models/policy/workers/megatron_policy_worker.py +++ b/nemo_rl/models/policy/workers/megatron_policy_worker.py @@ -368,6 +368,7 @@ def __init__( init_reference_model: bool = True, *, worker_sharding_annotations: NamedSharding, + skip_weight_load: bool = False, **kwargs: Any, ): """Initialize the MegatronPolicyWorker.""" @@ -489,11 +490,16 @@ def __init__( self.megatron_cfg.validate() # Step 4: Setup Megatron model and components + assert not (skip_weight_load and (init_optimizer or init_reference_model)), ( + "skip_weight_load is only valid for inference-only policies " + "(init_optimizer=False, init_reference_model=False)." + ) model_and_optimizer_state = setup_model_and_optimizer( config, self.megatron_cfg, init_optimizer, pre_load_checkpoint_hook=getattr(self, "_pre_load_checkpoint_hook", None), + load_weights=not skip_weight_load, ) self.mcore_state = model_and_optimizer_state.state diff --git a/tests/unit/models/generation/test_megatron_generation.py b/tests/unit/models/generation/test_megatron_generation.py index 787a051fa95..873e4d419b3 100644 --- a/tests/unit/models/generation/test_megatron_generation.py +++ b/tests/unit/models/generation/test_megatron_generation.py @@ -430,10 +430,15 @@ def test_megatron_generation_colocated(cluster, test_input_data, tokenizer): @pytest.mark.mcore @pytest.mark.timeout(900) +@pytest.mark.parametrize("skip_weight_load", [False, True]) def test_megatron_generation_non_colocated_refit( - policy_cluster_separate, test_input_data, tokenizer + policy_cluster_separate, test_input_data, tokenizer, skip_weight_load ): - """Non-colocated Megatron generation.""" + """Non-colocated Megatron generation. + + With skip_weight_load the inference engine builds without loading the + checkpoint and must still generate correctly once refit delivers weights. + """ generation_cluster = RayVirtualCluster( bundle_ct_per_node_list=[1], use_gpus=True, @@ -455,8 +460,22 @@ def test_megatron_generation_non_colocated_refit( policy = Policy( cluster=policy_cluster_separate, config=config, tokenizer=tokenizer ) + + # construction guard: skip_weight_load requires a dedicated inference + # policy; wrapping an existing (colocated) policy must be rejected. + with pytest.raises(AssertionError): + MegatronGeneration( + config=config, + tokenizer=tokenizer, + policy=policy, + skip_weight_load=True, + ) + mg = MegatronGeneration( - config=config, tokenizer=tokenizer, cluster=generation_cluster + config=config, + tokenizer=tokenizer, + cluster=generation_cluster, + skip_weight_load=skip_weight_load, ) # init the refit collective on both sides. @@ -478,12 +497,36 @@ def test_megatron_generation_non_colocated_refit( # refit the inference engine from the training weights, then generate refit_policy_generation(policy, mg, False) - outputs = mg.generate(test_input_data, greedy=True) + # Greedy needs to be false because processed logprobs doesn't handle it well. + outputs = mg.generate(test_input_data, greedy=False) _assert_valid_generation_output(outputs, test_input_data) generated_texts = tokenizer.batch_decode( outputs["output_ids"], skip_special_tokens=True ) assert all(len(t) > 0 for t in generated_texts), "Some texts are empty" + + # Training-policy logprobs must match generation-policy logprobs. + # A broken refit would fail this test. + fprop_data = BatchedDataDict( + { + "input_ids": outputs["output_ids"], + "input_lengths": outputs["unpadded_sequence_lengths"], + } + ) + policy.prepare_for_lp_inference() + train_logprobs = policy.get_logprobs(fprop_data)["logprobs"] + gen_mask = torch.zeros_like(outputs["logprobs"], dtype=torch.bool) + for i, (start, end) in enumerate( + zip(test_input_data["input_lengths"], outputs["unpadded_sequence_lengths"]) + ): + gen_mask[i, start:end] = True + abs_diff = (outputs["logprobs"] - train_logprobs).abs().masked_select(gen_mask) + avg_prob_mult_error = torch.exp(abs_diff).mean() + assert avg_prob_mult_error <= 1.05, ( + f"generation logprobs diverge from training-policy logprobs " + f"(avg prob mult error {avg_prob_mult_error:.4f}); inference weights " + f"do not match training weights after refit" + ) finally: if mg is not None: mg.shutdown() From 10e8b4cc002c58d52f5200cfb75a0b01f508a0ac Mon Sep 17 00:00:00 2001 From: alexchiu Date: Thu, 6 Aug 2026 09:06:21 +0800 Subject: [PATCH 02/10] fix(nemo-gym): refresh final-token route across turns (#3485) Signed-off-by: alexchiu <7390474+zpqiu@users.noreply.github.com> Co-authored-by: alexchiu <7390474+zpqiu@users.noreply.github.com> --- nemo_rl/environments/nemo_gym.py | 7 +++++++ .../environments/test_nemo_gym_router_replay.py | 16 ++++++++++------ 2 files changed, 17 insertions(+), 6 deletions(-) diff --git a/nemo_rl/environments/nemo_gym.py b/nemo_rl/environments/nemo_gym.py index 30fb04954b2..1b3daba4612 100644 --- a/nemo_rl/environments/nemo_gym.py +++ b/nemo_rl/environments/nemo_gym.py @@ -417,6 +417,13 @@ def _postprocess_nemo_gym_to_nemo_rl_result( "with enable_return_routed_experts." ) + # The next prompt prefill supplies the real route for the previous + # turn's final token, whose decode route was padded. + if routed_experts is not None and seen_token_ids: + previous_routes = nemo_rl_message_log[-1].get("routed_experts") + if isinstance(previous_routes, torch.Tensor): + previous_routes[-1] = routed_experts[len(seen_token_ids) - 1] + prompt_start = len(seen_token_ids) prompt_end = len(prompt_token_ids) generation_start = prompt_end diff --git a/tests/unit/environments/test_nemo_gym_router_replay.py b/tests/unit/environments/test_nemo_gym_router_replay.py index d1171856213..15383d8af5b 100644 --- a/tests/unit/environments/test_nemo_gym_router_replay.py +++ b/tests/unit/environments/test_nemo_gym_router_replay.py @@ -27,6 +27,10 @@ def _routes(num_tokens: int) -> list[list[list[int]]]: def test_nemo_gym_postprocess_slices_routed_experts(): + first_turn_routes = _routes(3) + first_turn_routes[-1] = [[0, 1]] + second_turn_routes = _routes(7) + second_turn_routes[2] = [[30, 31]] nemo_gym_result = { "response": { "output": [ @@ -34,13 +38,13 @@ def test_nemo_gym_postprocess_slices_routed_experts(): "prompt_token_ids": [1, 2], "generation_token_ids": [3], "generation_log_probs": [-0.1], - "routed_experts": _routes(3), + "routed_experts": first_turn_routes, }, { "prompt_token_ids": [1, 2, 3, 4, 5], "generation_token_ids": [6, 7], "generation_log_probs": [-0.2, -0.3], - "routed_experts": _routes(7), + "routed_experts": second_turn_routes, }, ] }, @@ -58,13 +62,13 @@ class _MockSelf: 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[0]["routed_experts"].tolist() == first_turn_routes[:2] assert message_log[1]["token_ids"].tolist() == [3] - assert message_log[1]["routed_experts"].tolist() == _routes(3)[2:3] + assert message_log[1]["routed_experts"].tolist() == second_turn_routes[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[2]["routed_experts"].tolist() == second_turn_routes[3:5] assert message_log[3]["token_ids"].tolist() == [6, 7] - assert message_log[3]["routed_experts"].tolist() == _routes(7)[5:7] + assert message_log[3]["routed_experts"].tolist() == second_turn_routes[5:7] def test_nemo_gym_postprocess_requires_routed_experts_when_configured(): From ae07eafe8035b5b2e84efa7234e70e7fd7e493c1 Mon Sep 17 00:00:00 2001 From: Yi-Fu Wu Date: Wed, 5 Aug 2026 21:42:04 -0700 Subject: [PATCH 03/10] feat(slurm): support external vLLM services in one allocation (#3386) Signed-off-by: Yi-Fu Wu Co-authored-by: Jiaqi Zeng Co-authored-by: Gerald Shen --- docs/guides/nemotron-3-ultra.md | 17 +- pyrefly.toml | 2 + tests/unit/tools/test_external_gym_vllm.py | 833 ++++++++++++++++++ tools/external_gym_vllm/README.md | 203 +++++ tools/external_gym_vllm/lb_watchdog.sh | 62 ++ tools/external_gym_vllm/pool_config.sh | 300 +++++++ tools/external_gym_vllm/run_in_allocation.sh | 702 +++++++++++++++ tools/external_gym_vllm/serve_vllm_on_ray.py | 41 + .../vllm_backend_registry.sh | 64 ++ tools/external_gym_vllm/vllm_pool_lb.py | 592 +++++++++++++ 10 files changed, 2809 insertions(+), 7 deletions(-) create mode 100644 tests/unit/tools/test_external_gym_vllm.py create mode 100644 tools/external_gym_vllm/README.md create mode 100755 tools/external_gym_vllm/lb_watchdog.sh create mode 100755 tools/external_gym_vllm/pool_config.sh create mode 100755 tools/external_gym_vllm/run_in_allocation.sh create mode 100755 tools/external_gym_vllm/serve_vllm_on_ray.py create mode 100755 tools/external_gym_vllm/vllm_backend_registry.sh create mode 100755 tools/external_gym_vllm/vllm_pool_lb.py diff --git a/docs/guides/nemotron-3-ultra.md b/docs/guides/nemotron-3-ultra.md index 11f4c432cce..8d401348e2e 100644 --- a/docs/guides/nemotron-3-ultra.md +++ b/docs/guides/nemotron-3-ultra.md @@ -203,13 +203,16 @@ Set the following before each `bash examples/nemo_gym/nemotron-3-ultra/ultra_lau | `NL2BASH_JUDGE_MODEL` | NL2Bash / general-purpose judge: HF repo id or local path. Default judge is `Qwen/Qwen3-235B-A22B-Instruct-2507-FP8`. | | `SAFETY_JUDGE_MODEL` | Content-safety judge: HF repo id or local path. Default is [`nvidia/Nemotron-Content-Safety-Reasoning-4B`](https://huggingface.co/nvidia/Nemotron-Content-Safety-Reasoning-4B). | -> **Serving GenRM as a standalone service.** The GenRM judge does not have to run -> inside the training job. You can bring it up separately — any OpenAI-compatible -> vLLM endpoint, or the external-GenRM service launcher under `tools/external_genrm/` -> (runs the judge fleet in a second Slurm hetgroup behind a load balancer) — and -> point the run at it with `GENRM_BASE_URL=http://:/v1`. Judging is then -> routed to that endpoint instead of being served from the gym pool, which frees -> those GPUs and lets one GenRM deployment back many training runs. +> **Serving GenRM outside Gym.** For a separately deployed OpenAI-compatible +> endpoint, set `GENRM_BASE_URL=http://:/v1`. Judging is then routed +> to that endpoint instead of being served from the Gym pool, allowing one +> deployment to back multiple training runs. +> +> To co-schedule dedicated model servers with training in one Slurm +> heterogeneous allocation, use the +> [external Gym vLLM pool helpers](https://github.com/NVIDIA-NeMo/RL/blob/main/tools/external_gym_vllm/README.md). +> They place the server fleet in a second hetgroup, start load balancers on the +> training component, and inject the resolved endpoints into the driver command. Optional knobs: diff --git a/pyrefly.toml b/pyrefly.toml index 27b1999d900..5f55643e09b 100644 --- a/pyrefly.toml +++ b/pyrefly.toml @@ -239,6 +239,8 @@ project-includes = [ "nemo_rl/weight_sync/nccl_reshard_weight_synchronizer.py", "nemo_rl/weight_sync/xferdtensor.py", "nemo_rl/weight_sync/xferdtensor_python.py", + "tools/external_gym_vllm/vllm_pool_lb.py", + "tools/external_gym_vllm/serve_vllm_on_ray.py", "tools/model_diagnostics/1.max_model_len_respected.py", "tools/model_diagnostics/2.long_generation_decode_vs_prefill.py", "tools/model_diagnostics/3.check_and_reinit_hf_model_embeddings_untrained.py", diff --git a/tests/unit/tools/test_external_gym_vllm.py b/tests/unit/tools/test_external_gym_vllm.py new file mode 100644 index 00000000000..96fb891d4c5 --- /dev/null +++ b/tests/unit/tools/test_external_gym_vllm.py @@ -0,0 +1,833 @@ +# 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 json +import os +import signal +import subprocess +import textwrap +import time +from pathlib import Path +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest +from aiohttp import ClientPayloadError, web + +from tools.external_gym_vllm.vllm_pool_lb import ( + SHUTDOWN_TIMEOUT_SECONDS, + Backend, + BackendPool, + LoadBalancer, + UpstreamRetryableStatus, + _read_current_rss_mb, +) + +REPO_ROOT = Path(__file__).parents[3] + + +def test_shutdown_timeout_bounds_watchdog_restart_outage(): + assert 0 < SHUTDOWN_TIMEOUT_SECONDS <= 120 + + +def test_read_current_rss_uses_vmrss_instead_of_process_high_water_mark(): + status = textwrap.dedent( + """\ + Name: python + VmHWM: 8388608 kB + VmRSS: 315392 kB + """ + ) + + with patch( + "tools.external_gym_vllm.vllm_pool_lb.Path.read_text", + return_value=status, + ): + assert _read_current_rss_mb() == 308 + + +@pytest.mark.asyncio +async def test_memory_watchdog_requests_graceful_shutdown(): + pool = BackendPool("/tmp", "test") + pool._running = True + + with ( + patch( + "tools.external_gym_vllm.vllm_pool_lb._read_current_rss_mb", + return_value=4097, + ), + patch("tools.external_gym_vllm.vllm_pool_lb.os.kill") as kill, + ): + await pool._health_check_loop() + + assert pool._running is False + kill.assert_called_once_with(os.getpid(), signal.SIGTERM) + + +def test_backend_pool_reads_only_ready_registry_entries(tmp_path): + registry = tmp_path / ".registry_test" + registry.write_text( + "\n".join( + [ + "ready-backend 10.0.0.1 8000 123 ready", + "starting-backend 10.0.0.2 8001 124 starting", + "malformed", + ] + ) + ) + + pool = BackendPool(str(tmp_path), "test") + + assert pool._read_registry() == {"ready-backend": ("10.0.0.1", 8000)} + + +def test_read_registry_skips_bad_line_without_dropping_later_entries(tmp_path): + registry = tmp_path / ".registry_test" + registry.write_text( + "\n".join( + [ + "good-1 10.0.0.1 8000 123 ready", + "bad-port 10.0.0.2 not-a-port 124 ready", + "good-2 10.0.0.3 8002 125 ready", + ] + ) + ) + + pool = BackendPool(str(tmp_path), "test") + + assert pool._read_registry() == { + "good-1": ("10.0.0.1", 8000), + "good-2": ("10.0.0.3", 8002), + } + + +def test_backend_pool_picks_least_loaded_healthy_backend(): + pool = BackendPool("/tmp", "test") + first = Backend("first", "10.0.0.1", 8000) + second = Backend("second", "10.0.0.2", 8000) + first.inflight = 4 + second.inflight = 1 + pool.backends = {first.job_id: first, second.job_id: second} + + assert pool.pick() is second + assert pool.pick(exclude={"second"}) is first + + first.healthy = False + assert pool.pick(exclude={"second"}) is None + + +def test_affinity_key_is_stable_and_ignores_invalid_json(): + body = json.dumps({"messages": [{"role": "user", "content": "prompt"}]}).encode() + + assert LoadBalancer._extract_affinity_key(body) == ( + LoadBalancer._extract_affinity_key(body) + ) + assert LoadBalancer._extract_affinity_key(b"not-json") is None + + +def test_extract_affinity_key_handles_json_that_is_not_an_object(): + assert LoadBalancer._extract_affinity_key(b"[1, 2]") is None + assert LoadBalancer._extract_affinity_key(b"null") is None + assert LoadBalancer._extract_affinity_key(b"123") is None + + +def test_pick_prefers_affinity_backend_until_it_becomes_a_hotspot(): + pool = BackendPool("/tmp", "test") + first = Backend("first", "10.0.0.1", 8000) + second = Backend("second", "10.0.0.2", 8000) + pool.backends = {first.job_id: first, second.job_id: second} + affinity_key = LoadBalancer._extract_affinity_key( + json.dumps({"messages": [{"role": "user", "content": "prompt"}]}).encode() + ) + + preferred = pool.pick(affinity_key=affinity_key) + assert pool.pick(affinity_key=affinity_key) is preferred + + other = next( + backend for backend in pool.backends.values() if backend is not preferred + ) + preferred.inflight = 2 * other.inflight + 11 + assert pool.pick(affinity_key=affinity_key) is other + + +@pytest.mark.asyncio +async def test_proxy_retries_a_5xx_on_another_backend(): + pool = BackendPool("/tmp", "test") + first = Backend("first", "10.0.0.1", 8000) + second = Backend("second", "10.0.0.2", 8000) + pool.backends = {first.job_id: first, second.job_id: second} + load_balancer = LoadBalancer(pool, 9213) + + expected_response = web.Response(status=200, body=b"ok") + load_balancer._proxy_once = AsyncMock( + side_effect=[ + UpstreamRetryableStatus(500, b"engine failed", {}), + expected_response, + ] + ) + request = MagicMock(spec=web.Request) + request.read = AsyncMock(return_value=b"{}") + request.method = "POST" + request.path_qs = "/v1/chat/completions" + request.headers = {} + + response = await load_balancer.handle_proxy(request) + + assert response is expected_response + assert load_balancer._proxy_once.await_count == 2 + assert first.healthy is True + assert second.healthy is True + + +@pytest.mark.asyncio +async def test_stream_failure_after_prepare_does_not_escape_for_retry(): + class FailingStreamContent: + async def _iterate(self): + yield b"first chunk" + raise ClientPayloadError("upstream disconnected") + + def iter_any(self): + return self._iterate() + + backend = Backend("first", "10.0.0.1", 8000) + load_balancer = LoadBalancer(BackendPool("/tmp", "test"), 9213) + upstream_response = MagicMock() + upstream_response.status = 200 + upstream_response.headers = {"Content-Type": "text/event-stream"} + upstream_response.content = FailingStreamContent() + request_context = MagicMock() + request_context.__aenter__ = AsyncMock(return_value=upstream_response) + request_context.__aexit__ = AsyncMock(return_value=None) + proxy_session = MagicMock() + proxy_session.request.return_value = request_context + load_balancer._proxy_session = proxy_session + + stream_response = MagicMock(spec=web.StreamResponse) + stream_response.prepare = AsyncMock() + stream_response.write = AsyncMock() + stream_response.write_eof = AsyncMock() + request = MagicMock(spec=web.Request) + + with patch( + "tools.external_gym_vllm.vllm_pool_lb.web.StreamResponse", + return_value=stream_response, + ): + result = await load_balancer._proxy_once( + backend, + "POST", + "/v1/responses", + {}, + b"{}", + request, + ) + + assert result is stream_response + stream_response.prepare.assert_awaited_once_with(request) + stream_response.write.assert_awaited_once_with(b"first chunk") + stream_response.write_eof.assert_awaited_once() + assert backend.healthy is False + assert backend.inflight == 0 + + +@pytest.mark.asyncio +async def test_proxy_drops_stale_length_after_upstream_decompression(): + backend = Backend("first", "10.0.0.1", 8000) + load_balancer = LoadBalancer(BackendPool("/tmp", "test"), 9213) + upstream_response = MagicMock() + upstream_response.status = 200 + upstream_response.headers = { + "Content-Type": "application/json", + "Content-Encoding": "gzip", + "Content-Length": "3", + "X-Request-Id": "request-1", + } + upstream_response.read = AsyncMock(return_value=b"decompressed") + request_context = MagicMock() + request_context.__aenter__ = AsyncMock(return_value=upstream_response) + request_context.__aexit__ = AsyncMock(return_value=None) + proxy_session = MagicMock() + proxy_session.request.return_value = request_context + load_balancer._proxy_session = proxy_session + request = MagicMock(spec=web.Request) + expected_response = MagicMock(spec=web.Response) + + with patch( + "tools.external_gym_vllm.vllm_pool_lb.web.Response", + return_value=expected_response, + ) as response_class: + result = await load_balancer._proxy_once( + backend, + "POST", + "/v1/responses", + {}, + b"{}", + request, + ) + + assert result is expected_response + assert response_class.call_args.kwargs["headers"] == { + "Content-Type": "application/json", + "X-Request-Id": "request-1", + } + assert backend.inflight == 0 + + +@pytest.mark.asyncio +async def test_proxy_forwards_last_upstream_5xx_after_exhausting_backends(): + pool = BackendPool("/tmp", "test") + first = Backend("first", "10.0.0.1", 8000) + second = Backend("second", "10.0.0.2", 8000) + pool.backends = {first.job_id: first, second.job_id: second} + load_balancer = LoadBalancer(pool, 9213) + load_balancer._proxy_once = AsyncMock( + side_effect=UpstreamRetryableStatus(503, b"engine dead", {"X-Request-Id": "1"}) + ) + request = MagicMock(spec=web.Request) + request.read = AsyncMock(return_value=b"{}") + request.method = "POST" + request.path_qs = "/v1/chat/completions" + request.headers = {} + + response = await load_balancer.handle_proxy(request) + + assert response.status == 503 + assert response.body == b"engine dead" + assert response.headers["X-Request-Id"] == "1" + assert load_balancer._proxy_once.await_count == 2 + assert first.healthy and second.healthy + + +def test_load_balancer_accepts_payloads_larger_than_aiohttp_default(): + app = LoadBalancer(BackendPool("/tmp", "test"), 9213).make_app() + + assert app._client_max_size == 0 + + +@pytest.mark.asyncio +async def test_proxy_returns_503_when_no_backend_is_available(): + load_balancer = LoadBalancer(BackendPool("/tmp", "test"), 9213) + request = MagicMock(spec=web.Request) + request.read = AsyncMock(return_value=b"{}") + request.method = "POST" + request.path_qs = "/v1/chat/completions" + request.headers = {} + + response = await load_balancer.handle_proxy(request) + + assert response.status == 503 + + +@pytest.mark.asyncio +async def test_health_reports_backend_counts(): + pool = BackendPool("/tmp", "test") + healthy = Backend("healthy", "10.0.0.1", 8000) + sick = Backend("sick", "10.0.0.2", 8000) + sick.healthy = False + pool.backends = {healthy.job_id: healthy, sick.job_id: sick} + + response = await LoadBalancer(pool, 9213).handle_health(MagicMock(spec=web.Request)) + + assert isinstance(response.body, bytes) + payload = json.loads(response.body) + assert payload["status"] == "ok" + assert payload["healthy_backends"] == 1 + assert payload["total_backends"] == 2 + + +def test_registry_shell_helpers_add_replace_remove(tmp_path): + script = REPO_ROOT / "tools/external_gym_vllm/vllm_backend_registry.sh" + program = textwrap.dedent( + f""" + set -euo pipefail + export EXTERNAL_VLLM_STATE_DIR={tmp_path} + export EXTERNAL_VLLM_GROUP_ID=test + source {script} + registry_add job-a 10.0.0.1 8000 + registry_add job-b 10.0.0.2 8001 + echo "count=$(registry_count_ready)" + registry_add job-a 10.0.0.9 8009 + echo "count=$(registry_count_ready)" + echo "ready=$(registry_list_ready | tr '\\n' ',')" + registry_remove job-b + echo "count=$(registry_count_ready)" + """ + ) + result = subprocess.run( + ["bash", "-c", program], + capture_output=True, + text=True, + check=True, + ) + + assert result.stdout.splitlines() == [ + "count=2", + "count=2", + "ready=10.0.0.2:8001,10.0.0.9:8009,", + "count=1", + ] + + +def test_load_balancer_watchdog_forwards_term_to_child(tmp_path): + watchdog = REPO_ROOT / "tools/external_gym_vllm/lb_watchdog.sh" + fake_python = tmp_path / "fake-python" + child_started = tmp_path / "child-started" + child_stopped = tmp_path / "child-stopped" + fake_python.write_text( + textwrap.dedent( + f"""\ + #!/bin/bash + touch {child_started} + trap 'touch {child_stopped}; exit 0' TERM INT + while true; do sleep 0.1; done + """ + ) + ) + fake_python.chmod(0o755) + process = subprocess.Popen( + ["bash", str(watchdog), "9213", str(tmp_path), "test"], + env={"PATH": os.environ["PATH"], "PYTHON": str(fake_python)}, + start_new_session=True, + ) + + try: + deadline = time.monotonic() + 5 + while not child_started.exists() and time.monotonic() < deadline: + time.sleep(0.05) + assert child_started.exists() + + process.send_signal(signal.SIGTERM) + assert process.wait(timeout=5) == 0 + assert child_stopped.exists() + finally: + if process.poll() is None: + os.killpg(process.pid, signal.SIGKILL) + process.wait(timeout=5) + + +def test_launcher_requires_a_heterogeneous_job(): + script = REPO_ROOT / "tools/external_gym_vllm/run_in_allocation.sh" + + result = subprocess.run( + ["bash", str(script)], + env={"PATH": os.environ["PATH"], "SLURM_JOB_ID": "123"}, + capture_output=True, + text=True, + ) + + assert result.returncode != 0 + assert "This script requires a Slurm heterogeneous job" in result.stderr + + +def test_launcher_rejects_more_than_two_hetgroups(): + script = REPO_ROOT / "tools/external_gym_vllm/run_in_allocation.sh" + env = { + "PATH": os.environ["PATH"], + "SLURM_JOB_ID": "123", + "SLURM_HET_SIZE": "3", + "SLURM_JOB_NODELIST_HET_GROUP_0": "ray[01-02]", + "SLURM_JOB_NODELIST_HET_GROUP_1": "genrm[01-02]", + "SLURM_JOB_ACCOUNT": "account", + "SLURM_JOB_PARTITION": "partition", + "SLURM_SUBMIT_DIR": "/tmp", + "BASE_LOG_DIR": "/lustre/logs", + "CONTAINER": "training.sqsh", + "MOUNTS": "/lustre:/lustre", + "COMMAND": "run __GENRM_BASE_URL__", + "EXTERNAL_VLLM_POOLS": "GENRM", + "EXTERNAL_VLLM_TOOLS_DIR_HOST": "/lustre/tools", + } + + result = subprocess.run( + ["bash", str(script)], + env=env, + capture_output=True, + text=True, + ) + + assert result.returncode != 0 + assert "Expected exactly two Slurm hetgroups, got 3" in result.stderr + + +def test_launcher_requires_nl2bash_placeholder_when_pool_is_enabled(): + script = REPO_ROOT / "tools/external_gym_vllm/run_in_allocation.sh" + env = { + "PATH": os.environ["PATH"], + "SLURM_JOB_ID": "123", + "SLURM_HET_SIZE": "2", + "SLURM_JOB_NODELIST_HET_GROUP_0": "ray[01-02]", + "SLURM_JOB_NODELIST_HET_GROUP_1": "judge[01-02]", + "SLURM_JOB_ACCOUNT": "account", + "SLURM_JOB_PARTITION": "partition", + "SLURM_SUBMIT_DIR": str(REPO_ROOT), + "BASE_LOG_DIR": "/lustre/logs", + "CONTAINER": "training.sqsh", + "MOUNTS": "/lustre:/lustre", + "COMMAND": "run __GENRM_BASE_URL__", + "EXTERNAL_VLLM_POOLS": "GENRM NL2BASH", + "EXTERNAL_VLLM_TOOLS_DIR_HOST": str(REPO_ROOT / "tools/external_gym_vllm"), + "GENRM_CONTAINER": "genrm.sqsh", + "GENRM_MODEL": "model-id", + "GENRM_VLLM_PYTHON": "/opt/python", + "GENRM_REPLICAS": "1", + "GENRM_TENSOR_PARALLEL_SIZE": "4", + "GENRM_LB_PORT": "9213", + "GENRM_URL_PLACEHOLDER": "__GENRM_BASE_URL__", + "NL2BASH_CONTAINER": "judge.sqsh", + "NL2BASH_MODEL": "judge-model-id", + "NL2BASH_VLLM_PYTHON": "/opt/python", + "NL2BASH_REPLICAS": "4", + "NL2BASH_TENSOR_PARALLEL_SIZE": "4", + "NL2BASH_LB_PORT": "9214", + "NL2BASH_URL_PLACEHOLDER": "__NL2BASH_BASE_URL__", + "RAY_SUB": str(REPO_ROOT / "ray.sub"), + } + + result = subprocess.run( + ["bash", str(script)], + env=env, + capture_output=True, + text=True, + ) + + assert result.returncode != 0 + assert "Driver command is missing __NL2BASH_BASE_URL__" in result.stderr + + +def test_launcher_rejects_duplicate_url_placeholders(): + script = REPO_ROOT / "tools/external_gym_vllm/run_in_allocation.sh" + env = { + "PATH": os.environ["PATH"], + "SLURM_JOB_ID": "123", + "SLURM_HET_SIZE": "2", + "SLURM_JOB_NODELIST_HET_GROUP_0": "ray[01-02]", + "SLURM_JOB_NODELIST_HET_GROUP_1": "judge[01-02]", + "SLURM_JOB_ACCOUNT": "account", + "SLURM_JOB_PARTITION": "partition", + "SLURM_SUBMIT_DIR": str(REPO_ROOT), + "BASE_LOG_DIR": "/lustre/logs", + "CONTAINER": "training.sqsh", + "MOUNTS": "/lustre:/lustre", + "COMMAND": "run __SHARED_BASE_URL__", + "EXTERNAL_VLLM_POOLS": "GENRM NL2BASH", + "EXTERNAL_VLLM_TOOLS_DIR_HOST": str(REPO_ROOT / "tools/external_gym_vllm"), + "GENRM_CONTAINER": "genrm.sqsh", + "GENRM_MODEL": "model-id", + "GENRM_VLLM_PYTHON": "/opt/python", + "GENRM_REPLICAS": "1", + "GENRM_TENSOR_PARALLEL_SIZE": "4", + "GENRM_LB_PORT": "9213", + "GENRM_URL_PLACEHOLDER": "__SHARED_BASE_URL__", + "NL2BASH_CONTAINER": "judge.sqsh", + "NL2BASH_MODEL": "judge-model-id", + "NL2BASH_VLLM_PYTHON": "/opt/python", + "NL2BASH_REPLICAS": "4", + "NL2BASH_TENSOR_PARALLEL_SIZE": "4", + "NL2BASH_LB_PORT": "9214", + "NL2BASH_URL_PLACEHOLDER": "__SHARED_BASE_URL__", + "RAY_SUB": str(REPO_ROOT / "ray.sub"), + } + + result = subprocess.run( + ["bash", str(script)], + env=env, + capture_output=True, + text=True, + ) + + assert result.returncode != 0 + assert "Multiple pools use URL placeholder __SHARED_BASE_URL__" in result.stderr + + +def test_launcher_routes_generic_pools_to_explicit_hetgroups(): + script = REPO_ROOT / "tools/external_gym_vllm/run_in_allocation.sh" + source = script.read_text() + + srun_blocks = [] + current_block = [] + for line in source.splitlines(): + if line.lstrip() == "srun \\": + current_block = [line] + elif current_block: + current_block.append(line) + if line.rstrip().endswith("&"): + srun_blocks.append("\n".join(current_block)) + current_block = [] + + assert len(srun_blocks) == 4 + pool_preflight = next(block for block in srun_blocks if "import ray, vllm" in block) + lb_preflight = next(block for block in srun_blocks if "import aiohttp" in block) + replica_launch = next(block for block in srun_blocks if "VLLM_SERVER_BODY" in block) + lb_launch = next( + block + for block in srun_blocks + if "lb_watchdog.sh" in block and "--output=" in block + ) + + for block in (pool_preflight, replica_launch): + assert "--het-group=1" in block + assert '-A "${SLURM_JOB_ACCOUNT}"' not in block + assert '-p "${SLURM_JOB_PARTITION}"' not in block + for block in (lb_preflight, lb_launch): + assert "--het-group=0" in block + assert '-A "${SLURM_JOB_ACCOUNT}"' in block + assert '-p "${SLURM_JOB_PARTITION}"' in block + + assert 'preflight_labels+=("${display_names[${pool}]} container")' in source + assert 'preflight_labels+=("load balancer container")' in source + assert "${preflight_labels[${preflight_index}]}" in source + assert 'export "${pool}_ENV_VARS=$(pool_value "${pool}" ENV_VARS)"' in source + assert 'export "${pool}_VLLM_ARGS=$(pool_value "${pool}" VLLM_ARGS)"' in source + assert 'SLURM_JOB_NODELIST="${SLURM_JOB_NODELIST_HET_GROUP_0}"' in source + assert 'scontrol show hostnames "${SLURM_JOB_NODELIST_HET_GROUP_1}"' in source + assert 'for pool in "${pool_names[@]}"' in source + assert "POOL_PREFIX=${pool}" in source + assert ( + 'COMMAND="${COMMAND//${placeholders[${pool}]}/${pool_urls[${pool}]}}"' in source + ) + assert "genrm" not in source.lower() + assert "nl2bash" not in source.lower() + assert "safety" not in source.lower() + assert "RAY_NODELIST" not in source + assert "external-vllm-lb-preflight" not in source + assert "if ! ready=$(" in source + assert 'env \\\n SLURM_JOB_NODELIST="${SLURM_JOB_NODELIST_HET_GROUP_0}"' in source + assert 'if [[ -n "${SLURM_RESTART_COUNT:-}" ]]; then' in source + assert ( + 'LOG_DIR="${BASE_LOG_DIR}/${SLURM_JOB_ID}-${SLURM_RESTART_COUNT}-logs"' + in source + ) + assert 'rm -f "${pool_log_dirs[${pool}]}"/head_ip_*' in source + assert ( + 'echo "[${REPLICA_ID}] ERROR: vLLM exited with status ${vllm_status}"' in source + ) + assert "if (( vllm_status == 0 )); then" in source + + +def test_private_ray_and_vllm_ports_match_sub_ephemeral_layout(): + script = REPO_ROOT / "tools/external_gym_vllm/run_in_allocation.sh" + source = script.read_text() + + assert "RAY_PORT=1200" in source + assert "RAY_CLIENT_SERVER_PORT=1201" in source + assert "MIN_WORKER_PORT=2000" in source + assert "MAX_WORKER_PORT=2999" in source + assert "VLLM_ENGINE_PORT=7000" in source + assert source.count('--min-worker-port="${MIN_WORKER_PORT}"') == 2 + assert source.count('--max-worker-port="${MAX_WORKER_PORT}"') == 2 + assert 'export VLLM_PORT="${VLLM_ENGINE_PORT}"' in source + assert '--port "${VLLM_HTTP_PORT}"' in source + + +def test_pool_config_interface_registers_an_arbitrary_third_pool(): + script = REPO_ROOT / "tools/external_gym_vllm/pool_config.sh" + program = textwrap.dedent( + f""" + set -euo pipefail + source {script} + register_external_vllm_pool SAFETY \\ + --display-name "Safety judge" \\ + --model safety-model \\ + --container service.sqsh \\ + --python /opt/vllm/bin/python \\ + --replicas 2 \\ + --tensor-parallel-size 4 \\ + --lb-port 9215 \\ + --url-placeholder __SAFETY_BASE_URL__ \\ + --group-id safety-pool + external_vllm_pool_env SAFETY NCCL_MNNVL_ENABLE=0 + external_vllm_pool_args SAFETY \\ + --dtype bfloat16 \\ + --attention-backend FLASH_ATTN + printf 'pools=%s\n' "$EXTERNAL_VLLM_POOLS" + printf 'name=%s\n' "$SAFETY_DISPLAY_NAME" + printf 'env=%s\n' "$SAFETY_ENV_VARS" + printf 'args=%s\n' "$(tr '\n' ',' <<< "$SAFETY_VLLM_ARGS")" + printf 'group=%s\n' "$SAFETY_GROUP_ID" + printf 'nodes=%s\n' "$EXTERNAL_VLLM_NUM_NODES" + """ + ) + result = subprocess.run( + ["bash", "-c", program], + capture_output=True, + text=True, + check=True, + ) + + assert result.stdout.splitlines() == [ + "pools=SAFETY", + "name=Safety judge", + "env=NCCL_MNNVL_ENABLE=0", + "args=--dtype,bfloat16,--attention-backend,FLASH_ATTN,", + "group=safety-pool", + "nodes=2", + ] + + +@pytest.mark.parametrize( + ("option", "value", "expected_error"), + [ + ("--replicas", "-5", "TEST_REPLICAS must be a positive integer"), + ( + "--tensor-parallel-size", + "0", + "TEST_TENSOR_PARALLEL_SIZE must be a positive integer", + ), + ("--lb-port", "99999999", "TEST_LB_PORT must be at most 65535"), + ("--vllm-port", "0", "TEST_VLLM_PORT must be a positive integer"), + ( + "--startup-timeout", + "nope", + "TEST_STARTUP_TIMEOUT must be a positive integer", + ), + ], +) +def test_pool_registration_rejects_invalid_numeric_values( + option, value, expected_error +): + script = REPO_ROOT / "tools/external_gym_vllm/pool_config.sh" + program = textwrap.dedent( + f""" + source {script} + register_external_vllm_pool TEST \\ + --model model \\ + --container image \\ + --python /opt/python \\ + --replicas 1 \\ + --tensor-parallel-size 4 \\ + --lb-port 9213 \\ + --url-placeholder __TEST_URL__ \\ + {option} {value} + """ + ) + + result = subprocess.run(["bash", "-c", program], capture_output=True, text=True) + + assert result.returncode == 2 + assert expected_error in result.stderr + + +@pytest.mark.parametrize( + ("second_pool_args", "expected_error"), + [ + ("--lb-port 9213 --url-placeholder __SECOND_URL__", "use LB port 9213"), + ( + "--lb-port 9214 --url-placeholder __FIRST_URL__", + "use URL placeholder __FIRST_URL__", + ), + ], +) +def test_pool_registration_rejects_duplicate_routing_keys( + second_pool_args, expected_error +): + script = REPO_ROOT / "tools/external_gym_vllm/pool_config.sh" + program = textwrap.dedent( + f""" + source {script} + register_external_vllm_pool FIRST \\ + --model model --container image --python /opt/python \\ + --replicas 1 --tensor-parallel-size 4 \\ + --lb-port 9213 --url-placeholder __FIRST_URL__ + register_external_vllm_pool SECOND \\ + --model model --container image --python /opt/python \\ + --replicas 1 --tensor-parallel-size 4 {second_pool_args} + """ + ) + + result = subprocess.run(["bash", "-c", program], capture_output=True, text=True) + + assert result.returncode == 2 + assert expected_error in result.stderr + + +def test_pool_registration_rejects_partial_nodes_and_unsafe_group_id(): + script = REPO_ROOT / "tools/external_gym_vllm/pool_config.sh" + command = textwrap.dedent( + f""" + source {script} + register_external_vllm_pool TEST \\ + --model model --container image --python /opt/python \\ + --replicas 1 --tensor-parallel-size 2 \\ + --lb-port 9213 --url-placeholder __TEST_URL__ + """ + ) + unsafe_group_command = command.replace( + "--tensor-parallel-size 2", + "--tensor-parallel-size 4 --group-id bad/id", + ) + + partial = subprocess.run(["bash", "-c", command], capture_output=True, text=True) + unsafe_group = subprocess.run( + ["bash", "-c", unsafe_group_command], capture_output=True, text=True + ) + + assert partial.returncode == 2 + assert "must be divisible by GPUS_PER_NODE=4" in partial.stderr + assert unsafe_group.returncode == 2 + assert "TEST_GROUP_ID may contain only" in unsafe_group.stderr + + +def test_submission_validation_checks_placeholders_paths_and_node_total(): + script = REPO_ROOT / "tools/external_gym_vllm/pool_config.sh" + tools_dir = REPO_ROOT / "tools/external_gym_vllm" + program = textwrap.dedent( + f""" + set -euo pipefail + source {script} + EXTERNAL_VLLM_SHARED_ROOT={REPO_ROOT} + BASE_LOG_DIR={REPO_ROOT}/logs + EXTERNAL_VLLM_TOOLS_DIR_HOST={tools_dir} + register_external_vllm_pool TEST \\ + --model model --container image --python /opt/python \\ + --replicas 2 --tensor-parallel-size 4 \\ + --lb-port 9213 --url-placeholder __TEST_URL__ + validate_external_vllm_submission 'run __TEST_URL__' 2 + """ + ) + + valid = subprocess.run(["bash", "-c", program], capture_output=True, text=True) + wrong_nodes = subprocess.run( + ["bash", "-c", program.replace("'run __TEST_URL__' 2", "'run __TEST_URL__' 3")], + capture_output=True, + text=True, + ) + missing_placeholder = subprocess.run( + [ + "bash", + "-c", + program.replace("'run __TEST_URL__' 2", "'run without endpoint' 2"), + ], + capture_output=True, + text=True, + ) + missing_node_count = subprocess.run( + [ + "bash", + "-c", + program.replace( + "validate_external_vllm_submission 'run __TEST_URL__' 2", + "validate_external_vllm_submission 'run __TEST_URL__'", + ), + ], + capture_output=True, + text=True, + ) + + assert valid.returncode == 0, valid.stderr + assert wrong_nodes.returncode == 2 + assert "expected 2 from registered pools" in wrong_nodes.stderr + assert missing_placeholder.returncode == 2 + assert "submission command is missing __TEST_URL__" in missing_placeholder.stderr + assert missing_node_count.returncode == 0, missing_node_count.stderr + assert ( + "skipping external hetgroup node-count validation" in missing_node_count.stderr + ) diff --git a/tools/external_gym_vllm/README.md b/tools/external_gym_vllm/README.md new file mode 100644 index 00000000000..4a66207ad75 --- /dev/null +++ b/tools/external_gym_vllm/README.md @@ -0,0 +1,203 @@ +# Heterogeneous-job external Gym vLLM pools + +These helpers run arbitrary fixed-model vLLM pools beside NeMo RL in one +two-component Slurm heterogeneous job: + +```text +one Slurm heterogeneous job +├── hetgroup 0: NeMo RL Ray cluster and per-pool load balancers +└── hetgroup 1: independent private Ray/vLLM replicas +``` + +`run_in_allocation.sh` is service-agnostic. A recipe launcher defines named +pools such as GenRM, NL2Bash, or safety, including each model's environment and +vLLM arguments. The wrapper then: + +1. validates every pool and splits hetgroup 1 into disjoint node slices; +2. starts every replica in its own private Ray cluster; +3. starts one OpenAI-compatible load balancer per pool; +4. waits for every backend and load balancer to become healthy; +5. replaces each pool's URL placeholder in `COMMAND`; +6. starts `ray.sub` on hetgroup 0 only; and +7. stops training if any required external service exits. + +Recipe launchers keep their concrete model and serving settings outside these +helpers. A launcher can add another service without adding another server body +or lifecycle path to this wrapper. + +## Pool registration interface + +Launchers source `pool_config.sh` and use the same three functions for every +model. `register_external_vllm_pool` declares topology and endpoint metadata; +the other functions append environment assignments and argv entries: + +```bash +source "$PROJECT_ROOT/tools/external_gym_vllm/pool_config.sh" +EXTERNAL_VLLM_POOLS="" +export GPUS_PER_NODE=4 +NUM_EXTERNAL_SERVICE_NODES=2 + +register_external_vllm_pool SAFETY \ + --display-name "Safety judge" \ + --model "$SAFETY_MODEL" \ + --container "$SERVICE_CONTAINER" \ + --python /opt/vllm/bin/python \ + --replicas 2 \ + --tensor-parallel-size 4 \ + --lb-port 9215 \ + --served-model-name "$SAFETY_GYM_MODEL_NAME" \ + --url-placeholder __SAFETY_BASE_URL__ \ + --shared-path "$SAFETY_REASONING_PARSER" + +external_vllm_pool_env SAFETY \ + FLASHINFER_WORKSPACE_BASE=/tmp \ + NCCL_MNNVL_ENABLE=0 + +external_vllm_pool_args SAFETY \ + --dtype bfloat16 \ + --attention-backend FLASH_ATTN + +validate_external_vllm_submission "$COMMAND" "$NUM_EXTERNAL_SERVICE_NODES" +``` + +Registration appends `SAFETY` to `EXTERNAL_VLLM_POOLS` and exports the +normalized `SAFETY_*` contract inherited by `sbatch`. Adding this pool requires +no change to `run_in_allocation.sh`. + +The generated fields are: + +| Variable | Required | Default | Purpose | +|---|---:|---|---| +| `POOL_MODEL` | yes | — | Checkpoint path under `EXTERNAL_VLLM_SHARED_ROOT` or Hugging Face model ID. | +| `POOL_CONTAINER` | yes | — | Container used by this pool's replicas. | +| `POOL_VLLM_PYTHON` | yes | — | Python executable containing vLLM, Ray, and NeMo RL's compatibility patch. | +| `POOL_REPLICAS` | yes | — | Number of independent DP=1 servers. | +| `POOL_TENSOR_PARALLEL_SIZE` | yes | — | Tensor parallel size per server. | +| `POOL_LB_PORT` | yes | — | Unique load-balancer port on the Ray head node. | +| `POOL_URL_PLACEHOLDER` | yes | — | Token in `COMMAND` replaced by this pool's `/v1` URL. | +| `POOL_GROUP_ID` | no | `inline--` | Registry namespace; set with `--group-id` only when an explicit stable namespace is needed. | +| `POOL_DISPLAY_NAME` | no | pool name | Human-readable log label. | +| `POOL_SERVED_MODEL_NAME` | no | `model` | OpenAI API model name. Must equal the calling Gym server's `model` field because Gym overwrites the request model. | +| `POOL_VLLM_PORT` | no | `8000` | Backend HTTP port; replicas use disjoint private clusters. | +| `POOL_STARTUP_TIMEOUT` | no | `3600` | Seconds allowed for startup. | +| `POOL_SHARED_PATHS` | no | empty | Newline-separated absolute paths under `EXTERNAL_VLLM_SHARED_ROOT` that the pool container must access. | +| `POOL_ENV_VARS` | no | empty | Newline-separated `NAME=value` assignments applied before vLLM starts. | +| `POOL_VLLM_ARGS` | no | empty | Newline-separated vLLM CLI arguments, one argv entry per line. | + +Registration validates required fields, positive topology values, TCP port +ranges, TP divisibility, and duplicate ports/placeholders before `sbatch`. +`EXTERNAL_VLLM_NUM_NODES` is exported as the node total computed from all +registered pools. Call `validate_external_vllm_submission` after constructing +`COMMAND` to check its placeholders, shared paths, tool files, and requested +external node count before submitting the allocation. + +The interface uses one-argument-per-line encoding internally, preserving JSON +configs and paths containing spaces without `eval`. The wrapper itself supplies `--tensor-parallel-size`, +`--distributed-executor-backend ray`, `--port`, and `--served-model-name`. +Everything model-specific—including attention, reasoning/tool parsers, expert +parallelism, MoE backend, cache settings, and loader settings—belongs in the +launcher's pool definition. A pool's reasoning-parser setting must also agree +with the consuming Gym server's `uses_reasoning_parser` setting; in particular, +do not pass `--reasoning-parser` when Gym explicitly disables it. + +## Global contract + +Required variables: + +| Variable | Purpose | +|---|---| +| `BASE_LOG_DIR` | Parent under `EXTERNAL_VLLM_SHARED_ROOT` for `-logs`. | +| `COMMAND` | NeMo RL command containing every pool's URL placeholder. | +| `CONTAINER` | NeMo RL and load-balancer container. | +| `MOUNTS` | Mount list required by `ray.sub` and `COMMAND`. | +| `EXTERNAL_VLLM_POOLS` | Ordered pool names; this order determines node slicing. | +| `EXTERNAL_VLLM_TOOLS_DIR_HOST` | Path to this directory under `EXTERNAL_VLLM_SHARED_ROOT`. | + +Optional globals: + +| Variable | Default | Purpose | +|---|---|---| +| `GPUS_PER_NODE` | `4` | GPUs claimed per node. Set this before registering pools; it sizes hetgroup 1 and is exported to `ray.sub`, so it must match the physical GPUs per node for both components. | +| `NUM_EXTERNAL_SERVICE_NODES` | empty | Expected hetgroup 1 node count. Pass it to `validate_external_vllm_submission` to fail before `sbatch` on a topology mismatch; validation warns and skips this check when unset. | +| `EXTERNAL_VLLM_LB_PYTHON` | `/opt/nemo_rl_venv/bin/python` | Python with `aiohttp` in `CONTAINER`. | +| `RAY_SUB` | `$SLURM_SUBMIT_DIR/ray.sub` | Normal NeMo RL Slurm entrypoint. | +| `EXTERNAL_VLLM_SHARED_ROOT` | `/lustre` | Shared host path mounted at the same path in every external-service container. | +| `DEDICATED_RAY_HEAD` | unset | Passed through to `ray.sub`. With `1`, include one extra node in hetgroup 0 while keeping `cluster.num_nodes` equal to the GPU worker-node count; the head node's GPUs remain allocated but idle. | + +The number of nodes in hetgroup 1 must equal: + +```text +sum(POOL_REPLICAS * POOL_TENSOR_PARALLEL_SIZE / GPUS_PER_NODE) +``` + +Every TP value must be divisible by `GPUS_PER_NODE`. The wrapper records the +actual split in `$BASE_LOG_DIR/[-]-logs/node-allocation.txt` +and writes each resolved URL to `_url` in the same +directory. + +### Private-cluster port layout + +Each replica uses the same fixed port layout on its disjoint node set. This +matches `ray.sub` and stays below the `9000` ephemeral-port floor observed on +GB200 nodes: + +| Range | Purpose | +|---|---| +| `1200-1201` | Private Ray GCS and client server. | +| `1301-1312` | Private Ray management services. | +| `2000-2999` | Private Ray worker gRPC ports. | +| `7000-7999` | vLLM engine rendezvous, anchored by `VLLM_PORT=7000`. | +| `8000` by default | Per-pool vLLM HTTP endpoint (`POOL_VLLM_PORT`). | + +The old `10002-19999` Ray worker default overlaps the `9000-65000` ephemeral +range on these nodes. A vLLM TCPStore probe could therefore select a Ray worker +port and later fail with `EADDRINUSE`. Keeping both Ray and vLLM internal ports +in disjoint sub-ephemeral bands removes that race; ports above `19999` are not +reserved for this helper. + +## Filesystem and container requirements + +External replicas mount `EXTERNAL_VLLM_SHARED_ROOT` at the same path inside the +container. Therefore `BASE_LOG_DIR`, `EXTERNAL_VLLM_TOOLS_DIR_HOST`, and +absolute local model paths must be under that root. A Hugging Face model ID is +also accepted. + +Each pool container must provide: + +- its configured `POOL_VLLM_PYTHON`; +- importable `nemo_rl`, `ray`, and `vllm` packages in that environment; and +- the `ray` command on `PATH`. + +`serve_vllm_on_ray.py` applies NeMo RL's vLLM compatibility patches before it +imports the vLLM API server. `CONTAINER` must provide +`EXTERNAL_VLLM_LB_PYTHON` with `aiohttp` installed. + +## Slurm submission + +The wrapper requires exactly two hetgroups. Submit the NeMo RL nodes first and +the sum of all external-pool nodes second: + +```bash +sbatch \ + --account= \ + --partition= \ + --nodes= \ + --exclusive \ + --gres=gpu:4 \ + --time=04:00:00 \ + --export=ALL \ + : \ + --account= \ + --partition= \ + --nodes= \ + --exclusive \ + --gres=gpu:4 \ + --time=04:00:00 \ + tools/external_gym_vllm/run_in_allocation.sh +``` + +Slurm gang-schedules the two components. Replica `srun` steps explicitly use +hetgroup 1 and load-balancer steps explicitly use hetgroup 0. Before starting +`ray.sub`, the wrapper scopes its unsuffixed Slurm nodelist and node-count +variables to component 0; Slurm then uses component 0 by default for its steps. +External nodes therefore cannot accidentally join the training Ray cluster. diff --git a/tools/external_gym_vllm/lb_watchdog.sh b/tools/external_gym_vllm/lb_watchdog.sh new file mode 100755 index 00000000000..10216b96537 --- /dev/null +++ b/tools/external_gym_vllm/lb_watchdog.sh @@ -0,0 +1,62 @@ +#!/bin/bash +# 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. + +# Restart an external Gym vLLM load balancer after an unexpected exit. + +set -uo pipefail + +PYTHON="${PYTHON:-python3}" +LB_SCRIPT="$(dirname "$0")/vllm_pool_lb.py" +PORT="${1:?port required}" +REGISTRY_DIR="${2:?registry directory required}" +GROUP_ID="${3:?group ID required}" + +lb_pid="" +shutdown() { + trap - TERM INT + if [[ -n "${lb_pid}" ]] && kill -0 "${lb_pid}" 2>/dev/null; then + kill "${lb_pid}" 2>/dev/null || true + wait "${lb_pid}" 2>/dev/null || true + fi + exit 0 +} +trap shutdown TERM INT + +fast_failures=0 +while true; do + started_at=$(date +%s) + echo "$(date) [WATCHDOG] Starting load balancer on port ${PORT}" + "${PYTHON}" "${LB_SCRIPT}" \ + --port "${PORT}" \ + --registry-dir "${REGISTRY_DIR}" \ + --group-id "${GROUP_ID}" & + lb_pid=$! + wait "${lb_pid}" + status=$? + lb_pid="" + elapsed=$(( $(date +%s) - started_at )) + echo "$(date) [WATCHDOG] Load balancer exited (${status}) after ${elapsed}s" + + if (( elapsed < 5 )); then + fast_failures=$((fast_failures + 1)) + if (( fast_failures >= 3 )); then + echo "$(date) [WATCHDOG] ERROR: load balancer failed three times during startup" + exit 1 + fi + else + fast_failures=0 + fi + sleep 2 +done diff --git a/tools/external_gym_vllm/pool_config.sh b/tools/external_gym_vllm/pool_config.sh new file mode 100755 index 00000000000..96d4363f9db --- /dev/null +++ b/tools/external_gym_vllm/pool_config.sh @@ -0,0 +1,300 @@ +#!/bin/bash +# 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. + +# Public shell interface for registering external Gym vLLM pools before sbatch. + +_external_vllm_set() { + local pool="$1" suffix="$2" value="$3" + printf -v "${pool}_${suffix}" '%s' "${value}" + export "${pool}_${suffix}" +} + +_external_vllm_is_registered() { + local requested="$1" existing + local -a existing_pools=() + read -r -a existing_pools <<< "${EXTERNAL_VLLM_POOLS:-}" + for existing in "${existing_pools[@]}"; do + [[ "${existing}" == "${requested}" ]] && return 0 + done + return 1 +} + +_external_vllm_append_lines() { + local pool="$1" suffix="$2" + shift 2 + local variable_name="${pool}_${suffix}" + local current_value="${!variable_name-}" + local new_value + printf -v new_value '%s\n' "$@" + new_value="${new_value%$'\n'}" + if [[ -n "${current_value}" && -n "${new_value}" ]]; then + new_value="${current_value}"$'\n'"${new_value}" + elif [[ -n "${current_value}" ]]; then + new_value="${current_value}" + fi + _external_vllm_set "${pool}" "${suffix}" "${new_value}" +} + +_external_vllm_require_positive_integer() { + local field="$1" value="$2" + if [[ ! "${value}" =~ ^[0-9]+$ ]] || (( value <= 0 )); then + echo "ERROR: ${field} must be a positive integer (got '${value}')" >&2 + return 2 + fi +} + +_external_vllm_require_port() { + local field="$1" value="$2" + _external_vllm_require_positive_integer "${field}" "${value}" || return + if (( value > 65535 )); then + echo "ERROR: ${field} must be at most 65535 (got '${value}')" >&2 + return 2 + fi +} + +_external_vllm_recompute_node_count() { + local gpus_per_node="${GPUS_PER_NODE:-4}" + local pool replicas_var tensor_parallel_size_var + local total=0 + local -a pools=() + + read -r -a pools <<< "${EXTERNAL_VLLM_POOLS:-}" + for pool in "${pools[@]}"; do + replicas_var="${pool}_REPLICAS" + tensor_parallel_size_var="${pool}_TENSOR_PARALLEL_SIZE" + total=$((total + ${!replicas_var} * ${!tensor_parallel_size_var} / gpus_per_node)) + done + EXTERNAL_VLLM_NUM_NODES="${total}" + export EXTERNAL_VLLM_NUM_NODES +} + +_external_vllm_require_shared_path() { + local field="$1" path="$2" + local allow_model_id="${3:-0}" + local shared_root="${EXTERNAL_VLLM_SHARED_ROOT:-/lustre}" + if [[ "${shared_root}" != /* ]]; then + echo "ERROR: EXTERNAL_VLLM_SHARED_ROOT must be absolute (got '${shared_root}')" >&2 + return 2 + fi + if [[ "${path}" != /* ]]; then + if [[ "${allow_model_id}" == "1" ]]; then + return 0 + fi + echo "ERROR: ${field} must be an absolute path (got '${path}')" >&2 + return 2 + fi + if [[ "${path}" != "${shared_root}" && "${path}" != "${shared_root}"/* ]]; then + echo "ERROR: ${field} must be under ${shared_root} (got '${path}')" >&2 + return 2 + fi +} + +register_external_vllm_pool() { + local pool="${1:?pool name required}" + shift + if [[ ! "${pool}" =~ ^[A-Z][A-Z0-9_]*$ ]]; then + echo "ERROR: invalid external vLLM pool name '${pool}'; use an uppercase shell identifier" >&2 + return 2 + fi + if _external_vllm_is_registered "${pool}"; then + echo "ERROR: external vLLM pool '${pool}' is already registered" >&2 + return 2 + fi + + local display_name="${pool}" + local model="" + local container="" + local python="" + local replicas="" + local tensor_parallel_size="" + local lb_port="" + local url_placeholder="" + local group_id="" + local served_model_name="model" + local vllm_port="8000" + local startup_timeout="3600" + local -a shared_paths=() + + while (( $# > 0 )); do + case "$1" in + --display-name) display_name="${2:?value required for $1}"; shift 2 ;; + --model) model="${2:?value required for $1}"; shift 2 ;; + --container) container="${2:?value required for $1}"; shift 2 ;; + --python) python="${2:?value required for $1}"; shift 2 ;; + --replicas) replicas="${2:?value required for $1}"; shift 2 ;; + --tensor-parallel-size) tensor_parallel_size="${2:?value required for $1}"; shift 2 ;; + --lb-port) lb_port="${2:?value required for $1}"; shift 2 ;; + --url-placeholder) url_placeholder="${2:?value required for $1}"; shift 2 ;; + --group-id) group_id="${2:?value required for $1}"; shift 2 ;; + --served-model-name) served_model_name="${2:?value required for $1}"; shift 2 ;; + --vllm-port) vllm_port="${2:?value required for $1}"; shift 2 ;; + --startup-timeout) startup_timeout="${2:?value required for $1}"; shift 2 ;; + --shared-path) shared_paths+=("${2:?value required for $1}"); shift 2 ;; + *) + echo "ERROR: unknown register_external_vllm_pool option: $1" >&2 + return 2 + ;; + esac + done + + local field value + for field in model container python replicas tensor_parallel_size lb_port url_placeholder; do + value="${!field}" + if [[ -z "${value}" ]]; then + echo "ERROR: ${field//_/-} is required for external vLLM pool ${pool}" >&2 + return 2 + fi + done + + local gpus_per_node="${GPUS_PER_NODE:-4}" + _external_vllm_require_positive_integer "GPUS_PER_NODE" "${gpus_per_node}" || return + _external_vllm_require_positive_integer "${pool}_REPLICAS" "${replicas}" || return + _external_vllm_require_positive_integer \ + "${pool}_TENSOR_PARALLEL_SIZE" "${tensor_parallel_size}" || return + _external_vllm_require_port "${pool}_LB_PORT" "${lb_port}" || return + _external_vllm_require_port "${pool}_VLLM_PORT" "${vllm_port}" || return + _external_vllm_require_positive_integer \ + "${pool}_STARTUP_TIMEOUT" "${startup_timeout}" || return + if (( tensor_parallel_size % gpus_per_node != 0 )); then + echo "ERROR: ${pool}_TENSOR_PARALLEL_SIZE must be divisible by GPUS_PER_NODE=${gpus_per_node}" >&2 + return 2 + fi + if [[ -n "${group_id}" && ! "${group_id}" =~ ^[A-Za-z0-9._-]+$ ]]; then + echo "ERROR: ${pool}_GROUP_ID may contain only letters, digits, '.', '_', and '-'" >&2 + return 2 + fi + _external_vllm_require_shared_path "${pool}_MODEL" "${model}" 1 || return + local shared_path + for shared_path in "${shared_paths[@]}"; do + _external_vllm_require_shared_path "${pool}_SHARED_PATHS" "${shared_path}" || return + done + + local existing existing_lb_var existing_placeholder_var + local -a existing_pools=() + read -r -a existing_pools <<< "${EXTERNAL_VLLM_POOLS:-}" + for existing in "${existing_pools[@]}"; do + existing_lb_var="${existing}_LB_PORT" + existing_placeholder_var="${existing}_URL_PLACEHOLDER" + if [[ "${!existing_lb_var}" == "${lb_port}" ]]; then + echo "ERROR: external vLLM pools ${existing} and ${pool} use LB port ${lb_port}" >&2 + return 2 + fi + if [[ "${!existing_placeholder_var}" == "${url_placeholder}" ]]; then + echo "ERROR: external vLLM pools ${existing} and ${pool} use URL placeholder ${url_placeholder}" >&2 + return 2 + fi + done + + _external_vllm_set "${pool}" DISPLAY_NAME "${display_name}" + _external_vllm_set "${pool}" MODEL "${model}" + _external_vllm_set "${pool}" CONTAINER "${container}" + _external_vllm_set "${pool}" VLLM_PYTHON "${python}" + _external_vllm_set "${pool}" REPLICAS "${replicas}" + _external_vllm_set "${pool}" TENSOR_PARALLEL_SIZE "${tensor_parallel_size}" + _external_vllm_set "${pool}" LB_PORT "${lb_port}" + _external_vllm_set "${pool}" URL_PLACEHOLDER "${url_placeholder}" + _external_vllm_set "${pool}" SERVED_MODEL_NAME "${served_model_name}" + _external_vllm_set "${pool}" VLLM_PORT "${vllm_port}" + _external_vllm_set "${pool}" STARTUP_TIMEOUT "${startup_timeout}" + if [[ -n "${group_id}" ]]; then + _external_vllm_set "${pool}" GROUP_ID "${group_id}" + fi + if (( ${#shared_paths[@]} > 0 )); then + _external_vllm_append_lines "${pool}" SHARED_PATHS "${shared_paths[@]}" + else + _external_vllm_set "${pool}" SHARED_PATHS "" + fi + _external_vllm_set "${pool}" ENV_VARS "" + _external_vllm_set "${pool}" VLLM_ARGS "" + + EXTERNAL_VLLM_POOLS="${EXTERNAL_VLLM_POOLS:+${EXTERNAL_VLLM_POOLS} }${pool}" + export EXTERNAL_VLLM_POOLS + _external_vllm_recompute_node_count +} + +validate_external_vllm_submission() { + local command="${1:-${COMMAND:-}}" + local expected_nodes="${2:-${NUM_EXTERNAL_SERVICE_NODES:-}}" + local shared_root="${EXTERNAL_VLLM_SHARED_ROOT:-/lustre}" + local pool placeholder_var path variable_name required_file + local -a pools=() + + if [[ -z "${command}" ]]; then + echo "ERROR: external vLLM submission command is empty" >&2 + return 2 + fi + if [[ -z "${EXTERNAL_VLLM_POOLS:-}" ]]; then + echo "ERROR: no external vLLM pools are registered" >&2 + return 2 + fi + if [[ -n "${expected_nodes}" ]]; then + _external_vllm_require_positive_integer \ + "NUM_EXTERNAL_SERVICE_NODES" "${expected_nodes}" || return + if (( expected_nodes != EXTERNAL_VLLM_NUM_NODES )); then + echo "ERROR: NUM_EXTERNAL_SERVICE_NODES=${expected_nodes}, expected ${EXTERNAL_VLLM_NUM_NODES} from registered pools" >&2 + return 2 + fi + else + echo "WARNING: NUM_EXTERNAL_SERVICE_NODES is unset; skipping external hetgroup node-count validation" >&2 + fi + + read -r -a pools <<< "${EXTERNAL_VLLM_POOLS}" + for pool in "${pools[@]}"; do + placeholder_var="${pool}_URL_PLACEHOLDER" + if [[ "${command}" != *"${!placeholder_var}"* ]]; then + echo "ERROR: submission command is missing ${!placeholder_var} for pool ${pool}" >&2 + return 2 + fi + done + + for variable_name in BASE_LOG_DIR EXTERNAL_VLLM_TOOLS_DIR_HOST; do + path="${!variable_name-}" + if [[ -z "${path}" ]]; then + echo "ERROR: ${variable_name} is required for external vLLM submission" >&2 + return 2 + fi + _external_vllm_require_shared_path "${variable_name}" "${path}" || return + done + for required_file in vllm_backend_registry.sh vllm_pool_lb.py lb_watchdog.sh serve_vllm_on_ray.py; do + if [[ ! -f "${EXTERNAL_VLLM_TOOLS_DIR_HOST}/${required_file}" ]]; then + echo "ERROR: missing ${EXTERNAL_VLLM_TOOLS_DIR_HOST}/${required_file}" >&2 + return 2 + fi + done + if [[ "${shared_root}" != /* ]]; then + echo "ERROR: EXTERNAL_VLLM_SHARED_ROOT must be absolute (got '${shared_root}')" >&2 + return 2 + fi +} + +external_vllm_pool_env() { + local pool="${1:?pool name required}" + shift + if ! _external_vllm_is_registered "${pool}"; then + echo "ERROR: external vLLM pool '${pool}' is not registered" >&2 + return 2 + fi + _external_vllm_append_lines "${pool}" ENV_VARS "$@" +} + +external_vllm_pool_args() { + local pool="${1:?pool name required}" + shift + if ! _external_vllm_is_registered "${pool}"; then + echo "ERROR: external vLLM pool '${pool}' is not registered" >&2 + return 2 + fi + _external_vllm_append_lines "${pool}" VLLM_ARGS "$@" +} diff --git a/tools/external_gym_vllm/run_in_allocation.sh b/tools/external_gym_vllm/run_in_allocation.sh new file mode 100755 index 00000000000..d07c6d42dc8 --- /dev/null +++ b/tools/external_gym_vllm/run_in_allocation.sh @@ -0,0 +1,702 @@ +#!/bin/bash +# 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. + +# Run external Gym vLLM pools beside NeMo RL in a single, two-component +# Slurm heterogeneous job. Pool definitions are supplied by the caller through +# EXTERNAL_VLLM_POOLS and consistently named, exported environment variables. + +set -euo pipefail + +: "${SLURM_JOB_ID:?This script must run inside a Slurm allocation}" +: "${SLURM_HET_SIZE:?This script requires a Slurm heterogeneous job}" +: "${SLURM_JOB_NODELIST_HET_GROUP_0:?Hetgroup 0 nodelist is required}" +: "${SLURM_JOB_NODELIST_HET_GROUP_1:?Hetgroup 1 nodelist is required}" +: "${SLURM_JOB_ACCOUNT:?SLURM_JOB_ACCOUNT is required}" +: "${SLURM_JOB_PARTITION:?SLURM_JOB_PARTITION is required}" +: "${SLURM_SUBMIT_DIR:?SLURM_SUBMIT_DIR is required}" +: "${BASE_LOG_DIR:?BASE_LOG_DIR is required}" +: "${CONTAINER:?CONTAINER is required}" +: "${MOUNTS:?MOUNTS is required}" +: "${COMMAND:?COMMAND is required}" +: "${EXTERNAL_VLLM_POOLS:?EXTERNAL_VLLM_POOLS is required}" +: "${EXTERNAL_VLLM_TOOLS_DIR_HOST:?EXTERNAL_VLLM_TOOLS_DIR_HOST is required}" + +if [[ "${SLURM_HET_SIZE}" != "2" ]]; then + echo "[FATAL] Expected exactly two Slurm hetgroups, got ${SLURM_HET_SIZE}" >&2 + exit 1 +fi + +RAY_SUB="${RAY_SUB:-${SLURM_SUBMIT_DIR}/ray.sub}" +export GPUS_PER_NODE="${GPUS_PER_NODE:-4}" +EXTERNAL_VLLM_LB_PYTHON="${EXTERNAL_VLLM_LB_PYTHON:-/opt/nemo_rl_venv/bin/python}" +EXTERNAL_VLLM_SHARED_ROOT="${EXTERNAL_VLLM_SHARED_ROOT:-/lustre}" + +if [[ ! -f "${RAY_SUB}" ]]; then + echo "[FATAL] ray.sub does not exist: ${RAY_SUB}" >&2 + exit 1 +fi +for required_file in vllm_backend_registry.sh vllm_pool_lb.py lb_watchdog.sh serve_vllm_on_ray.py; do + if [[ ! -f "${EXTERNAL_VLLM_TOOLS_DIR_HOST}/${required_file}" ]]; then + echo "[FATAL] Missing ${EXTERNAL_VLLM_TOOLS_DIR_HOST}/${required_file}" >&2 + exit 1 + fi +done +if [[ ! "${GPUS_PER_NODE}" =~ ^[0-9]+$ ]] || (( GPUS_PER_NODE <= 0 )); then + echo "[FATAL] GPUS_PER_NODE must be a positive integer" >&2 + exit 1 +fi +if [[ "${EXTERNAL_VLLM_SHARED_ROOT}" != /* ]]; then + echo "[FATAL] EXTERNAL_VLLM_SHARED_ROOT must be absolute" >&2 + exit 1 +fi + +read -r -a pool_names <<< "${EXTERNAL_VLLM_POOLS}" +if (( ${#pool_names[@]} == 0 )); then + echo "[FATAL] EXTERNAL_VLLM_POOLS must name at least one pool" >&2 + exit 1 +fi + +pool_value() { + local pool="$1" + local suffix="$2" + local default_value="${3-}" + local variable_name="${pool}_${suffix}" + printf '%s' "${!variable_name-${default_value}}" +} + +require_pool_value() { + local pool="$1" + local suffix="$2" + local value + value=$(pool_value "${pool}" "${suffix}") + if [[ -z "${value}" ]]; then + echo "[FATAL] ${pool}_${suffix} is required" >&2 + return 1 + fi + printf '%s' "${value}" +} + +declare -A display_names=() +declare -A models=() +declare -A containers=() +declare -A vllm_pythons=() +declare -A replicas=() +declare -A tensor_parallel_sizes=() +declare -A nodes_per_replica=() +declare -A node_offsets=() +declare -A node_counts=() +declare -A served_model_names=() +declare -A backend_ports=() +declare -A lb_ports=() +declare -A startup_timeouts=() +declare -A placeholders=() +declare -A group_ids=() +declare -A pool_log_dirs=() +declare -A state_dirs=() +declare -A lb_state_dirs=() +declare -A pool_urls=() + +total_external_nodes=0 +max_startup_timeout=0 +declare -A seen_pool_names=() +declare -A seen_lb_ports=() +declare -A seen_placeholders=() + +for pool in "${pool_names[@]}"; do + if [[ ! "${pool}" =~ ^[A-Z][A-Z0-9_]*$ ]]; then + echo "[FATAL] Invalid pool name '${pool}'; use uppercase shell identifiers" >&2 + exit 1 + fi + if [[ -n "${seen_pool_names[${pool}]-}" ]]; then + echo "[FATAL] Duplicate external vLLM pool: ${pool}" >&2 + exit 1 + fi + seen_pool_names["${pool}"]=1 + + display_names["${pool}"]=$(pool_value "${pool}" DISPLAY_NAME "${pool}") + models["${pool}"]=$(require_pool_value "${pool}" MODEL) + containers["${pool}"]=$(require_pool_value "${pool}" CONTAINER) + vllm_pythons["${pool}"]=$(require_pool_value "${pool}" VLLM_PYTHON) + replicas["${pool}"]=$(require_pool_value "${pool}" REPLICAS) + tensor_parallel_sizes["${pool}"]=$(require_pool_value "${pool}" TENSOR_PARALLEL_SIZE) + served_model_names["${pool}"]=$(pool_value "${pool}" SERVED_MODEL_NAME model) + backend_ports["${pool}"]=$(pool_value "${pool}" VLLM_PORT 8000) + lb_ports["${pool}"]=$(require_pool_value "${pool}" LB_PORT) + startup_timeouts["${pool}"]=$(pool_value "${pool}" STARTUP_TIMEOUT 3600) + placeholders["${pool}"]=$(require_pool_value "${pool}" URL_PLACEHOLDER) + group_ids["${pool}"]=$(pool_value "${pool}" GROUP_ID "inline-${pool,,}-${SLURM_JOB_ID}") + if [[ ! "${group_ids[${pool}]}" =~ ^[A-Za-z0-9._-]+$ ]]; then + echo "[FATAL] ${pool}_GROUP_ID may contain only letters, digits, '.', '_', and '-'" >&2 + exit 1 + fi + + # srun inherits the pool contract through exported variables. Re-export the + # normalized values so optional defaults are also visible inside replicas. + export "${pool}_DISPLAY_NAME=${display_names[${pool}]}" + export "${pool}_MODEL=${models[${pool}]}" + export "${pool}_CONTAINER=${containers[${pool}]}" + export "${pool}_VLLM_PYTHON=${vllm_pythons[${pool}]}" + export "${pool}_REPLICAS=${replicas[${pool}]}" + export "${pool}_TENSOR_PARALLEL_SIZE=${tensor_parallel_sizes[${pool}]}" + export "${pool}_SERVED_MODEL_NAME=${served_model_names[${pool}]}" + export "${pool}_VLLM_PORT=${backend_ports[${pool}]}" + export "${pool}_ENV_VARS=$(pool_value "${pool}" ENV_VARS)" + export "${pool}_VLLM_ARGS=$(pool_value "${pool}" VLLM_ARGS)" + + for numeric_suffix in REPLICAS TENSOR_PARALLEL_SIZE VLLM_PORT LB_PORT STARTUP_TIMEOUT; do + case "${numeric_suffix}" in + REPLICAS) numeric_value="${replicas[${pool}]}" ;; + TENSOR_PARALLEL_SIZE) numeric_value="${tensor_parallel_sizes[${pool}]}" ;; + VLLM_PORT) numeric_value="${backend_ports[${pool}]}" ;; + LB_PORT) numeric_value="${lb_ports[${pool}]}" ;; + STARTUP_TIMEOUT) numeric_value="${startup_timeouts[${pool}]}" ;; + esac + if [[ ! "${numeric_value}" =~ ^[0-9]+$ ]] || (( numeric_value <= 0 )); then + echo "[FATAL] ${pool}_${numeric_suffix} must be a positive integer" >&2 + exit 1 + fi + if [[ "${numeric_suffix}" == "VLLM_PORT" || "${numeric_suffix}" == "LB_PORT" ]] && (( numeric_value > 65535 )); then + echo "[FATAL] ${pool}_${numeric_suffix} must be at most 65535" >&2 + exit 1 + fi + done + if (( tensor_parallel_sizes[${pool}] % GPUS_PER_NODE != 0 )); then + echo "[FATAL] ${pool}_TENSOR_PARALLEL_SIZE must be divisible by GPUS_PER_NODE" >&2 + exit 1 + fi + if [[ -n "${seen_lb_ports[${lb_ports[${pool}]}]-}" ]]; then + echo "[FATAL] Multiple pools use load-balancer port ${lb_ports[${pool}]}" >&2 + exit 1 + fi + seen_lb_ports["${lb_ports[${pool}]}"]="${pool}" + if [[ -n "${seen_placeholders[${placeholders[${pool}]}]-}" ]]; then + echo "[FATAL] Multiple pools use URL placeholder ${placeholders[${pool}]}" >&2 + exit 1 + fi + seen_placeholders["${placeholders[${pool}]}"]="${pool}" + if [[ "${COMMAND}" != *"${placeholders[${pool}]}"* ]]; then + echo "[FATAL] Driver command is missing ${placeholders[${pool}]} for ${display_names[${pool}]}" >&2 + exit 1 + fi + + # Each private Ray cluster owns whole nodes. This makes its fixed Ray port safe + # to reuse across replicas because no two replicas ever share a host. + nodes_per_replica["${pool}"]=$((tensor_parallel_sizes[${pool}] / GPUS_PER_NODE)) + node_counts["${pool}"]=$((replicas[${pool}] * nodes_per_replica[${pool}])) + node_offsets["${pool}"]="${total_external_nodes}" + total_external_nodes=$((total_external_nodes + node_counts[${pool}])) + if (( startup_timeouts[${pool}] > max_startup_timeout )); then + max_startup_timeout="${startup_timeouts[${pool}]}" + fi +done + +shared_paths=("${BASE_LOG_DIR}" "${EXTERNAL_VLLM_TOOLS_DIR_HOST}") +for pool in "${pool_names[@]}"; do + if [[ "${models[${pool}]}" == /* ]]; then + shared_paths+=("${models[${pool}]}") + fi + while IFS= read -r shared_path; do + [[ -n "${shared_path}" ]] && shared_paths+=("${shared_path}") + done <<< "$(pool_value "${pool}" SHARED_PATHS)" +done +for shared_path in "${shared_paths[@]}"; do + if [[ "${shared_path}" != "${EXTERNAL_VLLM_SHARED_ROOT}" && "${shared_path}" != "${EXTERNAL_VLLM_SHARED_ROOT}"/* ]]; then + echo "[FATAL] Path must be under ${EXTERNAL_VLLM_SHARED_ROOT} for the external-service container mount: ${shared_path}" >&2 + exit 1 + fi +done + +mapfile -t ray_nodes < <( + scontrol show hostnames "${SLURM_JOB_NODELIST_HET_GROUP_0}" | sort +) +mapfile -t external_nodes < <( + scontrol show hostnames "${SLURM_JOB_NODELIST_HET_GROUP_1}" | sort +) +if (( ${#ray_nodes[@]} == 0 )); then + echo "[FATAL] Slurm hetgroup 0 contains no NeMo RL nodes" >&2 + exit 1 +fi +if (( ${#external_nodes[@]} != total_external_nodes )); then + echo "[FATAL] Slurm hetgroup 1 has ${#external_nodes[@]} nodes, expected ${total_external_nodes}" >&2 + for pool in "${pool_names[@]}"; do + echo "[FATAL] ${display_names[${pool}]}: ${node_counts[${pool}]} nodes" >&2 + done + exit 1 +fi + +# Must match ray.sub's LOG_DIR: ENDED is the teardown channel between them. +if [[ -n "${SLURM_RESTART_COUNT:-}" ]]; then + LOG_DIR="${BASE_LOG_DIR}/${SLURM_JOB_ID}-${SLURM_RESTART_COUNT}-logs" +else + LOG_DIR="${BASE_LOG_DIR}/${SLURM_JOB_ID}-logs" +fi +mkdir -p "${LOG_DIR}" +rm_args=() +for pool in "${pool_names[@]}"; do + pool_key="${pool,,}" + pool_log_dirs["${pool}"]="${LOG_DIR}/external_${pool_key}" + state_dirs["${pool}"]="${pool_log_dirs[${pool}]}/state" + lb_state_dirs["${pool}"]="/tmp/external-vllm-state-${pool_key}" + mkdir -p "${pool_log_dirs[${pool}]}" "${state_dirs[${pool}]}" + rm_args+=( + "${state_dirs[${pool}]}/.registry_${group_ids[${pool}]}" + "${state_dirs[${pool}]}/.registry_${group_ids[${pool}]}.lock" + "${LOG_DIR}/${pool_key}_url" + ) +done +rm -f "${rm_args[@]}" +for pool in "${pool_names[@]}"; do + rm -f "${pool_log_dirs[${pool}]}"/head_ip_* +done + +{ + for pool in "${pool_names[@]}"; do + echo "[external_${pool,,}]" + offset="${node_offsets[${pool}]}" + count="${node_counts[${pool}]}" + printf '%s\n' "${external_nodes[@]:offset:count}" + done + echo "[nemo_rl_ray]" + printf '%s\n' "${ray_nodes[@]}" +} > "${LOG_DIR}/node-allocation.txt" + +echo "[INFO] Heterogeneous-job external-vLLM topology" +echo "[INFO] Hetgroup 0, NeMo RL Ray: ${#ray_nodes[@]} nodes (${SLURM_JOB_NODELIST_HET_GROUP_0})" +for pool in "${pool_names[@]}"; do + echo "[INFO] Hetgroup 1, ${display_names[${pool}]}: ${node_counts[${pool}]} nodes, ${replicas[${pool}]} TP=${tensor_parallel_sizes[${pool}]} replicas" +done + +declare -a service_step_pids=() +declare -a service_step_labels=() +declare -a lb_step_pids=() +declare -a lb_step_labels=() +declare -a preflight_pids=() +declare -a preflight_labels=() +ray_sub_pid="" + +cleanup() { + local status=$? + trap - EXIT TERM INT + + touch "${LOG_DIR}/ENDED" 2>/dev/null || true + if [[ -n "${ray_sub_pid}" ]] && kill -0 "${ray_sub_pid}" 2>/dev/null; then + kill "${ray_sub_pid}" 2>/dev/null || true + fi + for pid in "${lb_step_pids[@]}" "${service_step_pids[@]}" "${preflight_pids[@]}"; do + if [[ -n "${pid}" ]] && kill -0 "${pid}" 2>/dev/null; then + kill "${pid}" 2>/dev/null || true + fi + done + + wait 2>/dev/null || true + exit "${status}" +} +trap cleanup EXIT +trap 'exit 143' TERM INT + +check_service_steps() { + local index + for index in "${!service_step_pids[@]}"; do + if ! kill -0 "${service_step_pids[${index}]}" 2>/dev/null; then + echo "[FATAL] ${service_step_labels[${index}]} exited unexpectedly" >&2 + return 1 + fi + done + for index in "${!lb_step_pids[@]}"; do + if ! kill -0 "${lb_step_pids[${index}]}" 2>/dev/null; then + echo "[FATAL] ${lb_step_labels[${index}]} exited unexpectedly" >&2 + return 1 + fi + done +} + +resolve_node_ip() { + local node="$1" ip + ip=$(getent ahostsv4 "${node}" 2>/dev/null | awk 'NR == 1 { print $1 }' || true) + if [[ -z "${ip}" ]]; then + ip=$(host "${node}" 2>/dev/null | awk '/has address/ { print $4; exit }' || true) + fi + if [[ -z "${ip}" ]]; then + echo "[FATAL] Could not resolve an IPv4 address for ${node}" >&2 + return 1 + fi + echo "${ip}" +} + +VLLM_SERVER_BODY=$(cat <<'VLLM_SERVER_BODY_EOF' +set -euo pipefail + +: "${POOL_PREFIX:?POOL_PREFIX is required}" +: "${REPLICA_ID:?REPLICA_ID is required}" +: "${EXTERNAL_VLLM_TOOLS_DIR:?EXTERNAL_VLLM_TOOLS_DIR is required}" +: "${EXTERNAL_VLLM_STATE_DIR:?EXTERNAL_VLLM_STATE_DIR is required}" +: "${EXTERNAL_VLLM_GROUP_ID:?EXTERNAL_VLLM_GROUP_ID is required}" +: "${HEAD_IP_FILE:?HEAD_IP_FILE is required}" +: "${LOG_FILE:?LOG_FILE is required}" + +pool_value() { + local variable_name="${POOL_PREFIX}_$1" + printf '%s' "${!variable_name-}" +} + +MODEL=$(pool_value MODEL) +VLLM_PYTHON=$(pool_value VLLM_PYTHON) +VLLM_HTTP_PORT=$(pool_value VLLM_PORT) +TENSOR_PARALLEL_SIZE=$(pool_value TENSOR_PARALLEL_SIZE) +SERVED_MODEL_NAME=$(pool_value SERVED_MODEL_NAME) +DISPLAY_NAME=$(pool_value DISPLAY_NAME) +[[ -n "${SERVED_MODEL_NAME}" ]] || SERVED_MODEL_NAME=model +[[ -n "${DISPLAY_NAME}" ]] || DISPLAY_NAME="${POOL_PREFIX}" + +source "${EXTERNAL_VLLM_TOOLS_DIR}/vllm_backend_registry.sh" + +# Match ray.sub's sub-ephemeral port layout. Some GB200 nodes use +# 9000-65000 for ephemeral ports, so Ray's former 10002-19999 worker range +# could race vLLM's bind-probe-and-release TCPStore allocation. Every replica +# owns disjoint nodes, so these fixed ports can be reused across replicas. +RAY_PORT=1200 +RAY_CLIENT_SERVER_PORT=1201 +NODE_MANAGER_PORT=1301 +OBJECT_MANAGER_PORT=1303 +RUNTIME_ENV_AGENT_PORT=1305 +DASHBOARD_AGENT_GRPC_PORT=1307 +METRICS_EXPORT_PORT=1309 +DASHBOARD_AGENT_LISTEN_PORT=1311 +MIN_WORKER_PORT=2000 +MAX_WORKER_PORT=2999 +VLLM_ENGINE_PORT=7000 + +cleanup_replica() { + if [[ "${SLURM_PROCID:-0}" -eq 0 ]]; then + registry_remove "${REPLICA_ID}" || true + fi + ray stop 2>/dev/null || true +} +trap cleanup_replica EXIT +trap 'trap - EXIT; cleanup_replica; exit 143' TERM INT + +if [[ "${SLURM_PROCID:-0}" -eq 0 ]]; then + rm -f "${HEAD_IP_FILE}" + HEAD_IP=$(hostname -I | awk '{ print $1 }') + if [[ -z "${HEAD_IP}" ]]; then + HEAD_IP=$(getent ahostsv4 "$(hostname)" | awk 'NR == 1 { print $1 }') + fi + if [[ -z "${HEAD_IP}" ]]; then + echo "[${REPLICA_ID}] ERROR: could not determine the head-node IP" >&2 + exit 1 + fi + + echo "${HEAD_IP}" > "${HEAD_IP_FILE}" + echo "[${REPLICA_ID}] Starting private Ray head at ${HEAD_IP}:${RAY_PORT}" + ray start \ + --head \ + --node-ip-address="${HEAD_IP}" \ + --port="${RAY_PORT}" \ + --ray-client-server-port="${RAY_CLIENT_SERVER_PORT}" \ + --min-worker-port="${MIN_WORKER_PORT}" \ + --max-worker-port="${MAX_WORKER_PORT}" \ + --node-manager-port="$((NODE_MANAGER_PORT + 1))" \ + --object-manager-port="$((OBJECT_MANAGER_PORT + 1))" \ + --runtime-env-agent-port="$((RUNTIME_ENV_AGENT_PORT + 1))" \ + --dashboard-agent-grpc-port="$((DASHBOARD_AGENT_GRPC_PORT + 1))" \ + --dashboard-agent-listen-port="$((DASHBOARD_AGENT_LISTEN_PORT + 1))" \ + --metrics-export-port="$((METRICS_EXPORT_PORT + 1))" \ + --disable-usage-stats + + while IFS= read -r assignment; do + [[ -n "${assignment}" ]] || continue + variable_name="${assignment%%=*}" + if [[ ! "${variable_name}" =~ ^[A-Za-z_][A-Za-z0-9_]*$ ]] || [[ "${assignment}" != *=* ]]; then + echo "[${REPLICA_ID}] ERROR: invalid environment assignment: ${assignment}" >&2 + exit 1 + fi + export "${assignment}" + done <<< "$(pool_value ENV_VARS)" + + # Keep vLLM's TCPStore and MessageQueue ports inside the reserved + # 7000-7999 band. serve_vllm_on_ray.py applies the NeMo RL compatibility + # patch that offsets the TCPStore search within this per-engine window. + export VLLM_PORT="${VLLM_ENGINE_PORT}" + + vllm_args=() + while IFS= read -r argument; do + [[ -n "${argument}" ]] && vllm_args+=("${argument}") + done <<< "$(pool_value VLLM_ARGS)" + + echo "[${REPLICA_ID}] Starting ${DISPLAY_NAME} vLLM server at TP=${TENSOR_PARALLEL_SIZE}/DP=1" + "${VLLM_PYTHON}" "${EXTERNAL_VLLM_TOOLS_DIR}/serve_vllm_on_ray.py" serve "${MODEL}" \ + --tensor-parallel-size "${TENSOR_PARALLEL_SIZE}" \ + --distributed-executor-backend ray \ + --port "${VLLM_HTTP_PORT}" \ + --served-model-name "${SERVED_MODEL_NAME}" \ + "${vllm_args[@]}" \ + > "${LOG_FILE}" 2>&1 & + VLLM_PID=$! + + while ! "${VLLM_PYTHON}" -c \ + 'import sys, urllib.request; urllib.request.urlopen(sys.argv[1], timeout=2).close()' \ + "http://${HEAD_IP}:${VLLM_HTTP_PORT}/health" >/dev/null 2>&1; do + if ! kill -0 "${VLLM_PID}" 2>/dev/null; then + echo "[${REPLICA_ID}] ERROR: vLLM exited before becoming healthy" >&2 + exit 1 + fi + sleep 5 + done + + registry_add "${REPLICA_ID}" "${HEAD_IP}" "${VLLM_HTTP_PORT}" + echo "[${REPLICA_ID}] Registered healthy backend ${HEAD_IP}:${VLLM_HTTP_PORT}" + if wait "${VLLM_PID}"; then + vllm_status=0 + else + vllm_status=$? + fi + echo "[${REPLICA_ID}] ERROR: vLLM exited with status ${vllm_status}" >&2 + if (( vllm_status == 0 )); then + exit 1 + fi + exit "${vllm_status}" +else + for _ in $(seq 1 120); do + [[ -s "${HEAD_IP_FILE}" ]] && break + sleep 1 + done + if [[ ! -s "${HEAD_IP_FILE}" ]]; then + echo "[${REPLICA_ID}] ERROR: private Ray head IP was not published" >&2 + exit 1 + fi + + HEAD_IP=$(cat "${HEAD_IP_FILE}") + joined=0 + for _ in $(seq 1 120); do + if ray start \ + --address="${HEAD_IP}:${RAY_PORT}" \ + --min-worker-port="${MIN_WORKER_PORT}" \ + --max-worker-port="${MAX_WORKER_PORT}" \ + --node-manager-port="${NODE_MANAGER_PORT}" \ + --object-manager-port="${OBJECT_MANAGER_PORT}" \ + --runtime-env-agent-port="${RUNTIME_ENV_AGENT_PORT}" \ + --dashboard-agent-grpc-port="${DASHBOARD_AGENT_GRPC_PORT}" \ + --dashboard-agent-listen-port="${DASHBOARD_AGENT_LISTEN_PORT}" \ + --metrics-export-port="${METRICS_EXPORT_PORT}" \ + --disable-usage-stats; then + joined=1 + break + fi + sleep 2 + done + if (( joined == 0 )); then + echo "[${REPLICA_ID}] ERROR: failed to join private Ray cluster" >&2 + exit 1 + fi + tail -f /dev/null +fi +VLLM_SERVER_BODY_EOF +) +bash -n <(printf '%s' "${VLLM_SERVER_BODY}") || { + echo "[FATAL] Generated vLLM server script has syntax errors" >&2 + exit 1 +} + +ray_head_node="${ray_nodes[0]}" +lb_mounts="${MOUNTS},${EXTERNAL_VLLM_TOOLS_DIR_HOST}:/opt/external-vllm-tools:ro" +external_service_mount="${EXTERNAL_VLLM_SHARED_ROOT}:${EXTERNAL_VLLM_SHARED_ROOT}" +for pool in "${pool_names[@]}"; do + lb_mounts+=",${state_dirs[${pool}]}:${lb_state_dirs[${pool}]}" +done + +for pool in "${pool_names[@]}"; do + offset="${node_offsets[${pool}]}" + first_pool_node="${external_nodes[${offset}]}" + echo "[INFO] Validating the ${display_names[${pool}]} container and Python environment" + srun \ + --het-group=1 \ + --no-container-mount-home \ + --container-image="${containers[${pool}]}" \ + --container-mounts="${external_service_mount}" \ + --mpi=pmix \ + --gres="gpu:${GPUS_PER_NODE}" \ + --overlap \ + --kill-on-bad-exit=1 \ + --nodelist="${first_pool_node}" \ + --nodes=1 \ + --ntasks=1 \ + bash -lc \ + "command -v ray >/dev/null && '${vllm_pythons[${pool}]}' -c 'import ray, vllm; from nemo_rl.models.generation.vllm.patches import _apply_vllm_patches'" \ + >/dev/null & + preflight_pids+=("$!") + preflight_labels+=("${display_names[${pool}]} container") +done + +echo "[INFO] Validating the load-balancer container and Python environment" +srun \ + --het-group=0 \ + --no-container-mount-home \ + --container-image="${CONTAINER}" \ + --container-mounts="${lb_mounts}" \ + --container-workdir="${SLURM_SUBMIT_DIR}" \ + --mpi=pmix \ + -A "${SLURM_JOB_ACCOUNT}" \ + -p "${SLURM_JOB_PARTITION}" \ + --overlap \ + --kill-on-bad-exit=1 \ + --nodelist="${ray_head_node}" \ + --nodes=1 \ + --ntasks=1 \ + --cpus-per-task=1 \ + bash -lc \ + "test -x /opt/external-vllm-tools/lb_watchdog.sh && '${EXTERNAL_VLLM_LB_PYTHON}' -c 'import aiohttp'" \ + >/dev/null & +preflight_pids+=("$!") +preflight_labels+=("load balancer container") + +preflight_failed=0 +for preflight_index in "${!preflight_pids[@]}"; do + preflight_pid="${preflight_pids[${preflight_index}]}" + if ! wait "${preflight_pid}"; then + echo "[FATAL] External-vLLM preflight failed: ${preflight_labels[${preflight_index}]} (pid=${preflight_pid})" >&2 + preflight_failed=1 + fi +done +preflight_pids=() +preflight_labels=() +if (( preflight_failed != 0 )); then + exit 1 +fi + +for pool in "${pool_names[@]}"; do + echo "[INFO] Launching ${display_names[${pool}]} replicas" + for (( replica_index = 0; replica_index < replicas[${pool}]; replica_index++ )); do + first_node_index=$((node_offsets[${pool}] + replica_index * nodes_per_replica[${pool}])) + replica_node_count="${nodes_per_replica[${pool}]}" + replica_nodes=("${external_nodes[@]:first_node_index:replica_node_count}") + replica_nodelist=$(IFS=,; echo "${replica_nodes[*]}") + replica_id="${SLURM_JOB_ID}-${pool,,}-${replica_index}" + head_ip_file="${pool_log_dirs[${pool}]}/head_ip_${replica_index}" + vllm_log="${pool_log_dirs[${pool}]}/vllm_${replica_index}.log" + + echo "[INFO] ${display_names[${pool}]} replica ${replica_index}: ${replica_nodelist}" + srun \ + --het-group=1 \ + --no-container-mount-home \ + --container-image="${containers[${pool}]}" \ + --container-mounts="${external_service_mount}" \ + --mpi=pmix \ + --gres="gpu:${GPUS_PER_NODE}" \ + --overlap \ + --kill-on-bad-exit=1 \ + --nodelist="${replica_nodelist}" \ + --nodes="${nodes_per_replica[${pool}]}" \ + --ntasks="${nodes_per_replica[${pool}]}" \ + --ntasks-per-node=1 \ + --export="ALL,POOL_PREFIX=${pool},REPLICA_ID=${replica_id},EXTERNAL_VLLM_TOOLS_DIR=${EXTERNAL_VLLM_TOOLS_DIR_HOST},EXTERNAL_VLLM_STATE_DIR=${state_dirs[${pool}]},EXTERNAL_VLLM_GROUP_ID=${group_ids[${pool}]},HEAD_IP_FILE=${head_ip_file},LOG_FILE=${vllm_log}" \ + --output="${pool_log_dirs[${pool}]}/replica_${replica_index}_%t.log" \ + bash -c "${VLLM_SERVER_BODY}" & + service_step_pids+=("$!") + service_step_labels+=("${display_names[${pool}]} replica ${replica_index}") + done +done + +ray_head_ip=$(resolve_node_ip "${ray_head_node}") +for pool in "${pool_names[@]}"; do + pool_urls["${pool}"]="http://${ray_head_ip}:${lb_ports[${pool}]}/v1" + echo "[INFO] Starting ${display_names[${pool}]} load balancer at ${pool_urls[${pool}]}" + srun \ + --het-group=0 \ + --no-container-mount-home \ + --container-name="external-vllm-lb-${pool,,}-${SLURM_JOB_ID}" \ + --container-image="${CONTAINER}" \ + --container-mounts="${lb_mounts}" \ + --container-workdir="${SLURM_SUBMIT_DIR}" \ + --mpi=pmix \ + -A "${SLURM_JOB_ACCOUNT}" \ + -p "${SLURM_JOB_PARTITION}" \ + --overlap \ + --nodelist="${ray_head_node}" \ + --nodes=1 \ + --ntasks=1 \ + --cpus-per-task=2 \ + --output="${pool_log_dirs[${pool}]}/load_balancer.log" \ + bash -lc "PYTHON='${EXTERNAL_VLLM_LB_PYTHON}' /opt/external-vllm-tools/lb_watchdog.sh '${lb_ports[${pool}]}' '${lb_state_dirs[${pool}]}' '${group_ids[${pool}]}'" & + lb_step_pids+=("$!") + lb_step_labels+=("${display_names[${pool}]} load balancer") +done + +deadline=$((SECONDS + max_startup_timeout)) +while true; do + all_ready=1 + for pool in "${pool_names[@]}"; do + if ! ready=$( + EXTERNAL_VLLM_STATE_DIR="${state_dirs[${pool}]}" \ + EXTERNAL_VLLM_TOOLS_DIR="${EXTERNAL_VLLM_TOOLS_DIR_HOST}" \ + EXTERNAL_VLLM_GROUP_ID="${group_ids[${pool}]}" \ + bash -c 'source "${EXTERNAL_VLLM_TOOLS_DIR}/vllm_backend_registry.sh"; registry_count_ready' + ); then + echo "[WARN] Could not read ${display_names[${pool}]} registry; retrying" >&2 + ready=0 + fi + echo "[INFO] ${display_names[${pool}]} ready: ${ready}/${replicas[${pool}]}" + if (( ready != replicas[${pool}] )); then + all_ready=0 + fi + done + (( all_ready == 1 )) && break + check_service_steps + if (( SECONDS >= deadline )); then + echo "[FATAL] Timed out waiting for all external vLLM pools" >&2 + exit 1 + fi + sleep 15 +done + +for pool in "${pool_names[@]}"; do + until curl -sfm 10 "${pool_urls[${pool}]}/models" >/dev/null 2>&1; do + check_service_steps + if (( SECONDS >= deadline )); then + echo "[FATAL] ${display_names[${pool}]} load balancer failed its end-to-end /models probe" >&2 + exit 1 + fi + sleep 5 + done + echo "${pool_urls[${pool}]}" > "${LOG_DIR}/${pool,,}_url" + COMMAND="${COMMAND//${placeholders[${pool}]}/${pool_urls[${pool}]}}" +done +export COMMAND + +echo "[INFO] External vLLM pools are healthy; starting NeMo RL" +# ray.sub predates hetjobs and consumes the unsuffixed allocation variables. +# Restrict those variables to component 0; srun also defaults to hetgroup 0. +# `env` execs bash directly, so ray_sub_pid is the process that owns its traps. +env \ + SLURM_JOB_NODELIST="${SLURM_JOB_NODELIST_HET_GROUP_0}" \ + SLURM_JOB_NUM_NODES="${#ray_nodes[@]}" \ + bash "${RAY_SUB}" & +ray_sub_pid=$! + +while kill -0 "${ray_sub_pid}" 2>/dev/null; do + if ! check_service_steps; then + touch "${LOG_DIR}/ENDED" + kill "${ray_sub_pid}" 2>/dev/null || true + wait "${ray_sub_pid}" 2>/dev/null || true + exit 1 + fi + sleep 5 +done + +set +e +wait "${ray_sub_pid}" +status=$? +set -e +ray_sub_pid="" +exit "${status}" diff --git a/tools/external_gym_vllm/serve_vllm_on_ray.py b/tools/external_gym_vllm/serve_vllm_on_ray.py new file mode 100755 index 00000000000..265a189d9d9 --- /dev/null +++ b/tools/external_gym_vllm/serve_vllm_on_ray.py @@ -0,0 +1,41 @@ +#!/usr/bin/env python3 +# 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. + +"""Run an external Gym vLLM CLI with its Python propagated to Ray workers.""" + +from __future__ import annotations + +import sys + +import ray + +from nemo_rl.models.generation.vllm.patches import _apply_vllm_patches + + +def main() -> None: + """Connect to the private cluster and start the requested vLLM command.""" + worker_python = sys.executable + _apply_vllm_patches(worker_python) + ray.init(address="auto", runtime_env={"py_executable": worker_python}) + + # vLLM is available only in the generation worker environment and must be + # imported after NeMo RL applies its runtime patches. + from vllm.entrypoints.cli.main import main as vllm_main + + vllm_main() + + +if __name__ == "__main__": + main() diff --git a/tools/external_gym_vllm/vllm_backend_registry.sh b/tools/external_gym_vllm/vllm_backend_registry.sh new file mode 100755 index 00000000000..134841ef202 --- /dev/null +++ b/tools/external_gym_vllm/vllm_backend_registry.sh @@ -0,0 +1,64 @@ +#!/bin/bash +# 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. + +# File-backed registry helpers for an in-allocation external Gym vLLM pool. + +set -euo pipefail + +EXTERNAL_VLLM_STATE_DIR="${EXTERNAL_VLLM_STATE_DIR:-$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)}" +EXTERNAL_VLLM_GROUP_ID="${EXTERNAL_VLLM_GROUP_ID:-default}" +REGISTRY_FILE="${EXTERNAL_VLLM_STATE_DIR}/.registry_${EXTERNAL_VLLM_GROUP_ID}" +REGISTRY_LOCK="${REGISTRY_FILE}.lock" + +_ensure_registry() { + touch "${REGISTRY_FILE}" "${REGISTRY_LOCK}" +} + +registry_add() { + local backend_id="$1" ip="$2" port="$3" + _ensure_registry + ( + flock -w 10 200 + grep -v "^${backend_id} " "${REGISTRY_FILE}" > "${REGISTRY_FILE}.tmp" 2>/dev/null || true + echo "${backend_id} ${ip} ${port} $(date +%s) ready" >> "${REGISTRY_FILE}.tmp" + mv "${REGISTRY_FILE}.tmp" "${REGISTRY_FILE}" + ) 200>"${REGISTRY_LOCK}" +} + +registry_remove() { + local backend_id="$1" + _ensure_registry + ( + flock -w 10 200 + grep -v "^${backend_id} " "${REGISTRY_FILE}" > "${REGISTRY_FILE}.tmp" 2>/dev/null || true + mv "${REGISTRY_FILE}.tmp" "${REGISTRY_FILE}" + ) 200>"${REGISTRY_LOCK}" +} + +registry_list() { + _ensure_registry + ( + flock -s -w 10 200 + cat "${REGISTRY_FILE}" + ) 200>"${REGISTRY_LOCK}" +} + +registry_list_ready() { + registry_list | awk '$5 == "ready" { print $2 ":" $3 }' +} + +registry_count_ready() { + registry_list_ready | wc -l +} diff --git a/tools/external_gym_vllm/vllm_pool_lb.py b/tools/external_gym_vllm/vllm_pool_lb.py new file mode 100755 index 00000000000..5f6f8d659fe --- /dev/null +++ b/tools/external_gym_vllm/vllm_pool_lb.py @@ -0,0 +1,592 @@ +#!/usr/bin/env python3 +# 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. + +"""Lightweight async reverse proxy for an external Gym vLLM backend pool. + +Reads the backend registry file and forwards OpenAI-compatible requests +to healthy backends using least-outstanding-requests routing. + +Supports: + - Dynamic backend discovery (re-reads registry every few seconds) + - Health checks (GET /health on each backend) + - Automatic failover when a backend goes down + - Single URL for Gym integration + +Usage: + python vllm_pool_lb.py --port 8080 --registry-dir /path/to/state --group-id default + +The load balancer exposes: + http://:8080/v1/... → proxied to backends + http://:8080/health → LB health + backend summary +""" + +import argparse +import asyncio +import hashlib +import json +import logging +import os +import signal +import sys +import time +from pathlib import Path + +MAX_RSS_MB = 4096 +SHUTDOWN_TIMEOUT_SECONDS = 120 + +print(f"[LB] Python: {sys.executable}", flush=True) +try: + import aiohttp + from aiohttp import web +except Exception as e: + print( + f"ERROR: Failed to import aiohttp: {type(e).__name__}: {e} (python={sys.executable})", + flush=True, + ) + import traceback + + traceback.print_exc() + sys.exit(1) + +logging.basicConfig( + level=logging.INFO, + format="%(asctime)s [%(levelname)s] %(message)s", +) +log = logging.getLogger("vllm_pool_lb") + + +# HTTP statuses that indicate the upstream backend is sick (crashed EngineCore, +# overloaded, etc). We fail over to a different backend and quarantine the sick +# one until the next health probe clears it. +RETRYABLE_UPSTREAM_STATUSES = {500, 502, 503, 504} + + +def _read_current_rss_mb() -> float | None: + """Read the process's current resident memory from procfs.""" + try: + for line in Path("/proc/self/status").read_text().splitlines(): + if line.startswith("VmRSS:"): + return int(line.split()[1]) / 1024 + except (OSError, ValueError): + log.exception("Failed to read current RSS from /proc/self/status") + return None + + +class UpstreamRetryableStatus(Exception): + """Raised from _proxy_once when a non-streaming upstream returned a 5xx we want to retry.""" + + def __init__(self, status: int, body: bytes, headers: dict[str, str]) -> None: + super().__init__(f"upstream status {status}") + self.status = status + self.body = body + self.headers = headers + + +class Backend: + __slots__ = ("job_id", "host", "port", "healthy", "inflight", "last_check") + + def __init__(self, job_id: str, host: str, port: int) -> None: + self.job_id = job_id + self.host = host + self.port = port + self.healthy = True + self.inflight = 0 + self.last_check = 0.0 + + @property + def base_url(self) -> str: + return f"http://{self.host}:{self.port}" + + def __repr__(self) -> str: + status = "UP" if self.healthy else "DOWN" + return f"Backend({self.job_id}, {self.host}:{self.port}, {status}, inflight={self.inflight})" + + +class BackendPool: + """Manages the set of backends, reads registry, performs health checks.""" + + def __init__( + self, registry_dir: str, group_id: str, health_interval: float = 5.0 + ) -> None: + self.registry_file = Path(registry_dir) / f".registry_{group_id}" + self.health_interval = health_interval + self.backends: dict[str, Backend] = {} # job_id -> Backend + self._session: aiohttp.ClientSession | None = None + self._running = False + + async def start(self) -> None: + self._session = aiohttp.ClientSession(timeout=aiohttp.ClientTimeout(total=5)) + self._running = True + asyncio.create_task(self._refresh_loop()) + asyncio.create_task(self._health_check_loop()) + + async def stop(self) -> None: + self._running = False + if self._session: + await self._session.close() + + def _read_registry(self) -> dict[str, tuple[str, int]] | None: + """Read registry file. Returns {job_id: (host, port)}.""" + result: dict[str, tuple[str, int]] = {} + if not self.registry_file.exists(): + return result + try: + lines = self.registry_file.read_text().strip().splitlines() + except (OSError, UnicodeError) as e: + log.warning("Failed to read registry: %s", e) + return None + + for line in lines: + parts = line.split() + if len(parts) >= 5 and parts[4] == "ready": + try: + result[parts[0]] = (parts[1], int(parts[2])) + except ValueError: + log.warning("Skipping malformed registry entry: %s", line) + return result + + async def _refresh_loop(self) -> None: + """Periodically re-read the registry to discover new/removed backends.""" + while self._running: + try: + registered = self._read_registry() + if registered is not None: + # Add new backends + for job_id, (host, port) in registered.items(): + if job_id not in self.backends: + b = Backend(job_id, host, port) + self.backends[job_id] = b + log.info("Discovered new backend: %s", b) + # Remove gone backends + gone = set(self.backends) - set(registered) + for job_id in gone: + b = self.backends.pop(job_id) + log.info("Removed backend: %s", b) + except Exception as e: + log.warning("Refresh error: %s", e) + await asyncio.sleep(self.health_interval) + + async def _check_backend_health(self, backend: Backend) -> None: + session = self._session + if session is None: + return + try: + async with session.get(f"{backend.base_url}/health") as response: + backend.healthy = response.status == 200 + except Exception: + backend.healthy = False + backend.last_check = time.time() + + async def _health_check_loop(self) -> None: + """Periodically check /health on each backend.""" + while self._running: + await asyncio.gather( + *( + self._check_backend_health(backend) + for backend in list(self.backends.values()) + ) + ) + rss_mb = _read_current_rss_mb() + if rss_mb is not None and rss_mb > MAX_RSS_MB: + log.error( + "Current RSS %.0f MB exceeds %d MB; stopping new requests and " + "draining in-flight requests before restart", + rss_mb, + MAX_RSS_MB, + ) + self._running = False + os.kill(os.getpid(), signal.SIGTERM) + return + await asyncio.sleep(self.health_interval) + + def pick( + self, exclude: set[str] | None = None, affinity_key: str | None = None + ) -> Backend | None: + """Pick a healthy backend. + + If affinity_key is set, use consistent hashing to prefer the same backend + for requests with the same prefix (enables vLLM prefix caching). + Falls back to least-outstanding-requests if the preferred backend is + excluded or unhealthy. + """ + exclude = exclude or set() + healthy = [ + b for b in self.backends.values() if b.healthy and b.job_id not in exclude + ] + if not healthy: + return None + + if affinity_key and len(healthy) > 1: + # Consistent hash: sort by hash(affinity_key + job_id) to get a + # stable preference order. Pick the first one (preferred), but if + # it's heavily loaded compared to the least-loaded, fall back. + h = hashlib.md5(affinity_key.encode()).hexdigest() + ranked = sorted( + healthy, key=lambda b: hashlib.md5((h + b.job_id).encode()).hexdigest() + ) + preferred = ranked[0] + least_loaded = min(healthy, key=lambda b: b.inflight) + # Use preferred backend unless it has 2x+ more inflight than the + # least loaded — avoids hotspots when one prefix dominates. + if preferred.inflight <= least_loaded.inflight * 2 + 10: + return preferred + return least_loaded + + return min(healthy, key=lambda b: b.inflight) + + def summary(self) -> list[dict[str, str | bool | int]]: + return [ + { + "job_id": b.job_id, + "url": b.base_url, + "healthy": b.healthy, + "inflight": b.inflight, + } + for b in self.backends.values() + ] + + +class LoadBalancer: + def __init__(self, pool: BackendPool, port: int) -> None: + self.pool = pool + self.port = port + self._proxy_session: aiohttp.ClientSession | None = None + + async def start(self) -> None: + # limit=5000: enough for 26 backends × ~50 concurrent reqs, but capped + # to prevent unbounded memory growth that caused OOM at 31GB with limit=0. + self._proxy_session = aiohttp.ClientSession( + connector=aiohttp.TCPConnector(limit=5000, limit_per_host=150), + timeout=aiohttp.ClientTimeout(total=1800), + ) + await self.pool.start() + + async def stop(self) -> None: + await self.pool.stop() + if self._proxy_session: + await self._proxy_session.close() + + async def handle_health(self, request: web.Request) -> web.Response: + backends = self.pool.summary() + healthy_count = sum(1 for b in backends if b["healthy"]) + return web.json_response( + { + "status": "ok" if healthy_count > 0 else "no_healthy_backends", + "healthy_backends": healthy_count, + "total_backends": len(backends), + "backends": backends, + } + ) + + async def _proxy_once( + self, + backend: Backend, + request_method: str, + path_qs: str, + headers: dict[str, str], + body: bytes, + request: web.Request, + ) -> web.StreamResponse: + """Attempt to proxy a single request to one backend.""" + target_url = f"{backend.base_url}{path_qs}" + backend.inflight += 1 + try: + if self._proxy_session is None: + raise RuntimeError("Load balancer has not been started") + async with self._proxy_session.request( + method=request_method, + url=target_url, + headers=headers, + data=body, + ) as upstream_resp: + content_type = upstream_resp.headers.get("Content-Type", "") + is_streaming = "text/event-stream" in content_type + + if is_streaming: + response = web.StreamResponse( + status=upstream_resp.status, + headers={ + "Content-Type": "text/event-stream", + "Cache-Control": "no-cache", + "Connection": "keep-alive", + }, + ) + await response.prepare(request) + iterator = upstream_resp.content.iter_any().__aiter__() + while True: + try: + chunk = await anext(iterator) + except StopAsyncIteration: + break + except Exception as e: + backend.healthy = False + log.warning( + "Upstream stream from backend %s failed after response " + "headers were sent; returning the partial response without " + "retrying: %s: %s", + backend, + type(e).__name__, + e, + ) + break + try: + await response.write(chunk) + except Exception as e: + log.info( + "Downstream disconnected while streaming from backend %s: " + "%s: %s", + backend, + type(e).__name__, + e, + ) + return response + try: + await response.write_eof() + except Exception as e: + log.info( + "Could not finish downstream stream from backend %s: %s: %s", + backend, + type(e).__name__, + e, + ) + return response + resp_body = await upstream_resp.read() + resp_headers = { + k: v + for k, v in upstream_resp.headers.items() + if k.lower() + not in ("transfer-encoding", "content-encoding", "content-length") + } + if upstream_resp.status in RETRYABLE_UPSTREAM_STATUSES: + raise UpstreamRetryableStatus( + status=upstream_resp.status, + body=resp_body, + headers=resp_headers, + ) + return web.Response( + status=upstream_resp.status, + headers=resp_headers, + body=resp_body, + ) + finally: + backend.inflight -= 1 + raise RuntimeError("Upstream request exited without producing a response") + + @staticmethod + def _extract_affinity_key(body: bytes) -> str | None: + """Extract a prefix-affinity key from the request body. + + For chat completions, hash the messages array (the shared prompt prefix). + This ensures requests with the same prompt go to the same backend, + maximizing vLLM prefix cache hits. + """ + try: + data = json.loads(body) + if not isinstance(data, dict): + return None + messages = data.get("messages") + if messages: + # Hash messages content only (not metadata/response pairs which vary) + return hashlib.md5( + json.dumps(messages, sort_keys=True).encode() + ).hexdigest() + except (json.JSONDecodeError, TypeError, UnicodeError): + pass + return None + + async def handle_proxy(self, request: web.Request) -> web.StreamResponse: + """Forward request to a backend, retrying on a different backend if one fails.""" + body = await request.read() + headers = { + k: v + for k, v in request.headers.items() + if k.lower() not in ("host", "transfer-encoding") + } + + # Extract affinity key for prefix-cache-aware routing + affinity_key = ( + self._extract_affinity_key(body) if request.method == "POST" else None + ) + + tried: set[str] = set() + last_error: Exception | None = None + last_upstream_5xx: UpstreamRetryableStatus | None = None + + # Try up to MAX_RETRIES different healthy backends. `pool.pick` filters + # on `healthy` and excludes anything in `tried`, so the loop also + # terminates early if every non-quarantined backend has been attempted. + # The per-attempt aiohttp ClientTimeout bounds the worst-case latency + # of a wedged backend. + MAX_RETRIES = 5 + for attempt in range(1, MAX_RETRIES + 1): + backend = self.pool.pick(exclude=tried, affinity_key=affinity_key) + if backend is None: + log.warning( + "[proxy %s %s] no more healthy untried backends after %d attempt(s); giving up", + request.method, + request.path_qs, + attempt - 1, + ) + break + + tried.add(backend.job_id) + try: + resp = await self._proxy_once( + backend, + request.method, + request.path_qs, + headers, + body, + request, + ) + if attempt > 1: + log.warning( + "[proxy %s %s] succeeded on attempt %d/%d via backend %s after failing on %s", + request.method, + request.path_qs, + attempt, + MAX_RETRIES, + backend, + sorted(tried - {backend.job_id}), + ) + return resp + except UpstreamRetryableStatus as e: + log.warning( + "[proxy %s %s] attempt %d/%d: backend %s returned %d, " + "failing over without changing backend health. Body: %s", + request.method, + request.path_qs, + attempt, + MAX_RETRIES, + backend, + e.status, + e.body[:500], + ) + last_upstream_5xx = e + except Exception as e: + log.warning( + "[proxy %s %s] attempt %d/%d: backend %s raised %s: %s, " + "quarantining and failing over", + request.method, + request.path_qs, + attempt, + MAX_RETRIES, + backend, + type(e).__name__, + e, + ) + backend.healthy = False + last_error = e + + # All retries exhausted. If the last failure was an upstream 5xx with a real + # response body, forward that body verbatim so callers see the true error. + # Otherwise synthesize a 502/503. + if last_upstream_5xx is not None: + log.error( + "[proxy %s %s] all %d attempt(s) failed; returning last upstream status %d. " + "Tried backends: %s", + request.method, + request.path_qs, + len(tried), + last_upstream_5xx.status, + sorted(tried), + ) + return web.Response( + status=last_upstream_5xx.status, + headers=last_upstream_5xx.headers, + body=last_upstream_5xx.body, + ) + log.error( + "[proxy %s %s] all %d attempt(s) failed with connection errors; " + "last_error=%s. Tried backends: %s", + request.method, + request.path_qs, + len(tried), + last_error, + sorted(tried), + ) + return web.json_response( + {"error": f"All backends failed. Last error: {last_error}"}, + status=502 if last_error else 503, + ) + + def make_app(self) -> web.Application: + # Judge payloads can include many long candidate responses. Let each + # upstream model enforce its own request/model-length limits. + app = web.Application(client_max_size=0) + app.router.add_get("/health", self.handle_health) + # Catch-all: proxy everything else + app.router.add_route("*", "/{path:.*}", self.handle_proxy) + + async def on_startup(_app: web.Application) -> None: + await self.start() + + async def on_cleanup(_app: web.Application) -> None: + await self.stop() + + app.on_startup.append(on_startup) + # Cleanup runs after aiohttp has stopped accepting requests and drained + # active handlers, so the upstream client session remains usable while + # long generations finish. + app.on_cleanup.append(on_cleanup) + return app + + +def main() -> None: + parser = argparse.ArgumentParser(description="External vLLM Pool Load Balancer") + parser.add_argument("--port", type=int, default=8080, help="LB listen port") + parser.add_argument( + "--registry-dir", + default=os.environ.get( + "EXTERNAL_VLLM_STATE_DIR", os.path.dirname(os.path.abspath(__file__)) + ), + help="Directory containing the registry file", + ) + parser.add_argument( + "--group-id", + default=os.environ.get("EXTERNAL_VLLM_GROUP_ID", "default"), + help="Server group ID", + ) + parser.add_argument( + "--health-interval", + type=float, + default=5.0, + help="Seconds between health checks / registry refresh", + ) + args = parser.parse_args() + + pool = BackendPool(args.registry_dir, args.group_id, args.health_interval) + lb = LoadBalancer(pool, args.port) + app = lb.make_app() + + log.info( + "Starting external vLLM load balancer on port %d (group=%s)", + args.port, + args.group_id, + ) + log.info("Registry: %s", pool.registry_file) + + # Stop accepting new connections immediately, but give in-flight requests + # a bounded window to finish before the watchdog restarts the process. + web.run_app( + app, + port=args.port, + print=log.info, + shutdown_timeout=SHUTDOWN_TIMEOUT_SECONDS, + ) + + +if __name__ == "__main__": + main() From d4c0d4d755e160b093288bc7b4ea630c19b0c8a1 Mon Sep 17 00:00:00 2001 From: bg51717 <91998467+bg51717@users.noreply.github.com> Date: Thu, 6 Aug 2026 19:09:14 +0800 Subject: [PATCH 04/10] feat(ppo): support non-colocated generation (#3262) Signed-off-by: bg51717 --- docs/guides/ppo.md | 32 +- examples/configs/ppo_math_1B.yaml | 9 + ....5b-gsm8k-1n8g-automodel-noncolocated.yaml | 59 + ...tron-valuetp2sp-dynbatch-noncolocated.yaml | 83 ++ nemo_rl/algorithms/ppo.py | 345 ++++- nemo_rl/data/datasets/utils.py | 15 +- tests/functional/L1_Functional_Tests_PPO.sh | 2 + .../functional/ppo_megatron_non_colocated.sh | 68 + tests/functional/ppo_non_colocated.sh | 63 + ...-1.5b-gsm8k-1n8g-automodel-noncolocated.sh | 49 + ...gatron-valuetp2sp-dynbatch-noncolocated.sh | 49 + tests/test_suites/nightly.txt | 6 + tests/unit/algorithms/test_ppo.py | 1130 +++++++++++++++++ tests/unit/data/test_utils.py | 40 +- .../ppo_math_1B_megatron.yaml | 4 + 15 files changed, 1895 insertions(+), 59 deletions(-) create mode 100644 examples/configs/recipes/llm/ppo-qwen2.5-1.5b-gsm8k-1n8g-automodel-noncolocated.yaml create mode 100644 examples/configs/recipes/llm/ppo-qwen2.5-1.5b-gsm8k-2n8g-megatron-valuetp2sp-dynbatch-noncolocated.yaml create mode 100755 tests/functional/ppo_megatron_non_colocated.sh create mode 100755 tests/functional/ppo_non_colocated.sh create mode 100755 tests/test_suites/llm/ppo-qwen2.5-1.5b-gsm8k-1n8g-automodel-noncolocated.sh create mode 100755 tests/test_suites/llm/ppo-qwen2.5-1.5b-gsm8k-2n8g-megatron-valuetp2sp-dynbatch-noncolocated.sh diff --git a/docs/guides/ppo.md b/docs/guides/ppo.md index c7cfee08b56..d99396ec1da 100644 --- a/docs/guides/ppo.md +++ b/docs/guides/ppo.md @@ -38,9 +38,24 @@ We define a [ValueInterface](../../nemo_rl/models/value/interfaces.py) that cont The value model supports the **Megatron-Core backend** (`value.megatron_cfg.enabled: true`) and the **DTensor backend** (`value.dtensor_cfg.enabled: true`). It uses the same architecture and tokenizer as the policy (configured via `value.model_name`), but is trained with a separate MSE loss on GAE returns. -### Colocated Architecture +### Deployment Architectures -PPO uses a colocated architecture where the **policy**, **value model**, and **vLLM generation engine** share the same set of GPUs. GPU memory is managed by offloading models to CPU between stages: the value model is loaded to GPU only during its inference and training phases, then offloaded to make room for other components. +By default, PPO uses a colocated architecture where the **policy**, **value model**, and **generation engine** share one `RayVirtualCluster`. GPU memory is managed by offloading models to CPU between stages: the value model is loaded to GPU only during its inference and training phases, then offloaded to make room for the other components. + +PPO also supports non-colocated vLLM generation. In this mode, the policy and value model continue to time-share one training `RayVirtualCluster`, while vLLM runs on a separate inference `RayVirtualCluster` in the same Ray cluster. Updated policy weights are transferred to vLLM through the cross-cluster collective refit path. + +```yaml +policy: + generation: + backend: vllm + colocated: + enabled: false + resources: + gpus_per_node: 2 + num_nodes: null +``` + +When only one node remains for policy and generation after other resources are reserved, `gpus_per_node` reserves that many GPUs for generation and `num_nodes` must be `null` or `1`. When more than one node remains for training and generation, generation uses complete nodes: set `num_nodes` to the number of inference nodes and `gpus_per_node` equal to `cluster.gpus_per_node`. Non-colocated SGLang generation is not currently supported by PPO. ### Value Model Configuration @@ -217,6 +232,8 @@ ppo: seed: 42 use_dynamic_sampling: false overlong_filtering: false + # null logs mismatch metrics without masking; set a threshold to mask sequences. + seq_logprob_error_threshold: null adv_estimator: name: "gae" @@ -257,6 +274,7 @@ value_loss_fn: **PPO-specific parameters:** - **`ppo.ppo_epochs`**: Number of training updates per rollout batch - **`ppo.policy_training_start_step`**: Number of critic-only warmup steps before policy training begins +- **`ppo.seq_logprob_error_threshold`**: Nullable sequence-level multiplicative probability-error threshold. PPO always logs sequence-level train/generation mismatch metrics; when this is set, sequences above the threshold are excluded from advantage and loss computation. - **`ppo.adv_estimator.name`**: Set to `"gae"` for GAE advantage estimation (PPO default) - **`ppo.adv_estimator.gae_lambda`**: GAE $\lambda$ parameter (bias-variance tradeoff, typically 0.95) - **`ppo.adv_estimator.gae_gamma`**: Discount factor $\gamma$ (typically 1.0 for outcome-supervised tasks) @@ -268,7 +286,7 @@ All other parameters (clipping, KL, importance sampling, dynamic sampling, rewar ## Metrics -PPO logs all the same metrics as GRPO (see [GRPO Metrics](grpo.md#metrics)). In addition, the following critic-specific metrics are logged: +PPO logs all the same metrics as GRPO (see [GRPO Metrics](grpo.md#metrics)). It also logs the following PPO-specific metrics: | Metric | Description | |--------|-------------| @@ -279,6 +297,14 @@ PPO logs all the same metrics as GRPO (see [GRPO Metrics](grpo.md#metrics)). In | `critic/values_max` | Maximum predicted value | | `critic/returns_mean` | Mean of GAE returns | | `critic/explained_var` | Explained variance: $1 - \text{Var}(R - V) / \text{Var}(R)$. Higher is better; values near 1.0 indicate the critic accurately predicts returns. | +| `max_seq_mult_prob_error` | Maximum sequence-level multiplicative probability error between generation and training logprobs before optional masking. | +| `mean_seq_mult_prob_error` | Mean sequence-level multiplicative probability error before optional masking. | +| `min_seq_mult_prob_error` | Minimum sequence-level multiplicative probability error before optional masking. | +| `max_seq_mult_prob_error_after_mask` | Maximum sequence-level multiplicative probability error among sequences retained after optional masking. | +| `mean_seq_mult_prob_error_after_mask` | Mean sequence-level multiplicative probability error among sequences retained after optional masking. | +| `min_seq_mult_prob_error_after_mask` | Minimum sequence-level multiplicative probability error among sequences retained after optional masking. | +| `num_masked_seqs_by_logprob_error` | Number of sequences excluded by `ppo.seq_logprob_error_threshold`. | +| `masked_correct_pct` | Fraction of sequences excluded by `ppo.seq_logprob_error_threshold` that received a reward of 1. | ## Evaluate the Trained Model diff --git a/examples/configs/ppo_math_1B.yaml b/examples/configs/ppo_math_1B.yaml index 1d79963534b..1bd173aa132 100644 --- a/examples/configs/ppo_math_1B.yaml +++ b/examples/configs/ppo_math_1B.yaml @@ -47,6 +47,9 @@ ppo: source_max: 1.0 target_min: -1.0 # DAPO: scale rewards to [-1, 1] target_max: 1.0 + # Log train/generation probability mismatch. Set a threshold to mask + # high-error sequences; null keeps metrics-only behavior. + seq_logprob_error_threshold: null loss_fn: disable_ppo_ratio: false @@ -459,3 +462,9 @@ logger: cluster: gpus_per_node: 1 num_nodes: 1 + # Port range for the distributed master address (TCPStore / NCCL rendezvous) + # and per-worker available ports. Kept below the OS ephemeral range + # (32768-60999 on stock Linux). See ray.sub for the full port layout. + master_port_range_low: 1400 + master_port_range_high: 1999 + segment_size: null # Nodes per NVLink domain segment for topology-aware alignment; null to disable diff --git a/examples/configs/recipes/llm/ppo-qwen2.5-1.5b-gsm8k-1n8g-automodel-noncolocated.yaml b/examples/configs/recipes/llm/ppo-qwen2.5-1.5b-gsm8k-1n8g-automodel-noncolocated.yaml new file mode 100644 index 00000000000..c62d2c3e8da --- /dev/null +++ b/examples/configs/recipes/llm/ppo-qwen2.5-1.5b-gsm8k-1n8g-automodel-noncolocated.yaml @@ -0,0 +1,59 @@ +defaults: ../../ppo_math_1B.yaml +ppo: + num_prompts_per_step: 1024 + num_generations_per_prompt: 1 + max_num_epochs: 15 + ppo_epochs: 1 + val_period: 1 + overlong_filtering: true + reward_shaping: + enabled: false + adv_estimator: + gae_lambda_value: 1.0 + gae_lambda_policy: 1 + reward_scaling: + enabled: false +loss_fn: + ratio_clip_max: 0.2 + ratio_clip_c: 3 +value_loss_fn: + scale: 1.0 + cliprange: 0.5 +checkpointing: + checkpoint_dir: results/ppo-qwen2.5-1.5b-gsm8k-1n8g-automodel-noncolocated +policy: + model_name: Qwen/Qwen2.5-1.5B-Instruct + train_global_batch_size: 256 + max_total_sequence_length: 1024 + generation: + max_new_tokens: 512 + vllm_cfg: + gpu_memory_utilization: 0.4 + max_model_len: 1024 + colocated: + enabled: false + resources: + gpus_per_node: 4 +value: + model_name: Qwen/Qwen2.5-1.5B-Instruct + train_micro_batch_size: 4 +data: + max_input_seq_length: 512 + train: + dataset_name: gsm8k + split: train + validation: + dataset_name: gsm8k + split: test + default: + system_prompt_file: examples/prompts/gsm8k.txt +env: + math: + math_verify_impl: hf_math_verify +logger: + log_dir: logs/ppo-qwen2.5-1.5b-gsm8k-1n8g-automodel-noncolocated + wandb: + project: nemo-rl + name: ppo-qwen2.5-1.5b-gsm8k-1n8g-automodel-noncolocated +cluster: + gpus_per_node: 8 diff --git a/examples/configs/recipes/llm/ppo-qwen2.5-1.5b-gsm8k-2n8g-megatron-valuetp2sp-dynbatch-noncolocated.yaml b/examples/configs/recipes/llm/ppo-qwen2.5-1.5b-gsm8k-2n8g-megatron-valuetp2sp-dynbatch-noncolocated.yaml new file mode 100644 index 00000000000..9fc719a2b65 --- /dev/null +++ b/examples/configs/recipes/llm/ppo-qwen2.5-1.5b-gsm8k-2n8g-megatron-valuetp2sp-dynbatch-noncolocated.yaml @@ -0,0 +1,83 @@ +defaults: ../../ppo_math_1B_megatron.yaml +ppo: + num_prompts_per_step: 1024 + num_generations_per_prompt: 1 + max_num_epochs: 15 + ppo_epochs: 1 + val_period: 1 + overlong_filtering: true + reward_shaping: + enabled: false + adv_estimator: + gae_lambda_value: 1.0 + gae_lambda_policy: 1 + reward_scaling: + enabled: false +loss_fn: + ratio_clip_max: 0.2 + ratio_clip_c: 3 +value_loss_fn: + scale: 1.0 + cliprange: 0.5 +checkpointing: + checkpoint_dir: results/ppo-qwen2.5-1.5b-gsm8k-2n8g-megatron-valuetp2sp-dynbatch-noncolocated +policy: + model_name: Qwen/Qwen2.5-1.5B-Instruct + train_global_batch_size: 256 + max_total_sequence_length: 1024 + megatron_cfg: + tensor_model_parallel_size: 2 + context_parallel_size: 2 + sequence_parallel: true + optimizer: + weight_decay: 0.01 + scheduler: + start_weight_decay: 0.01 + end_weight_decay: 0.01 + lr_warmup_iters: 0 + lr_warmup_init: 0 + make_sequence_length_divisible_by: 8 + generation: + max_new_tokens: 512 + colocated: + enabled: false + resources: + gpus_per_node: 8 + num_nodes: 1 + vllm_cfg: + gpu_memory_utilization: 0.4 + max_model_len: 1024 +value: + model_name: Qwen/Qwen2.5-1.5B-Instruct + train_micro_batch_size: 4 + megatron_cfg: + tensor_model_parallel_size: 2 + sequence_parallel: true + optimizer: + lr: 1.0e-05 + weight_decay: 0.01 + scheduler: + lr_warmup_iters: 0 + dynamic_batching: + enabled: true +data: + max_input_seq_length: 512 + train: + dataset_name: gsm8k + split: train + validation: + dataset_name: gsm8k + split: test + default: + system_prompt_file: examples/prompts/gsm8k.txt +env: + math: + math_verify_impl: hf_math_verify +logger: + log_dir: logs/ppo-qwen2.5-1.5b-gsm8k-2n8g-megatron-valuetp2sp-dynbatch-noncolocated + wandb: + project: nemo-rl + name: ppo-qwen2.5-1.5b-gsm8k-2n8g-megatron-valuetp2sp-dynbatch-noncolocated +cluster: + num_nodes: 2 + gpus_per_node: 8 diff --git a/nemo_rl/algorithms/ppo.py b/nemo_rl/algorithms/ppo.py index f12aaabfb2c..30374bbb313 100644 --- a/nemo_rl/algorithms/ppo.py +++ b/nemo_rl/algorithms/ppo.py @@ -18,6 +18,7 @@ from typing import Any, NotRequired, Optional, TypedDict, TypeVar, cast import numpy as np +import ray import torch from pydantic import BaseModel from torchdata.stateful_dataloader import StatefulDataLoader @@ -32,6 +33,7 @@ RewardScalingConfig, _should_use_async_rollouts, _should_use_nemo_gym, + compute_and_apply_seq_logprob_error_masking, extract_initial_prompt_messages, refit_policy_generation, scale_rewards, @@ -56,9 +58,15 @@ batched_message_log_to_flat_message, get_keys_from_message_log, ) -from nemo_rl.data.utils import load_dataloader_state +from nemo_rl.data.utils import extract_necessary_env_names, load_dataloader_state from nemo_rl.distributed.batched_data_dict import BatchedDataDict -from nemo_rl.distributed.virtual_cluster import ClusterConfig, RayVirtualCluster +from nemo_rl.distributed.virtual_cluster import ( + TOPO_RANK_UNKNOWN, + ClusterConfig, + RayVirtualCluster, + get_ray_cluster_topology, + prepare_segment_topology, +) from nemo_rl.environments.interfaces import EnvironmentInterface from nemo_rl.experience.rollouts import ( run_async_multi_turn_rollout, @@ -146,6 +154,9 @@ class PPOConfig(TypedDict): # Value model trains from step 0; policy training is skipped for # total_steps < this value. Default 0 (train from start). policy_training_start_step: NotRequired[int] + # Nullable sequence-level multiplicative probability-error threshold. + # None logs metrics without masking; values above the threshold are excluded. + seq_logprob_error_threshold: float | None class PPOSaveState(TypedDict): @@ -170,6 +181,30 @@ def _default_ppo_save_state() -> PPOSaveState: } +def _apply_ppo_seq_logprob_error_masking( + train_data: BatchedDataDict, + rewards: torch.Tensor, + seq_logprob_error_threshold: float | None, +) -> tuple[torch.Tensor, dict[str, float | int]]: + """Apply optional mismatch masking and return the advantage mask and metrics.""" + metrics = compute_and_apply_seq_logprob_error_masking( + train_data=train_data, + rewards=rewards, + seq_logprob_error_threshold=seq_logprob_error_threshold, + ) + metrics["num_masked_seqs_by_logprob_error"] = metrics.pop("num_masked_seqs") + + advantage_mask = train_data["token_mask"] * train_data["sample_mask"].unsqueeze(-1) + if not advantage_mask.bool().any(): + raise RuntimeError( + "PPO has no valid response tokens after filtering. Check overlong " + "filtering and ppo.seq_logprob_error_threshold to avoid an optimizer " + "step with an empty batch." + ) + + return advantage_mask, metrics + + class PPOLoggerConfig(LoggerConfig): num_val_samples_to_print: int # number of val samples to print to stdout @@ -248,10 +283,10 @@ def setup( ) if refit_transport is not None: raise ValueError( - "Checkpoint-engine refit requires non-colocated generation, but " - "PPO currently requires colocated generation. Non-colocated PPO " - "support is tracked in " - "https://github.com/NVIDIA-NeMo/RL/issues/3275." + f"policy.generation.refit_transport={refit_transport!r} is not yet " + "supported by PPO (colocated or non-colocated); PPO refits over the " + "default collective path. Set policy.generation.refit_transport=null. " + "Tracked in https://github.com/NVIDIA-NeMo/RL/issues/3275." ) if "megatron_cfg" in policy_config and policy_config["megatron_cfg"]["enabled"]: @@ -393,14 +428,16 @@ def setup( # ========================== print("\n▶ Setting up compute cluster...", flush=True) colocated_inference = generation_config["colocated"]["enabled"] - assert colocated_inference, ( - "PPO currently requires colocated generation (vLLM / SGLang sharing GPUs " - "with the policy worker). Set policy.generation.colocated.enabled=true. " - "Non-colocated PPO is not yet supported." - ) - reward_model_enabled = ( - "env_name" in data_config and data_config["env_name"] == "reward_model" - ) + backend = generation_config["backend"] + if not colocated_inference and backend != "vllm": + raise NotImplementedError( + "Non-colocated PPO generation currently supports only vLLM; " + f"got backend={backend!r}. SGLang does not yet implement the " + "cross-cluster collective weight update path." + ) + + reward_model_enabled = "reward_model" in extract_necessary_env_names(data_config) + segment_size = cluster_config.get("segment_size") total_nodes = cluster_config["num_nodes"] if reward_model_enabled: @@ -420,39 +457,206 @@ def setup( f"policy_nodes:{policy_nodes} + rm_nodes:{rm_nodes} = total_nodes:{total_nodes}" ) - if total_nodes == 1: - policy_gpus_per_node = cluster_config["gpus_per_node"] - rm_gpus_per_node - assert policy_gpus_per_node > 0, ( - "policy.generation.colocated.resources.gpus_per_node must be > 0 " - "when cluster.num_nodes = 1, " - f"but got {policy_gpus_per_node}." + if colocated_inference: + if total_nodes == 1: + policy_gpus_per_node = cluster_config["gpus_per_node"] - rm_gpus_per_node + assert policy_gpus_per_node > 0, ( + "policy.generation.colocated.resources.gpus_per_node must be > 0 " + "when cluster.num_nodes = 1, " + f"but got {policy_gpus_per_node}." + ) + else: + policy_gpus_per_node = cluster_config["gpus_per_node"] + + cluster = RayVirtualCluster( + name="ppo_policy_cluster", + bundle_ct_per_node_list=[policy_gpus_per_node] * policy_nodes, + use_gpus=True, + num_gpus_per_node=policy_gpus_per_node, + max_colocated_worker_groups=1 if backend == "megatron" else 3, + ) + train_cluster = cluster + inference_cluster = cluster + print( + f" ✓ Ray cluster for policy initialized with {policy_nodes} nodes", + flush=True, ) else: - policy_gpus_per_node = cluster_config["gpus_per_node"] - - cluster = RayVirtualCluster( - name="grpo_policy_cluster", - bundle_ct_per_node_list=[policy_gpus_per_node] * policy_nodes, - use_gpus=True, - num_gpus_per_node=policy_gpus_per_node, - max_colocated_worker_groups=1 - if generation_config["backend"] == "megatron" - else 3, - ) - train_cluster = cluster - inference_cluster = cluster - print( - f" ✓ Ray cluster for policy initialized with {policy_nodes} nodes", - flush=True, - ) + train_gpus_per_node = cluster_config["gpus_per_node"] + train_nodes = policy_nodes + + inference_resources = generation_config["colocated"]["resources"] + inference_gpus_per_node = inference_resources["gpus_per_node"] + inference_nodes = inference_resources["num_nodes"] + shared_node_inference = policy_nodes == 1 + + if shared_node_inference: + assert ( + inference_gpus_per_node is not None and inference_gpus_per_node > 0 + ), ( + "policy.generation.colocated.resources.gpus_per_node must be explicitly set to a value > 0 " + "when policy_nodes = 1 and inference is non-colocated, " + f"but got {inference_gpus_per_node}." + ) + assert inference_nodes is None or inference_nodes == 1, ( + "policy.generation.colocated.resources.num_nodes must be 1 or set to null " + "when policy_nodes = 1 and inference is non-colocated, " + f"but got {inference_nodes}." + ) + + inference_nodes = 1 + reward_gpus_to_subtract = rm_gpus_per_node if total_nodes == 1 else 0 + train_gpus_per_node -= inference_gpus_per_node + reward_gpus_to_subtract + assert train_gpus_per_node > 0, ( + "Not enough GPUs for PPO training after reserving non-colocated " + "generation resources: " + f"train_gpus_per_node={train_gpus_per_node}, " + f"cluster.gpus_per_node={cluster_config['gpus_per_node']}, " + f"inference_gpus_per_node={inference_gpus_per_node}, " + f"reward_gpus_per_node={reward_gpus_to_subtract}." + ) + else: + assert inference_nodes is not None and inference_nodes > 0, ( + "policy.generation.colocated.resources.num_nodes must be > 0 " + "when cluster.num_nodes > 1 and inference is non-colocated, " + f"but got {inference_nodes}." + ) + assert ( + inference_gpus_per_node is not None + and inference_gpus_per_node == cluster_config["gpus_per_node"] + ), ( + "policy.generation.colocated.resources.gpus_per_node must be explicitly set and equal to cluster.gpus_per_node " + "when cluster.num_nodes > 1 and inference is non-colocated, " + f"but got inference_gpus_per_node={inference_gpus_per_node}, " + f"cluster.gpus_per_node={cluster_config['gpus_per_node']}." + ) + train_nodes -= inference_nodes + + assert train_nodes > 0 and inference_nodes > 0, ( + "Non-colocated PPO requires both training and inference resources, " + f"but got train_nodes={train_nodes}, inference_nodes={inference_nodes}." + ) + assert inference_gpus_per_node is not None + + node_resource_constraints = None + inference_node_resource_constraints = None + inference_segment_size = None + if segment_size is not None: + topology = get_ray_cluster_topology() + num_alive_nodes = len(topology) + required_nodes = ( + train_nodes if shared_node_inference else train_nodes + inference_nodes + ) + assert num_alive_nodes >= required_nodes, ( + "Not enough alive Ray nodes for all PPO roles: " + f"need {required_nodes} " + f"(train={train_nodes}, inference={inference_nodes}, " + f"shared_node={shared_node_inference}), " + f"but only {num_alive_nodes} alive nodes found" + ) + node_resource_constraints, remaining_node_ids, topology = ( + prepare_segment_topology( + segment_size, + train_nodes, + topology=topology, + role="training", + ) + ) + if node_resource_constraints is not None: + training_node_ids = set(topology) - set(remaining_node_ids) + nodes_missing_topo_rank = [ + nid + for nid in training_node_ids + if topology[nid][1] == TOPO_RANK_UNKNOWN + ] + if nodes_missing_topo_rank: + print( + f" ⚠ {len(nodes_missing_topo_rank)} selected training nodes have NVLink domain " + f"info but no topo_rank; intra-domain rank ordering may be suboptimal", + flush=True, + ) + + if generation_config["backend"] == "vllm": + vllm_cfg = generation_config.get("vllm_cfg", {}) + gpus_per_instance = vllm_cfg["tensor_parallel_size"] * vllm_cfg.get( + "pipeline_parallel_size", 1 + ) + elif generation_config["backend"] == "trtllm": + trtllm_cfg = generation_config.get("trtllm_cfg", {}) + gpus_per_instance = trtllm_cfg[ + "tensor_parallel_size" + ] * trtllm_cfg.get("pipeline_parallel_size", 1) + else: + sglang_cfg = generation_config.get("sglang_cfg", {}) + gpus_per_instance = sglang_cfg.get("gpus_per_server", 1) + nodes_per_instance = ( + gpus_per_instance + inference_gpus_per_node - 1 + ) // inference_gpus_per_node + if nodes_per_instance > 1 and inference_nodes % nodes_per_instance == 0: + remaining_topology = { + node_id: topology[node_id] for node_id in remaining_node_ids + } + ( + inference_node_resource_constraints, + _, + _, + ) = prepare_segment_topology( + nodes_per_instance, + inference_nodes, + topology=remaining_topology, + role="inference", + ) + inference_segment_size = nodes_per_instance + elif nodes_per_instance > 1: + print( + f" ⚠ inference_nodes={inference_nodes} is not divisible by " + f"nodes_per_instance={nodes_per_instance}; skipping inference " + "topology constraints", + flush=True, + ) + + train_cluster = RayVirtualCluster( + name="ppo_train_cluster", + bundle_ct_per_node_list=[train_gpus_per_node] * train_nodes, + use_gpus=True, + num_gpus_per_node=train_gpus_per_node, + max_colocated_worker_groups=2, + port_range_low=cluster_config.get("master_port_range_low"), + port_range_high=cluster_config.get("master_port_range_high"), + segment_size=segment_size, + node_resource_constraints=node_resource_constraints, + ) + if node_resource_constraints is not None: + train_cluster.get_placement_groups() + + inference_cluster = RayVirtualCluster( + name="ppo_inference_cluster", + bundle_ct_per_node_list=[inference_gpus_per_node] * inference_nodes, + use_gpus=True, + num_gpus_per_node=inference_gpus_per_node, + max_colocated_worker_groups=1, + port_range_low=cluster_config.get("master_port_range_low"), + port_range_high=cluster_config.get("master_port_range_high"), + segment_size=inference_segment_size, + node_resource_constraints=inference_node_resource_constraints, + ) + if inference_node_resource_constraints is not None: + VllmGeneration.init_cluster_placement_groups( + inference_cluster, generation_config + ) + + print( + " ✓ Separate PPO clusters initialized: " + f"train={train_nodes}x{train_gpus_per_node} GPUs for policy/value, " + f"inference={inference_nodes}x{inference_gpus_per_node} GPUs for vLLM", + flush=True, + ) # ========================== # Training and Inference # ========================== print("\n▶ Setting up model and training...", flush=True) - # vllm model loading prefers clean environment, initialize policy_generation before policy in colocated mode - backend = generation_config["backend"] generation_config["model_name"] = policy_config["model_name"] # Needed for vLLM # Dictionary to store worker initialization timing stats for logging @@ -552,12 +756,12 @@ def initialize_generation_with_policy( worker_init_timing_metrics: Dictionary to store timing metrics Returns: - Tuple of (policy_generation, policy) + Tuple of (policy_generation, policy, value_model) """ - # Initialize generation engine first so it claims its GPU memory - # before policy/value workers are constructed; then policy, then value. - print(" ⚙️ Initializing workers (colocated mode)", flush=True) + mode = "colocated" if colocated_inference else "non-colocated" + print(f" ⚙️ Initializing workers ({mode} mode)", flush=True) + # Policy and value initialize serially because they share training GPUs. policy_generation, generation_time = init_generation_fn() worker_init_timing_metrics[init_time_key] = generation_time @@ -655,6 +859,26 @@ def initialize_generation_with_policy( # during setup to free GPU for value model initialization). policy.prepare_for_training() + if not colocated_inference: + assert policy_generation is not None + t0 = time.perf_counter() + ip, port = train_cluster.get_master_address_and_port() + print( + f"Using ip: {ip}, port: {port} for collective communication", + flush=True, + ) + train_world_size = train_cluster.world_size() + world_size = train_world_size + inference_cluster.world_size() + + futures_train = policy.init_collective( + ip, port, world_size, train_world_size=train_world_size + ) + futures_inference = policy_generation.init_collective( + ip, port, world_size, train_world_size=train_world_size + ) # type: ignore[call-arg] + ray.get(futures_train + futures_inference) + worker_init_timing_metrics["collective_init_time_s"] = time.perf_counter() - t0 + # prepare refit info state_dict_info = policy.prepare_refit_info() if policy_generation is not None: @@ -962,6 +1186,11 @@ def ppo_train( if NEED_REFIT and POLICY_GENERATION_STALE: refit_policy_generation(policy, policy_generation, colocated_inference) + if not colocated_inference: + # Colocated refit offloads policy inside + # `refit_policy_generation`. Do it here so the value + # model can reuse the training GPUs. + policy.offload_to_cpu() POLICY_GENERATION_STALE = False else: policy_generation.prepare_for_generation() @@ -1055,6 +1284,12 @@ def ppo_train( timer=timer, kv_scales=kv_scales_cache if sync_kv_scales else None, ) + if not colocated_inference: + # Colocated refit offloads policy inside + # `refit_policy_generation`. Do it here so the value + # model can reuse the training GPUs. + with timer.time("policy_offload_after_refit"): + policy.offload_to_cpu() POLICY_GENERATION_STALE = False else: if colocated_inference: @@ -1234,6 +1469,17 @@ def ppo_train( policy.finish_inference() + ( + advantage_mask, + seq_logprob_error_metrics, + ) = _apply_ppo_seq_logprob_error_masking( + train_data=train_data, + rewards=rewards, + seq_logprob_error_threshold=master_config.ppo[ + "seq_logprob_error_threshold" + ], + ) + # Build prompt IDs for advantage estimation (groups responses from same prompt). # Use the token-length-based extractor so multi-turn prompts containing # assistant messages still resolve to the original prompt only. @@ -1254,7 +1500,7 @@ def ppo_train( adv_kwargs = dict( prompt_ids=prompt_ids_for_adv, rewards=train_data["rewards"], - mask=train_data["token_mask"], + mask=advantage_mask, reference_logprobs=train_data.get("reference_policy_logprobs"), logprobs=train_data["prev_logprobs"], ) @@ -1357,6 +1603,12 @@ def ppo_train( colocated_inference, kv_scales=kv_scales_cache if sync_kv_scales else None, ) + if not colocated_inference: + # Colocated refit offloads policy inside + # `refit_policy_generation`. Do it here so the value + # model can reuse the training GPUs. + with timer.time("policy_offload_after_refit"): + policy.offload_to_cpu() POLICY_GENERATION_STALE = False else: if colocated_inference: @@ -1381,11 +1633,10 @@ def ppo_train( # Metrics flat_advantages = train_data["advantages"] - flat_token_mask = flat_messages["token_loss_mask"] del flat_messages response_advantages = torch.masked_select( - flat_advantages, flat_token_mask.bool() + flat_advantages, advantage_mask.bool() ) memory_tracker.snapshot_start_of_stage("Metrics", dir()) @@ -1494,6 +1745,7 @@ def ppo_train( metrics.update(rollout_metrics) metrics["generation_logger_metrics"] = generation_logger_metrics + metrics.update(seq_logprob_error_metrics) if "global_valid_toks" in metrics: total_valid_tokens += metrics["global_valid_toks"] @@ -1627,6 +1879,7 @@ def ppo_train( print("\n📊 Training Results:") if train_results is not None: print(f" • Policy Loss: {metrics.get('loss', 'N/A')}") + print(f" • Generation KL Error: {metrics.get('gen_kl_error', 'N/A')}") if value_results is not None: print(f" • Critic Loss: {metrics.get('critic/loss', 'N/A')}") print(f" • Critic Grad Norm: {metrics.get('critic/grad_norm', 'N/A')}") diff --git a/nemo_rl/data/datasets/utils.py b/nemo_rl/data/datasets/utils.py index 4e924d1ffb6..ef9b8fd4c6a 100644 --- a/nemo_rl/data/datasets/utils.py +++ b/nemo_rl/data/datasets/utils.py @@ -301,14 +301,13 @@ def extract_necessary_env_names(data_config: dict) -> list[str]: The necessary environment names. """ necessary_env_names = set() - keys = ["train", "validation", "default"] - for key in keys: - if ( - key in data_config - and data_config[key] is not None - and "env_name" in data_config[key] - ): - necessary_env_names.add(data_config[key]["env_name"]) + for key in ("train", "validation", "default"): + configs = data_config.get(key) + if not isinstance(configs, list): + configs = [configs] + for config in configs: + if isinstance(config, dict) and "env_name" in config: + necessary_env_names.add(config["env_name"]) return list(necessary_env_names) diff --git a/tests/functional/L1_Functional_Tests_PPO.sh b/tests/functional/L1_Functional_Tests_PPO.sh index c398722dbc5..c330c6b07ec 100755 --- a/tests/functional/L1_Functional_Tests_PPO.sh +++ b/tests/functional/L1_Functional_Tests_PPO.sh @@ -36,6 +36,8 @@ run_test() { run_test fast uv run --no-sync bash ./tests/functional/ppo_automodel.sh run_test fast uv run --no-sync bash ./tests/functional/ppo_megatron.sh +run_test fast uv run --no-sync bash ./tests/functional/ppo_non_colocated.sh +run_test fast uv run --no-sync bash ./tests/functional/ppo_megatron_non_colocated.sh cd ${PROJECT_ROOT}/tests if compgen -G ".coverage*" > /dev/null; then diff --git a/tests/functional/ppo_megatron_non_colocated.sh b/tests/functional/ppo_megatron_non_colocated.sh new file mode 100755 index 00000000000..b1d5a8e67da --- /dev/null +++ b/tests/functional/ppo_megatron_non_colocated.sh @@ -0,0 +1,68 @@ +#!/bin/bash + +SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd) +PROJECT_ROOT=$(realpath $SCRIPT_DIR/../..) +# Mark the current repo as safe, since wandb fetches metadata about the repo +git config --global --add safe.directory $PROJECT_ROOT + +set -eou pipefail + +EXP_NAME=$(basename $0 .sh) +EXP_DIR=$SCRIPT_DIR/$EXP_NAME +LOG_DIR=$EXP_DIR/logs +JSON_METRICS=$EXP_DIR/metrics.json +RUN_LOG=$EXP_DIR/run.log +export PYTHONPATH=${PROJECT_ROOT}:${PYTHONPATH:-} + +rm -rf $EXP_DIR $LOG_DIR +mkdir -p $EXP_DIR $LOG_DIR + +cd $PROJECT_ROOT +uv run coverage run -a --data-file=$PROJECT_ROOT/tests/.coverage --source=$PROJECT_ROOT/nemo_rl \ + $PROJECT_ROOT/examples/run_ppo.py \ + --config $PROJECT_ROOT/examples/configs/ppo_math_1B_megatron.yaml \ + policy.model_name=Qwen/Qwen2.5-0.5B \ + value.model_name=Qwen/Qwen2.5-0.5B \ + ppo.num_prompts_per_step=2 \ + ppo.num_generations_per_prompt=4 \ + ppo.ppo_epochs=2 \ + ppo.policy_training_start_step=1 \ + ppo.seq_logprob_error_threshold=1000 \ + ppo.val_at_end=true \ + policy.train_global_batch_size=4 \ + policy.logprob_batch_size=4 \ + policy.train_micro_batch_size=1 \ + value.train_global_batch_size=4 \ + value.train_micro_batch_size=1 \ + policy.generation.colocated.enabled=false \ + policy.generation.colocated.resources.gpus_per_node=1 \ + cluster.gpus_per_node=2 \ + ppo.max_num_steps=2 \ + logger.tensorboard_enabled=true \ + logger.log_dir=$LOG_DIR \ + logger.wandb_enabled=false \ + logger.monitor_gpus=true \ + checkpointing.enabled=false \ + $@ \ + 2>&1 | tee $RUN_LOG + +# Prove that this run exercised the separate-cluster collective path. +grep -q "Separate PPO clusters initialized" $RUN_LOG +grep -q "collective communication" $RUN_LOG + +uv run tests/json_dump_tb_logs.py $LOG_DIR --output_path $JSON_METRICS + +uv run tests/check_metrics.py $JSON_METRICS \ + 'len(data["train/loss"]) == 1' \ + 'len(data["train/critic/loss"]) == 2' \ + 'max(data["train/num_masked_seqs_by_logprob_error"]) == 0' \ + 'max(data["train/token_mult_prob_error"]) < 1.05' \ + 'max(data["train/max_seq_mult_prob_error"]) < 1.1' \ + 'min(data["train/probs_ratio_clamped_min"]) > 0.79' \ + 'max(data["train/probs_ratio_clamped_min"]) < 1.21' \ + 'min(data["train/probs_ratio_clamped_max"]) > 0.79' \ + 'max(data["train/probs_ratio_clamped_max"]) < 1.29' \ + 'max(data["train/critic/loss"]) < 6.0' \ + 'min(data["train/critic/loss"]) >= 0' \ + 'max(data["train/critic/explained_var"]) <= 1.0001' \ + 'max(data["train/critic/grad_norm"]) < 350' diff --git a/tests/functional/ppo_non_colocated.sh b/tests/functional/ppo_non_colocated.sh new file mode 100755 index 00000000000..82fd897c052 --- /dev/null +++ b/tests/functional/ppo_non_colocated.sh @@ -0,0 +1,63 @@ +#!/bin/bash + +SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd) +PROJECT_ROOT=$(realpath $SCRIPT_DIR/../..) +# Mark the current repo as safe, since wandb fetches metadata about the repo +git config --global --add safe.directory $PROJECT_ROOT + +set -eou pipefail + +EXP_NAME=$(basename $0 .sh) +EXP_DIR=$SCRIPT_DIR/$EXP_NAME +LOG_DIR=$EXP_DIR/logs +JSON_METRICS=$EXP_DIR/metrics.json +RUN_LOG=$EXP_DIR/run.log +export PYTHONPATH=${PROJECT_ROOT}:${PYTHONPATH:-} + +rm -rf $EXP_DIR $LOG_DIR +mkdir -p $EXP_DIR $LOG_DIR + +cd $PROJECT_ROOT +uv run coverage run -a --data-file=$PROJECT_ROOT/tests/.coverage --source=$PROJECT_ROOT/nemo_rl \ + $PROJECT_ROOT/examples/run_ppo.py \ + policy.model_name=Qwen/Qwen2.5-0.5B \ + value.model_name=Qwen/Qwen2.5-0.5B \ + ppo.num_prompts_per_step=2 \ + ppo.num_generations_per_prompt=4 \ + ppo.ppo_epochs=2 \ + policy.train_global_batch_size=4 \ + policy.logprob_batch_size=4 \ + policy.train_micro_batch_size=1 \ + policy.generation.colocated.enabled=false \ + policy.generation.colocated.resources.gpus_per_node=1 \ + policy.generation.vllm_cfg.async_engine=false \ + value.train_global_batch_size=4 \ + value.train_micro_batch_size=1 \ + cluster.gpus_per_node=2 \ + ppo.max_num_steps=2 \ + logger.tensorboard_enabled=true \ + logger.log_dir=$LOG_DIR \ + logger.wandb_enabled=false \ + logger.monitor_gpus=true \ + checkpointing.enabled=false \ + $@ \ + 2>&1 | tee $RUN_LOG + +# Prove that this run exercised the separate-cluster collective path. +grep -q "Separate PPO clusters initialized" $RUN_LOG +grep -q "collective communication" $RUN_LOG + +uv run tests/json_dump_tb_logs.py $LOG_DIR --output_path $JSON_METRICS + +uv run tests/check_metrics.py $JSON_METRICS \ + 'len(data["train/loss"]) == 2' \ + 'max(data["train/token_mult_prob_error"]) < 1.05' \ + 'max(data["train/max_seq_mult_prob_error"]) < 1.1' \ + 'min(data["train/probs_ratio_clamped_min"]) > 0.79' \ + 'max(data["train/probs_ratio_clamped_min"]) < 1.21' \ + 'min(data["train/probs_ratio_clamped_max"]) > 0.79' \ + 'max(data["train/probs_ratio_clamped_max"]) < 1.29' \ + 'max(data["train/critic/loss"]) < 6.0' \ + 'min(data["train/critic/loss"]) >= 0' \ + 'max(data["train/critic/explained_var"]) <= 1.0001' \ + 'max(data["train/critic/grad_norm"]) < 350' diff --git a/tests/test_suites/llm/ppo-qwen2.5-1.5b-gsm8k-1n8g-automodel-noncolocated.sh b/tests/test_suites/llm/ppo-qwen2.5-1.5b-gsm8k-1n8g-automodel-noncolocated.sh new file mode 100755 index 00000000000..76243c398fc --- /dev/null +++ b/tests/test_suites/llm/ppo-qwen2.5-1.5b-gsm8k-1n8g-automodel-noncolocated.sh @@ -0,0 +1,49 @@ +#!/bin/bash +SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd) +source $SCRIPT_DIR/common.env + +# ===== BEGIN CONFIG ===== +NUM_NODES=1 +STEPS_PER_RUN=40 +MAX_STEPS=40 +NUM_RUNS=$(( (MAX_STEPS + STEPS_PER_RUN - 1) / STEPS_PER_RUN )) # Round up +NUM_MINUTES=60 +# ===== END CONFIG ===== + +exit_if_max_steps_reached + +# Run the experiment +cd $PROJECT_ROOT +uv run examples/run_ppo.py \ + --config $CONFIG_PATH \ + ppo.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 + +# Prove that this run exercised the separate-cluster collective path. +grep -q "Separate PPO clusters initialized" $RUN_LOG +grep -q "collective communication" $RUN_LOG + +# Convert tensorboard logs to json +uv run tests/json_dump_tb_logs.py $LOG_DIR --output_path $JSON_METRICS + +# Only run metrics if the target step is reached +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.1' \ + 'data["train/token_mult_prob_error"]["40"] < 1.1' \ + 'median(data["train/max_seq_mult_prob_error"]) < 1.2' \ + 'data["train/reward"]["40"] > 0.75' \ + 'data["validation/accuracy"]["40"] > 0.65' + + # Clean up checkpoint directory after successful run to save space. + rm -rf "$CKPT_DIR" +fi diff --git a/tests/test_suites/llm/ppo-qwen2.5-1.5b-gsm8k-2n8g-megatron-valuetp2sp-dynbatch-noncolocated.sh b/tests/test_suites/llm/ppo-qwen2.5-1.5b-gsm8k-2n8g-megatron-valuetp2sp-dynbatch-noncolocated.sh new file mode 100755 index 00000000000..4c89ade3c26 --- /dev/null +++ b/tests/test_suites/llm/ppo-qwen2.5-1.5b-gsm8k-2n8g-megatron-valuetp2sp-dynbatch-noncolocated.sh @@ -0,0 +1,49 @@ +#!/bin/bash +SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd) +source $SCRIPT_DIR/common.env + +# ===== BEGIN CONFIG ===== +NUM_NODES=2 +STEPS_PER_RUN=40 +MAX_STEPS=40 +NUM_RUNS=$(( (MAX_STEPS + STEPS_PER_RUN - 1) / STEPS_PER_RUN )) # Round up +NUM_MINUTES=60 +# ===== END CONFIG ===== + +exit_if_max_steps_reached + +# Run the experiment +cd $PROJECT_ROOT +uv run examples/run_ppo.py \ + --config $CONFIG_PATH \ + ppo.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 + +# Prove that this run exercised the separate-cluster collective path. +grep -q "Separate PPO clusters initialized" $RUN_LOG +grep -q "collective communication" $RUN_LOG + +# Convert tensorboard logs to json +uv run tests/json_dump_tb_logs.py $LOG_DIR --output_path $JSON_METRICS + +# Only run metrics if the target step is reached +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.1' \ + 'data["train/token_mult_prob_error"]["40"] < 1.1' \ + 'median(data["train/max_seq_mult_prob_error"]) < 1.2' \ + 'data["train/reward"]["40"] > 0.75' \ + 'data["validation/accuracy"]["40"] > 0.65' + + # Clean up checkpoint directory after successful run to save space. + rm -rf "$CKPT_DIR" +fi diff --git a/tests/test_suites/nightly.txt b/tests/test_suites/nightly.txt index dea2565dbcc..b8280bd740b 100644 --- a/tests/test_suites/nightly.txt +++ b/tests/test_suites/nightly.txt @@ -276,9 +276,15 @@ tests/test_suites/llm/distillation-xtoken-off-policy-multiteacher-qwen3-4b-llama # DTensor PPO with value sequence parallelism (TP2+SP) (Qwen2.5-1.5B, GSM8K) tests/test_suites/llm/ppo-qwen2.5-1.5b-gsm8k-1n8g-automodel-valuetp2sp.sh +# Non-colocated AutoModel PPO with vLLM on a dedicated GPU split (Qwen2.5-1.5B, GSM8K) +tests/test_suites/llm/ppo-qwen2.5-1.5b-gsm8k-1n8g-automodel-noncolocated.sh + # Megatron PPO with value sequence parallelism (TP2+SP) and dynamic batching (Qwen2.5-1.5B, GSM8K) tests/test_suites/llm/ppo-qwen2.5-1.5b-gsm8k-1n8g-megatron-valuetp2sp-dynbatch.sh +# Non-colocated Megatron PPO with dedicated training and vLLM nodes (Qwen2.5-1.5B, GSM8K) +tests/test_suites/llm/ppo-qwen2.5-1.5b-gsm8k-2n8g-megatron-valuetp2sp-dynbatch-noncolocated.sh + # Megatron PPO with value SP + pipeline + context parallelism + sequence packing (TP2+SP+PP2+CP2) (Qwen2.5-1.5B, GSM8K) tests/test_suites/llm/ppo-qwen2.5-1.5b-gsm8k-1n8g-megatron-valuetp2sp-pp2cp2-pack.sh diff --git a/tests/unit/algorithms/test_ppo.py b/tests/unit/algorithms/test_ppo.py index 5682f574bfd..0d6a9c9388b 100644 --- a/tests/unit/algorithms/test_ppo.py +++ b/tests/unit/algorithms/test_ppo.py @@ -12,6 +12,10 @@ # See the License for the specific language governing permissions and # limitations under the License. +from contextlib import nullcontext +from types import SimpleNamespace +from unittest.mock import MagicMock + import pytest import torch @@ -24,6 +28,8 @@ MseValueLossConfig, MseValueLossFn, ) +from nemo_rl.algorithms.reward_functions import RewardShapingConfig +from nemo_rl.data import DataConfig from nemo_rl.distributed.batched_data_dict import BatchedDataDict @@ -693,3 +699,1127 @@ def test_create_advantage_estimator_requires_adv_estimator_key(): with pytest.raises(KeyError): _create_advantage_estimator(master_config) + + +def _make_ppo_loop_batch( + truncated_samples: tuple[bool, bool] = (False, False), +) -> BatchedDataDict: + return BatchedDataDict( + { + "message_log": [ + [ + { + "role": "user", + "content": "prompt-0", + "token_ids": torch.tensor([1]), + }, + { + "role": "assistant", + "content": "answer-0", + "token_ids": torch.tensor([2, 3]), + }, + ], + [ + { + "role": "user", + "content": "prompt-1", + "token_ids": torch.tensor([1]), + }, + { + "role": "assistant", + "content": "answer-1", + "token_ids": torch.tensor([4, 5]), + }, + ], + ], + "total_reward": torch.tensor([0.0, 1.0]), + "loss_multiplier": torch.ones(2), + "truncated": torch.tensor(truncated_samples, dtype=torch.bool), + "length": torch.ones(2, dtype=torch.int32), + } + ) + + +def _run_mock_ppo_train( + monkeypatch, + *, + max_num_steps: int, + ppo_epochs: int, + seq_logprob_error_threshold: float | None, + policy_training_start_step: int = 0, + overlong_filtering: bool = False, + truncated_samples: tuple[bool, bool] = (False, False), +): + """Run the real PPO loop with deterministic in-process collaborators.""" + from nemo_rl.algorithms import ppo as ppo_mod + + events: list[str] = [] + generation_logprobs = torch.tensor( + [ + [0.0, 0.0, 0.0], + [0.0, 0.6931472, 0.6931472], + ] + ) + + def fake_flatten(message_logs, *_args, **_kwargs): + batch_size = len(message_logs) + return ( + BatchedDataDict( + { + "token_ids": torch.tensor([[1, 2, 3], [1, 4, 5]])[:batch_size], + "generation_logprobs": generation_logprobs[:batch_size].clone(), + "token_loss_mask": torch.tensor([[0.0, 1.0, 1.0], [0.0, 1.0, 1.0]])[ + :batch_size + ], + "content": [["prompt", "answer"] for _ in range(batch_size)], + } + ), + torch.full((batch_size,), 3, dtype=torch.int32), + ) + + class DummyAdvantageEstimator: + def __init__(self): + self.masks = [] + + def compute_advantage(self, **kwargs): + mask = kwargs["mask"].clone() + self.masks.append(mask) + return mask.clone(), mask.clone() + + class DummyTimer: + def time(self, *_args, **_kwargs): + return nullcontext() + + def get_timing_metrics(self, **_kwargs): + return {"total_step_time": 1.0} + + def reset(self): + pass + + class DummyTimeoutChecker: + def __init__(self, *_args, **_kwargs): + pass + + def start_iterations(self): + pass + + def mark_iteration(self): + pass + + def check_save(self): + return False + + class DummyMemoryTracker: + def snapshot_start_of_stage(self, *_args, **_kwargs): + pass + + class DummyLoader: + def __init__(self, batches): + self.batches = batches + + def __iter__(self): + return iter(self.batches) + + def __len__(self): + return len(self.batches) + + train_result = { + "loss": torch.tensor([0.1]), + "grad_norm": torch.tensor([1.0]), + "all_mb_metrics": {}, + } + value_result = { + "loss": torch.tensor([0.2]), + "grad_norm": torch.tensor([2.0]), + "all_mb_metrics": {}, + } + + policy = MagicMock() + policy.prepare_for_lp_inference.side_effect = lambda: events.append("policy_lp") + policy.offload_to_cpu.side_effect = lambda: events.append("policy_offload") + policy.prepare_for_training.side_effect = lambda: events.append("policy_train_prep") + policy.train.side_effect = lambda *_args, **_kwargs: ( + events.append("policy_train") or train_result + ) + policy.get_logprobs.return_value = {"logprobs": torch.zeros(2, 3)} + policy.get_reference_policy_logprobs.return_value = { + "reference_logprobs": torch.zeros(2, 3) + } + + value_model = MagicMock() + value_model.finish_training.side_effect = lambda: events.append("value_finish") + value_model.get_values.return_value = {"values": torch.zeros(2, 3, 1)} + value_model.train.side_effect = lambda *_args, **_kwargs: ( + events.append("value_train") or value_result + ) + + policy_generation = MagicMock() + policy_generation.requires_kv_scale_sync = False + policy_generation.prepare_for_generation.side_effect = lambda: events.append( + "generation_prepare" + ) + policy_generation.get_logger_metrics.return_value = {} + policy_generation.get_step_metrics.return_value = {} + + def fake_rollout(*_args, input_batch, **_kwargs): + events.append("rollout") + return input_batch, {"mean_gen_tokens_per_sample": 2.0} + + refit = MagicMock(side_effect=lambda *_args, **_kwargs: events.append("refit")) + advantage_estimator = DummyAdvantageEstimator() + + monkeypatch.setattr(ppo_mod, "Timer", DummyTimer) + monkeypatch.setattr(ppo_mod, "TimeoutChecker", DummyTimeoutChecker) + monkeypatch.setattr(ppo_mod, "MemoryTracker", DummyMemoryTracker) + monkeypatch.setattr(ppo_mod, "maybe_gpu_profile_step", lambda *_args: None) + monkeypatch.setattr(ppo_mod, "print_performance_metrics", lambda *_args: {}) + monkeypatch.setattr(ppo_mod, "scale_rewards", lambda batch, _config: batch) + monkeypatch.setattr(ppo_mod, "_should_use_nemo_gym", lambda _config: False) + monkeypatch.setattr(ppo_mod, "_should_use_async_rollouts", lambda _config: False) + monkeypatch.setattr(ppo_mod, "run_multi_turn_rollout", fake_rollout) + monkeypatch.setattr(ppo_mod, "refit_policy_generation", refit) + monkeypatch.setattr(ppo_mod, "batched_message_log_to_flat_message", fake_flatten) + monkeypatch.setattr( + ppo_mod, + "extract_initial_prompt_messages", + lambda message_logs, _lengths: message_logs, + ) + monkeypatch.setattr( + ppo_mod, + "_create_advantage_estimator", + lambda _config: advantage_estimator, + ) + + master_config = SimpleNamespace( + ppo={ + "max_num_steps": max_num_steps, + "max_num_epochs": 1, + "max_rollout_turns": 1, + "num_prompts_per_step": 2, + "num_generations_per_prompt": 1, + "overlong_filtering": overlong_filtering, + "policy_training_start_step": policy_training_start_step, + "ppo_epochs": ppo_epochs, + "reward_scaling": {"enabled": False}, + "reward_shaping": RewardShapingConfig(enabled=False), + "seq_logprob_error_threshold": seq_logprob_error_threshold, + "val_at_start": False, + "val_at_end": False, + "val_period": 0, + }, + policy={ + "generation": { + "backend": "vllm", + "colocated": {"enabled": False}, + "vllm_cfg": {"async_engine": False}, + }, + "max_total_sequence_length": 3, + "make_sequence_length_divisible_by": 1, + }, + loss_fn=_make_loss_config(), + checkpointing={ + "enabled": False, + "checkpoint_must_save_by": None, + "save_period": 100, + "metric_name": None, + }, + cluster={"num_nodes": 1, "gpus_per_node": 2}, + ) + + logger = MagicMock() + checkpointer = MagicMock() + checkpointer.save_optimizer = False + dataloader = DummyLoader( + [_make_ppo_loop_batch(truncated_samples) for _ in range(max_num_steps)] + ) + tokenizer = SimpleNamespace(pad_token_id=0) + + ppo_mod.ppo_train( + policy, + policy_generation, + value_model, + dataloader, + None, + tokenizer, + MagicMock(), + MagicMock(), + {}, + None, + logger, + checkpointer, + ppo_mod._default_ppo_save_state(), + master_config, + ) + + return SimpleNamespace( + policy=policy, + policy_generation=policy_generation, + value_model=value_model, + logger=logger, + checkpointer=checkpointer, + advantage_estimator=advantage_estimator, + refit=refit, + events=events, + ) + + +def test_ppo_train_noncolocated_refit_offload_lifecycle(monkeypatch): + harness = _run_mock_ppo_train( + monkeypatch, + max_num_steps=2, + ppo_epochs=2, + seq_logprob_error_threshold=None, + ) + + assert harness.refit.call_count == 2 + assert harness.policy.train.call_count == 4 + assert harness.value_model.train.call_count == 4 + assert harness.policy.offload_to_cpu.call_count == 4 + harness.policy_generation.prepare_for_generation.assert_not_called() + assert harness.policy_generation.finish_generation.call_count == 2 + + for call in harness.refit.call_args_list: + assert call.args[0] is harness.policy + assert call.args[1] is harness.policy_generation + assert call.args[2] is False + + refit_indices = [ + index for index, event in enumerate(harness.events) if event == "refit" + ] + for index in refit_indices: + assert harness.events[index - 2 : index + 3] == [ + "value_finish", + "policy_lp", + "refit", + "policy_offload", + "rollout", + ] + + +def test_ppo_train_critic_warmup_reuses_generation_until_policy_update(monkeypatch): + harness = _run_mock_ppo_train( + monkeypatch, + max_num_steps=2, + ppo_epochs=1, + seq_logprob_error_threshold=None, + policy_training_start_step=1, + ) + + assert harness.refit.call_count == 1 + assert harness.policy_generation.prepare_for_generation.call_count == 1 + assert harness.policy.train.call_count == 1 + assert harness.value_model.train.call_count == 2 + assert harness.policy_generation.finish_generation.call_count == 2 + + rollout_indices = [ + index for index, event in enumerate(harness.events) if event == "rollout" + ] + assert harness.events[rollout_indices[1] - 1 : rollout_indices[1] + 1] == [ + "generation_prepare", + "rollout", + ] + + +def test_ppo_train_excludes_overlong_samples_from_advantage(monkeypatch): + harness = _run_mock_ppo_train( + monkeypatch, + max_num_steps=1, + ppo_epochs=1, + seq_logprob_error_threshold=None, + overlong_filtering=True, + truncated_samples=(False, True), + ) + + expected_advantage_mask = torch.tensor([[0.0, 1.0, 1.0], [0.0, 0.0, 0.0]]) + torch.testing.assert_close( + harness.advantage_estimator.masks[0], expected_advantage_mask + ) + + +def test_ppo_train_rejects_all_masked_batch(monkeypatch): + with pytest.raises( + RuntimeError, + match="no valid response tokens after filtering", + ): + _run_mock_ppo_train( + monkeypatch, + max_num_steps=1, + ppo_epochs=1, + seq_logprob_error_threshold=None, + overlong_filtering=True, + truncated_samples=(True, True), + ) + + +def test_ppo_train_wires_logprob_mask_to_advantage_training_and_metrics(monkeypatch): + harness = _run_mock_ppo_train( + monkeypatch, + max_num_steps=1, + ppo_epochs=1, + seq_logprob_error_threshold=1.5, + ) + + expected_sample_mask = torch.tensor([1.0, 0.0]) + expected_advantage_mask = torch.tensor([[0.0, 1.0, 1.0], [0.0, 0.0, 0.0]]) + torch.testing.assert_close( + harness.advantage_estimator.masks[0], expected_advantage_mask + ) + + policy_train_data = harness.policy.train.call_args.args[0] + value_train_data = harness.value_model.train.call_args.args[0] + torch.testing.assert_close(policy_train_data["sample_mask"], expected_sample_mask) + torch.testing.assert_close(value_train_data["sample_mask"], expected_sample_mask) + + final_train_metrics = [ + call.args[0] + for call in harness.logger.log_metrics.call_args_list + if call.kwargs.get("prefix") == "train" + and "num_masked_seqs_by_logprob_error" in call.args[0] + ] + assert len(final_train_metrics) == 1 + assert final_train_metrics[0]["num_masked_seqs_by_logprob_error"] == 1 + assert final_train_metrics[0]["max_seq_mult_prob_error"] == pytest.approx(2.0) + assert final_train_metrics[0]["advantages/mean"] == pytest.approx(1.0) + assert final_train_metrics[0]["advantages/min"] == pytest.approx(1.0) + assert final_train_metrics[0]["advantages/max"] == pytest.approx(1.0) + + +# ============================================================================ +# Tests for non-colocated setup +# ============================================================================ + + +def _make_noncolocated_setup_config( + *, + backend: str = "vllm", + total_nodes: int = 1, + total_gpus_per_node: int = 8, + inference_nodes: int | None = None, + inference_gpus_per_node: int | None = 2, + reward_model_gpus_per_node: int | None = None, + segment_size: int | None = None, + tensor_parallel_size: int = 1, +): + """Build the minimal config needed to exercise PPO cluster setup.""" + from nemo_rl.algorithms.ppo import MasterConfig + + data_config: DataConfig = { + "max_input_seq_length": 1, + "shuffle": False, + "num_workers": 0, + "train": {"dataset_name": "fake-dataset"}, + } + env_config = {} + if reward_model_gpus_per_node is not None: + data_config["train"] = [ + {"dataset_name": "fake-dataset", "env_name": "reward_model"} + ] + env_config["reward_model"] = { + "resources": { + "num_nodes": 1, + "gpus_per_node": reward_model_gpus_per_node, + } + } + + return MasterConfig.model_construct( + policy={ + "model_name": "fake-model", + "train_global_batch_size": 1, + "train_micro_batch_size": 1, + "dtensor_cfg": {"enabled": True}, + "megatron_cfg": {"enabled": False}, + "generation": { + "backend": backend, + "colocated": { + "enabled": False, + "resources": { + "num_nodes": inference_nodes, + "gpus_per_node": inference_gpus_per_node, + }, + }, + "vllm_cfg": { + "precision": "bf16", + "kv_cache_dtype": "auto", + "tensor_parallel_size": tensor_parallel_size, + "pipeline_parallel_size": 1, + }, + "vllm_kwargs": {}, + "sglang_cfg": {}, + }, + }, + value={ + "megatron_cfg": { + "enabled": True, + "context_parallel_size": 1, + }, + "sequence_packing": {"enabled": False}, + }, + loss_fn=ClippedPGLossConfig(), + value_loss_fn=MseValueLossConfig(), + env=env_config, + data=data_config, + ppo={ + "max_num_steps": 1, + "max_num_epochs": 1, + "num_prompts_per_step": 1, + "num_generations_per_prompt": 1, + "max_rollout_turns": 1, + "val_period": 0, + "val_batch_size": 1, + "val_at_start": False, + "val_at_end": False, + "max_val_samples": 1, + "seed": 42, + "overlong_filtering": False, + "use_dynamic_sampling": False, + "batch_multiplier": 1, + "ppo_epochs": 1, + "policy_training_start_step": 0, + "reward_shaping": {"enabled": False}, + "reward_scaling": {"enabled": False}, + "adv_estimator": {"name": "raw_reward"}, + }, + logger={"num_val_samples_to_print": 0}, + cluster={ + "num_nodes": total_nodes, + "gpus_per_node": total_gpus_per_node, + "segment_size": segment_size, + }, + checkpointing={ + "enabled": False, + "save_optimizer": False, + }, + ) + + +def _patch_ppo_setup_prerequisites(monkeypatch): + """Replace setup dependencies that are unrelated to resource validation.""" + from nemo_rl.algorithms import ppo as ppo_mod + + class DummyLogger: + def log_hyperparams(self, *_args, **_kwargs): + pass + + def log_metrics(self, *_args, **_kwargs): + pass + + class DummyCheckpointer: + def get_latest_checkpoint_path(self): + return None + + def load_training_info(self, _path): + return None + + def get_resume_paths(self, _path, *, model_component="policy"): + return None, None + + class DummyLoader: + def __init__(self, *_args, **_kwargs): + pass + + def __len__(self): + return 1 + + monkeypatch.setattr(ppo_mod, "Logger", lambda *_args, **_kwargs: DummyLogger()) + monkeypatch.setattr( + ppo_mod, + "CheckpointManager", + lambda *_args, **_kwargs: DummyCheckpointer(), + ) + monkeypatch.setattr(ppo_mod, "StatefulDataLoader", DummyLoader) + return ppo_mod + + +def _setup_dataset(): + from unittest.mock import MagicMock + + dataset = MagicMock() + dataset.__len__.return_value = 1 + return dataset + + +def _run_noncolocated_setup(monkeypatch, config): + """Run setup with lightweight workers and return the observable topology.""" + from unittest.mock import MagicMock + + ppo_mod = _patch_ppo_setup_prerequisites(monkeypatch) + cluster_calls = [] + + class DummyCluster: + def __init__(self, **kwargs): + self.kwargs = kwargs + self.get_placement_groups_called = False + cluster_calls.append(self) + + def world_size(self): + return sum(self.kwargs["bundle_ct_per_node_list"]) + + def get_master_address_and_port(self): + return "127.0.0.1", 1234 + + def get_placement_groups(self): + self.get_placement_groups_called = True + return [] + + policy = MagicMock() + policy.prepare_refit_info.return_value = {"state": "dict"} + policy.init_collective.return_value = ["policy-future"] + value_model = MagicMock() + generation = MagicMock() + generation.init_collective.return_value = ["generation-future"] + policy_factory = MagicMock(return_value=policy) + value_factory = MagicMock(return_value=value_model) + generation_factory = MagicMock(return_value=generation) + ray_get = MagicMock(side_effect=lambda futures: futures) + + monkeypatch.setattr(ppo_mod, "RayVirtualCluster", DummyCluster) + monkeypatch.setattr(ppo_mod, "Policy", policy_factory) + monkeypatch.setattr(ppo_mod, "Value", value_factory) + monkeypatch.setattr(ppo_mod, "VllmGeneration", generation_factory) + monkeypatch.setattr(ppo_mod.ray, "get", ray_get) + + result = ppo_mod.setup(config, MagicMock(), _setup_dataset(), None) + return ( + result, + cluster_calls, + policy, + generation, + policy_factory, + value_factory, + generation_factory, + ray_get, + ) + + +@pytest.mark.parametrize( + ("refit_transport", "error_match"), + [ + ( + "vllm_s3_sparse", + "Remote sparse refit is currently supported only by GRPO", + ), + ( + "nixl", + "PPO refits over the default collective path.*Set " + "policy.generation.refit_transport=null", + ), + ], +) +def test_ppo_rejects_explicit_vllm_refit_transport_before_cluster_creation( + monkeypatch, + refit_transport, + error_match, +): + """PPO uses its default collective refit path for both cluster layouts.""" + from unittest.mock import MagicMock + + ppo_mod = _patch_ppo_setup_prerequisites(monkeypatch) + cluster_cls = MagicMock() + monkeypatch.setattr(ppo_mod, "RayVirtualCluster", cluster_cls) + config = _make_noncolocated_setup_config() + config.policy["generation"]["refit_transport"] = refit_transport + + with pytest.raises(ValueError, match=error_match): + ppo_mod.setup(config, MagicMock(), _setup_dataset(), None) + + cluster_cls.assert_not_called() + + +def test_noncolocated_sglang_is_rejected_before_cluster_creation(monkeypatch): + """SGLang has no cross-cluster refit path, so setup must reject it early.""" + from unittest.mock import MagicMock + + ppo_mod = _patch_ppo_setup_prerequisites(monkeypatch) + cluster_cls = MagicMock() + monkeypatch.setattr(ppo_mod, "RayVirtualCluster", cluster_cls) + config = _make_noncolocated_setup_config(backend="sglang") + + with pytest.raises( + NotImplementedError, + match="Non-colocated PPO generation currently supports only vLLM", + ): + ppo_mod.setup(config, MagicMock(), _setup_dataset(), None) + + cluster_cls.assert_not_called() + + +def test_noncolocated_inference_must_leave_training_gpus_single_node(monkeypatch): + """A single-node split must reserve at least one GPU for policy and value.""" + from unittest.mock import MagicMock + + ppo_mod = _patch_ppo_setup_prerequisites(monkeypatch) + cluster_cls = MagicMock() + monkeypatch.setattr(ppo_mod, "RayVirtualCluster", cluster_cls) + config = _make_noncolocated_setup_config( + total_gpus_per_node=8, + inference_gpus_per_node=8, + ) + + with pytest.raises( + AssertionError, + match="Not enough GPUs for PPO training after reserving", + ): + ppo_mod.setup(config, MagicMock(), _setup_dataset(), None) + + cluster_cls.assert_not_called() + + +def test_noncolocated_inference_requires_explicit_gpus_per_node_single_node( + monkeypatch, +): + """A single-node split cannot infer how many GPUs belong to rollout.""" + from unittest.mock import MagicMock + + ppo_mod = _patch_ppo_setup_prerequisites(monkeypatch) + cluster_cls = MagicMock() + monkeypatch.setattr(ppo_mod, "RayVirtualCluster", cluster_cls) + config = _make_noncolocated_setup_config(inference_gpus_per_node=None) + + with pytest.raises( + AssertionError, + match=( + "policy.generation.colocated.resources.gpus_per_node must be explicitly set" + ), + ): + ppo_mod.setup(config, MagicMock(), _setup_dataset(), None) + + cluster_cls.assert_not_called() + + +def test_noncolocated_inference_rejects_multiple_nodes_single_node(monkeypatch): + """A single-node split cannot allocate multiple rollout nodes.""" + from unittest.mock import MagicMock + + ppo_mod = _patch_ppo_setup_prerequisites(monkeypatch) + cluster_cls = MagicMock() + monkeypatch.setattr(ppo_mod, "RayVirtualCluster", cluster_cls) + config = _make_noncolocated_setup_config( + total_nodes=1, + total_gpus_per_node=8, + inference_nodes=2, + inference_gpus_per_node=2, + ) + + with pytest.raises( + AssertionError, + match="policy.generation.colocated.resources.num_nodes must be 1 or set to null", + ): + ppo_mod.setup(config, MagicMock(), _setup_dataset(), None) + + cluster_cls.assert_not_called() + + +def test_noncolocated_inference_requires_explicit_num_nodes_multi_node(monkeypatch): + """A multi-node split must state how many full nodes belong to rollout.""" + from unittest.mock import MagicMock + + ppo_mod = _patch_ppo_setup_prerequisites(monkeypatch) + cluster_cls = MagicMock() + monkeypatch.setattr(ppo_mod, "RayVirtualCluster", cluster_cls) + config = _make_noncolocated_setup_config( + total_nodes=2, + inference_nodes=None, + inference_gpus_per_node=8, + ) + + with pytest.raises( + AssertionError, + match=( + "policy.generation.colocated.resources.num_nodes must be > 0 " + "when cluster.num_nodes > 1" + ), + ): + ppo_mod.setup(config, MagicMock(), _setup_dataset(), None) + + cluster_cls.assert_not_called() + + +def test_noncolocated_inference_requires_explicit_gpus_per_node_multi_node( + monkeypatch, +): + """A multi-node split requires full-node GPU allocation for rollout.""" + from unittest.mock import MagicMock + + ppo_mod = _patch_ppo_setup_prerequisites(monkeypatch) + cluster_cls = MagicMock() + monkeypatch.setattr(ppo_mod, "RayVirtualCluster", cluster_cls) + config = _make_noncolocated_setup_config( + total_nodes=2, + inference_nodes=1, + inference_gpus_per_node=None, + ) + + with pytest.raises( + AssertionError, + match=( + "policy.generation.colocated.resources.gpus_per_node must be " + "explicitly set and equal to cluster.gpus_per_node" + ), + ): + ppo_mod.setup(config, MagicMock(), _setup_dataset(), None) + + cluster_cls.assert_not_called() + + +def test_noncolocated_inference_requires_full_node_gpus_multi_node(monkeypatch): + """A multi-node rollout allocation cannot reserve a partial physical node.""" + from unittest.mock import MagicMock + + ppo_mod = _patch_ppo_setup_prerequisites(monkeypatch) + cluster_cls = MagicMock() + monkeypatch.setattr(ppo_mod, "RayVirtualCluster", cluster_cls) + config = _make_noncolocated_setup_config( + total_nodes=2, + total_gpus_per_node=8, + inference_nodes=1, + inference_gpus_per_node=4, + ) + + with pytest.raises( + AssertionError, + match=( + "policy.generation.colocated.resources.gpus_per_node must be " + "explicitly set and equal to cluster.gpus_per_node" + ), + ): + ppo_mod.setup(config, MagicMock(), _setup_dataset(), None) + + cluster_cls.assert_not_called() + + +def test_noncolocated_inference_must_leave_training_nodes_multi_node(monkeypatch): + """A multi-node split must leave at least one node for policy and value.""" + from unittest.mock import MagicMock + + ppo_mod = _patch_ppo_setup_prerequisites(monkeypatch) + cluster_cls = MagicMock() + monkeypatch.setattr(ppo_mod, "RayVirtualCluster", cluster_cls) + config = _make_noncolocated_setup_config( + total_nodes=2, + total_gpus_per_node=8, + inference_nodes=2, + inference_gpus_per_node=8, + ) + + with pytest.raises( + AssertionError, + match="Non-colocated PPO requires both training and inference resources", + ): + ppo_mod.setup(config, MagicMock(), _setup_dataset(), None) + + cluster_cls.assert_not_called() + + +def test_noncolocated_topology_requires_enough_alive_nodes(monkeypatch): + """Topology placement fails before creating partially schedulable clusters.""" + from unittest.mock import MagicMock + + ppo_mod = _patch_ppo_setup_prerequisites(monkeypatch) + cluster_cls = MagicMock() + monkeypatch.setattr(ppo_mod, "RayVirtualCluster", cluster_cls) + monkeypatch.setattr( + ppo_mod, + "get_ray_cluster_topology", + lambda: { + "node-0": ("domain-0", 0), + "node-1": ("domain-0", 1), + }, + ) + config = _make_noncolocated_setup_config( + total_nodes=3, + inference_nodes=1, + inference_gpus_per_node=8, + segment_size=1, + ) + + with pytest.raises( + AssertionError, + match="Not enough alive Ray nodes for all PPO roles", + ): + ppo_mod.setup(config, MagicMock(), _setup_dataset(), None) + + cluster_cls.assert_not_called() + + +def test_noncolocated_topology_counts_shared_node_once(monkeypatch): + """A one-node GPU split must not require two physical Ray nodes.""" + from nemo_rl.algorithms import ppo as ppo_mod + + monkeypatch.setattr( + ppo_mod, + "get_ray_cluster_topology", + lambda: {"node-0": ("domain-0", 0)}, + ) + config = _make_noncolocated_setup_config( + total_nodes=1, + total_gpus_per_node=8, + inference_gpus_per_node=2, + segment_size=1, + ) + + result, cluster_calls, *_ = _run_noncolocated_setup(monkeypatch, config) + + train_cluster, inference_cluster = result[3] + assert train_cluster.kwargs["node_resource_constraints"] == [{"domain-0": 0.001}] + assert inference_cluster.kwargs["node_resource_constraints"] is None + assert train_cluster.get_placement_groups_called + assert len(cluster_calls) == 2 + + +def test_noncolocated_multi_node_topology_constraints(monkeypatch): + """Topology constraints are applied to a successful multi-node train split.""" + from nemo_rl.algorithms import ppo as ppo_mod + + monkeypatch.setattr( + ppo_mod, + "get_ray_cluster_topology", + lambda: { + "node-0": ("domain-0", 0), + "node-1": ("domain-0", 1), + "node-2": ("domain-0", 2), + }, + ) + config = _make_noncolocated_setup_config( + total_nodes=3, + total_gpus_per_node=8, + inference_nodes=1, + inference_gpus_per_node=8, + segment_size=1, + ) + + result, _, *_ = _run_noncolocated_setup(monkeypatch, config) + + train_cluster, inference_cluster = result[3] + assert train_cluster.kwargs["node_resource_constraints"] == [ + {"domain-0": 0.001}, + {"domain-0": 0.001}, + ] + assert inference_cluster.kwargs["node_resource_constraints"] is None + assert train_cluster.get_placement_groups_called + + +def test_noncolocated_multi_node_inference_topology_constraints(monkeypatch): + """Multi-node vLLM instances are pinned and initialized eagerly.""" + from nemo_rl.algorithms import ppo as ppo_mod + + monkeypatch.setattr( + ppo_mod, + "get_ray_cluster_topology", + lambda: { + "node-0": ("domain-0", 0), + "node-1": ("domain-0", 1), + "node-2": ("domain-1", 0), + "node-3": ("domain-1", 1), + }, + ) + config = _make_noncolocated_setup_config( + total_nodes=4, + total_gpus_per_node=8, + inference_nodes=2, + inference_gpus_per_node=8, + segment_size=2, + tensor_parallel_size=16, + ) + + result, _, _, _, _, _, generation_factory, _ = _run_noncolocated_setup( + monkeypatch, config + ) + + train_cluster, inference_cluster = result[3] + assert train_cluster.kwargs["node_resource_constraints"] == [ + {"domain-0": 0.001}, + {"domain-0": 0.001}, + ] + assert inference_cluster.kwargs["node_resource_constraints"] == [ + {"domain-1": 0.001}, + {"domain-1": 0.001}, + ] + assert inference_cluster.kwargs["segment_size"] == 2 + generation_factory.init_cluster_placement_groups.assert_called_once_with( + inference_cluster, config.policy["generation"] + ) + + +def test_noncolocated_skips_nondivisible_inference_topology_constraints( + monkeypatch, + capsys, +): + """Non-divisible inference instances fall back to unconstrained placement.""" + from nemo_rl.algorithms import ppo as ppo_mod + + monkeypatch.setattr( + ppo_mod, + "get_ray_cluster_topology", + lambda: { + "node-0": ("domain-0", 0), + "node-1": ("domain-0", 1), + "node-2": ("domain-1", 0), + "node-3": ("domain-1", 1), + "node-4": ("domain-1", 2), + }, + ) + config = _make_noncolocated_setup_config( + total_nodes=5, + total_gpus_per_node=8, + inference_nodes=3, + inference_gpus_per_node=8, + segment_size=2, + tensor_parallel_size=16, + ) + + result, _, _, _, _, _, generation_factory, _ = _run_noncolocated_setup( + monkeypatch, config + ) + + train_cluster, inference_cluster = result[3] + assert train_cluster.kwargs["node_resource_constraints"] == [ + {"domain-0": 0.001}, + {"domain-0": 0.001}, + ] + assert inference_cluster.kwargs["node_resource_constraints"] is None + assert inference_cluster.kwargs["segment_size"] is None + generation_factory.init_cluster_placement_groups.assert_not_called() + assert ( + "inference_nodes=3 is not divisible by nodes_per_instance=2" + in capsys.readouterr().out + ) + + +def test_noncolocated_vllm_builds_separate_clusters_and_collective(monkeypatch): + """Verify separate-cluster construction and collective call wiring only. + + The AutoModel and Megatron functional tests cover the real collective handshake. + """ + config = _make_noncolocated_setup_config( + total_gpus_per_node=8, + inference_gpus_per_node=2, + ) + config.cluster["master_port_range_low"] = 1400 + config.cluster["master_port_range_high"] = 1999 + ( + result, + cluster_calls, + policy, + generation, + policy_factory, + value_factory, + generation_factory, + ray_get, + ) = _run_noncolocated_setup(monkeypatch, config) + + train_cluster, inference_cluster = result[3] + assert [cluster.kwargs["name"] for cluster in cluster_calls] == [ + "ppo_train_cluster", + "ppo_inference_cluster", + ] + assert train_cluster.kwargs["bundle_ct_per_node_list"] == [6] + assert train_cluster.kwargs["max_colocated_worker_groups"] == 2 + assert train_cluster.kwargs["port_range_low"] == 1400 + assert train_cluster.kwargs["port_range_high"] == 1999 + assert inference_cluster.kwargs["bundle_ct_per_node_list"] == [2] + assert inference_cluster.kwargs["max_colocated_worker_groups"] == 1 + assert inference_cluster.kwargs["port_range_low"] == 1400 + assert inference_cluster.kwargs["port_range_high"] == 1999 + assert policy_factory.call_args.kwargs["cluster"] is train_cluster + assert value_factory.call_args.kwargs["cluster"] is train_cluster + assert generation_factory.call_args.kwargs["cluster"] is inference_cluster + + policy.init_collective.assert_called_once_with( + "127.0.0.1", 1234, 8, train_world_size=6 + ) + generation.init_collective.assert_called_once_with( + "127.0.0.1", 1234, 8, train_world_size=6 + ) + ray_get.assert_called_once_with(["policy-future", "generation-future"]) + + policy.offload_to_cpu.assert_called_once_with() + value_model = result[2] + value_model.finish_training.assert_called_once_with() + policy.prepare_for_training.assert_called_once_with() + policy.prepare_refit_info.assert_called_once_with() + generation.prepare_refit_info.assert_called_once_with({"state": "dict"}) + + +def test_colocated_setup_keeps_single_cluster_and_skips_collective(monkeypatch): + """The default colocated setup remains unchanged by the cluster split.""" + config = _make_noncolocated_setup_config() + config.policy["generation"]["colocated"] = { + "enabled": True, + "resources": {"num_nodes": None, "gpus_per_node": None}, + } + + result, cluster_calls, policy, generation, *_, ray_get = _run_noncolocated_setup( + monkeypatch, config + ) + + train_cluster, inference_cluster = result[3] + assert train_cluster is inference_cluster + assert len(cluster_calls) == 1 + assert train_cluster.kwargs["name"] == "ppo_policy_cluster" + assert train_cluster.kwargs["max_colocated_worker_groups"] == 3 + policy.init_collective.assert_not_called() + generation.init_collective.assert_not_called() + ray_get.assert_not_called() + + +def test_noncolocated_vllm_multi_node_cluster_and_collective_sizes(monkeypatch): + """A full inference node is carved out of a three-node PPO allocation.""" + config = _make_noncolocated_setup_config( + total_nodes=3, + total_gpus_per_node=8, + inference_nodes=1, + inference_gpus_per_node=8, + ) + result, _, policy, generation, *_ = _run_noncolocated_setup(monkeypatch, config) + + train_cluster, inference_cluster = result[3] + assert train_cluster.kwargs["bundle_ct_per_node_list"] == [8, 8] + assert train_cluster.kwargs["max_colocated_worker_groups"] == 2 + assert inference_cluster.kwargs["bundle_ct_per_node_list"] == [8] + assert inference_cluster.kwargs["max_colocated_worker_groups"] == 1 + policy.init_collective.assert_called_once_with( + "127.0.0.1", 1234, 24, train_world_size=16 + ) + generation.init_collective.assert_called_once_with( + "127.0.0.1", 1234, 24, train_world_size=16 + ) + + +def test_noncolocated_vllm_single_node_reserves_reward_model_gpu(monkeypatch): + """Training receives GPUs left after rollout and reward-model reservations.""" + config = _make_noncolocated_setup_config( + total_gpus_per_node=8, + inference_gpus_per_node=2, + reward_model_gpus_per_node=1, + ) + result, *_ = _run_noncolocated_setup(monkeypatch, config) + + train_cluster, inference_cluster = result[3] + assert train_cluster.kwargs["bundle_ct_per_node_list"] == [5] + assert train_cluster.kwargs["max_colocated_worker_groups"] == 2 + assert inference_cluster.kwargs["bundle_ct_per_node_list"] == [2] + assert inference_cluster.kwargs["max_colocated_worker_groups"] == 1 + + +def test_noncolocated_reward_model_node_leaves_shared_train_inference_node( + monkeypatch, +): + """A dedicated RM node does not consume GPUs on the shared PPO node.""" + config = _make_noncolocated_setup_config( + total_nodes=2, + total_gpus_per_node=8, + inference_gpus_per_node=2, + reward_model_gpus_per_node=1, + ) + + result, _, policy, generation, *_ = _run_noncolocated_setup(monkeypatch, config) + + train_cluster, inference_cluster = result[3] + assert train_cluster.kwargs["bundle_ct_per_node_list"] == [6] + assert inference_cluster.kwargs["bundle_ct_per_node_list"] == [2] + policy.init_collective.assert_called_once_with( + "127.0.0.1", 1234, 8, train_world_size=6 + ) + generation.init_collective.assert_called_once_with( + "127.0.0.1", 1234, 8, train_world_size=6 + ) diff --git a/tests/unit/data/test_utils.py b/tests/unit/data/test_utils.py index 54580fd4cd1..82e54e9e45b 100644 --- a/tests/unit/data/test_utils.py +++ b/tests/unit/data/test_utils.py @@ -11,8 +11,7 @@ # 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. -"""Unit tests for ``nemo_rl.data.utils.get_train_dataset_name`` and -``nemo_rl.data.utils.load_dataloader_state``. +"""Unit tests for shared data configuration and checkpoint helpers. These helpers underpin the dataset-swap-aware checkpoint resume logic: when the saved ``dataset_name`` (read from the ``config.yaml`` written alongside @@ -29,6 +28,7 @@ import yaml from torchdata.stateful_dataloader import StatefulDataLoader +from nemo_rl.data.datasets import extract_necessary_env_names from nemo_rl.data.utils import get_train_dataset_name, load_dataloader_state # --------------------------------------------------------------------------- @@ -100,6 +100,42 @@ def _write_checkpoint( return dir_path +# --------------------------------------------------------------------------- +# extract_necessary_env_names +# --------------------------------------------------------------------------- + + +def test_extract_necessary_env_names_handles_dataset_lists(): + data_config = { + "train": [ + {"dataset_name": "train-a", "env_name": "math"}, + {"dataset_name": "train-b", "env_name": "reward_model"}, + ], + "validation": [ + {"dataset_name": "val-a", "env_name": "math"}, + {"dataset_name": "val-b", "env_name": "code"}, + ], + "default": {"env_name": "default-env"}, + } + + assert set(extract_necessary_env_names(data_config)) == { + "math", + "reward_model", + "code", + "default-env", + } + + +def test_extract_necessary_env_names_ignores_missing_or_none_entries(): + data_config = { + "train": [{"dataset_name": "train"}, None], + "validation": None, + "default": None, + } + + assert extract_necessary_env_names(data_config) == [] + + # --------------------------------------------------------------------------- # get_train_dataset_name # --------------------------------------------------------------------------- diff --git a/tests/unit/reference_configs/ppo_math_1B_megatron.yaml b/tests/unit/reference_configs/ppo_math_1B_megatron.yaml index 3610bfd09e8..37a291f6d93 100644 --- a/tests/unit/reference_configs/ppo_math_1B_megatron.yaml +++ b/tests/unit/reference_configs/ppo_math_1B_megatron.yaml @@ -47,6 +47,7 @@ ppo: source_max: 1.0 target_min: -1.0 # DAPO: scale rewards to [-1, 1] target_max: 1.0 + seq_logprob_error_threshold: null loss_fn: disable_ppo_ratio: false @@ -430,3 +431,6 @@ logger: cluster: gpus_per_node: 1 num_nodes: 1 + master_port_range_low: 1400 + master_port_range_high: 1999 + segment_size: null From daf46ff37a64832d8c5637bdfbb0bc3252b47052 Mon Sep 17 00:00:00 2001 From: Michal Futrega Date: Thu, 6 Aug 2026 21:24:39 +0200 Subject: [PATCH 05/10] feat(grpo): validation-only sampling params and grouped pass@k validation (#3401) Signed-off-by: Michal Futrega Co-authored-by: Claude Fable 5 --- examples/configs/distillation_math.yaml | 3 + examples/configs/evals/eval.yaml | 3 + examples/configs/evals/mmau.yaml | 3 + examples/configs/grpo_math_1B.yaml | 7 + examples/configs/ppo_math_1B.yaml | 3 + .../mopd-qwen3-1.7b-3n8g-megatron-pack.yaml | 1 - examples/nemo_gym/grpo_nanov3.yaml | 5 +- .../grpo_qwen3_30ba3b_thinking_swe1.yaml | 2 +- .../grpo_qwen3_30ba3b_thinking_swe2.yaml | 2 +- ...rkplace_assistant_nemotron_nano_v2_9b.yaml | 4 + .../nemotron-3-super/stage1_rlvr.yaml | 5 +- .../nemotron-3-super/stage2_swe1.yaml | 5 +- .../nemotron-3-super/stage2_swe2.yaml | 5 +- .../nemotron-3-super/stage3_rlhf.yaml | 5 +- .../nemotron-3-ultra/ifbench_teacher.yaml | 5 +- examples/nemo_gym/nemotron-3-ultra/mopd.yaml | 5 +- .../nemotron-3-ultra/reasoning_teacher.yaml | 5 +- .../nemotron-3-ultra/rlhf_teacher.yaml | 5 +- .../nemotron-3-ultra/student_rlvr1.yaml | 5 +- .../nemotron-3-ultra/student_rlvr2.yaml | 5 +- .../nemotron-3-ultra/swe_teacher.yaml | 5 +- nemo_rl/algorithms/grpo.py | 76 +++++++++- nemo_rl/experience/rollouts.py | 28 +++- nemo_rl/models/generation/interfaces.py | 35 +++++ .../generation/vllm/vllm_worker_async.py | 35 ++++- .../configs/grpo_math_1B.yaml | 5 + tests/unit/algorithms/test_distillation.py | 15 ++ tests/unit/algorithms/test_grpo.py | 140 ++++++++++++++++++ .../environments/test_code_environment.py | 3 + tests/unit/environments/test_retriever.py | 3 + tests/unit/experience/test_rollouts.py | 6 + .../models/generation/test_vllm_generation.py | 3 + .../generation/test_vllm_large_model.py | 3 + .../generation/test_vllm_quant_backend.py | 3 + .../trtllm/test_trtllm_generation.py | 3 + .../reference_configs/distillation_math.yaml | 3 + tests/unit/reference_configs/eval.yaml | 3 + .../unit/reference_configs/grpo_math_1B.yaml | 4 + .../ppo_math_1B_megatron.yaml | 3 + tests/unit/utils/test_native_checkpoint.py | 3 + 40 files changed, 437 insertions(+), 25 deletions(-) diff --git a/examples/configs/distillation_math.yaml b/examples/configs/distillation_math.yaml index 90e16000aa8..76ccf33e0ce 100644 --- a/examples/configs/distillation_math.yaml +++ b/examples/configs/distillation_math.yaml @@ -201,6 +201,9 @@ policy: &POLICY_BASE temperature: 1.0 top_p: 1.0 top_k: null + val_temperature: ${.temperature} + val_top_p: ${.top_p} + val_top_k: ${.top_k} stop_token_ids: null stop_strings: null vllm_cfg: diff --git a/examples/configs/evals/eval.yaml b/examples/configs/evals/eval.yaml index c492ac34cd1..820d503f27e 100644 --- a/examples/configs/evals/eval.yaml +++ b/examples/configs/evals/eval.yaml @@ -12,6 +12,9 @@ generation: temperature: 0.0 top_p: 1.0 top_k: -1 # -1 means disable + val_temperature: ${.temperature} + val_top_p: ${.top_p} + val_top_k: ${.top_k} num_prompts_per_step: -1 # -1 means pass all prompts at once model_name: "Qwen/Qwen2.5-Math-1.5B-Instruct" stop_token_ids: null diff --git a/examples/configs/evals/mmau.yaml b/examples/configs/evals/mmau.yaml index e12c3ea0aec..7bc324d492f 100644 --- a/examples/configs/evals/mmau.yaml +++ b/examples/configs/evals/mmau.yaml @@ -11,6 +11,9 @@ generation: temperature: 0.0 top_p: 1.0 top_k: -1 + val_temperature: ${.temperature} + val_top_p: ${.top_p} + val_top_k: ${.top_k} num_prompts_per_step: -1 model_name: "Qwen/Qwen2.5-Omni-3B" stop_token_ids: null diff --git a/examples/configs/grpo_math_1B.yaml b/examples/configs/grpo_math_1B.yaml index c4682a2c976..062abddaa4c 100644 --- a/examples/configs/grpo_math_1B.yaml +++ b/examples/configs/grpo_math_1B.yaml @@ -15,6 +15,9 @@ grpo: advantage_clip_low: null advantage_clip_high: null max_val_samples: 256 + # Validation rollouts per prompt; k > 1 also reports the pass_k metric. + # max_val_samples counts PROMPTS: total validation rollouts = max_val_samples * k. + val_num_generations_per_prompt: 1 # Early stop once this metric (e.g. accuracy or pass_k) reaches the threshold; null disables. stop_at_validation_metric: null # Required when stop_at_validation_metric is set. @@ -344,6 +347,10 @@ policy: temperature: 1.0 top_p: 1.0 top_k: null + # Validation-only sampling; defaults follow the train values above. + val_temperature: ${.temperature} + val_top_p: ${.top_p} + val_top_k: ${.top_k} stop_token_ids: null stop_strings: null # null = topology default (IPC colocated, NCCL non-colocated). diff --git a/examples/configs/ppo_math_1B.yaml b/examples/configs/ppo_math_1B.yaml index 1bd173aa132..366cf08cf31 100644 --- a/examples/configs/ppo_math_1B.yaml +++ b/examples/configs/ppo_math_1B.yaml @@ -237,6 +237,9 @@ policy: temperature: 1.0 top_p: 1.0 top_k: null + val_temperature: ${.temperature} + val_top_p: ${.top_p} + val_top_k: ${.top_k} stop_token_ids: null stop_strings: null mcore_generation_config: diff --git a/examples/configs/recipes/llm/mopd-qwen3-1.7b-3n8g-megatron-pack.yaml b/examples/configs/recipes/llm/mopd-qwen3-1.7b-3n8g-megatron-pack.yaml index d967ea86519..e8be20b666c 100644 --- a/examples/configs/recipes/llm/mopd-qwen3-1.7b-3n8g-megatron-pack.yaml +++ b/examples/configs/recipes/llm/mopd-qwen3-1.7b-3n8g-megatron-pack.yaml @@ -2,7 +2,6 @@ defaults: ../../grpo_math_1B.yaml grpo: num_prompts_per_step: 8 num_generations_per_prompt: 4 - num_val_generations_per_prompt: 1 max_num_steps: 5 val_period: 1000 overlong_filtering: true diff --git a/examples/nemo_gym/grpo_nanov3.yaml b/examples/nemo_gym/grpo_nanov3.yaml index 8fa3175d9e4..103d470f610 100644 --- a/examples/nemo_gym/grpo_nanov3.yaml +++ b/examples/nemo_gym/grpo_nanov3.yaml @@ -2,7 +2,7 @@ grpo: num_prompts_per_step: 128 num_generations_per_prompt: 16 - num_val_generations_per_prompt: 4 + val_num_generations_per_prompt: 4 max_rollout_turns: 1 # for multi-turn rollouts. Math Environments just have 1 turn (answering the question) max_num_epochs: 1 max_num_steps: 1000000 @@ -210,6 +210,9 @@ policy: temperature: 1.0 top_p: 1.0 top_k: null + val_temperature: ${.temperature} + val_top_p: ${.top_p} + val_top_k: ${.top_k} stop_token_ids: null stop_strings: null mcore_generation_config: diff --git a/examples/nemo_gym/grpo_qwen3_30ba3b_thinking_swe1.yaml b/examples/nemo_gym/grpo_qwen3_30ba3b_thinking_swe1.yaml index 48f284cb711..c19fb351b65 100644 --- a/examples/nemo_gym/grpo_qwen3_30ba3b_thinking_swe1.yaml +++ b/examples/nemo_gym/grpo_qwen3_30ba3b_thinking_swe1.yaml @@ -13,7 +13,7 @@ checkpointing: grpo: num_prompts_per_step: 64 num_generations_per_prompt: 8 - num_val_generations_per_prompt: 1 + val_num_generations_per_prompt: 1 max_num_epochs: 100 advantage_clip_low: -100 advantage_clip_high: 100 diff --git a/examples/nemo_gym/grpo_qwen3_30ba3b_thinking_swe2.yaml b/examples/nemo_gym/grpo_qwen3_30ba3b_thinking_swe2.yaml index 6764b43126d..449b873e5f2 100644 --- a/examples/nemo_gym/grpo_qwen3_30ba3b_thinking_swe2.yaml +++ b/examples/nemo_gym/grpo_qwen3_30ba3b_thinking_swe2.yaml @@ -13,7 +13,7 @@ checkpointing: grpo: num_prompts_per_step: 8 num_generations_per_prompt: 8 - num_val_generations_per_prompt: 1 + val_num_generations_per_prompt: 1 max_num_epochs: 100 advantage_clip_low: -100 advantage_clip_high: 100 diff --git a/examples/nemo_gym/grpo_workplace_assistant_nemotron_nano_v2_9b.yaml b/examples/nemo_gym/grpo_workplace_assistant_nemotron_nano_v2_9b.yaml index e7b78207ff4..b3641dd5437 100644 --- a/examples/nemo_gym/grpo_workplace_assistant_nemotron_nano_v2_9b.yaml +++ b/examples/nemo_gym/grpo_workplace_assistant_nemotron_nano_v2_9b.yaml @@ -14,6 +14,7 @@ grpo: advantage_clip_low: null advantage_clip_high: null max_val_samples: null # inferred from size of val dataset. for multi evals, repeat val ds via `num_repeats` in `ng_prepare_data`. + val_num_generations_per_prompt: 1 # Early stop once this metric (e.g. accuracy or pass_k) reaches the threshold; null disables. stop_at_validation_metric: null # Required when stop_at_validation_metric is set. @@ -226,6 +227,9 @@ policy: temperature: 1.0 top_p: 1.0 top_k: null + val_temperature: ${.temperature} + val_top_p: ${.top_p} + val_top_k: ${.top_k} stop_token_ids: null stop_strings: null vllm_cfg: diff --git a/examples/nemo_gym/nemotron-3-super/stage1_rlvr.yaml b/examples/nemo_gym/nemotron-3-super/stage1_rlvr.yaml index 1b2d300dc26..03317acd30e 100644 --- a/examples/nemo_gym/nemotron-3-super/stage1_rlvr.yaml +++ b/examples/nemo_gym/nemotron-3-super/stage1_rlvr.yaml @@ -13,7 +13,7 @@ checkpointing: grpo: num_prompts_per_step: 256 num_generations_per_prompt: 16 - num_val_generations_per_prompt: 2 + val_num_generations_per_prompt: 2 max_rollout_turns: 1 max_num_epochs: 1 max_num_steps: 1000000 @@ -226,6 +226,9 @@ policy: temperature: 1.0 top_p: 1.0 top_k: null + val_temperature: ${.temperature} + val_top_p: ${.top_p} + val_top_k: ${.top_k} stop_token_ids: null stop_strings: null vllm_cfg: diff --git a/examples/nemo_gym/nemotron-3-super/stage2_swe1.yaml b/examples/nemo_gym/nemotron-3-super/stage2_swe1.yaml index 59c41c3d0b2..8a8d52ca042 100644 --- a/examples/nemo_gym/nemotron-3-super/stage2_swe1.yaml +++ b/examples/nemo_gym/nemotron-3-super/stage2_swe1.yaml @@ -13,7 +13,7 @@ checkpointing: grpo: num_prompts_per_step: 64 num_generations_per_prompt: 16 - num_val_generations_per_prompt: 1 + val_num_generations_per_prompt: 1 max_rollout_turns: 1 max_num_epochs: 100 max_num_steps: 1000000 @@ -226,6 +226,9 @@ policy: temperature: 1.0 top_p: 1.0 top_k: null + val_temperature: ${.temperature} + val_top_p: ${.top_p} + val_top_k: ${.top_k} stop_token_ids: null stop_strings: null vllm_cfg: diff --git a/examples/nemo_gym/nemotron-3-super/stage2_swe2.yaml b/examples/nemo_gym/nemotron-3-super/stage2_swe2.yaml index 72a22c7913d..4c5284f5aae 100644 --- a/examples/nemo_gym/nemotron-3-super/stage2_swe2.yaml +++ b/examples/nemo_gym/nemotron-3-super/stage2_swe2.yaml @@ -13,7 +13,7 @@ checkpointing: grpo: num_prompts_per_step: 16 num_generations_per_prompt: 32 - num_val_generations_per_prompt: 1 + val_num_generations_per_prompt: 1 max_rollout_turns: 1 max_num_epochs: 100 max_num_steps: 1000000 @@ -219,6 +219,9 @@ policy: temperature: 1.0 top_p: 1.0 top_k: null + val_temperature: ${.temperature} + val_top_p: ${.top_p} + val_top_k: ${.top_k} stop_token_ids: null stop_strings: null vllm_cfg: diff --git a/examples/nemo_gym/nemotron-3-super/stage3_rlhf.yaml b/examples/nemo_gym/nemotron-3-super/stage3_rlhf.yaml index e56c770a2b4..f2d2b38930f 100644 --- a/examples/nemo_gym/nemotron-3-super/stage3_rlhf.yaml +++ b/examples/nemo_gym/nemotron-3-super/stage3_rlhf.yaml @@ -13,7 +13,7 @@ checkpointing: grpo: num_prompts_per_step: 128 num_generations_per_prompt: 16 - num_val_generations_per_prompt: 2 + val_num_generations_per_prompt: 2 max_rollout_turns: 1 max_num_epochs: 1 max_num_steps: 1000000 @@ -226,6 +226,9 @@ policy: temperature: 1.0 top_p: 1.0 top_k: null + val_temperature: ${.temperature} + val_top_p: ${.top_p} + val_top_k: ${.top_k} stop_token_ids: null stop_strings: null vllm_cfg: diff --git a/examples/nemo_gym/nemotron-3-ultra/ifbench_teacher.yaml b/examples/nemo_gym/nemotron-3-ultra/ifbench_teacher.yaml index 026824cf562..18dc57404a5 100644 --- a/examples/nemo_gym/nemotron-3-ultra/ifbench_teacher.yaml +++ b/examples/nemo_gym/nemotron-3-ultra/ifbench_teacher.yaml @@ -47,7 +47,7 @@ checkpointing: grpo: num_prompts_per_step: 128 num_generations_per_prompt: 16 - num_val_generations_per_prompt: 2 + val_num_generations_per_prompt: 2 max_rollout_turns: 1 max_num_epochs: 1 max_num_steps: 1000000 @@ -305,6 +305,9 @@ policy: temperature: 1.0 top_p: 1.0 top_k: null + val_temperature: ${.temperature} + val_top_p: ${.top_p} + val_top_k: ${.top_k} stop_token_ids: null stop_strings: null # TP=8 EP=8: EP=TP so vllm_dp_size=1, async_engine=true works with NeMo Gym. diff --git a/examples/nemo_gym/nemotron-3-ultra/mopd.yaml b/examples/nemo_gym/nemotron-3-ultra/mopd.yaml index fe85213fdd4..203bf03e83f 100644 --- a/examples/nemo_gym/nemotron-3-ultra/mopd.yaml +++ b/examples/nemo_gym/nemotron-3-ultra/mopd.yaml @@ -62,7 +62,7 @@ checkpointing: grpo: num_prompts_per_step: 1024 num_generations_per_prompt: 1 - num_val_generations_per_prompt: 2 + val_num_generations_per_prompt: 2 max_rollout_turns: 1 max_num_epochs: 1 max_num_steps: 1000000 @@ -326,6 +326,9 @@ policy: temperature: 1.0 top_p: 1.0 top_k: null + val_temperature: ${.temperature} + val_top_p: ${.top_p} + val_top_k: ${.top_k} stop_token_ids: null stop_strings: null # TP=8 EP=8: EP=TP so vllm_dp_size=1, async_engine=true works with NeMo Gym. diff --git a/examples/nemo_gym/nemotron-3-ultra/reasoning_teacher.yaml b/examples/nemo_gym/nemotron-3-ultra/reasoning_teacher.yaml index b32979ce39a..e62f261e867 100644 --- a/examples/nemo_gym/nemotron-3-ultra/reasoning_teacher.yaml +++ b/examples/nemo_gym/nemotron-3-ultra/reasoning_teacher.yaml @@ -51,7 +51,7 @@ checkpointing: grpo: num_prompts_per_step: 128 num_generations_per_prompt: 16 - num_val_generations_per_prompt: 2 + val_num_generations_per_prompt: 2 max_rollout_turns: 1 max_num_epochs: 10 max_num_steps: 1000000 @@ -308,6 +308,9 @@ policy: temperature: 1.0 top_p: 1.0 top_k: null + val_temperature: ${.temperature} + val_top_p: ${.top_p} + val_top_k: ${.top_k} stop_token_ids: null stop_strings: null # TP=8 EP=8: EP=TP so vllm_dp_size=1, async_engine=true works with NeMo Gym. diff --git a/examples/nemo_gym/nemotron-3-ultra/rlhf_teacher.yaml b/examples/nemo_gym/nemotron-3-ultra/rlhf_teacher.yaml index 32fc72deb4b..d757cde7a0e 100644 --- a/examples/nemo_gym/nemotron-3-ultra/rlhf_teacher.yaml +++ b/examples/nemo_gym/nemotron-3-ultra/rlhf_teacher.yaml @@ -48,7 +48,7 @@ checkpointing: grpo: num_prompts_per_step: 128 num_generations_per_prompt: 16 - num_val_generations_per_prompt: 2 + val_num_generations_per_prompt: 2 max_rollout_turns: 1 max_num_epochs: 1 max_num_steps: 1000000 @@ -306,6 +306,9 @@ policy: temperature: 1.0 top_p: 1.0 top_k: null + val_temperature: ${.temperature} + val_top_p: ${.top_p} + val_top_k: ${.top_k} stop_token_ids: null stop_strings: null # TP=8 EP=8: EP=TP so vllm_dp_size=1, async_engine=true works with NeMo Gym. diff --git a/examples/nemo_gym/nemotron-3-ultra/student_rlvr1.yaml b/examples/nemo_gym/nemotron-3-ultra/student_rlvr1.yaml index 03e43b694dc..f6c3b308ad5 100644 --- a/examples/nemo_gym/nemotron-3-ultra/student_rlvr1.yaml +++ b/examples/nemo_gym/nemotron-3-ultra/student_rlvr1.yaml @@ -44,7 +44,7 @@ checkpointing: grpo: num_prompts_per_step: 512 num_generations_per_prompt: 16 - num_val_generations_per_prompt: 2 + val_num_generations_per_prompt: 2 max_rollout_turns: 1 max_num_epochs: 1 max_num_steps: 1000000 @@ -302,6 +302,9 @@ policy: temperature: 1.0 top_p: 1.0 top_k: null + val_temperature: ${.temperature} + val_top_p: ${.top_p} + val_top_k: ${.top_k} stop_token_ids: null stop_strings: null # TP=8 EP=8: EP=TP so vllm_dp_size=1, async_engine=true works with NeMo Gym. diff --git a/examples/nemo_gym/nemotron-3-ultra/student_rlvr2.yaml b/examples/nemo_gym/nemotron-3-ultra/student_rlvr2.yaml index 3a9fbae3712..262d1679cd8 100644 --- a/examples/nemo_gym/nemotron-3-ultra/student_rlvr2.yaml +++ b/examples/nemo_gym/nemotron-3-ultra/student_rlvr2.yaml @@ -45,7 +45,7 @@ checkpointing: grpo: num_prompts_per_step: 512 num_generations_per_prompt: 16 - num_val_generations_per_prompt: 2 + val_num_generations_per_prompt: 2 max_rollout_turns: 1 max_num_epochs: 1 max_num_steps: 1000000 @@ -303,6 +303,9 @@ policy: temperature: 1.0 top_p: 1.0 top_k: null + val_temperature: ${.temperature} + val_top_p: ${.top_p} + val_top_k: ${.top_k} stop_token_ids: null stop_strings: null # TP=8 EP=8: EP=TP so vllm_dp_size=1, async_engine=true works with NeMo Gym. diff --git a/examples/nemo_gym/nemotron-3-ultra/swe_teacher.yaml b/examples/nemo_gym/nemotron-3-ultra/swe_teacher.yaml index c8db114be15..75224b0e556 100644 --- a/examples/nemo_gym/nemotron-3-ultra/swe_teacher.yaml +++ b/examples/nemo_gym/nemotron-3-ultra/swe_teacher.yaml @@ -67,7 +67,7 @@ checkpointing: grpo: num_prompts_per_step: 32 num_generations_per_prompt: 16 - num_val_generations_per_prompt: 2 + val_num_generations_per_prompt: 2 max_rollout_turns: 1 max_num_epochs: 4 max_num_steps: 1000000 @@ -325,6 +325,9 @@ policy: temperature: 1.0 top_p: 1.0 top_k: null + val_temperature: ${.temperature} + val_top_p: ${.top_p} + val_top_k: ${.top_k} stop_token_ids: null stop_strings: null vllm_cfg: diff --git a/nemo_rl/algorithms/grpo.py b/nemo_rl/algorithms/grpo.py index a646d7357b7..278abeab7da 100644 --- a/nemo_rl/algorithms/grpo.py +++ b/nemo_rl/algorithms/grpo.py @@ -98,6 +98,7 @@ from nemo_rl.models.generation.interfaces import ( GenerationConfig, GenerationInterface, + GenerationSamplingParams, resolve_routed_experts_dtype_name_for_model, ) from nemo_rl.models.generation.megatron import MegatronGeneration @@ -246,7 +247,13 @@ class GRPOConfig(BaseModel, extra="allow"): # Whether to run validation on the last training step. Setting this to True ensures the # final checkpoint has validation metrics, which is required for get_best_checkpoint_path(). val_at_end: bool = False + # Counts PROMPTS, not rollouts: with val_num_generations_per_prompt = k, + # total validation rollouts = max_val_samples * k. max_val_samples: int | None = 256 # None for NeMo-Gym compatibility + # Number of independent validation rollouts generated for each prompt; + # k > 1 additionally reports pass@k over each prompt's k rollouts as the + # pass_k metric. + val_num_generations_per_prompt: int = 1 # Early stop: end training once this validation metric (e.g. accuracy, # always reported, or pass_k with grouped validation) reaches # stop_at_validation_threshold; null disables early stopping. @@ -398,6 +405,40 @@ def setup( if generation_config["backend"] == "vllm": normalize_vllm_refit_config(cast(VllmConfig, generation_config)) + # Validation-only sampling is honored only on the NeMo-Gym vLLM rollout + # path; everywhere else validation must sample exactly like training. + val_sampling_overridden = ( + generation_config["val_temperature"] != generation_config["temperature"] + or generation_config["val_top_p"] != generation_config["top_p"] + or generation_config["val_top_k"] != generation_config["top_k"] + ) + if val_sampling_overridden: + assert generation_config["backend"] == "vllm" and _should_use_nemo_gym( + master_config + ), ( + "generation.val_temperature/val_top_p/val_top_k differing from the " + "train sampling params is only supported for vLLM NeMo-Gym rollouts." + ) + # The NeMo-Gym path only stamps temperature/top_p onto requests and + # rejects any top_k at rollout time, so a val_top_k override can never + # be honored — fail here instead of at the first validation step. + assert not generation_config["val_top_k"], ( + "generation.val_top_k is not supported: the NeMo-Gym rollout path " + "only honors val_temperature/val_top_p. Leave val_top_k null." + ) + assert grpo_config.val_num_generations_per_prompt >= 1, ( + "grpo.val_num_generations_per_prompt must be >= 1" + ) + # pass_k is only reported when k > 1; catch the mismatch here instead of + # at the first validation step. + assert not ( + grpo_config.stop_at_validation_metric == "pass_k" + and grpo_config.val_num_generations_per_prompt <= 1 + ), ( + "grpo.stop_at_validation_metric='pass_k' requires " + "grpo.val_num_generations_per_prompt > 1" + ) + # Set seed for all random number generators set_seed(grpo_config.seed) @@ -3682,6 +3723,10 @@ def validate( timer = Timer(context={"worker": "validator"}) with timer.time("total_validation_time"): print(f"▶ Starting validation at step {step}...", flush=True) + # >= 1 is validated in setup(). + val_num_generations_per_prompt = ( + master_config.grpo.val_num_generations_per_prompt + ) total_rewards = [] total_lengths = [] @@ -3694,12 +3739,23 @@ def validate( if batch_idx >= max_batches: break + if val_num_generations_per_prompt > 1: + val_batch = val_batch.repeat_interleave(val_num_generations_per_prompt) + additional_metrics_to_report = dict() # Generate responses (updates the LLMMessageLogType in batch_with_msg_logs) # Use async rollouts when enabled by config/backend defaults. # We cascade NeMo-Gym first since NeMo-Gym also uses async rollouts. if _should_use_nemo_gym(master_config): generation_config = master_config.policy["generation"] + # Validation-only sampling (e.g. near-greedy validation); + # defaults to the train profile via the exemplar YAML + # interpolations. Training rollouts keep policy.generation. + val_sampling_params = GenerationSamplingParams( + temperature=generation_config["val_temperature"], + top_p=generation_config["val_top_p"], + top_k=generation_config["val_top_k"], + ) nemo_gym_rollout_result = run_nemo_gym_rollout_sync( policy_generation=policy_generation, input_batch=val_batch, @@ -3707,6 +3763,7 @@ def validate( task_to_env=val_task_to_env, max_seq_len=master_config.policy["max_total_sequence_length"], generation_config=generation_config, + sampling_params=val_sampling_params, log_full_result_tables=should_log_nemo_gym_full_result_tables( wandb_enabled=master_config.logger["wandb_enabled"], wandb_config=master_config.logger["wandb"], @@ -3757,11 +3814,26 @@ def validate( all_message_logs.extend(to_env) - # Calculate validation metrics + # Calculate validation metrics. accuracy is the mean reward over all + # rollouts; grouped validation (val_num_generations_per_prompt > 1) + # additionally reports pass@k over each prompt's k rollouts as pass_k. num_samples = len(total_rewards) + pass_k = None if num_samples > 0: rewards_t = torch.tensor(total_rewards, dtype=torch.float32) accuracy = rewards_t.mean().item() + if val_num_generations_per_prompt > 1: + assert num_samples % val_num_generations_per_prompt == 0, ( + "Validation rewards must be divisible by " + "grpo.val_num_generations_per_prompt" + ) + pass_k = ( + (rewards_t.view(-1, val_num_generations_per_prompt) > 0) + .any(dim=1) + .float() + .mean() + .item() + ) else: accuracy = 0.0 @@ -3774,6 +3846,8 @@ def validate( "avg_length": avg_length, **additional_metrics_to_report, } + if pass_k is not None: + val_metrics["pass_k"] = pass_k # Print sample conversations only once at the end of validation try: diff --git a/nemo_rl/experience/rollouts.py b/nemo_rl/experience/rollouts.py index 1a6145160b8..84ff46b13d3 100644 --- a/nemo_rl/experience/rollouts.py +++ b/nemo_rl/experience/rollouts.py @@ -55,6 +55,7 @@ GenerationDatumSpec, GenerationInterface, GenerationOutputSpec, + GenerationSamplingParams, ) from nemo_rl.utils.timer import Timer @@ -2007,7 +2008,9 @@ def apply_reward_penalties( def _prepare_nemo_gym_rows( - rows: list[dict], generation_config: GenerationConfig + rows: list[dict], + generation_config: GenerationConfig, + sampling_params: GenerationSamplingParams, ) -> None: """Apply NeMo-RL sampling parameters and stable row indices in place.""" for row_index, row in enumerate(rows): @@ -2017,8 +2020,8 @@ def _prepare_nemo_gym_rows( "Each NeMo-Gym row must contain a responses_create_params dict" ) - responses_create_params["temperature"] = generation_config["temperature"] - responses_create_params["top_p"] = generation_config["top_p"] + responses_create_params["temperature"] = sampling_params.temperature + responses_create_params["top_p"] = sampling_params.top_p configured_max_tokens = generation_config["max_new_tokens"] row_max_tokens = responses_create_params.get("max_output_tokens") responses_create_params["max_output_tokens"] = ( @@ -2059,6 +2062,7 @@ async def run_async_nemo_gym_rollout( thinking_tags: list[str] | tuple[str, ...] | None = None, mask_env_flagged_samples: bool = True, returns_entire_batch: bool = False, + sampling_params: Optional[GenerationSamplingParams] = None, ) -> AsyncGenerator[NemoGymRolloutResult, None]: """Stream complete NeMo-Gym prompt groups in group-completion order. @@ -2089,6 +2093,9 @@ async def run_async_nemo_gym_rollout( returns_entire_batch: Whether to treat the input as one potentially heterogeneous group. This requires ``num_generations`` to equal the batch size and is used by synchronous callers. + sampling_params: Sampling profile stamped onto every NeMo-Gym row. + ``None`` uses the train profile from ``generation_config``; + validation passes its own profile explicitly. Yields: ``NemoGymRolloutResult`` objects in prompt-group completion order. Rows @@ -2139,9 +2146,13 @@ async def run_async_nemo_gym_rollout( assert not generation_config["stop_token_ids"], ( "Stop strings is not supported in the generation config in NeMo-Gym path!" ) + if sampling_params is None: + sampling_params = GenerationSamplingParams.from_generation_config( + generation_config + ) # Top k is not OpenAI compatible, so NeMo-Gym does not guarantee support over it. - assert not generation_config["top_k"], ( - "Top k is not supported in the generation config in NeMo-Gym path!" + assert not sampling_params.top_k, ( + "Top k is not supported in the sampling params in NeMo-Gym path!" ) if num_generations <= 0: raise ValueError("num_generations must be greater than zero") @@ -2162,7 +2173,7 @@ async def run_async_nemo_gym_rollout( run_rollouts_timer_label = f"{timer_prefix}/run_rollouts" with timer.time(total_timer_label): - _prepare_nemo_gym_rows(nemo_gym_rows, generation_config) + _prepare_nemo_gym_rows(nemo_gym_rows, generation_config, sampling_params) accumulator = _NemoGymStreamAccumulator( rows=nemo_gym_rows, num_generations=num_generations, @@ -2248,6 +2259,7 @@ def run_nemo_gym_rollout_sync( effort_config: Optional[EffortLevelsConfig] = None, reward_penalty_config: dict[str, Any] | BaseModel | None = None, thinking_tags: list[str] | tuple[str, ...] | None = None, + sampling_params: Optional[GenerationSamplingParams] = None, mask_env_flagged_samples: bool = True, ) -> NemoGymRolloutResult: """Run and return one complete NeMo-Gym batch synchronously. @@ -2271,6 +2283,9 @@ def run_nemo_gym_rollout_sync( effort_config: Optional configuration for effort-based reward shaping. reward_penalty_config: Optional reward-penalty configuration. thinking_tags: Optional opening and closing tags used by thinking penalties. + sampling_params: Sampling profile stamped onto every NeMo-Gym row. + ``None`` uses the train profile from ``generation_config``; + validation passes its own profile explicitly. mask_env_flagged_samples: Whether to carry env-driven ``mask_sample`` flags in the rollout batch for loss masking. @@ -2304,6 +2319,7 @@ async def _consume_rollout() -> NemoGymRolloutResult: thinking_tags=thinking_tags, mask_env_flagged_samples=mask_env_flagged_samples, returns_entire_batch=True, + sampling_params=sampling_params, ): pass if rollout_result is None: diff --git a/nemo_rl/models/generation/interfaces.py b/nemo_rl/models/generation/interfaces.py index a757ad3a8ff..791657c394f 100644 --- a/nemo_rl/models/generation/interfaces.py +++ b/nemo_rl/models/generation/interfaces.py @@ -12,6 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. from abc import ABC, abstractmethod +from dataclasses import dataclass from typing import Any, NotRequired, Optional, TypedDict, Union import ray @@ -201,6 +202,13 @@ class GenerationConfig(TypedDict): temperature: float top_p: float top_k: int | None + # Validation-only sampling. The exemplar YAMLs default these to the train + # values above via interpolation (${.temperature}, ...), so validation + # samples exactly like training unless overridden. Only honored on the + # NeMo-Gym vLLM rollout path (guarded in grpo.setup()). + val_temperature: float + val_top_p: float + val_top_k: int | None model_name: NotRequired[str] # Not Required b/c GRPO writes this stop_token_ids: list[int] | None stop_strings: list[str] | None @@ -215,6 +223,33 @@ class GenerationConfig(TypedDict): _mtp_weights_from_refit: NotRequired[bool] +@dataclass +class GenerationSamplingParams: + """Sampling profile threaded explicitly through rollout entry points. + + Rollout callers construct one from the relevant ``GenerationConfig`` + fields (train or validation) so the sampling used for a rollout is + visible at the call site instead of flowing through config side-channels. + Named to distinguish it from ``TrainingSamplingParams`` (train-time logit + filtering) and vLLM's own ``SamplingParams``. + """ + + temperature: float + top_p: float + top_k: int | None + + @classmethod + def from_generation_config( + cls, generation_config: "GenerationConfig" + ) -> "GenerationSamplingParams": + """Build the train-time sampling profile from a generation config.""" + return cls( + temperature=generation_config["temperature"], + top_p=generation_config["top_p"], + top_k=generation_config["top_k"], + ) + + class GenerationDatumSpec(TypedDict): """Specification for input data required by generation models. diff --git a/nemo_rl/models/generation/vllm/vllm_worker_async.py b/nemo_rl/models/generation/vllm/vllm_worker_async.py index 201d0790044..d1da8a3ec1a 100644 --- a/nemo_rl/models/generation/vllm/vllm_worker_async.py +++ b/nemo_rl/models/generation/vllm/vllm_worker_async.py @@ -678,9 +678,38 @@ async def create_chat_completion( request.top_k = -1 # The request sampling params need to exactly match those as are set in NeMo RL. - # If they do not match, the inference will be off policy and destroy training stability. - assert request.temperature == generation_config["temperature"] - assert request.top_p == generation_config["top_p"] + # If they do not match, the inference will be off policy and destroy training + # stability. Validation rollouts are the one exception: they are stamped with + # the validation sampling profile (generation.val_temperature / val_top_p), + # which is metric-only and safe to serve — grpo.validate() is the only + # caller that constructs a non-train GenerationSamplingParams. Multi-turn + # agents issue their own requests, so this server-side check is the one + # chokepoint they all pass. + # vLLM resolves an unset top_p from the model's generation_config.json + # (ModelConfig.generation_config defaults to "auto"), NOT to 1.0, so a + # request omitting it would sample off-policy while passing this check. + assert request.top_p is not None, ( + "top_p must be set explicitly on NeMo-RL requests; an unset top_p is " + "resolved by vLLM from the model's generation_config.json and would " + "bypass the on-policy sampling check." + ) + request_top_p = request.top_p + is_train_sampling = ( + request.temperature == generation_config["temperature"] + and request_top_p == generation_config["top_p"] + ) + is_val_sampling = ( + request.temperature == generation_config["val_temperature"] + and request_top_p == generation_config["val_top_p"] + ) + assert is_train_sampling or is_val_sampling, ( + f"request sampling (temperature={request.temperature}, " + f"top_p={request.top_p}) matches neither the train sampling params " + f"(temperature={generation_config['temperature']}, " + f"top_p={generation_config['top_p']}) nor the validation sampling " + f"params (val_temperature={generation_config['val_temperature']}, " + f"val_top_p={generation_config['val_top_p']})" + ) try: generator = await openai_serving_chat.create_chat_completion( diff --git a/research/template_project/configs/grpo_math_1B.yaml b/research/template_project/configs/grpo_math_1B.yaml index 2862ad17e02..9d7b2b86296 100644 --- a/research/template_project/configs/grpo_math_1B.yaml +++ b/research/template_project/configs/grpo_math_1B.yaml @@ -13,6 +13,8 @@ grpo: val_at_end: false overlong_filtering: false max_val_samples: 256 + # Validation rollouts per prompt; k > 1 also reports the pass_k metric. + val_num_generations_per_prompt: 1 # Early stop once this metric (e.g. accuracy or pass_k) reaches the threshold; null disables. stop_at_validation_metric: null # Required when stop_at_validation_metric is set. @@ -286,6 +288,9 @@ policy: temperature: 1.0 top_p: 1.0 top_k: null + val_temperature: ${.temperature} + val_top_p: ${.top_p} + val_top_k: ${.top_k} stop_token_ids: null stop_strings: null mcore_generation_config: diff --git a/tests/unit/algorithms/test_distillation.py b/tests/unit/algorithms/test_distillation.py index dfda02093b8..df9188c02b2 100644 --- a/tests/unit/algorithms/test_distillation.py +++ b/tests/unit/algorithms/test_distillation.py @@ -150,6 +150,9 @@ def val_iter(self): "temperature": 1.0, "top_p": 1.0, "top_k": None, + "val_temperature": 1.0, + "val_top_p": 1.0, + "val_top_k": None, "colocated": { "enabled": False, }, @@ -851,6 +854,9 @@ def test_noncolocated_inference_requires_explicit_gpus_per_node_single_node(): "temperature": 1.0, "top_p": 1.0, "top_k": None, + "val_temperature": 1.0, + "val_top_p": 1.0, + "val_top_k": None, "backend": "vllm", "colocated": { "enabled": False, # Non-colocated @@ -926,6 +932,9 @@ def test_distillation_setup_non_colocated_smoke(monkeypatch, refit_transport): "temperature": 1.0, "top_p": 1.0, "top_k": None, + "val_temperature": 1.0, + "val_top_p": 1.0, + "val_top_k": None, "backend": "vllm", "refit_transport": refit_transport, "refit_cfg": None, @@ -1110,6 +1119,9 @@ def test_distillation_setup_nemo_gym_uses_deferred_vllm( "temperature": 1.0, "top_p": 1.0, "top_k": None, + "val_temperature": 1.0, + "val_top_p": 1.0, + "val_top_k": None, "backend": "vllm", "vllm_kwargs": {}, "vllm_cfg": { @@ -1363,6 +1375,9 @@ def test_noncolocated_inference_requires_explicit_gpus_per_node_multi_node(): "temperature": 1.0, "top_p": 1.0, "top_k": None, + "val_temperature": 1.0, + "val_top_p": 1.0, + "val_top_k": None, "backend": "vllm", "colocated": { "enabled": False, # Non-colocated diff --git a/tests/unit/algorithms/test_grpo.py b/tests/unit/algorithms/test_grpo.py index dee6cdb260b..479e5479919 100644 --- a/tests/unit/algorithms/test_grpo.py +++ b/tests/unit/algorithms/test_grpo.py @@ -50,6 +50,7 @@ dynamic_sampling, grpo_train, refit_policy_generation, + setup, validate, ) from nemo_rl.algorithms.grpo_sync import _train_fields_for_step, grpo_train_sync @@ -269,6 +270,7 @@ def val_iter(self): max_rollout_turns=1, val_period=100, val_start_at=-1, + val_num_generations_per_prompt=1, val_batch_size=1, val_at_start=False, val_at_end=False, @@ -305,6 +307,9 @@ def val_iter(self): "temperature": 1.0, "top_p": 1.0, "top_k": None, + "val_temperature": 1.0, + "val_top_p": 1.0, + "val_top_k": None, "backend": "vllm", "colocated": {"enabled": True}, "vllm_cfg": {"async_engine": True}, # Support async mode @@ -3773,6 +3778,141 @@ def test_validate_works_without_logger(self, mock_grpo_components): assert "accuracy" in val_metrics assert "avg_length" in val_metrics + def test_grouped_validation_reports_pass_k(self, mock_grpo_components): + mock_batch = BatchedDataDict[DatumSpec]( + { + "message_log": [ + [{"role": "user", "content": "a", "token_ids": torch.tensor([1])}], + [{"role": "user", "content": "b", "token_ids": torch.tensor([2])}], + ], + "task_name": ["math", "math"], + "extra_env_info": [{}, {}], + "loss_multiplier": torch.tensor([1.0, 1.0]), + "idx": torch.tensor([0, 1]), + "length": torch.tensor([1, 1]), + "total_reward": torch.tensor([0.0, 0.0]), + } + ) + mock_dataloader = MagicMock(spec=StatefulDataLoader) + mock_dataloader.__iter__ = MagicMock(return_value=iter([mock_batch])) + mock_config = mock_grpo_components["master_config"] + mock_config.grpo.max_val_samples = 2 + mock_config.grpo.val_batch_size = 2 + mock_config.grpo.val_num_generations_per_prompt = 4 + + def run_rollout(_policy, repeated_batch, *_args, **_kwargs): + # Each prompt is repeated k=4 times, contiguously. + assert repeated_batch["idx"].tolist() == [0, 0, 0, 0, 1, 1, 1, 1] + # Prompt 0 passes once out of 4; prompt 1 never passes. + repeated_batch["total_reward"] = torch.tensor( + [1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0] + ) + return repeated_batch, {"mean_gen_tokens_per_sample": 1.0} + + with ( + patch( + "nemo_rl.algorithms.grpo.run_multi_turn_rollout", + side_effect=run_rollout, + ), + patch("nemo_rl.algorithms.grpo._should_use_nemo_gym", return_value=False), + patch( + "nemo_rl.algorithms.grpo._should_use_async_rollouts", + return_value=False, + ), + patch("nemo_rl.algorithms.grpo.print_message_log_samples"), + ): + val_metrics, _ = validate( + MagicMock(), + mock_dataloader, + MagicMock(), + {"math": MagicMock(spec=EnvironmentInterface)}, + step=0, + master_config=mock_config, + ) + + # accuracy stays the plain mean over all 8 rollouts; pass@4 counts + # prompts with at least one passing rollout (1 of 2). + assert val_metrics["accuracy"] == pytest.approx(0.125) + assert val_metrics["pass_k"] == pytest.approx(0.5) + + def test_validation_uses_val_sampling_params_on_gym_path( + self, mock_grpo_components + ): + mock_batch = BatchedDataDict[DatumSpec]( + { + "message_log": [ + [{"role": "user", "content": "a", "token_ids": torch.tensor([1])}], + [{"role": "user", "content": "b", "token_ids": torch.tensor([2])}], + ], + "task_name": ["math", "math"], + "extra_env_info": [{}, {}], + "loss_multiplier": torch.tensor([1.0, 1.0]), + "idx": torch.tensor([0, 1]), + "length": torch.tensor([1, 1]), + "total_reward": torch.tensor([0.0, 0.0]), + } + ) + mock_dataloader = MagicMock(spec=StatefulDataLoader) + mock_dataloader.__iter__ = MagicMock(return_value=iter([mock_batch])) + mock_config = mock_grpo_components["master_config"] + mock_config.grpo.max_val_samples = 2 + mock_config.grpo.val_batch_size = 2 + mock_config.grpo.val_num_generations_per_prompt = 2 + # Train samples at 1.0/1.0; validation runs near-greedy. + mock_config.policy["generation"].update( + {"val_temperature": 0.1, "val_top_p": 0.9, "val_top_k": None} + ) + mock_config.logger.update({"wandb_enabled": False, "wandb": {}}) + mock_config.env = {} + + def run_gym_rollout(**kwargs): + repeated_batch = kwargs["input_batch"] + # 2 prompts x k=2 validation rollouts, contiguous per prompt. + assert repeated_batch["idx"].tolist() == [0, 0, 1, 1] + repeated_batch["total_reward"] = torch.tensor([1.0, 0.0, 0.0, 0.0]) + return MagicMock( + final_batch=repeated_batch, + rollout_metrics={"mean_gen_tokens_per_sample": 1.0}, + ) + + with ( + patch( + "nemo_rl.algorithms.grpo.run_nemo_gym_rollout_sync", + side_effect=run_gym_rollout, + ) as mock_rollout, + patch("nemo_rl.algorithms.grpo._should_use_nemo_gym", return_value=True), + patch("nemo_rl.algorithms.grpo.print_message_log_samples"), + ): + val_metrics, _ = validate( + MagicMock(), + mock_dataloader, + MagicMock(), + {"math": MagicMock(spec=EnvironmentInterface)}, + step=0, + master_config=mock_config, + ) + + sampling_params = mock_rollout.call_args.kwargs["sampling_params"] + assert sampling_params.temperature == pytest.approx(0.1) + assert sampling_params.top_p == pytest.approx(0.9) + assert sampling_params.top_k is None + assert val_metrics["accuracy"] == pytest.approx(0.25) + assert val_metrics["pass_k"] == pytest.approx(0.5) + + def test_setup_rejects_val_sampling_outside_gym_vllm_path( + self, mock_grpo_components + ): + master_config = mock_grpo_components["master_config"] + # Non-gym rollouts (env has no nemo_gym) with validation sampling + # different from training must be rejected at setup time. + master_config.policy["generation"].update( + {"backend": "megatron", "val_temperature": 0.1} + ) + master_config.env = {} + + with pytest.raises(AssertionError, match="only supported for vLLM NeMo-Gym"): + setup(master_config, MagicMock(), MagicMock(), None) + def test_validate_returns_empty_when_no_dataloader(self, mock_grpo_components): """Test that validate returns empty dicts when no dataloader is provided.""" mock_policy_gen = MagicMock() diff --git a/tests/unit/environments/test_code_environment.py b/tests/unit/environments/test_code_environment.py index d32550aba1e..cfdb96cc089 100644 --- a/tests/unit/environments/test_code_environment.py +++ b/tests/unit/environments/test_code_environment.py @@ -46,6 +46,9 @@ "temperature": 1.0, "top_p": 1.0, "top_k": None, + "val_temperature": 1.0, + "val_top_p": 1.0, + "val_top_k": None, "stop_token_ids": None, "stop_strings": None, "vllm_cfg": { diff --git a/tests/unit/environments/test_retriever.py b/tests/unit/environments/test_retriever.py index c9413e67590..9e0c27e62f1 100644 --- a/tests/unit/environments/test_retriever.py +++ b/tests/unit/environments/test_retriever.py @@ -45,6 +45,9 @@ "temperature": 1.0, "top_p": 1.0, "top_k": None, + "val_temperature": 1.0, + "val_top_p": 1.0, + "val_top_k": None, "stop_token_ids": None, "stop_strings": None, "vllm_cfg": { diff --git a/tests/unit/experience/test_rollouts.py b/tests/unit/experience/test_rollouts.py index 038384a5958..40676612602 100644 --- a/tests/unit/experience/test_rollouts.py +++ b/tests/unit/experience/test_rollouts.py @@ -357,6 +357,9 @@ def initial_multi_step_calculator_batch(rollout_tokenizer): "temperature": 0.01, # Near-greedy "top_p": 1.0, "top_k": None, + "val_temperature": 0.01, + "val_top_p": 1.0, + "val_top_k": None, "stop_token_ids": None, "stop_strings": None, "vllm_cfg": { @@ -1188,6 +1191,9 @@ async def _collect(): "top_k": None, "temperature": 1.0, "top_p": 1.0, + "val_temperature": 1.0, + "val_top_p": 1.0, + "val_top_k": None, "max_new_tokens": 32, }, num_generations=2, diff --git a/tests/unit/models/generation/test_vllm_generation.py b/tests/unit/models/generation/test_vllm_generation.py index 0821f0865b6..93429ccfb5e 100644 --- a/tests/unit/models/generation/test_vllm_generation.py +++ b/tests/unit/models/generation/test_vllm_generation.py @@ -67,6 +67,9 @@ "temperature": 1.0, "top_p": 1.0, "top_k": None, + "val_temperature": 1.0, + "val_top_p": 1.0, + "val_top_k": None, "stop_token_ids": None, "stop_strings": None, "vllm_cfg": { diff --git a/tests/unit/models/generation/test_vllm_large_model.py b/tests/unit/models/generation/test_vllm_large_model.py index 89eaece234c..9b18a446ebf 100644 --- a/tests/unit/models/generation/test_vllm_large_model.py +++ b/tests/unit/models/generation/test_vllm_large_model.py @@ -38,6 +38,9 @@ "temperature": 0.8, "top_p": 1.0, "top_k": None, + "val_temperature": 0.8, + "val_top_p": 1.0, + "val_top_k": None, "stop_token_ids": None, "stop_strings": None, "vllm_cfg": { diff --git a/tests/unit/models/generation/test_vllm_quant_backend.py b/tests/unit/models/generation/test_vllm_quant_backend.py index 6b8bd6487c4..2cb80268661 100644 --- a/tests/unit/models/generation/test_vllm_quant_backend.py +++ b/tests/unit/models/generation/test_vllm_quant_backend.py @@ -57,6 +57,9 @@ def _make_vllm_config(tokenizer, *, async_engine=False, is_eval=True): "temperature": 0.0, "top_p": 1.0, "top_k": None, + "val_temperature": 0.0, + "val_top_p": 1.0, + "val_top_k": None, "stop_token_ids": None, "stop_strings": None, "quant_cfg": _QUANT_CFG, diff --git a/tests/unit/models/generation/trtllm/test_trtllm_generation.py b/tests/unit/models/generation/trtllm/test_trtllm_generation.py index d9dc1eaea70..1123ebae58a 100644 --- a/tests/unit/models/generation/trtllm/test_trtllm_generation.py +++ b/tests/unit/models/generation/trtllm/test_trtllm_generation.py @@ -43,6 +43,9 @@ def _config(**trtllm_overrides): "temperature": 1.0, "top_p": 1.0, "top_k": None, + "val_temperature": 1.0, + "val_top_p": 1.0, + "val_top_k": None, "stop_token_ids": None, "stop_strings": None, "_pad_token_id": 0, diff --git a/tests/unit/reference_configs/distillation_math.yaml b/tests/unit/reference_configs/distillation_math.yaml index 5b14bd9bb8b..9fe04653393 100644 --- a/tests/unit/reference_configs/distillation_math.yaml +++ b/tests/unit/reference_configs/distillation_math.yaml @@ -191,6 +191,9 @@ policy: &POLICY_BASE temperature: 1.0 top_p: 1.0 top_k: null + val_temperature: ${.temperature} + val_top_p: ${.top_p} + val_top_k: ${.top_k} stop_token_ids: null stop_strings: null vllm_cfg: diff --git a/tests/unit/reference_configs/eval.yaml b/tests/unit/reference_configs/eval.yaml index abe20f4d74d..95813a0a433 100644 --- a/tests/unit/reference_configs/eval.yaml +++ b/tests/unit/reference_configs/eval.yaml @@ -12,6 +12,9 @@ generation: temperature: 0.0 top_p: 1.0 top_k: -1 # -1 means disable + val_temperature: ${.temperature} + val_top_p: ${.top_p} + val_top_k: ${.top_k} num_prompts_per_step: -1 # -1 means pass all prompts at once model_name: "Qwen/Qwen2.5-Math-1.5B-Instruct" stop_token_ids: null diff --git a/tests/unit/reference_configs/grpo_math_1B.yaml b/tests/unit/reference_configs/grpo_math_1B.yaml index 5f4ffb2270d..7f026fca478 100644 --- a/tests/unit/reference_configs/grpo_math_1B.yaml +++ b/tests/unit/reference_configs/grpo_math_1B.yaml @@ -15,6 +15,7 @@ grpo: advantage_clip_low: null advantage_clip_high: null max_val_samples: 256 + val_num_generations_per_prompt: 1 # Early stop once this metric (e.g. accuracy or pass_k) reaches the threshold; null disables. stop_at_validation_metric: null # Required when stop_at_validation_metric is set. @@ -340,6 +341,9 @@ policy: temperature: 1.0 top_p: 1.0 top_k: null + val_temperature: ${.temperature} + val_top_p: ${.top_p} + val_top_k: ${.top_k} stop_token_ids: null stop_strings: null # null = topology default (IPC colocated, NCCL non-colocated). diff --git a/tests/unit/reference_configs/ppo_math_1B_megatron.yaml b/tests/unit/reference_configs/ppo_math_1B_megatron.yaml index 37a291f6d93..6af4416c2f5 100644 --- a/tests/unit/reference_configs/ppo_math_1B_megatron.yaml +++ b/tests/unit/reference_configs/ppo_math_1B_megatron.yaml @@ -223,6 +223,9 @@ policy: temperature: 1.0 top_p: 1.0 top_k: null + val_temperature: ${.temperature} + val_top_p: ${.top_p} + val_top_k: ${.top_k} stop_token_ids: null stop_strings: null mcore_generation_config: diff --git a/tests/unit/utils/test_native_checkpoint.py b/tests/unit/utils/test_native_checkpoint.py index 33240d92888..ba7b1151d64 100755 --- a/tests/unit/utils/test_native_checkpoint.py +++ b/tests/unit/utils/test_native_checkpoint.py @@ -73,6 +73,9 @@ "temperature": 1.0, "top_p": 1.0, "top_k": None, + "val_temperature": 1.0, + "val_top_p": 1.0, + "val_top_k": None, "backend": "vllm", "colocated": {"enabled": True}, }, From 154aebad23f8c2678cba8926ea978eadbaebf004 Mon Sep 17 00:00:00 2001 From: Rohit Jena Date: Thu, 6 Aug 2026 20:51:08 -0700 Subject: [PATCH 06/10] feat: Multimodal nemo gym compatible grpo pipeline (#3414) Signed-off-by: rohitrango Signed-off-by: Yi-Fu Wu Co-authored-by: Claude Opus 4.7 (1M context) Co-authored-by: Yi-Fu Wu --- ...-circle-click-2n8g-megatron-tp2ep8.v1.yaml | 160 ++++++++++++ examples/nemo_gym/grpo_nanov3.yaml | 2 +- examples/nemo_gym/run_grpo_nemo_gym.py | 26 +- nemo_rl/algorithms/grpo.py | 4 + nemo_rl/data/llm_message_utils.py | 23 +- nemo_rl/data/multimodal_utils.py | 88 +++++++ nemo_rl/data/processors.py | 28 +-- nemo_rl/environments/nemo_gym.py | 235 +++++++++++++++++- nemo_rl/models/megatron/setup.py | 8 + nemo_rl/models/policy/__init__.py | 2 + nemo_rl/utils/packed_tensor.py | 24 +- tests/test_suites/disabled.txt | 8 + ...3b-circle-click-2n8g-megatron-tp2ep8.v1.sh | 70 ++++++ tests/unit/data/datasets/test_mmpr_tiny.py | 18 +- tests/unit/data/test_llm_message_utils.py | 35 +++ .../data/test_multimodal_image_encoding.py | 100 ++++++++ .../environments/test_nemo_gym_mm_utils.py | 178 +++++++++++++ 17 files changed, 975 insertions(+), 34 deletions(-) create mode 100644 examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-circle-click-2n8g-megatron-tp2ep8.v1.yaml create mode 100755 tests/test_suites/vlm/vlm_grpo-nemotron-omni-30ba3b-circle-click-2n8g-megatron-tp2ep8.v1.sh create mode 100644 tests/unit/data/test_multimodal_image_encoding.py create mode 100644 tests/unit/environments/test_nemo_gym_mm_utils.py diff --git a/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-circle-click-2n8g-megatron-tp2ep8.v1.yaml b/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-circle-click-2n8g-megatron-tp2ep8.v1.yaml new file mode 100644 index 00000000000..3695064845c --- /dev/null +++ b/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-circle-click-2n8g-megatron-tp2ep8.v1.yaml @@ -0,0 +1,160 @@ +defaults: ../../vlm_grpo_3B_megatron.yaml +grpo: + num_prompts_per_step: 1 + num_val_generations_per_prompt: 1 + max_num_steps: 100 + val_period: 500 + overlong_filtering: true + max_val_samples: null + val_batch_size: null + seq_logprob_error_threshold: 2 +loss_fn: + reference_policy_kl_penalty: 0 + kl_input_clamp_value: null + kl_output_clamp_value: null + ratio_clip_max: 0.28 + use_on_policy_kl_approximation: true + use_importance_sampling_correction: true +checkpointing: + enabled: false + checkpoint_dir: results/grpo-nemotron-omni-30ba3b-gymv-circle-click + metric_name: val:total_reward/mean + keep_top_k: 1000000 + checkpoint_must_save_by: 00:03:40:00 +policy: + model_name: nvidia/Nemotron-3-Nano-Omni-30B-A3B-Reasoning-BF16 + is_vlm: true + tokenizer: + chat_template_kwargs: + enable_thinking: true + truncate_history_thinking: false + train_global_batch_size: ${mul:${grpo.num_prompts_per_step}, ${grpo.num_generations_per_prompt}} + generation_batch_size: 64 + logprob_batch_size: 1 + max_total_sequence_length: 8192 + logprob_chunk_size: 2048 + megatron_cfg: + activation_checkpointing: true + bias_activation_fusion: false + tensor_model_parallel_size: 2 + expert_model_parallel_size: 8 + context_parallel_size: 2 + sequence_parallel: true + moe_router_dtype: fp32 + moe_router_bias_update_rate: 0.001 + moe_aux_loss_coeff: 0.0 + moe_router_enable_expert_bias: true + defer_fp32_logits: true + track_moe_metrics: true + moe_per_layer_logging: true + radio_force_cpe_eval_mode: true + clear_memory_caches_before_refit: true + optimizer: + lr: 3.0e-06 + min_lr: 3.0e-06 + weight_decay: 0.0 + optimizer_cpu_offload: true + optimizer_offload_fraction: 1.0 + scheduler: + lr_decay_iters: null + lr_warmup_iters: 0 + lr_warmup_init: 3.0e-07 + distributed_data_parallel_config: + overlap_param_gather: false + average_in_collective: false + sequence_packing: + enabled: true + make_sequence_length_divisible_by: 32 + optimizer: null + scheduler: null + generation: + max_new_tokens: ${policy.max_total_sequence_length} + bad_words: [] + mcore_generation_config: + transformer_impl: inference_optimized + activation_checkpointing: false + mamba_inference_ssm_states_dtype: float32 + inference_moe_token_dispatcher_type: nccl + inference_grouped_gemm_backend: vllm + moe_router_num_groups: null + moe_router_group_topk: null + pipeline_model_parallel_size: 1 + expert_tensor_parallel_size: 1 + expert_model_parallel_size: 8 + sequence_parallel: true + context_parallel_size: 1 + tensor_model_parallel_size: 2 + buffer_size_gb: 20 + num_cuda_graphs: -1 + max_tokens: ${policy.max_total_sequence_length} + async_engine: true + expose_http_server: true + enable_prefix_caching: true + parsers: + - deepseek-r1-reasoning + - qwen3-coder-tool + vllm_cfg: + async_engine: true + tensor_parallel_size: 8 + enforce_eager: true + enable_prefix_caching: false + expose_http_server: true + reasoning_parser_plugin: nemo_rl/models/generation/vllm/reasoning_parsers/nano_v3_reasoning_parser.py + http_server_serving_chat_kwargs: + enable_auto_tools: true + tool_parser: qwen3_coder + reasoning_parser: nano_v3 + chat_template_content_format: string + default_chat_template_kwargs: + enable_thinking: true + truncate_history_thinking: false + vllm_kwargs: + limit_mm_per_prompt: + image: 1 + max_num_batched_tokens: 16384 + mamba_ssm_cache_dtype: float32 + compilation_config: + backend: eager + colocated: + enabled: false + resources: + gpus_per_node: 8 + num_nodes: 1 +data: + max_input_seq_length: null + shuffle: false + train: + data_path: /path/to/train_dataset.jsonl + validation: + data_path: /path/to/eval_dataset.jsonl + default: + dataset_name: NemoGymDataset + env_name: nemo_gym + prompt_file: null + processor: nemo_gym_data_processor +env: + should_use_nemo_gym: true + should_log_nemo_gym_responses: true + nemo_gym: + is_trajectory_collection: false + port_range_low: 5000 + port_range_high: 5999 + config_paths: + - responses_api_models/vllm_model/configs/vllm_model_for_training.yaml + - resources_servers/circle_click/configs/circle_click.yaml + circle_click_simple_agent: + responses_api_agents: + simple_agent: + max_steps: 1 +logger: + log_dir: logs/grpo-nemotron-omni-30ba3b-gymv-circle-click + wandb_enabled: true + wandb: + project: grpo-nemotron-omni-gymv + name: grpo-nemotron-omni-30ba3b-gymv-circle-click + mlflow: + experiment_name: grpo-nemotron-omni-gymv + run_name: grpo-nemotron-omni-30ba3b-gymv-circle-click +cluster: + gpus_per_node: 8 + num_nodes: 2 diff --git a/examples/nemo_gym/grpo_nanov3.yaml b/examples/nemo_gym/grpo_nanov3.yaml index 103d470f610..3c371660e41 100644 --- a/examples/nemo_gym/grpo_nanov3.yaml +++ b/examples/nemo_gym/grpo_nanov3.yaml @@ -240,7 +240,7 @@ policy: block_size_tokens: 256 # Size of each KV cache block in tokens (affects memory granularity) use_cuda_graphs_for_non_decode_steps: true # Enable CUDA graphs for prefill/context processing enable_chunked_prefill: true - max_tokens: 16384 # Maximum number of tokens to use in a single step. Analogous to vllm's max_num_batched_tokens + max_tokens: ${policy.max_total_sequence_length} # Maximum number of tokens to use in a single step. Analogous to vllm's max_num_batched_tokens kv_cache_management_mode: "persist" # KV cache lifecycle across suspend/resume. Options: "persist", "offload". To select "recompute", set grpo.async_grpo.recompute_kv_cache_after_weight_updates=true. materialize_only_last_token_logits: true num_speculative_tokens: 0 diff --git a/examples/nemo_gym/run_grpo_nemo_gym.py b/examples/nemo_gym/run_grpo_nemo_gym.py index 89081a6f948..4d83b9c129f 100644 --- a/examples/nemo_gym/run_grpo_nemo_gym.py +++ b/examples/nemo_gym/run_grpo_nemo_gym.py @@ -152,8 +152,14 @@ def main() -> None: ) with rl_init_timer.time("tokenizer"): - # setup tokenizer - tokenizer = get_tokenizer(config.policy["tokenizer"]) + is_vlm = bool(config.policy.get("is_vlm")) + if is_vlm: + processor = get_tokenizer(config.policy["tokenizer"], get_processor=True) + tokenizer = processor.tokenizer + else: + processor = None + tokenizer = get_tokenizer(config.policy["tokenizer"]) + assert config.policy["generation"] is not None, ( "A generation config is required for GRPO" ) @@ -171,6 +177,11 @@ def main() -> None: has_refit_draft_weights=has_refit_draft_weights, trains_mtp=trains_mtp, ) + if is_vlm and "vllm_cfg" in config.policy["generation"]: + assert not config.policy["generation"]["vllm_cfg"]["skip_tokenizer_init"], ( + "VLMs require tokenizer to be initialized before generation, " + "so skip_tokenizer_init must be set to False." + ) # NeMo-Gym specific config setup. setup_nemo_gym_config(config, tokenizer) @@ -181,8 +192,9 @@ def main() -> None: # NeMo-Gym environment needs to get dp_openai_server_base_urls from policy_generation, so we don't setup env here. with rl_init_timer.time("data"): print("\n▶ Setting up data...") + data_tokenizer = processor if processor is not None else tokenizer train_dataset, val_dataset = setup_response_data( - tokenizer, config.data, env_configs=None + data_tokenizer, config.data, env_configs=None ) # Validation dataset config setup. @@ -231,7 +243,13 @@ def main() -> None: master_config, teacher_worker_groups, alias_to_group_alias, - ) = setup(config, tokenizer, train_dataset, val_dataset) + ) = setup( + config, + tokenizer, + train_dataset, + val_dataset, + processor=processor, + ) rl_init_timer.record("total", time.perf_counter() - main_start) rl_init_metrics = rl_init_timer.get_timing_metrics(reduction_op="sum") diff --git a/nemo_rl/algorithms/grpo.py b/nemo_rl/algorithms/grpo.py index 278abeab7da..cc85990b55d 100644 --- a/nemo_rl/algorithms/grpo.py +++ b/nemo_rl/algorithms/grpo.py @@ -2042,6 +2042,9 @@ def _build_async_grpo_train_data( } ) _preserve_router_replay_routed_experts(train_data, flat_messages, policy_config) + # update multimodal data unconditionally + extra_multimodal_data = flat_messages.get_multimodal_dict(as_tensors=False) + train_data.update(extra_multimodal_data) return train_data @@ -3988,6 +3991,7 @@ def async_grpo_train( assert master_config.loss_fn.use_importance_sampling_correction, ( "Importance sampling correction must be enabled for async GRPO for good convergence due to off-policy samples!" ) + if router_replay_enabled(master_config.policy) and ( master_config.data_plane or {} ).get("enabled", False): diff --git a/nemo_rl/data/llm_message_utils.py b/nemo_rl/data/llm_message_utils.py index d3e1a4df903..c3e6bf76586 100644 --- a/nemo_rl/data/llm_message_utils.py +++ b/nemo_rl/data/llm_message_utils.py @@ -363,9 +363,26 @@ def batched_message_log_to_flat_message( result = BatchedDataDict() for key in all_keys: values = [seq.get(key) for seq in sequenced_lists] - # if the values are PackedTensors, create a new PackedTensor from the list of values - if values and isinstance(values[0], PackedTensor): - result[key] = PackedTensor.flattened_concat(values) + packed_template = next( + (value for value in values if isinstance(value, PackedTensor)), None + ) + if packed_template is not None: + if any( + value is not None and not isinstance(value, PackedTensor) + for value in values + ): + raise TypeError( + f"Expected PackedTensor or None for {key=}, " + f"got {[type(value).__name__ for value in values]}" + ) + filled_packed_values = cast( + list[PackedTensor], + [ + PackedTensor.empty_like(packed_template) if value is None else value + for value in values + ], + ) + result[key] = PackedTensor.flattened_concat(filled_packed_values) continue if not values or not isinstance(values[0], Tensor): result[key] = values diff --git a/nemo_rl/data/multimodal_utils.py b/nemo_rl/data/multimodal_utils.py index 470618d1e74..368608f5909 100644 --- a/nemo_rl/data/multimodal_utils.py +++ b/nemo_rl/data/multimodal_utils.py @@ -45,6 +45,13 @@ "audio": ["wav", "flac", "mp3"], } +_PLACEHOLDER_STYLE_PROCESSOR_NAMES = frozenset( + { + "NemotronNanoVLV2Processor", + "NemotronH_Nano_Omni_Reasoning_V3Processor", + } +) + # different media namings maybe used in the raw dataset, # in which case, they need to be mapped to the allowed ones @@ -67,6 +74,19 @@ logger = logging.getLogger(__name__) +def uses_image_placeholder(processor: Any) -> bool: + """Return whether a processor requires explicit image placeholders. + + Args: + processor: Multimodal processor to classify. + + Returns: + Whether the processor expands image placeholders through ``__call__`` + rather than tokenized ``apply_chat_template``. + """ + return type(processor).__name__ in _PLACEHOLDER_STYLE_PROCESSOR_NAMES + + class PackedTensor: """Wrapper around a list of torch tensors and a dimension along which to pack the tensors. @@ -376,11 +396,79 @@ def resolve_to_image(image_path_or_image: str | Image.Image) -> Image.Image: header, encoded = image_path_or_image.split(",", 1) image_data = base64.b64decode(encoded) return Image.open(BytesIO(image_data)).convert("RGB") + elif image_path_or_image.startswith("file://"): + return Image.open(image_path_or_image.removeprefix("file://")).convert("RGB") else: # Handle local file path return Image.open(image_path_or_image).convert("RGB") +def image_to_data_url(image: Image.Image, fmt: str = "PNG") -> str: + """Encode a PIL Image as a base64 ``data:`` URL. + + Args: + image: PIL image to encode. + fmt: PIL image format used for serialization (e.g. ``"PNG"``, ``"JPEG"``). + The value is also lowercased and embedded in the MIME type of the + returned URL. + + Returns: + A ``data:image/;base64,`` URL suitable for embedding in + an OpenAI Responses ``input_image`` content part. + """ + buf = BytesIO() + image.save(buf, format=fmt) + encoded = base64.b64encode(buf.getvalue()).decode("utf-8") + return f"data:image/{fmt.lower()};base64,{encoded}" + + +def encode_images_in_examples(nemo_gym_examples: list[dict]) -> list[dict]: + """Replace local image paths in NeMo Gym examples with base64 data URLs. + + Walks each example's ``responses_create_params.input[].content[]`` items + and rewrites any ``input_image`` part whose ``image_url`` is a local path + (or ``file://`` URL) into a base64 ``data:`` URL via + :func:`image_to_data_url`. Parts whose URL already starts with ``http://``, + ``https://``, or ``data:`` are left untouched. Malformed items (non-dict + entries, missing/empty URLs, non-list ``input``/``content``) are skipped + without raising. + + The examples are mutated in place; the same list is also returned for + convenience so callers can chain the call. + + Args: + nemo_gym_examples: List of NeMo Gym example dicts. Each example is + expected to contain a ``responses_create_params`` mapping with an + ``input`` list of Responses API messages. + + Returns: + The same ``nemo_gym_examples`` list, with local image references + rewritten to base64 data URLs in place. + """ + for example in nemo_gym_examples: + input_items = example.get("responses_create_params", {}).get("input", []) + if not isinstance(input_items, list): + continue + for item in input_items: + if not isinstance(item, dict): + continue + content = item.get("content", []) + if not isinstance(content, list): + continue + for part in content: + if not isinstance(part, dict) or part.get("type") != "input_image": + continue + url = part.get("image_url", "") + if isinstance(url, dict): + url = url.get("url", "") + if not isinstance(url, str) or not url: + continue + if url.startswith(("http://", "https://", "data:")): + continue + part["image_url"] = image_to_data_url(resolve_to_image(url)) + return nemo_gym_examples + + def get_media_from_message(message: dict[str, Any]) -> dict[str, list[Any]]: """Get all media from a message log item.""" # Handle None or missing content (e.g., assistant messages with only tool_calls) diff --git a/nemo_rl/data/processors.py b/nemo_rl/data/processors.py index a6a072b0357..a6e8ef14727 100644 --- a/nemo_rl/data/processors.py +++ b/nemo_rl/data/processors.py @@ -464,6 +464,7 @@ def vlm_hf_data_processor( get_multimodal_default_settings_from_processor, get_multimodal_keys_from_processor, resolve_to_image, + uses_image_placeholder, ) # depending on the task, format the data differently @@ -555,13 +556,9 @@ def vlm_hf_data_processor( # vs OpenAI content list style (e.g., Qwen-VL, Gemma). # These processors expand tokens in __call__ but NOT in apply_chat_template, # so we must use processor(text=..., images=...) directly. - _PLACEHOLDER_STYLE_PROCESSORS = ( - "NemotronNanoVLV2Processor", - "NemotronH_Nano_Omni_Reasoning_V3Processor", - ) - _uses_image_placeholder = type(processor).__name__ in _PLACEHOLDER_STYLE_PROCESSORS + uses_placeholder = uses_image_placeholder(processor) - if _uses_image_placeholder and images: + if uses_placeholder and images: # Convert content list to placeholder text format image_token = getattr(processor, "image_token", "") text_parts = [] @@ -582,27 +579,25 @@ def vlm_hf_data_processor( else: user_message_for_chat_template = user_message_for_tokenize - # this is the string-tokenized conversation template for the generation policy (for vllm) string_formatted_dialog = processor.apply_chat_template( [user_message_for_chat_template], tokenize=False, add_generation_prompt=True, ) - # this is the id-tokenized and image processed conversation template for the policy - if _uses_image_placeholder and images: + if uses_placeholder and images: # Dynamic-resolution path: keep pixel_values in float32 to match vLLM's # DynamicResolutionImageTiler bit-for-bit. vLLM stores/normalizes in # float32 and only casts at the vision_model boundary; matching that # rounding order tightens rollout/train logprob agreement. The model # forward dispatches on imgs_sizes and handles the bf16 cast. - message: dict = processor( + message = processor( text=string_formatted_dialog, images=images, return_tensors="pt", ) else: - message: dict = processor.apply_chat_template( + message = processor.apply_chat_template( [user_message_for_tokenize], tokenize=True, add_generation_prompt=True, @@ -620,7 +615,7 @@ def vlm_hf_data_processor( # the Nemotron Omni path can patchify it and preserve the processor's exact # placeholder count. if ( - _uses_image_placeholder + uses_placeholder and "pixel_values" in message and "imgs_sizes" not in message and message["pixel_values"].ndim == 4 @@ -646,7 +641,7 @@ def vlm_hf_data_processor( user_message[key] = PackedTensor( message[key], dim_to_pack=get_dim_to_pack_along(processor, key), - pad_to_max_shape=_uses_image_placeholder and key == "pixel_values", + pad_to_max_shape=uses_placeholder and key == "pixel_values", ) # specifically for gemma, we need to add token_type_ids to the user message as a sequence-type value @@ -789,7 +784,12 @@ def nemo_gym_data_processor( max_seq_length: int | None, idx: int, ) -> DatumSpec: - """Process a datum dictionary (directly loaded from dataset) into a DatumSpec for Nemo Gym.""" + """Process a datum dictionary (directly loaded from dataset) into a DatumSpec for Nemo Gym. + + NeMo-Gym builds the real cumulative prompt server-side. Both LLM and VLM + rows therefore use a placeholder here; VLM inputs are processed once after + the complete rollout has been collected. + """ output: DatumSpec = { # load to dict format here since `Dataset` cannot handle nested structure well in `NemoGymDataset` "extra_env_info": json.loads(datum_dict["extra_env_info"]), diff --git a/nemo_rl/environments/nemo_gym.py b/nemo_rl/environments/nemo_gym.py index 1b3daba4612..375d3b2d41d 100644 --- a/nemo_rl/environments/nemo_gym.py +++ b/nemo_rl/environments/nemo_gym.py @@ -22,9 +22,18 @@ import ray import torch +from PIL import Image from ray.util.scheduling_strategies import NodeAffinitySchedulingStrategy from transformers import PreTrainedTokenizerBase +from nemo_rl.data.multimodal_utils import ( + PackedTensor, + encode_images_in_examples, + get_dim_to_pack_along, + get_multimodal_keys_from_processor, + resolve_to_image, + uses_image_placeholder, +) from nemo_rl.distributed.ray_actor_environment_registry import get_actor_python_env from nemo_rl.distributed.virtual_cluster import ( DEFAULT_GYM_PORT_RANGE_HIGH, @@ -33,6 +42,7 @@ _get_node_ip_local, ) from nemo_rl.environments.interfaces import EnvironmentInterface +from nemo_rl.models.policy import TokenizerConfig from nemo_rl.utils.routed_experts_codec import decode_routed_experts from nemo_rl.utils.timer import Timer from nemo_rl.utils.venvs import create_local_venv_on_each_node @@ -112,6 +122,10 @@ class NemoGymConfig(TypedDict): # Forwarded from policy.tokenizer.use_fastokens so rollout actors patch their # tokenizer consistently with the driver. Defaults to off when absent. use_fastokens: NotRequired[bool] + # Multimodal fields (populated by `setup_nemo_gym_config` when VLM is enabled). + tokenizer_config: NotRequired[ + Optional[TokenizerConfig] + ] # For processor reconstruction inside the actor def _detect_invalid_tool_call_and_malformed_thinking( @@ -174,12 +188,190 @@ def _detect_invalid_tool_call_and_malformed_thinking( return is_invalid_tool_call, has_malformed_thinking +######################################## +# Multimodal helpers +######################################## + + +# WARNING: A function-call output beginning with HTTP(S) is accepted here and +# passed to ``resolve_to_image``, which performs an outbound request during +# postprocessing even when the tool result is not actually an image. +_IMAGE_SRC_PREFIXES = ("data:image/", "http://", "https://", "file://") + + +def _looks_like_image_src(src: str) -> bool: + """True when ``src`` plausibly points at an image the loader can open. + + Guards against tool responses (e.g. ``{"x": 0.65, "y": 0.83}`` from a + click tool) that are strings but not image URLs. Without this, the + indexer forwards the JSON payload to ``resolve_to_image`` → PIL.open, + which treats it as a filesystem path and raises ``FileNotFoundError``. + """ + return src.startswith(_IMAGE_SRC_PREFIXES) + + +def _extract_input_images_from_message(item: dict) -> list[Image.Image]: + """Pull PIL images out of a non-assistant Responses-API item. + + Handles both content-list items (user / tool messages carrying + ``input_image``/``image``/``image_url`` parts) and ``function_call_output`` + items whose ``output`` field is an image data URL. Tool outputs that are + non-image strings (e.g. structured JSON returned by tools like + ``click(x, y)``) contribute zero images to the bucket. + """ + images: list[Image.Image] = [] + if item.get("type") == "function_call_output": + src = item.get("output") + if isinstance(src, str) and _looks_like_image_src(src): + images.append(resolve_to_image(src)) + return images + content = item.get("content") or [] + if not isinstance(content, list): + return images + for part in content: + if not isinstance(part, dict): + continue + if part.get("type") not in ("input_image", "image", "image_url"): + continue + src = part.get("image") or part.get("image_url") or part.get("url") + if src is None: + continue + if isinstance(src, dict): + src = src.get("url") + if src is None: + continue + images.append(resolve_to_image(src)) + return images + + +def _index_per_turn_images( + output: list[dict], + input_messages: list[dict] | None = None, +) -> list[list[Image.Image]]: + """Bin server-returned images by the trainable turn that saw them. + + Walks the Responses-API items in order and flushes ``pending`` into a + per-turn bucket each time it hits an item carrying truthy + ``generation_token_ids`` — matching the exact gate that + ``_postprocess_nemo_gym_to_nemo_rl_result`` uses to decide which items + become trainable turns. Every other item (user turns, tool messages, + ``function_call_output``, non-trainable reasoning) contributes its images + to ``pending`` for the next trainable turn. This ensures the returned list + has one entry per trainable turn, aligned with the postprocess loop's + ``turn_idx`` even when the trainable item's role is not ``assistant`` + (e.g. a reasoning-only response, or a ``function_call``). + + ``input_messages`` is the initial ``responses_create_params.input`` list — + images there (e.g. a single-shot user prompt for tool-based envs like + circle-click) are consumed by the first trainable turn's tokenized prompt + and must land in the first bucket. Agents like ``gym_v_agent`` that keep + ``input`` empty and inject observations as ``function_call_output`` items + are unaffected — the seed is a no-op when ``input_messages`` is empty. + """ + per_turn: list[list[Image.Image]] = [] + pending: list[Image.Image] = [] + for item in input_messages or (): + if isinstance(item, dict) and item.get("role") != "assistant": + pending.extend(_extract_input_images_from_message(item)) + for item in output: + if item.get( + "generation_token_ids" + ): # trainable turn; empty generation_token_ids is skipped by the postprocess loop and must not consume a bucket + per_turn.append(pending) + pending = [] + elif item.get("role") != "assistant": + pending.extend(_extract_input_images_from_message(item)) + return per_turn + + +def _attach_multimodal_data_to_user_message( + user_message: dict, + *, + images: list[Image.Image], + processor: Any, +) -> None: + """Attach per-turn multimodal tensors to ``user_message``. + + The processor is only invoked to extract multimodal tensors (pixel_values, + imgs_sizes, num_patches, etc.); its text output is discarded — vLLM's + tokens remain the trajectory. We therefore feed it the minimal placeholder + text it needs to count image regions: one ``processor.image_token`` per + image. Passing the vLLM-decoded text does not work because that text + already contains expanded ``...*N...`` regions, and the + processor would try to re-expand every embedded ````. + """ + if not images or processor is None: + return + image_token = getattr(processor, "image_token", "") + processed = processor( + text=image_token * len(images), + images=images, + return_tensors="pt", + ) + uses_placeholder = uses_image_placeholder(processor) + multimodal_keys = list(get_multimodal_keys_from_processor(processor)) + # Historical checkpoints may emit dynamic image tiles without imgs_sizes. + # Mirror the media-metadata handling in vlm_hf_data_processor. + if ( + uses_placeholder + and "pixel_values" in processed + and "imgs_sizes" not in processed + and processed["pixel_values"].ndim == 4 + ): + pixel_values = processed["pixel_values"] + num_tiles, _, height, width = pixel_values.shape + processed["imgs_sizes"] = torch.tensor( + [[height, width]] * num_tiles, dtype=torch.long + ) + + # imgs_sizes / num_frames are not always declared in model_input_names by + # bundled image processors. RADIO uses temporal patching even for still + # images and requires one num_frames=1 entry per image/tile. + if "imgs_sizes" in processed and "imgs_sizes" not in multimodal_keys: + multimodal_keys.append("imgs_sizes") + if "imgs_sizes" in processed and "num_frames" not in processed: + processed["num_frames"] = torch.ones( + len(processed["imgs_sizes"]), dtype=torch.long + ) + if "num_frames" in processed and "num_frames" not in multimodal_keys: + multimodal_keys.append("num_frames") + for key in multimodal_keys: + if key not in processed: + continue + value = processed[key] + if key == "imgs_sizes": + value = value.to(dtype=torch.int32) + user_message[key] = PackedTensor( + value, + dim_to_pack=get_dim_to_pack_along(processor, key), + pad_to_max_shape=uses_placeholder and key == "pixel_values", + ) + + @ray.remote(max_restarts=-1, max_task_retries=-1) # pragma: no cover class NemoGym(EnvironmentInterface): """This environment class isn't really used for training. It's really meant as an integration wrapper around NeMo-Gym that hooks into the existing NeMo RL resource management via ray. So there is still one source of truth for resource management in NeMo RL.""" def __init__(self, cfg: NemoGymConfig): self.cfg = cfg + # Reconstruct the processor inside the actor (rather than serializing it + # per rollout call) for full-trajectory multimodal postprocessing. + self._processor: Optional[Any] = None + tokenizer_config = cfg.get("tokenizer_config") + if tokenizer_config: + from nemo_rl.algorithms.utils import get_tokenizer + + self._processor = get_tokenizer(tokenizer_config, get_processor=True) + # _attach_multimodal_data_to_user_message assumes a placeholder-style + # processor (imgs_sizes / num_frames reconstruction + pad_to_max_shape + # PackedTensor build). A non-placeholder VLM would silently produce + # wrong multimodal tensors — fail at actor construction instead. + assert uses_image_placeholder(self._processor), ( + "NemoGym multimodal path assumes a placeholder-style processor " + "(see _PLACEHOLDER_STYLE_PROCESSOR_NAMES in nemo_rl/data/multimodal_utils.py); " + f"got {type(self._processor).__name__}. Update " + "_attach_multimodal_data_to_user_message before enabling." + ) def _spinup(self) -> None: """Start the NeMo-Gym head server and rollout collection helper. @@ -293,6 +485,11 @@ async def run_rollouts( timer = Timer() counts_left = Counter(row["agent_ref"]["name"] for row in nemo_gym_examples) + # For multimodal runs, replace local filesystem image paths in the + # examples with base64 data URLs before shipping to vLLM. No-op when + # examples carry no `input_image` items (text-only case). + encode_images_in_examples(nemo_gym_examples) + timer.start("_run_rollouts_total") nemo_gym_result_iterator = self.rch.run_examples( examples=nemo_gym_examples, head_server_config=self.head_server_config @@ -350,12 +547,27 @@ async def run_rollouts( yield nemo_gym_row["_rowidx"], nemo_rl_result, timing_metrics def _postprocess_nemo_gym_to_nemo_rl_result( - self, nemo_gym_result: dict, tokenizer: PreTrainedTokenizerBase + self, + nemo_gym_result: dict, + tokenizer: PreTrainedTokenizerBase, ) -> dict: assert isinstance(nemo_gym_result, dict), ( f"Hit a non-successful response when querying NeMo Gym for rollouts: {nemo_gym_result}" ) + processor = getattr(self, "_processor", None) + per_turn_images = ( + _index_per_turn_images( + nemo_gym_result["response"]["output"], + input_messages=nemo_gym_result.get("responses_create_params", {}).get( + "input" + ), + ) + if processor is not None + else [] + ) + turn_idx = 0 + nemo_rl_message_log = [] seen_token_ids: List[int] = [] batch_decode_items = [] @@ -378,6 +590,7 @@ def _postprocess_nemo_gym_to_nemo_rl_result( ), f"""Non-contiguous messages found! This may be a tokenization issue where certain tokens are combined when messages are concatenated, or it may be due to part of the chat history being truncated (like if super long history is truncated or if reasoning is stripped out). Seen token IDs: {seen_token_ids} Output prompt token IDs: {output_item_dict["prompt_token_ids"]} +output prompt token ids till seen: {output_item_dict["prompt_token_ids"][: len(seen_token_ids)]} """ prompt_token_ids = output_item_dict.pop("prompt_token_ids") @@ -437,6 +650,16 @@ def _postprocess_nemo_gym_to_nemo_rl_result( if routed_experts is not None: user_message["routed_experts"] = routed_experts[prompt_start:prompt_end] nemo_rl_message_log.append(user_message) + + if processor is not None: + images_this_turn = ( + per_turn_images[turn_idx] if turn_idx < len(per_turn_images) else [] + ) + _attach_multimodal_data_to_user_message( + user_message, + images=images_this_turn, + processor=processor, + ) # 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. @@ -471,6 +694,7 @@ def _postprocess_nemo_gym_to_nemo_rl_result( batch_decode_items.append( (output_item_dict, prompt_token_ids, generation_token_ids) ) + turn_idx += 1 if batch_decode_items: prompt_strs = tokenizer.batch_decode( @@ -616,6 +840,13 @@ def setup_nemo_gym_config(config, tokenizer) -> None: generation_config["stop_strings"] = None generation_config["stop_token_ids"] = None + # For VLM runs, plumb the tokenizer config into the gym env config so the + # NemoGym actor can reconstruct the processor inside itself (needed for + # multi-turn multimodal postprocessing). + if config.policy.get("is_vlm"): + env_cfg = config.env.setdefault("nemo_gym", {}) + env_cfg.setdefault("tokenizer_config", dict(config.policy["tokenizer"])) + def spinup_nemo_gym_actor( env_configs: dict[str, Any], @@ -653,6 +884,7 @@ def spinup_nemo_gym_actor( # (where the detector reads them), not part of Gym's global config. invalid_tool_call_patterns = nemo_gym_dict.pop("invalid_tool_call_patterns", None) thinking_tags = nemo_gym_dict.pop("thinking_tags", None) + tokenizer_config = nemo_gym_dict.pop("tokenizer_config", None) # Pass prebuilt cache + venv dirs through the global config so the gym reuses # image-baked venvs instead of rebuilding them. @@ -668,6 +900,7 @@ def spinup_nemo_gym_actor( base_urls=base_urls, invalid_tool_call_patterns=invalid_tool_call_patterns, thinking_tags=thinking_tags, + tokenizer_config=tokenizer_config, require_routed_experts=enable_router_replay, routed_experts_dtype=routed_experts_dtype, use_fastokens=use_fastokens, diff --git a/nemo_rl/models/megatron/setup.py b/nemo_rl/models/megatron/setup.py index 2776acffbef..e8faea4f9cf 100644 --- a/nemo_rl/models/megatron/setup.py +++ b/nemo_rl/models/megatron/setup.py @@ -1407,6 +1407,14 @@ def freeze_moe_router(megatron_model): # Handle VLM models if hasattr(model_module, "thinker"): model_module = model_module.thinker + # NemotronVLModel / NemotronOmniModel wrap the GPT under + # `.llava_model.language_model`; unwrap that layer first so the + # generic `.language_model.decoder.layers` walk below finds the + # MoE router. + if getattr(model_module, "llava_model", None) is not None and hasattr( + model_module.llava_model, "language_model" + ): + model_module = model_module.llava_model if hasattr(model_module, "language_model"): model_module = model_module.language_model for layer in model_module.decoder.layers: diff --git a/nemo_rl/models/policy/__init__.py b/nemo_rl/models/policy/__init__.py index 6248f4b2c71..7994f37d7b6 100644 --- a/nemo_rl/models/policy/__init__.py +++ b/nemo_rl/models/policy/__init__.py @@ -562,3 +562,5 @@ class PolicyConfig(TypedDict): # If true, use standard Megatron layer specs while keeping ModelOpt # quantization enabled. Useful for faster QARL runs and logged in configs. disable_modelopt_layer_spec: NotRequired[bool] + + is_vlm: NotRequired[bool] diff --git a/nemo_rl/utils/packed_tensor.py b/nemo_rl/utils/packed_tensor.py index 01f58c55a32..c0dd856e9d6 100644 --- a/nemo_rl/utils/packed_tensor.py +++ b/nemo_rl/utils/packed_tensor.py @@ -77,14 +77,16 @@ def packed_broadcast_producer(iterator, group, src, post_iter_func): # Apply backend specific post processing and then convert to linearized uint8 tensor. # contiguous() is required because the upstream iterator may # yield non-contiguous tensors that view(...) cannot handle. - tensor = ( - post_iter_func(next(iterator)) - .contiguous() - .view(torch.uint8) - .view(-1) - ) + # 0-D tensors (e.g. BN `num_batches_tracked` counters on + # Nemotron-Omni's sound encoder) must be reshape(1)-ed + # before `.view(torch.uint8)` — Long→Byte view is illegal + # on scalars. + tensor = post_iter_func(next(iterator)).contiguous() + if tensor.dim() == 0: + tensor = tensor.reshape(1) + tensor = tensor.view(torch.uint8).view(-1) packing_tensor_list[buffer_idx].append(tensor) - packing_tensor_sizes[buffer_idx] += tensor.view(torch.uint8).numel() + packing_tensor_sizes[buffer_idx] += tensor.numel() if packing_tensor_sizes[buffer_idx] > target_packed_tensor_size: break # Pack the tensors and call broadcast collective @@ -140,11 +142,15 @@ def unpack_tensor( packed_tensor_sizes = list(map(lambda x: x[4], meta_data_list)) unpacked_tensor = packed_tensor.split_with_sizes(packed_tensor_sizes) - # unpacked_list = List[(name, torch.Tensor.view(dtype).view(*shape))] + # unpacked_list = List[(name, torch.Tensor.view(dtype).reshape(shape))] + # reshape(tuple) accepts an empty tuple for 0-D targets, whereas + # view(*shape) would call view() with no args and raise. Producer + # side reshapes 0-D tensors to (1,) before packing, and this consumer + # must reshape back to the original 0-D shape stored in meta_data. unpacked_list = [ ( meta_data_list[i][0], - tensor.view(meta_data_list[i][2]).view(*meta_data_list[i][1]), + tensor.view(meta_data_list[i][2]).reshape(tuple(meta_data_list[i][1])), ) for i, tensor in enumerate(unpacked_tensor) ] diff --git a/tests/test_suites/disabled.txt b/tests/test_suites/disabled.txt index fe460c2185f..14c0c1e28df 100644 --- a/tests/test_suites/disabled.txt +++ b/tests/test_suites/disabled.txt @@ -5,3 +5,11 @@ # grpo-qwen3.5-35ba3b-2n8g-megatron-ep16tp2cp2 run hangs the same way on main, # so this is the pre-existing Qwen3.5 + Megatron + EP hang. tests/test_suites/vlm/vlm_grpo-qwen3.5-35ba3b-geo3k-2n8g-megatron-ep16.sh + +# First multimodal NeMo-Gym recipe (circle_click at the pinned Gym). Disabled on +# landing for two reasons: the recipe has not been run end to end yet, so its +# reward threshold is an unvalidated smoke bound; and nightly.txt is at 3755 of +# its 3800 GPU-hour budget, which this run's 32 GPU-hours would leave only 13 to +# spare. Move to nightly.txt once a real run confirms it converges and the +# budget has room. +tests/test_suites/vlm/vlm_grpo-nemotron-omni-30ba3b-circle-click-2n8g-megatron-tp2ep8.v1.sh diff --git a/tests/test_suites/vlm/vlm_grpo-nemotron-omni-30ba3b-circle-click-2n8g-megatron-tp2ep8.v1.sh b/tests/test_suites/vlm/vlm_grpo-nemotron-omni-30ba3b-circle-click-2n8g-megatron-tp2ep8.v1.sh new file mode 100755 index 00000000000..f7c4214e7cb --- /dev/null +++ b/tests/test_suites/vlm/vlm_grpo-nemotron-omni-30ba3b-circle-click-2n8g-megatron-tp2ep8.v1.sh @@ -0,0 +1,70 @@ +#!/bin/bash +SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd) +source $SCRIPT_DIR/common.env + +# ===== BEGIN CONFIG ===== +NUM_NODES=2 +GPUS_PER_NODE=8 +STEPS_PER_RUN=10 +MAX_STEPS=10 +NUM_RUNS=$(( (MAX_STEPS + STEPS_PER_RUN - 1) / STEPS_PER_RUN )) # Round up +# 30B MoE across 2 nodes, plus nemo_gym head-server startup and vLLM warmup on top +# of the 10 steps; 120 min leaves margin for teardown + metric dump. +NUM_MINUTES=120 +# ===== END CONFIG ===== + +exit_if_max_steps_reached + +cd $PROJECT_ROOT + +# circle_click generates its own data (no HF download). Regenerate rather than reuse the +# 5-row example.jsonl so the run never trains on stale committed data, and so train/eval +# are disjoint (distinct --seed-offset). +DATA_DIR=$EXP_DIR/data +mkdir -p $DATA_DIR +GYM_DIR=3rdparty/Gym-workspace/Gym +RAW_TRAIN=$DATA_DIR/circle_click_train_raw.jsonl +RAW_VALIDATION=$DATA_DIR/circle_click_validation_raw.jsonl +( cd $GYM_DIR && uv run python resources_servers/circle_click/generate_data.py \ + --n 512 --seed-offset 0 --out $PROJECT_ROOT/$RAW_TRAIN ) +( cd $GYM_DIR && uv run python resources_servers/circle_click/generate_data.py \ + --n 32 --seed-offset 100000 --out $PROJECT_ROOT/$RAW_VALIDATION ) + +# Attach `agent_ref` so rollouts are routed to the env's agent. The name must match the +# group registered in resources_servers/circle_click/configs/circle_click.yaml. +TRAIN_PATH=$DATA_DIR/circle_click_train.jsonl +VALIDATION_PATH=$DATA_DIR/circle_click_validation.jsonl +jq -c '. + {agent_ref: {name: "circle_click_simple_agent"}}' $RAW_TRAIN > $TRAIN_PATH +jq -c '. + {agent_ref: {name: "circle_click_simple_agent"}}' $RAW_VALIDATION > $VALIDATION_PATH + +# Run the experiment via the gym entrypoint (circle_click is a NeMo-Gym env, so this +# recipe runs through run_grpo_nemo_gym.py rather than run_vlm_grpo.py). +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 + +# Convert tensorboard logs to json +uv run tests/json_dump_tb_logs.py $LOG_DIR --output_path $JSON_METRICS + +# Only run metrics if the target step is reached +if [[ $(jq 'to_entries | .[] | select(.key == "train/loss") | .value | keys | map(tonumber) | max' $JSON_METRICS) -ge $MAX_STEPS ]]; then + # Smoke-level threshold: this recipe has not been run end to end yet, so assert only + # that the multimodal gym path produces non-zero reward. Tighten once real runs land. + uv run tests/check_metrics.py $JSON_METRICS \ + 'max(data["train/reward"]) > 0.0' + + # Clean up checkpoint directory after successful run to save space. + rm -rf "$CKPT_DIR" +fi diff --git a/tests/unit/data/datasets/test_mmpr_tiny.py b/tests/unit/data/datasets/test_mmpr_tiny.py index 74822398d7b..4b0c9e4fc28 100644 --- a/tests/unit/data/datasets/test_mmpr_tiny.py +++ b/tests/unit/data/datasets/test_mmpr_tiny.py @@ -194,7 +194,7 @@ def tiny_image_path(tmp_path): ) -def _run_processor(tiny_image_path): +def _run_processor(tiny_image_path, processor=None): """Helper: run vlm_hf_data_processor on an MMPR sample and return (result DatumSpec, stub processor with captured_call_text).""" from nemo_rl.data.interfaces import TaskDataSpec @@ -202,7 +202,7 @@ def _run_processor(tiny_image_path): task_data_spec = TaskDataSpec(task_name="mmpr-tiny") task_data_spec.prompt = _TEST_PROMPT_TEMPLATE - processor = _make_stub_nemotron_processor() + processor = processor or _make_stub_nemotron_processor() sample = { "images": [tiny_image_path], "question": _RAW_QUESTION, @@ -238,6 +238,20 @@ def test_processor_produces_valid_datum_spec(self, tiny_image_path): assert result["task_name"] == "mmpr-tiny" user_message = result["message_log"][0] assert torch.equal(user_message["num_frames"].as_tensor(), torch.tensor([1])) + assert user_message["pixel_values"].pad_to_max_shape is True + assert user_message["pixel_values"].as_tensor().dtype == torch.float32 + + def test_conversation_preprocessor_is_preserved(self, tiny_image_path): + processor = _make_stub_nemotron_processor() + processor.conversation_preprocessor = MagicMock( + return_value={"role": "user", "content": "preprocessed"} + ) + + result, _ = _run_processor(tiny_image_path, processor=processor) + + processor.conversation_preprocessor.assert_called_once() + assert result["vllm_content"] == "preprocessed" + assert processor.captured_call_text == "preprocessed" def test_historical_tiled_processor_gets_media_metadata(self, tiny_image_path): from nemo_rl.data.interfaces import TaskDataSpec diff --git a/tests/unit/data/test_llm_message_utils.py b/tests/unit/data/test_llm_message_utils.py index b39f1175934..113fd9ce0b9 100644 --- a/tests/unit/data/test_llm_message_utils.py +++ b/tests/unit/data/test_llm_message_utils.py @@ -772,6 +772,41 @@ def test_batched_message_log_to_flat_message_with_packed_images() -> None: assert torch.equal(input_lengths, torch.tensor([4, 5], dtype=torch.int32)) +@pytest.mark.parametrize("image_first", [True, False]) +def test_batched_message_log_to_flat_message_with_image_free_sample( + image_first: bool, +) -> None: + from nemo_rl.data.multimodal_utils import PackedTensor + + image = torch.randn(1, 3, 4, 4) + image_log: LLMMessageLogType = [ + { + "role": "user", + "token_ids": torch.tensor([1, 2]), + "pixel_values": PackedTensor(image, dim_to_pack=0), + } + ] + image_free_log: LLMMessageLogType = [ + {"role": "user", "token_ids": torch.tensor([3, 4])} + ] + batch_logs = ( + [image_log, image_free_log] if image_first else [image_free_log, image_log] + ) + + batched, _ = batched_message_log_to_flat_message(batch_logs) + + pixel_values = batched["pixel_values"] + assert isinstance(pixel_values, PackedTensor) + assert len(pixel_values) == 2 + expected = [image, None] if image_first else [None, image] + for actual, expected_value in zip(pixel_values.tensors, expected): + if expected_value is None: + assert actual is None + else: + assert torch.equal(actual, expected_value) + assert "pixel_values" in batched.get_multimodal_dict() + + @pytest.mark.hf_gated def test_get_formatted_message_log_multimodal_prompt_formatting() -> None: processor = AutoProcessor.from_pretrained("Qwen/Qwen2.5-VL-3B-Instruct") diff --git a/tests/unit/data/test_multimodal_image_encoding.py b/tests/unit/data/test_multimodal_image_encoding.py new file mode 100644 index 00000000000..e36c64135e0 --- /dev/null +++ b/tests/unit/data/test_multimodal_image_encoding.py @@ -0,0 +1,100 @@ +# 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. + +from PIL import Image + +from nemo_rl.data.multimodal_utils import ( + encode_images_in_examples, + image_to_data_url, + resolve_to_image, +) + + +def _example(*content_parts: dict) -> dict: + return { + "responses_create_params": { + "input": [{"role": "user", "content": list(content_parts)}] + } + } + + +def _write_png(tmp_path, name: str, size: tuple[int, int]) -> str: + path = tmp_path / name + Image.new("RGB", size, color=(10, 20, 30)).save(path, format="PNG") + return str(path) + + +def test_image_to_data_url_round_trips_through_resolve_to_image(): + url = image_to_data_url(Image.new("RGB", (4, 3))) + assert url.startswith("data:image/png;base64,") + assert resolve_to_image(url).size == (4, 3) + + +def test_resolve_to_image_accepts_file_scheme(tmp_path): + path = _write_png(tmp_path, "img.png", (5, 6)) + assert resolve_to_image(f"file://{path}").size == (5, 6) + assert resolve_to_image(path).size == (5, 6) + + +def test_encode_images_encodes_local_paths_and_file_urls(tmp_path): + plain = _write_png(tmp_path, "plain.png", (2, 2)) + file_url = "file://" + _write_png(tmp_path, "scheme.png", (3, 3)) + + examples = [ + _example( + {"type": "input_image", "image_url": plain}, + {"type": "input_image", "image_url": {"url": file_url}}, + {"type": "input_text", "text": "describe"}, + ) + ] + encode_images_in_examples(examples) + + parts = examples[0]["responses_create_params"]["input"][0]["content"] + assert parts[0]["image_url"].startswith("data:image/png;base64,") + assert parts[1]["image_url"].startswith("data:image/png;base64,") + assert resolve_to_image(parts[0]["image_url"]).size == (2, 2) + assert resolve_to_image(parts[1]["image_url"]).size == (3, 3) + # Non-image parts are untouched. + assert parts[2] == {"type": "input_text", "text": "describe"} + + +def test_encode_images_passes_through_http_and_data_urls(): + data_url = image_to_data_url(Image.new("RGB", (2, 2))) + examples = [ + _example( + {"type": "input_image", "image_url": "https://example.com/cat.png"}, + {"type": "input_image", "image_url": "http://example.com/dog.png"}, + {"type": "input_image", "image_url": data_url}, + ) + ] + encode_images_in_examples(examples) + + parts = examples[0]["responses_create_params"]["input"][0]["content"] + assert parts[0]["image_url"] == "https://example.com/cat.png" + assert parts[1]["image_url"] == "http://example.com/dog.png" + assert parts[2]["image_url"] == data_url + + +def test_encode_images_is_a_noop_for_text_only_examples(): + examples = [_example({"type": "input_text", "text": "no images here"})] + before = [ + dict(part) + for part in examples[0]["responses_create_params"]["input"][0]["content"] + ] + assert encode_images_in_examples(examples) is examples + assert examples[0]["responses_create_params"]["input"][0]["content"] == before + + # Missing/oddly-shaped payloads must not raise. + assert encode_images_in_examples([{}, {"responses_create_params": {}}]) is not None + assert encode_images_in_examples([{"responses_create_params": {"input": "nope"}}]) diff --git a/tests/unit/environments/test_nemo_gym_mm_utils.py b/tests/unit/environments/test_nemo_gym_mm_utils.py new file mode 100644 index 00000000000..3d8a4186e17 --- /dev/null +++ b/tests/unit/environments/test_nemo_gym_mm_utils.py @@ -0,0 +1,178 @@ +# 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. + +from PIL import Image + +from nemo_rl.data.multimodal_utils import image_to_data_url +from nemo_rl.environments.nemo_gym import ( + _extract_input_images_from_message, + _index_per_turn_images, +) + + +def _image(size: tuple[int, int]) -> str: + """Return a data URL for a solid RGB image of the given size.""" + return image_to_data_url(Image.new("RGB", size)) + + +def _user(*data_urls: str) -> dict: + return { + "role": "user", + "content": [{"type": "input_image", "image_url": url} for url in data_urls], + } + + +def _assistant(token_ids: list[int]) -> dict: + return {"role": "assistant", "generation_token_ids": token_ids} + + +def test_extract_input_images_handles_flat_and_dict_image_url(): + item = { + "role": "user", + "content": [ + {"type": "input_image", "image_url": _image((2, 2))}, + {"type": "input_image", "image_url": {"url": _image((3, 3))}}, + {"type": "input_text", "text": "ignore me"}, + ], + } + images = _extract_input_images_from_message(item) + assert [img.size for img in images] == [(2, 2), (3, 3)] + + +def test_extract_input_images_returns_empty_for_string_content(): + assert _extract_input_images_from_message({"role": "user", "content": "hi"}) == [] + assert _extract_input_images_from_message({"role": "user"}) == [] + + +def test_extract_input_images_ignores_text_function_call_output(): + item = { + "type": "function_call_output", + "call_id": "c1", + "output": '{"ok": true}', + } + assert _extract_input_images_from_message(item) == [] + + item["output"] = "Tool failed to create result.png" + assert _extract_input_images_from_message(item) == [] + + +def test_index_per_turn_images_bins_images(): + output = [ + _user(_image((2, 2))), + _assistant([1, 2]), + _user(_image((3, 3)), _image((4, 4))), + _assistant([3, 4]), + ] + per_turn = _index_per_turn_images(output) + + assert len(per_turn) == 2 + assert [img.size for img in per_turn[0]] == [(2, 2)] + assert [img.size for img in per_turn[1]] == [(3, 3), (4, 4)] + + +def test_index_per_turn_images_seeds_first_turn_from_input_messages(): + input_messages = [_user(_image((2, 2)))] + output = [_assistant([1, 2])] + + per_turn = _index_per_turn_images(output, input_messages=input_messages) + + assert len(per_turn) == 1 + assert [img.size for img in per_turn[0]] == [(2, 2)] + + +def test_index_per_turn_images_text_only_rollout_yields_empty_buckets(): + output = [ + {"role": "user", "content": "solve this"}, + _assistant([1, 2]), + {"role": "user", "content": "and this"}, + _assistant([3, 4]), + ] + assert _index_per_turn_images(output) == [[], []] + + +def test_index_per_turn_images_assigns_tool_result_image_to_next_turn(): + """A tool-result image contributes to the following assistant turn.""" + output = [ + _user(_image((2, 2))), + _assistant([1, 2]), + {"type": "function_call_output", "output": _image((5, 5))}, + _assistant([3, 4]), + ] + per_turn = _index_per_turn_images(output) + + assert len(per_turn) == 2 + assert [img.size for img in per_turn[0]] == [(2, 2)] + assert [img.size for img in per_turn[1]] == [(5, 5)] + + +def test_index_per_turn_images_aligns_with_postprocess_skip_of_empty_generations(): + """Turns skipped by the postprocess loop must not consume an image bucket. + + ``_postprocess_nemo_gym_to_nemo_rl_result`` skips output items whose + ``generation_token_ids`` is present but empty, so the bucket list must skip + them too or every later turn is attached to the wrong images. + """ + output = [ + _user(_image((2, 2))), + _assistant([]), # all-EOS generation, skipped by the postprocess loop + _user(_image((6, 6))), + _assistant([7, 8]), + ] + per_turn = _index_per_turn_images(output) + + assert len(per_turn) == 1 + assert [img.size for img in per_turn[0]] == [(2, 2), (6, 6)] + + +def test_index_per_turn_images_flushes_on_non_assistant_trainable_item(): + """Trainable items whose role is not ``assistant`` (reasoning-only responses, + function_call items) still carry ``generation_token_ids`` and are treated as + turns by the postprocess loop. The per-turn image bucket must flush for them + too, or the batched flatten path will see a ``PackedTensor`` for turns + where the model produced a normal assistant message and a missing key for + turns where it produced only reasoning — crashing + ``PackedTensor.flattened_concat`` on the None entry. + """ + reasoning_only = {"type": "reasoning", "generation_token_ids": [9, 10]} + output = [ + _user(_image((2, 2))), + reasoning_only, + ] + per_turn = _index_per_turn_images(output) + + assert len(per_turn) == 1 + assert [img.size for img in per_turn[0]] == [(2, 2)] + + +def test_index_per_turn_images_flushes_on_function_call_trainable_item(): + """Same as the reasoning-only case, but for tool-calling turns where the + model call's last output item is a ``function_call`` (no ``role`` field).""" + function_call = { + "type": "function_call", + "name": "tool", + "arguments": "{}", + "call_id": "c1", + "generation_token_ids": [11, 12], + } + output = [ + _user(_image((2, 2))), + function_call, + {"type": "function_call_output", "output": _image((5, 5)), "call_id": "c1"}, + _assistant([13, 14]), + ] + per_turn = _index_per_turn_images(output) + + assert len(per_turn) == 2 + assert [img.size for img in per_turn[0]] == [(2, 2)] + assert [img.size for img in per_turn[1]] == [(5, 5)] From e496258b0285d42d5d9af30671e81722bca916dc Mon Sep 17 00:00:00 2001 From: Teodor-Dumitru Ene Date: Fri, 7 Aug 2026 04:12:20 -0500 Subject: [PATCH 07/10] fix: Reconcile #2315 and #2612 (#2902) Signed-off-by: Teodor-Dumitru Ene Co-authored-by: Claude Fable 5 --- ...3B-2n8g-megatron_generation-async-gym.yaml | 2 - examples/nemo_gym/grpo_nanov3.yaml | 3 +- nemo_rl/algorithms/grpo.py | 36 ++++++-- .../megatron/megatron_generation.py | 41 +++++++-- .../L1_Functional_Tests_Megatron_4.sh | 1 + .../grpo_megatron_generation_topology.sh | 71 +++++++++++++++ ...BA3B-2n8g-megatron_generation-async-gym.sh | 8 +- .../test_megatron_placement_groups.py | 86 +++++++++++++++++++ 8 files changed, 222 insertions(+), 26 deletions(-) create mode 100644 tests/functional/grpo_megatron_generation_topology.sh create mode 100644 tests/unit/distributed/test_megatron_placement_groups.py diff --git a/examples/configs/recipes/llm/grpo-nanov3-30BA3B-2n8g-megatron_generation-async-gym.yaml b/examples/configs/recipes/llm/grpo-nanov3-30BA3B-2n8g-megatron_generation-async-gym.yaml index f076885eb5a..c2ba70d7b72 100644 --- a/examples/configs/recipes/llm/grpo-nanov3-30BA3B-2n8g-megatron_generation-async-gym.yaml +++ b/examples/configs/recipes/llm/grpo-nanov3-30BA3B-2n8g-megatron_generation-async-gym.yaml @@ -15,8 +15,6 @@ policy: megatron_cfg: pipeline_model_parallel_size: 1 context_parallel_size: 2 - sequence_packing: - enabled: false generation: backend: megatron colocated: diff --git a/examples/nemo_gym/grpo_nanov3.yaml b/examples/nemo_gym/grpo_nanov3.yaml index 3c371660e41..88d38dfb7b9 100644 --- a/examples/nemo_gym/grpo_nanov3.yaml +++ b/examples/nemo_gym/grpo_nanov3.yaml @@ -225,7 +225,6 @@ policy: inference_cuda_graph_scope: "block" activation_checkpointing: false mamba_inference_ssm_states_dtype: "float32" - inference_moe_token_dispatcher_type: "nccl" # Fall back to NCCL for now inference_grouped_gemm_backend: "vllm" moe_router_num_groups: null # InferenceTopKRouter requires num_groups=None moe_router_group_topk: null # paired with moe_router_num_groups=null @@ -244,7 +243,7 @@ policy: kv_cache_management_mode: "persist" # KV cache lifecycle across suspend/resume. Options: "persist", "offload". To select "recompute", set grpo.async_grpo.recompute_kv_cache_after_weight_updates=true. materialize_only_last_token_logits: true num_speculative_tokens: 0 - refit_backend: "nvshmem" # Copy-service backend for non-colocated megatron weight refit. Options: "gloo" or "nvshmem". + refit_backend: "nccl" # Copy-service backend for non-colocated megatron weight refit. Options: "gloo" or "nccl". async_engine: true expose_http_server: true enable_prefix_caching: true diff --git a/nemo_rl/algorithms/grpo.py b/nemo_rl/algorithms/grpo.py index cc85990b55d..6d612e50b0d 100644 --- a/nemo_rl/algorithms/grpo.py +++ b/nemo_rl/algorithms/grpo.py @@ -744,6 +744,12 @@ def _spinup_nemo_gym(base_urls, model_name): ) train_cluster = cluster inference_cluster = cluster + # Colocated generation reuses the policy's cluster; need to decide topology here. + if ( + node_resource_constraints is not None + and generation_config["backend"] == "megatron" + ): + MegatronGeneration.init_cluster_placement_groups(cluster, policy_config) print( f" ✓ Ray cluster for policy initialized with {policy_nodes} nodes", flush=True, @@ -858,7 +864,7 @@ def _spinup_nemo_gym(base_urls, model_name): flush=True, ) - # Inference topology: each vLLM/SGLang instance spans + # Inference topology: each inference instance spans # nodes_per_instance nodes; keep those within one domain # so cross-node all-reduce uses NVLink, not InfiniBand. # @@ -866,7 +872,13 @@ def _spinup_nemo_gym(base_urls, model_name): # For SGLang: gpus_per_server already includes all parallelism # dimensions (TP, DP-attention, PP are internal subdivisions), # so we use it directly without multiplying by pp_size. - if generation_config["backend"] == "vllm": + # For Megatron: the NVLink-domain span of the parallelism the + # generation workers actually run with. + if generation_config["backend"] == "megatron": + gpus_per_instance = MegatronGeneration.nvlink_domain_span( + policy_config + ) + elif generation_config["backend"] == "vllm": vllm_cfg = generation_config.get("vllm_cfg", {}) gpus_per_instance = vllm_cfg["tensor_parallel_size"] * vllm_cfg.get( "pipeline_parallel_size", 1 @@ -944,13 +956,19 @@ def _spinup_nemo_gym(base_urls, model_name): node_resource_constraints=inference_node_resource_constraints, ) if inference_node_resource_constraints is not None: - { - "vllm": VllmGeneration, - "trtllm": TrtllmGeneration, - }[generation_config["backend"]].init_cluster_placement_groups( - inference_cluster, - generation_config, - ) + if generation_config["backend"] == "megatron": + # Megatron inference reuses the training parallelism config. + MegatronGeneration.init_cluster_placement_groups( + inference_cluster, policy_config + ) + else: + { + "vllm": VllmGeneration, + "trtllm": TrtllmGeneration, + }[generation_config["backend"]].init_cluster_placement_groups( + inference_cluster, + generation_config, + ) print( f" ✓ Ray inference cluster initialized with {inference_nodes} nodes with {inference_gpus_per_node} GPUs per node", flush=True, diff --git a/nemo_rl/models/generation/megatron/megatron_generation.py b/nemo_rl/models/generation/megatron/megatron_generation.py index 707f2230c6c..866f3c6b8f6 100644 --- a/nemo_rl/models/generation/megatron/megatron_generation.py +++ b/nemo_rl/models/generation/megatron/megatron_generation.py @@ -36,7 +36,35 @@ class MegatronGeneration(GenerationInterface): """Generation interface backed by Megatron (colocated or non-colocated).""" @staticmethod + def effective_megatron_cfg(config: PolicyConfig) -> dict[str, Any]: + """The megatron_cfg the generation workers actually run with. + + Colocated generation shares the training model, so the training + values apply; non-colocated builds a dedicated policy with + mcore_generation_config merged on top. Always returns a fresh dict. + """ + megatron_cfg = config["megatron_cfg"] + if config["generation"]["colocated"]["enabled"]: + return dict(megatron_cfg) + return { + **megatron_cfg, + **config["generation"].get("mcore_generation_config", {}), + } + + @classmethod + def nvlink_domain_span(cls, config: PolicyConfig) -> int: + """Largest GPU group requiring full NVLink connectivity.""" + megatron_cfg = cls.effective_megatron_cfg(config) + return max( + megatron_cfg["tensor_model_parallel_size"] + * megatron_cfg["context_parallel_size"], + megatron_cfg.get("expert_tensor_parallel_size", 1) + * megatron_cfg.get("expert_model_parallel_size", 1), + ) + + @classmethod def init_cluster_placement_groups( + cls, cluster: RayVirtualCluster, config: PolicyConfig, ) -> None: @@ -46,16 +74,10 @@ def init_cluster_placement_groups( cluster: The inference `RayVirtualCluster`. config: The full `PolicyConfig` (megatron parallelism + colocation). """ - megatron_cfg = config["megatron_cfg"] - model_parallel_size = ( - megatron_cfg["tensor_model_parallel_size"] - * megatron_cfg["pipeline_model_parallel_size"] - * megatron_cfg["context_parallel_size"] - ) colocated = config["generation"]["colocated"]["enabled"] cluster._init_placement_groups( strategy=None if colocated else "PACK", - use_unified_pg=model_parallel_size > cluster.num_gpus_per_node, + use_unified_pg=cls.nvlink_domain_span(config) > cluster.num_gpus_per_node, ) def __init__( @@ -111,7 +133,10 @@ def __init__( # Stand up a dedicated inference-only policy. self._owns_policy = True - self._policy_config["megatron_cfg"].update(self.cfg["mcore_generation_config"]) + self._policy_config = { + **config, + "megatron_cfg": self.effective_megatron_cfg(config), + } # Activation checkpointing is not compatible or useful in inference. self._policy_config["megatron_cfg"]["activation_checkpointing"] = False # Reserve GPUs before Policy workers grab them, to prevent disjoint NVLS domains. diff --git a/tests/functional/L1_Functional_Tests_Megatron_4.sh b/tests/functional/L1_Functional_Tests_Megatron_4.sh index 0b4bf49bd2b..0c00d2d7f2f 100644 --- a/tests/functional/L1_Functional_Tests_Megatron_4.sh +++ b/tests/functional/L1_Functional_Tests_Megatron_4.sh @@ -50,6 +50,7 @@ megatron_generation_supported() { if megatron_generation_supported; then run_test uv run --no-sync bash ./tests/functional/grpo_megatron_generation.sh + run_test uv run --no-sync bash ./tests/functional/grpo_megatron_generation_topology.sh run_test uv run --no-sync bash ./tests/functional/grpo_megatron_generation_non_colocated.sh run_test uv run --no-sync bash ./tests/functional/grpo_megatron_generation_async.sh run_test fast uv run --no-sync bash ./tests/functional/grpo_megatron_generation_colocated_async.sh diff --git a/tests/functional/grpo_megatron_generation_topology.sh b/tests/functional/grpo_megatron_generation_topology.sh new file mode 100644 index 00000000000..b6ed44dc8c3 --- /dev/null +++ b/tests/functional/grpo_megatron_generation_topology.sh @@ -0,0 +1,71 @@ +#!/bin/bash + +SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd) +PROJECT_ROOT=$(realpath $SCRIPT_DIR/../..) +# Mark the current repo as safe, since wandb fetches metadata about the repo +git config --global --add safe.directory $PROJECT_ROOT + +set -eou pipefail + +EXP_NAME=$(basename $0 .sh) +EXP_DIR=$SCRIPT_DIR/$EXP_NAME +LOG_DIR=$EXP_DIR/logs +JSON_METRICS=$EXP_DIR/metrics.json +RUN_LOG=$EXP_DIR/run.log +export PYTHONPATH=${PROJECT_ROOT}:${PYTHONPATH:-} + +rm -rf $EXP_DIR $LOG_DIR +mkdir -p $EXP_DIR $LOG_DIR + +cd $PROJECT_ROOT + +# cluster.segment_size only engages when Ray nodes carry nvlink_domain_* labels, +# which ray.sub probes from `nvidia-smi -q` ClusterUUID on NVLink-fabric clusters +# (e.g. GB200 NVL72); CI runners have none. Pre-start a Ray head with a synthetic +# domain label so init_ray() attaches to it (externally managed cluster) and the +# topology-aware megatron placement path runs for real. +cleanup() { + uv run ray stop --force || true +} +trap cleanup EXIT +uv run ray stop --force || true # don't attach to a stale cluster +uv run ray start --head --disable-usage-stats \ + --resources='{"nvlink_domain_ci_synthetic": 1, "topo_rank": 1}' + +uv run coverage run -a --data-file=$PROJECT_ROOT/tests/.coverage --source=$PROJECT_ROOT/nemo_rl \ + $PROJECT_ROOT/examples/run_grpo.py \ + --config $PROJECT_ROOT/examples/configs/grpo_math_1B_megatron.yaml \ + policy.model_name=Qwen/Qwen2.5-0.5B \ + grpo.num_prompts_per_step=2 \ + grpo.num_generations_per_prompt=4 \ + policy.train_global_batch_size=4 \ + policy.logprob_batch_size=4 \ + policy.train_micro_batch_size=1 \ + policy.generation.backend=megatron \ + cluster.gpus_per_node=2 \ + cluster.segment_size=1 \ + grpo.max_num_steps=2 \ + logger.tensorboard_enabled=true \ + logger.log_dir=$LOG_DIR \ + logger.wandb_enabled=false \ + logger.monitor_gpus=true \ + checkpointing.enabled=false \ + $@ \ + 2>&1 | tee $RUN_LOG + +# Guard against the silent fallback: with no (or unreadable) domain labels the run +# would succeed without ever exercising the topology placement path under test. +grep -q "Topology-aware allocation" $RUN_LOG || { + echo "ERROR: topology-aware allocation did not engage (no segment selection logged)" >&2 + exit 1 +} +# NOTE: `! grep` is exempt from `set -e`, hence the explicit if. +if grep -q "no NVLink domain info" $RUN_LOG; then + echo "ERROR: segment_size fell back to unordered allocation" >&2 + exit 1 +fi + +uv run tests/json_dump_tb_logs.py $LOG_DIR --output_path $JSON_METRICS + +uv run tests/check_metrics.py $JSON_METRICS \ + 'max(data["train/token_mult_prob_error"]) < 1.05' diff --git a/tests/test_suites/llm/grpo-nanov3-30BA3B-2n8g-megatron_generation-async-gym.sh b/tests/test_suites/llm/grpo-nanov3-30BA3B-2n8g-megatron_generation-async-gym.sh index 437a5756ffc..1819c65df23 100755 --- a/tests/test_suites/llm/grpo-nanov3-30BA3B-2n8g-megatron_generation-async-gym.sh +++ b/tests/test_suites/llm/grpo-nanov3-30BA3B-2n8g-megatron_generation-async-gym.sh @@ -8,11 +8,9 @@ GPUS_PER_NODE=8 STEPS_PER_RUN=8 MAX_STEPS=8 NUM_RUNS=$(( (MAX_STEPS + STEPS_PER_RUN - 1) / STEPS_PER_RUN )) # Round up -# ~25 min startup (30B-MoE load + CUDA-graph warmup + nemo_gym servers) plus -# ~63 min for 8 async steps left no headroom at 90 min, so the driver finished -# but Slurm SIGKILLed teardown/metric-dump at the wall-clock limit (CI mislabels -# the exit-137 as OOM). 120 min leaves margin for teardown + metrics. -NUM_MINUTES=120 +# ~25 min startup (30B-MoE load + CUDA-graph warmup + nemo_gym servers) plus ~130 min for 8 steps +# 180 min leaves margin for teardown + metric-dump within the 4 h job allocation. +NUM_MINUTES=180 # ===== END CONFIG ===== exit_if_max_steps_reached diff --git a/tests/unit/distributed/test_megatron_placement_groups.py b/tests/unit/distributed/test_megatron_placement_groups.py new file mode 100644 index 00000000000..606a2d03b8a --- /dev/null +++ b/tests/unit/distributed/test_megatron_placement_groups.py @@ -0,0 +1,86 @@ +# 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. + +from unittest.mock import MagicMock + +import pytest + +from nemo_rl.models.generation.megatron import MegatronGeneration + + +def _placement_policy_config( + *, + tp: int = 1, + pp: int = 1, + cp: int = 1, + ep: int = 1, + etp: int = 1, + colocated: bool = False, + mcore_overrides: dict | None = None, +) -> dict: + """Minimal PolicyConfig slice consumed by init_cluster_placement_groups.""" + return { + "megatron_cfg": { + "tensor_model_parallel_size": tp, + "pipeline_model_parallel_size": pp, + "context_parallel_size": cp, + "expert_model_parallel_size": ep, + "expert_tensor_parallel_size": etp, + }, + "generation": { + "colocated": {"enabled": colocated}, + "mcore_generation_config": mcore_overrides or {}, + }, + } + + +@pytest.mark.parametrize( + "config_kwargs,expected_strategy,expected_unified", + [ + # cross-node span via TP alone -> one unified PG + (dict(tp=8), "PACK", True), + # PP is excluded from the NVLink span: TP*CP=4 fits a node even + # though the full TP*PP*CP instance would not + (dict(tp=2, pp=2, cp=2), "PACK", False), + # node-local span at the == boundary -> per-node PGs + (dict(tp=4), "PACK", False), + # cross-node MoE expert group (ETP*EP > TP*CP) -> one unified PG: + # the NVLS dispatcher needs the ep_group fully NVLink-connected + (dict(tp=2, ep=8), "PACK", True), + # non-colocated generation parallelism overrides the training values + # (mirrors MegatronGeneration's megatron_cfg merge) + ( + dict(tp=8, mcore_overrides={"tensor_model_parallel_size": 2}), + "PACK", + False, + ), + # colocated reuses the training layout (overrides do not apply): + # no PACK strategy, span from the training config incl. its EP + (dict(tp=2, ep=8, colocated=True), None, True), + ], +) +def test_megatron_init_cluster_placement_groups( + config_kwargs, expected_strategy, expected_unified +): + """The NVLink-domain span is max(TP*CP, ETP*EP) of the operative config.""" + cluster = MagicMock(num_gpus_per_node=4) + + MegatronGeneration.init_cluster_placement_groups( + cluster, _placement_policy_config(**config_kwargs) + ) + + cluster._init_placement_groups.assert_called_once_with( + strategy=expected_strategy, + use_unified_pg=expected_unified, + ) From 8a65e64664b94ca23f0249752b7f1bb642722881 Mon Sep 17 00:00:00 2001 From: Anish Mahishi <20884035+macandro96@users.noreply.github.com> Date: Fri, 7 Aug 2026 16:55:45 -0400 Subject: [PATCH 08/10] chore: bump TQ to v0.19 (#3423) Signed-off-by: Anish Mahishi --- nemo_rl/data_plane/adapters/transfer_queue.py | 128 ++++++---- nemo_rl/data_plane/schema.py | 19 ++ pyproject.toml | 16 +- tests/unit/data_plane/README.md | 12 +- tests/unit/data_plane/test_codec_mooncake.py | 238 +++++++++++++++++- tests/unit/data_plane/test_tq_lifecycle.py | 21 +- uv.lock | 7 +- 7 files changed, 366 insertions(+), 75 deletions(-) diff --git a/nemo_rl/data_plane/adapters/transfer_queue.py b/nemo_rl/data_plane/adapters/transfer_queue.py index ce293b4e269..995cfa24c37 100644 --- a/nemo_rl/data_plane/adapters/transfer_queue.py +++ b/nemo_rl/data_plane/adapters/transfer_queue.py @@ -27,8 +27,9 @@ import socket import subprocess import time +import warnings from importlib import resources -from typing import Any +from typing import Any, cast import torch import transfer_queue as tq @@ -39,6 +40,7 @@ DataPlaneConfig, KVBatchMeta, ) +from nemo_rl.data_plane.schema import PROMOTE_1D_FIELDS # ────────────────────────────────────────────────────────────────────────── # Backend init — lifted from rl-arena/arena/backends.py. @@ -171,27 +173,28 @@ def patched(*args, **kwargs): cls.options = patched # type: ignore[method-assign] return True - patched_any = False + unpatched_classes: list[str] = [] try: - from transfer_queue.storage.simple_backend import SimpleStorageUnit + from transfer_queue.storage.simple_storage import SimpleStorageUnit - patched_any |= _install(SimpleStorageUnit) + if not _install(SimpleStorageUnit): + unpatched_classes.append("SimpleStorageUnit") except ImportError: - pass + unpatched_classes.append("SimpleStorageUnit") try: from transfer_queue.controller import TransferQueueController - patched_any |= _install(TransferQueueController) + if not _install(TransferQueueController): + unpatched_classes.append("TransferQueueController") except ImportError: - pass + unpatched_classes.append("TransferQueueController") - if not patched_any: + if unpatched_classes: # Soft-fail: TQ may have moved its actor classes. The driver will # still work; multi-node TQ may need the per-node `uv sync` workaround. - import warnings - warnings.warn( - "Could not patch TQ actor classes for runtime_env injection. " + "Could not patch every TQ actor class for runtime_env injection: " + f"unpatched={unpatched_classes}. " "Multi-node TQ may fail with ModuleNotFoundError: 'transfer_queue' " "on worker nodes. Workaround: run `uv sync` inside each node's " "container before the driver runs.", @@ -322,19 +325,23 @@ def _assert_no_key_loss(src_dict: dict, new_td: TensorDict, fn: str) -> None: def _promote_1d_leaves(td: TensorDict) -> TensorDict: - """Unsqueeze 1D tensor leaves to ``(N, 1)`` — mooncake_cpu KV-path workaround. + """Promote declared scalar leaves to ``(N, 1)`` for Mooncake. - Works around TQ's ``KVStorageManager`` 1D schema/data mismatch; - :func:`_from_wire` squeezes the trailing 1 back on read. Symmetric - with `_from_wire` — callers gate on ``self._promote_1d``. - ``NonTensorStack`` / ``NonTensorData`` leaves pass through. + The authoritative field list lives in + :data:`nemo_rl.data_plane.schema.PROMOTE_1D_FIELDS`. Declared fields must + arrive as dense ``(N,)`` tensors. Any other dense 1D tensor is rejected so + it cannot silently encounter TQ v0.1.9's schema/data mismatch. + ``NonTensorStack`` and ``NonTensorData`` leaves pass through. Args: - td: ``TensorDict`` whose 1D tensor leaves should be promoted. + td: TensorDict to validate and encode for the Mooncake wire format. Returns: - ``TensorDict`` with 1D tensor leaves unsqueezed to ``(N, 1)``; - all other leaves pass through unchanged. + TensorDict with declared scalar leaves promoted to ``(N, 1)``. + + Raises: + ValueError: If a declared field is not a dense 1D tensor, or an + undeclared field is a dense 1D tensor. """ # td.keys() (top-level) includes NonTensorData / NonTensorStack leaves. # keys(include_nested=True, leaves_only=True) enumerates tensor leaves @@ -343,9 +350,23 @@ def _promote_1d_leaves(td: TensorDict) -> TensorDict: changed = False for k in td.keys(): v = td.get(k) - if isinstance(v, torch.Tensor) and not v.is_nested and v.dim() == 1: + field_name = str(k) + if field_name in PROMOTE_1D_FIELDS: + if not isinstance(v, torch.Tensor) or v.is_nested or v.dim() != 1: + shape = tuple(v.shape) if isinstance(v, torch.Tensor) else None + raise ValueError( + f"Mooncake scalar field {field_name!r} must be a dense " + f"1D tensor with shape (N,), got {type(v).__name__} " + f"with shape {shape}." + ) new_dict[str(k)] = v.unsqueeze(-1).contiguous() changed = True + elif isinstance(v, torch.Tensor) and not v.is_nested and v.dim() == 1: + raise ValueError( + f"Mooncake field {field_name!r} is a dense 1D tensor but is " + "not declared in data_plane.schema.PROMOTE_1D_FIELDS. Add " + "the field to the schema if it is a per-sample scalar." + ) else: new_dict[str(k)] = v if not changed: @@ -356,23 +377,45 @@ def _promote_1d_leaves(td: TensorDict) -> TensorDict: def _from_wire(td: TensorDict) -> TensorDict: - """Inverse of `_promote_1d_leaves`: squeeze trailing 1 back to (N,).""" + """Normalize TQ reads and invert :func:`_promote_1d_leaves` when needed. + + Both TQ v0.1.9 storage managers reconstruct every non-scalar field as a + nested tensor, including fields whose rows all have the same shape. + Densify those uniform nested tensors first so regular batched inputs retain + their dense representation. Truly ragged fields remain nested. Finally, + squeeze only singleton dimensions declared in + :data:`nemo_rl.data_plane.schema.PROMOTE_1D_FIELDS`. + """ # Same top-level iteration as `_promote_1d_leaves`: NonTensorData / # NonTensorStack leaves are only visible via td.keys(), not leaves_only. new_dict: dict[str, Any] = {} changed = False for k in td.keys(): v = td.get(k) - if ( - isinstance(v, torch.Tensor) - and not v.is_nested - and v.dim() >= 2 - and v.shape[-1] == 1 - ): - new_dict[str(k)] = v.squeeze(-1).contiguous() - changed = True + field_name = str(k) + if isinstance(v, torch.Tensor) and v.is_nested: + rows = list(v.unbind()) + if rows and all(row.shape == rows[0].shape for row in rows[1:]): + v = torch.stack(rows) + changed = True + if field_name in PROMOTE_1D_FIELDS: + if not isinstance(v, torch.Tensor) or v.is_nested: + raise ValueError( + f"Mooncake scalar field {field_name!r} could not be " + "restored as a dense tensor." + ) + if v.dim() == 1: + new_dict[field_name] = v + elif v.dim() == 2 and v.shape[-1] == 1: + new_dict[field_name] = v.squeeze(-1).contiguous() + changed = True + else: + raise ValueError( + f"Mooncake scalar field {field_name!r} must decode as " + f"(N,) or (N, 1), got shape {tuple(v.shape)}." + ) else: - new_dict[str(k)] = v + new_dict[field_name] = v if not changed: return td new_td = TensorDict(new_dict, batch_size=td.batch_size) @@ -583,9 +626,9 @@ def put_samples( return KVBatchMeta( partition_id=partition_id, task_name=None, sample_ids=[], fields=None ) - if tags is None: - tags = [{} for _ in sample_ids] - + user_tags = ( + [{} for _ in sample_ids] if tags is None else [dict(tag) for tag in tags] + ) wire_fields: TensorDict | None = None field_names: list[str] | None = None if fields is not None: @@ -594,17 +637,21 @@ def put_samples( # TDs. TQ's encoder forces ``.contiguous()`` per tensor leaf # itself, so the call here was redundant for tensors and # destructive for non-tensors. - wire_fields = fields.detach() # type: ignore[bad-assignment,missing-argument] + detached_fields = cast( + TensorDict, + fields.detach(), # type: ignore[missing-argument] + ) if self._promote_1d: - wire_fields = _promote_1d_leaves(wire_fields) # type: ignore[bad-argument-type] - field_names = list(wire_fields.keys()) + detached_fields = _promote_1d_leaves(detached_fields) + wire_fields = detached_fields + field_names = [str(key) for key in detached_fields.keys()] # TQ's wire vocabulary is `keys=` — translation point. tq.kv_batch_put( keys=list(sample_ids), partition_id=partition_id, fields=wire_fields, - tags=tags, + tags=user_tags, ) return KVBatchMeta( @@ -612,7 +659,7 @@ def put_samples( task_name=None, sample_ids=list(sample_ids), fields=field_names, - tags=[dict(t) for t in tags] if tags else None, + tags=user_tags if user_tags else None, ) def get_samples( @@ -623,15 +670,12 @@ def get_samples( ) -> TensorDict: if not sample_ids: return TensorDict({}, batch_size=(0,)) - # TQ's wire vocabulary is `keys=` — translation point. td = tq.kv_batch_get( keys=list(sample_ids), partition_id=partition_id, select_fields=select_fields, ) - if self._promote_1d: - td = _from_wire(td) - return td + return _from_wire(td) def clear_samples(self, sample_ids: list[str] | None, partition_id: str) -> None: cleared_via_none = sample_ids is None diff --git a/nemo_rl/data_plane/schema.py b/nemo_rl/data_plane/schema.py index e451fc361cc..49cf79422e7 100644 --- a/nemo_rl/data_plane/schema.py +++ b/nemo_rl/data_plane/schema.py @@ -63,6 +63,25 @@ ROUTED_EXPERTS_FIELD = "routed_experts" +# Per-sample 1D scalar fields. The TQ adapter promotes these to ``(N, 1)`` +# on write to work around TQ v0.1.9's KVStorageManager schema/data mismatch on +# the Mooncake backend, and squeezes them back to ``(N,)`` on read. This is the +# authoritative user-level schema; no per-row shape metadata is carried. +# +# Fields listed here must be dense ``(N,)`` tensors when written through the +# Mooncake adapter. Dense 1D fields not listed here are rejected on that path so +# a new field cannot silently reintroduce the upstream shape mismatch. +# +# Delete this set and the corresponding adapter transforms when upstream TQ +# fixes 1D field schema extraction. +PROMOTE_1D_FIELDS: frozenset[str] = frozenset( + { + INPUT_LENGTHS, + "total_reward", + SAMPLE_MASK, + } +) + def fields_with_optional_routed_experts( fields: Sequence[str], diff --git a/pyproject.toml b/pyproject.toml index 9958048b578..1b4244c5727 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -84,11 +84,10 @@ dependencies = [ # automatically include them. Removes the need for a `[data-plane]` # extra and the corresponding plumbing in the per-worker venv builder. "tensordict", - # Pinned to b266d39 (post-0.1.6, pre-0.1.7) for PR #77's MooncakeStore - # refactor: `clear` switched from unanchored `remove_by_regex` to - # exact-key `batch_remove`, which fixes a collateral-key-deletion bug - # that breaks DAPO + mooncake_cpu. Bump to the 0.1.7 tag when released. - "TransferQueue @ git+https://github.com/Ascend/TransferQueue.git@b266d39", + # TransferQueue v0.1.9 adds full-system checkpoint save/load APIs and + # retains the exact-key MooncakeStore clear behavior required by DAPO. + # Pin the immutable release commit rather than the mutable version tag. + "TransferQueue @ git+https://github.com/Ascend/TransferQueue.git@c51614308b68c8d7a87c9b3ef62d59e14c69bde2", # Backs data_plane.backend="mooncake_cpu". Default backend is "simple" # (in-process), but the mooncake_cpu path needs the `mooncake_master` # binary that ships in this wheel at /mooncake/. Bundled @@ -391,10 +390,9 @@ override-dependencies = [ "pytest>=9.0.3", "langchain>=0.3.28", # Address CVE-2025-65106 "langchain-core>=0.3.80", # Address CVE-2025-65106 - # TransferQueue (data-plane extra) pins numpy<2.0.0; megatron-core needs - # numpy>=2.1.0 via onnx → ml-dtypes. Override globally so the data-plane - # extra composes with mcore/automodel without version-mirroring TQ's - # requirements.txt. Forward-compatible across TQ minor bumps. + # Forces numpy past tensorrt-llm's `numpy>=2.0.0,<2.4` cap (resolves to + # 2.5.x). The original driver — TransferQueue pinning `numpy<2.0.0` — is + # gone as of TQ v0.1.9; drop this override once the trtllm cap lifts. "numpy>=2.1.0", # av (PyAV) carries CVE-bundled codec libs (libx264, libx265, libopenh264, libmp3lame). # It is only needed by megatron-bridge's optional WAN diffusion path, which installs it diff --git a/tests/unit/data_plane/README.md b/tests/unit/data_plane/README.md index f37b6b5852a..9ef9f60d417 100644 --- a/tests/unit/data_plane/README.md +++ b/tests/unit/data_plane/README.md @@ -30,10 +30,14 @@ Generated audit of every test function under `tests/unit/data_plane/` with a one - `test_materialize_default_pad_value_is_zero` — No `pad_value_dict` → pad with 0. - `test_response_from_nested_extracts_response_slice` — Worker write-back: jagged (prompt+response) → response only. -## `test_codec_mooncake.py` (4 tests) - -- `test_promote_1d_leaves_unsqueezes_1d` — `_promote_1d_leaves` turns 1D `(N,)` leaves into `(N, 1)` for mooncake wire. -- `test_promote_1d_roundtrip_via_from_wire` — `_promote_1d_leaves` + `_from_wire` restores original `(N,)` shape and values. +## `test_codec_mooncake.py` + +- `test_promote_1d_leaves_unsqueezes_1d` — `_promote_1d_leaves` turns schema-declared scalar fields from `(N,)` into `(N, 1)` for the Mooncake wire. +- `test_promote_1d_roundtrip_via_from_wire` — `_promote_1d_leaves` + `_from_wire` restores the schema-declared field's original `(N,)` shape and values. +- `test_from_wire_rejects_invalid_declared_field_shape` — Corrupt or incompatible scalar wire shapes fail at the data-plane boundary. +- `test_promote_1d_leaves_rejects_undeclared_1d_field` — Unknown dense 1D Mooncake fields fail loudly instead of silently hitting TQ's shape mismatch. +- `test_put_samples_uses_schema_without_private_shape_tags` — Promotion does not add per-row adapter metadata to user tags. +- `test_get_samples_uses_static_shape_schema` — Reads restore scalar fields by the shared schema while preserving genuine `(N, 1)` columns. - `test_pack_per_token_field_truncates_sp_padding` — pack_per_token_field slices each row to its own length, dropping SP padding. - `test_pack_per_token_field_exact_fit_matches_to_nested_by_length` — At exact fit, `pack_per_token_field` matches `to_nested_by_length`. diff --git a/tests/unit/data_plane/test_codec_mooncake.py b/tests/unit/data_plane/test_codec_mooncake.py index c9b71820d74..68752b5bb58 100644 --- a/tests/unit/data_plane/test_codec_mooncake.py +++ b/tests/unit/data_plane/test_codec_mooncake.py @@ -14,7 +14,7 @@ """Unit tests for the mooncake_cpu-specific wire workarounds. Covers: - P1 — `promote_1d` round-trip: writer unsqueezes 1D → (N,1), reader squeezes back. + P1 — schema-declared 1D scalar round-trip through the Mooncake workaround. P2 — pack_per_token_field: tolerates SP padding wider than max(lengths). No Ray, no GPU, no transfer_queue required. @@ -22,6 +22,7 @@ from __future__ import annotations +import pytest import torch from nemo_rl.data_plane.codec import pack_per_token_field, to_nested_by_length @@ -44,11 +45,12 @@ def test_promote_1d_leaves_unsqueezes_1d() -> None: n = 8 t = torch.arange(n, dtype=torch.float32) - td = TensorDict({"reward": t}, batch_size=[n]) + td = TensorDict({"input_lengths": t}, batch_size=[n]) out = _promote_1d_leaves(td) - assert out["reward"].shape == (n, 1), ( - f"Expected wire shape ({n}, 1) but got {tuple(out['reward'].shape)}." + assert out["input_lengths"].shape == (n, 1), ( + "Expected input_lengths to use the Mooncake wire shape " + f"({n}, 1), got {tuple(out['input_lengths'].shape)}." ) @@ -63,14 +65,234 @@ def test_promote_1d_roundtrip_via_from_wire() -> None: n = 6 original = torch.arange(n, dtype=torch.float32) - td = TensorDict({"reward": original}, batch_size=[n]) + td = TensorDict({"input_lengths": original}, batch_size=[n]) wire = _promote_1d_leaves(td) - assert wire["reward"].shape == (n, 1) + assert wire["input_lengths"].shape == (n, 1) back = _from_wire(wire) - assert back["reward"].shape == (n,) - assert torch.equal(back["reward"], original) + assert back["input_lengths"].shape == (n,) + assert torch.equal(back["input_lengths"], original) + + +def test_from_wire_densifies_uniform_nested_rows() -> None: + """TQ v0.1.9's uniform nested reads are restored to dense tensors.""" + from tensordict import TensorDict + + from nemo_rl.data_plane.adapters.transfer_queue import _from_wire + + rows = [torch.tensor([i, i + 1], dtype=torch.float32) for i in range(4)] + wire = TensorDict( + {"input_ids": torch.nested.as_nested_tensor(rows, layout=torch.jagged)}, + batch_size=[len(rows)], + ) + + back = _from_wire(wire) + + assert not back["input_ids"].is_nested + assert back["input_ids"].shape == (len(rows), 2) + assert torch.equal(back["input_ids"], torch.stack(rows)) + + +def test_from_wire_preserves_genuine_length_one_token_column() -> None: + """Only fields promoted from ``(N,)`` are squeezed after a TQ read.""" + from tensordict import TensorDict + + from nemo_rl.data_plane.adapters.transfer_queue import _from_wire + + n = 4 + wire = TensorDict( + { + "total_reward": torch.nested.as_nested_tensor( + [torch.tensor([float(i)]) for i in range(n)], layout=torch.jagged + ), + "input_ids": torch.nested.as_nested_tensor( + [torch.tensor([i]) for i in range(n)], layout=torch.jagged + ), + }, + batch_size=[n], + ) + + back = _from_wire(wire) + + assert back["total_reward"].shape == (n,) + assert back["input_ids"].shape == (n, 1) + assert torch.equal(back["input_ids"], torch.arange(n).unsqueeze(-1)) + + +def test_from_wire_rejects_invalid_declared_field_shape() -> None: + """A corrupted scalar wire shape fails at the data-plane boundary.""" + from tensordict import TensorDict + + from nemo_rl.data_plane.adapters.transfer_queue import _from_wire + + wire = TensorDict({"input_lengths": torch.ones(3, 2)}, batch_size=[3]) + + with pytest.raises(ValueError, match=r"input_lengths.*\(N, 1\)"): + _from_wire(wire) + + +def test_promote_1d_leaves_rejects_undeclared_1d_field() -> None: + """New scalar fields must be added to the authoritative schema.""" + from tensordict import TensorDict + + from nemo_rl.data_plane.adapters.transfer_queue import _promote_1d_leaves + + fields = TensorDict({"new_scalar": torch.arange(3)}, batch_size=[3]) + + with pytest.raises(ValueError, match="not declared.*PROMOTE_1D_FIELDS"): + _promote_1d_leaves(fields) + + +def test_promote_1d_leaves_rejects_invalid_declared_field_shape() -> None: + """A schema-declared scalar cannot silently change its user-level rank.""" + from tensordict import TensorDict + + from nemo_rl.data_plane.adapters.transfer_queue import _promote_1d_leaves + + fields = TensorDict({"input_lengths": torch.ones(3, 2)}, batch_size=[3]) + + with pytest.raises(ValueError, match=r"input_lengths.*shape \(N,\)"): + _promote_1d_leaves(fields) + + +def test_put_samples_uses_schema_without_private_shape_tags(monkeypatch) -> None: + """Mooncake promotion changes tensors but not user-provided TQ tags.""" + from tensordict import TensorDict + + import nemo_rl.data_plane.adapters.transfer_queue as tq_adapter + + n = 3 + original_fields = TensorDict( + { + "input_lengths": torch.arange(n), + "input_ids": torch.arange(n).unsqueeze(-1), + }, + batch_size=[n], + ) + user_tags = [{"weight_version": 7} for _ in range(n)] + + def fake_kv_batch_put( + *, + keys: list[str], + partition_id: str, + fields: TensorDict, + tags: list[dict[str, object]], + ) -> None: + assert keys == ["a", "b", "c"] + assert partition_id == "train" + assert fields["input_lengths"].shape == (n, 1) + assert fields["input_ids"].shape == (n, 1) + assert tags == user_tags + + monkeypatch.setattr(tq_adapter.tq, "kv_batch_put", fake_kv_batch_put) + client = object.__new__(tq_adapter.TQDataPlaneClient) + client._promote_1d = True + + meta = client.put_samples( + ["a", "b", "c"], "train", fields=original_fields, tags=user_tags + ) + + assert meta.tags == user_tags + + +def test_get_samples_uses_static_shape_schema(monkeypatch) -> None: + """The Mooncake adapter restores scalar ranks without row metadata.""" + from tensordict import TensorDict + + import nemo_rl.data_plane.adapters.transfer_queue as tq_adapter + + n = 3 + original = TensorDict( + { + "total_reward": torch.arange(n, dtype=torch.float32), + "input_ids": torch.arange(n).unsqueeze(-1), + }, + batch_size=[n], + ) + wire_data = TensorDict( + { + "total_reward": torch.nested.as_nested_tensor( + [row for row in original["total_reward"].unsqueeze(-1)], + layout=torch.jagged, + ), + "input_ids": torch.nested.as_nested_tensor( + [row for row in original["input_ids"]], layout=torch.jagged + ), + }, + batch_size=[n], + ) + + def fake_kv_batch_get( + *, keys: list[str], partition_id: str, select_fields: list[str] + ) -> TensorDict: + assert keys == ["a", "b", "c"] + assert partition_id == "train" + assert select_fields == ["total_reward", "input_ids"] + return wire_data + + monkeypatch.setattr(tq_adapter.tq, "kv_batch_get", fake_kv_batch_get) + client = object.__new__(tq_adapter.TQDataPlaneClient) + client._promote_1d = True + + restored = client.get_samples( + ["a", "b", "c"], "train", ["total_reward", "input_ids"] + ) + + assert restored["total_reward"].shape == (n,) + assert restored["input_ids"].shape == (n, 1) + assert torch.equal(restored["total_reward"], original["total_reward"]) + assert torch.equal(restored["input_ids"], original["input_ids"]) + + +def test_get_samples_densifies_uniform_rows_without_1d_promotion(monkeypatch) -> None: + """The simple backend normalizes uniform nested rows without squeezing.""" + from tensordict import TensorDict + + import nemo_rl.data_plane.adapters.transfer_queue as tq_adapter + + rows = [torch.tensor([1, 2]), torch.tensor([3, 4])] + wire_data = TensorDict( + {"input_ids": torch.nested.as_nested_tensor(rows, layout=torch.jagged)}, + batch_size=[len(rows)], + ) + + def fake_kv_batch_get( + *, keys: list[str], partition_id: str, select_fields: list[str] + ) -> TensorDict: + assert keys == ["a", "b"] + assert partition_id == "train" + assert select_fields == ["input_ids"] + return wire_data + + monkeypatch.setattr(tq_adapter.tq, "kv_batch_get", fake_kv_batch_get, raising=False) + client = object.__new__(tq_adapter.TQDataPlaneClient) + client._promote_1d = False + + restored = client.get_samples(["a", "b"], "train", ["input_ids"]) + + assert not restored["input_ids"].is_nested + assert restored["input_ids"].shape == (2, 2) + assert torch.equal(restored["input_ids"], torch.stack(rows)) + + +def test_from_wire_preserves_ragged_nested_rows() -> None: + """Variable-length rollout fields must remain nested.""" + from tensordict import TensorDict + + from nemo_rl.data_plane.adapters.transfer_queue import _from_wire + + rows = [torch.arange(i + 1) for i in range(3)] + nested = torch.nested.as_nested_tensor(rows, layout=torch.jagged) + wire = TensorDict({"token_ids": nested}, batch_size=[len(rows)]) + + back = _from_wire(wire) + + assert back["token_ids"].is_nested + assert all( + torch.equal(actual, expected) + for actual, expected in zip(back["token_ids"].unbind(), rows, strict=True) + ) # ── P2: pack_per_token_field — tolerates SP padding ────────────────────────── diff --git a/tests/unit/data_plane/test_tq_lifecycle.py b/tests/unit/data_plane/test_tq_lifecycle.py index a349c27c2e7..3e79a50ff84 100644 --- a/tests/unit/data_plane/test_tq_lifecycle.py +++ b/tests/unit/data_plane/test_tq_lifecycle.py @@ -134,10 +134,11 @@ def test_smoke_round_trip_backends(tq_client_backends) -> None: consumer_tasks=["read"], ) keys = ["a", "b", "c", "d"] + values = torch.arange(12).reshape(4, 3) client.put_samples( sample_ids=keys, partition_id="smoke-backend", - fields=TensorDict({"x": torch.arange(4)}, batch_size=[4]), + fields=TensorDict({"x": values}, batch_size=[4]), ) meta = client.claim_meta( @@ -150,7 +151,9 @@ def test_smoke_round_trip_backends(tq_client_backends) -> None: assert meta.size == 4 data = client.get_data(meta) - expected = torch.tensor([keys.index(k) for k in meta.sample_ids]) + expected = torch.stack([values[keys.index(k)] for k in meta.sample_ids]) + assert not data["x"].is_nested + assert data["x"].shape == expected.shape assert torch.equal(data["x"], expected) client.clear_samples(sample_ids=None, partition_id="smoke-backend") @@ -164,11 +167,11 @@ def test_smoke_round_trip_1d_fields(tq_client_backends) -> None: this for mooncake_cpu; simple passes the tensor through unchanged. """ n = 6 - reward = torch.arange(n, dtype=torch.float32) + total_reward = torch.arange(n, dtype=torch.float32) tq_client_backends.register_partition( partition_id="smoke-1d", - fields=["reward"], + fields=["total_reward"], num_samples=n, consumer_tasks=["read"], ) @@ -176,21 +179,21 @@ def test_smoke_round_trip_1d_fields(tq_client_backends) -> None: tq_client_backends.put_samples( sample_ids=keys, partition_id="smoke-1d", - fields=TensorDict({"reward": reward}, batch_size=[n]), + fields=TensorDict({"total_reward": total_reward}, batch_size=[n]), ) meta = tq_client_backends.claim_meta( partition_id="smoke-1d", task_name="read", - required_fields=["reward"], + required_fields=["total_reward"], batch_size=n, timeout_s=30.0, ) data = tq_client_backends.get_data(meta) - assert data["reward"].shape == reward.shape, ( - f"Expected shape {tuple(reward.shape)} for 1D field, " - f"got {tuple(data['reward'].shape)}. " + assert data["total_reward"].shape == total_reward.shape, ( + f"Expected shape {tuple(total_reward.shape)} for 1D field, " + f"got {tuple(data['total_reward'].shape)}. " "TQ must not unsqueeze 1D tensors silently (R-C2)." ) diff --git a/uv.lock b/uv.lock index 940c29c2ed9..ac5225635a8 100644 --- a/uv.lock +++ b/uv.lock @@ -4467,7 +4467,7 @@ requires-dist = [ { name = "torchdata" }, { name = "torchvision", marker = "sys_platform != 'darwin'", specifier = "==0.26.0", index = "https://download.pytorch.org/whl/cu130" }, { name = "torchvision", marker = "sys_platform == 'darwin'", specifier = "==0.26.0", index = "https://pypi.org/simple" }, - { name = "transferqueue", git = "https://github.com/Ascend/TransferQueue.git?rev=b266d39" }, + { name = "transferqueue", git = "https://github.com/Ascend/TransferQueue.git?rev=c51614308b68c8d7a87c9b3ef62d59e14c69bde2" }, { name = "transformer-engine", extras = ["core-cu13", "pytorch"], marker = "extra == 'automodel'", git = "https://github.com/NVIDIA/TransformerEngine.git?rev=v2.14.1" }, { name = "transformers", specifier = ">=5.5.0,<5.9.0" }, { name = "transformers", marker = "extra == 'automodel'", specifier = ">=5.5.0,<5.6.0" }, @@ -8004,14 +8004,15 @@ wheels = [ [[package]] name = "transferqueue" -version = "0.1.7.dev0" -source = { git = "https://github.com/Ascend/TransferQueue.git?rev=b266d39#b266d39a15aae114730de36cf8317b6285436f7f" } +version = "0.1.9" +source = { git = "https://github.com/Ascend/TransferQueue.git?rev=c51614308b68c8d7a87c9b3ef62d59e14c69bde2#c51614308b68c8d7a87c9b3ef62d59e14c69bde2" } dependencies = [ { name = "hydra-core", version = "1.3.2", source = { registry = "https://pypi.org/simple" }, marker = "(platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-fsdp') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-mcore') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-trtllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-fsdp' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-trtllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-trtllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-trtllm' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine == 'aarch64' and sys_platform == 'linux' and extra == 'extra-7-nemo-rl-mcore') or (platform_machine == 'x86_64' and sys_platform == 'linux' and extra == 'extra-7-nemo-rl-mcore') or (platform_machine == 'aarch64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-fsdp') or (platform_machine == 'aarch64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-mcore') or (platform_machine == 'aarch64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine == 'aarch64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-trtllm') or (platform_machine == 'aarch64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine == 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-fsdp') or (platform_machine == 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-mcore') or (platform_machine == 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine == 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-trtllm') or (platform_machine == 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine == 'aarch64' and extra == 'extra-7-nemo-rl-fsdp' and extra != 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine == 'aarch64' and extra != 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-trtllm') or (platform_machine == 'aarch64' and extra != 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine == 'aarch64' and extra != 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-trtllm' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine == 'x86_64' and extra == 'extra-7-nemo-rl-fsdp' and extra != 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine == 'x86_64' and extra != 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-trtllm') or (platform_machine == 'x86_64' and extra != 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine == 'x86_64' and extra != 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-trtllm' and extra == 'extra-7-nemo-rl-vllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-fsdp') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-mcore') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-trtllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-vllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-fsdp' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-trtllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-trtllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-trtllm' and extra == 'extra-7-nemo-rl-vllm')" }, { name = "hydra-core", version = "1.3.4", source = { registry = "https://pypi.org/simple" }, marker = "(platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-fsdp') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-mcore') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-trtllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-fsdp' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-trtllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-trtllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-trtllm' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine == 'aarch64' and sys_platform == 'linux' and extra == 'extra-7-nemo-rl-automodel') or (platform_machine == 'aarch64' and sys_platform == 'linux' and extra != 'extra-7-nemo-rl-mcore') or (platform_machine == 'x86_64' and sys_platform == 'linux' and extra == 'extra-7-nemo-rl-automodel') or (platform_machine == 'x86_64' and sys_platform == 'linux' and extra != 'extra-7-nemo-rl-mcore') or (platform_machine == 'aarch64' and extra != 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine == 'aarch64' and extra != 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-trtllm') or (platform_machine == 'aarch64' and extra != 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine == 'x86_64' and extra != 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine == 'x86_64' and extra != 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-trtllm') or (platform_machine == 'x86_64' and extra != 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-fsdp') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-mcore') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-trtllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-vllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-fsdp' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-trtllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-trtllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-trtllm' and extra == 'extra-7-nemo-rl-vllm')" }, { name = "msgspec", marker = "(platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-fsdp') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-mcore') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-trtllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-fsdp' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-trtllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-trtllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-trtllm' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-fsdp') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-mcore') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-trtllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-vllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-fsdp' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-trtllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-trtllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-trtllm' and extra == 'extra-7-nemo-rl-vllm')" }, { name = "numpy", marker = "(platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-fsdp') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-mcore') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-trtllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-fsdp' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-trtllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-trtllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-trtllm' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-fsdp') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-mcore') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-trtllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-vllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-fsdp' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-trtllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-trtllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-trtllm' and extra == 'extra-7-nemo-rl-vllm')" }, { name = "omegaconf", marker = "(platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-fsdp') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-mcore') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-trtllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-fsdp' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-trtllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-trtllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-trtllm' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-fsdp') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-mcore') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-trtllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-vllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-fsdp' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-trtllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-trtllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-trtllm' and extra == 'extra-7-nemo-rl-vllm')" }, + { name = "prometheus-client", marker = "(platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-fsdp') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-mcore') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-trtllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-fsdp' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-trtllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-trtllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-trtllm' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-fsdp') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-mcore') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-trtllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-vllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-fsdp' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-trtllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-trtllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-trtllm' and extra == 'extra-7-nemo-rl-vllm')" }, { name = "psutil", marker = "(platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-fsdp') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-mcore') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-trtllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-fsdp' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-trtllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-trtllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-trtllm' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-fsdp') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-mcore') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-trtllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-vllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-fsdp' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-trtllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-trtllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-trtllm' and extra == 'extra-7-nemo-rl-vllm')" }, { name = "pyzmq", marker = "(platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-fsdp') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-mcore') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-trtllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-fsdp' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-trtllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-trtllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-trtllm' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-fsdp') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-mcore') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-trtllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-vllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-fsdp' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-trtllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-trtllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-trtllm' and extra == 'extra-7-nemo-rl-vllm')" }, { name = "ray", extra = ["default"], marker = "(platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-fsdp') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-mcore') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-trtllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-fsdp' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-trtllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-trtllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-trtllm' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-fsdp') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-mcore') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-trtllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-vllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-fsdp' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-trtllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-trtllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-trtllm' and extra == 'extra-7-nemo-rl-vllm')" }, From 0e687e6d07623d780a4174310e92382ce738a8a2 Mon Sep 17 00:00:00 2001 From: Yuki Huang Date: Sat, 8 Aug 2026 07:10:15 +0800 Subject: [PATCH 09/10] feat(sc): overlap NeMo-Gym spinup with deferred vLLM load (#3499) Signed-off-by: Yuki Huang --- examples/run_grpo_single_controller.py | 8 +- nemo_rl/algorithms/grpo.py | 129 ++++----- nemo_rl/algorithms/metric_utils.py | 105 +++++++ nemo_rl/algorithms/single_controller.py | 6 + .../single_controller_utils/setup.py | 266 ++++++++++++++---- nemo_rl/environments/nemo_gym.py | 2 +- pyrefly.toml | 1 + ...async-1off-single-controller-streaming2.sh | 2 +- tests/unit/algorithms/test_metric_utils.py | 144 ++++++++++ .../single_controller/test_rollout_pump.py | 5 +- .../test_run_grpo_single_controller.py | 3 +- .../test_single_controller.py | 50 ++++ .../test_single_controller_setup.py | 216 ++++++++++++-- .../unit/single_controller/test_train_pump.py | 2 + tools/refit_verifier.py | 2 +- 15 files changed, 783 insertions(+), 158 deletions(-) create mode 100644 nemo_rl/algorithms/metric_utils.py create mode 100644 tests/unit/algorithms/test_metric_utils.py diff --git a/examples/run_grpo_single_controller.py b/examples/run_grpo_single_controller.py index 544124c9fa0..009f683cb36 100644 --- a/examples/run_grpo_single_controller.py +++ b/examples/run_grpo_single_controller.py @@ -122,10 +122,14 @@ def main() -> None: if bool(config.env.get("should_use_nemo_gym")): setup_nemo_gym_config(config, tokenizer) - actor_args = setup_single_controller(config, tokenizer) + actor_args, setup_timing_metrics = setup_single_controller(config, tokenizer) print("🚀 Launching SingleControllerActor") - sc = SingleControllerActor.remote(master_config=config, actor_args=actor_args) + sc = SingleControllerActor.remote( + master_config=config, + actor_args=actor_args, + setup_timing_metrics=setup_timing_metrics, + ) try: result = ray.get(sc.run.remote()) print(f"SC run complete: {result}") diff --git a/nemo_rl/algorithms/grpo.py b/nemo_rl/algorithms/grpo.py index 6d612e50b0d..2cc604ba2c4 100644 --- a/nemo_rl/algorithms/grpo.py +++ b/nemo_rl/algorithms/grpo.py @@ -47,6 +47,10 @@ ClippedPGLossFn, ) from nemo_rl.algorithms.loss.interfaces import LossFunction +from nemo_rl.algorithms.metric_utils import ( + SetupTimingMetrics, + print_setup_timing_summary, +) from nemo_rl.algorithms.opd import OnPolicyDistillationConfig from nemo_rl.algorithms.reward_functions import ( RewardShapingConfig, @@ -996,18 +1000,21 @@ def _spinup_nemo_gym(base_urls, model_name): # vllm model loading prefers clean environment, initialize policy_generation before policy in colocated mode backend = generation_config["backend"] + gen_init_time_key = ( + "megatron_generation_init_time_s" + if backend == "megatron" + else f"{backend}_init_time_s" + ) generation_config["model_name"] = policy_config["model_name"] # Needed for vLLM remote_transport = None remote_synchronizer_cls = None remote_baseline_init_refs: list[Any] = [] checkpoint_engine_config = None - # Dictionary to store worker initialization timing stats for logging - worker_init_timing_metrics = {} + # Worker initialization timing stats — populated as each phase completes. + setup_timing_metrics = SetupTimingMetrics() if teacher_reservation_time: - worker_init_timing_metrics["teacher_reservation_time_s"] = ( - teacher_reservation_time - ) + setup_timing_metrics.teacher_reservation_time_s = teacher_reservation_time weights_path, optimizer_path = checkpointer.get_resume_paths(last_checkpoint_path) @@ -1112,22 +1119,18 @@ def init_megatron_generation(policy=None): def initialize_generation_with_policy( init_generation_fn, - generation_name: str, - init_time_key: str, colocated_inference: bool, - worker_init_timing_metrics: dict, + setup_timing_metrics: SetupTimingMetrics, ): - """Generic function to initialize a generation engine (vLLM or SGLang) along with policy. + """Initialize a generation engine along with policy, sequentially or in parallel. Args: - init_generation_fn: Function that initializes the generation engine (init_vllm or init_sglang) - generation_name: Name of the generation engine ("vLLM" or "SGLang") - init_time_key: Key name for storing initialization time in metrics ("vllm_init_time_s" or "sglang_init_time_s") - colocated_inference: Whether inference is colocated with training - worker_init_timing_metrics: Dictionary to store timing metrics + init_generation_fn: Function that initializes the generation engine (init_vllm, ...). + colocated_inference: Whether inference is colocated with training. + setup_timing_metrics: SetupTimingMetrics to store timings on. Returns: - Tuple of (policy_generation, policy) + Tuple of (policy_generation, policy). """ # Determine if parallel initialization is possible (non-colocated mode) use_parallel_init = not colocated_inference @@ -1149,10 +1152,10 @@ def initialize_generation_with_policy( parallel_wall_time = time.perf_counter() - parallel_start_time # Store timing metrics - worker_init_timing_metrics[init_time_key] = generation_time - worker_init_timing_metrics["policy_init_time_s"] = policy_time - worker_init_timing_metrics["parallel_wall_time_s"] = parallel_wall_time - worker_init_timing_metrics["parallel_init_enabled"] = True + setattr(setup_timing_metrics, gen_init_time_key, generation_time) + setup_timing_metrics.policy_init_time_s = policy_time + setup_timing_metrics.parallel_wall_time_s = parallel_wall_time + setup_timing_metrics.parallel_init_enabled = 1.0 else: # Sequential initialization: colocated mode (GPU memory requires generation engine first) @@ -1163,11 +1166,11 @@ def initialize_generation_with_policy( # Initialize generation engine first (clean GPU memory), then policy policy_generation, generation_time = init_generation_fn() - worker_init_timing_metrics[init_time_key] = generation_time + setattr(setup_timing_metrics, gen_init_time_key, generation_time) policy, policy_time = init_policy() - worker_init_timing_metrics["policy_init_time_s"] = policy_time - worker_init_timing_metrics["parallel_init_enabled"] = 0.0 + setup_timing_metrics.policy_init_time_s = policy_time + setup_timing_metrics.parallel_init_enabled = 0.0 return policy_generation, policy @@ -1175,13 +1178,11 @@ def initialize_generation_with_policy( if backend == "megatron": # Initialize training first so checkpoint conversion completes before inference starts. policy, policy_time = init_policy() - worker_init_timing_metrics["policy_init_time_s"] = policy_time + setup_timing_metrics.policy_init_time_s = policy_time # Colocated wraps the training policy; non-colocated builds a dedicated inference policy. policy_generation, megatron_gen_time = init_megatron_generation(policy) - worker_init_timing_metrics["megatron_generation_init_time_s"] = ( - megatron_gen_time - ) + setup_timing_metrics.megatron_generation_init_time_s = megatron_gen_time if enable_nemo_gym: # The Megatron inference engine must be up before its server URLs exist. @@ -1189,7 +1190,7 @@ def initialize_generation_with_policy( policy_generation.dp_openai_server_base_urls, generation_config["model_name"], ) - worker_init_timing_metrics["nemo_gym_init_time_s"] = nemo_gym_time + setup_timing_metrics.nemo_gym_init_time_s = nemo_gym_time print( f" ✓ Using {backend} backend for generation with {policy_config['model_name']}", @@ -1254,11 +1255,13 @@ def initialize_generation_with_policy( " ⚡ Deferred model load: reserving vLLM ports for overlapped NeMo Gym init", flush=True, ) + vllm_reserve_t0 = time.perf_counter() deferred_vllm = VllmGeneration( cluster=inference_cluster, config=generation_config, defer_model_load=True, ) + vllm_reserve_time = time.perf_counter() - vllm_reserve_t0 print( f" ✓ Reserved {len(deferred_vllm.dp_openai_server_base_urls)} vLLM server URLs: " f"{deferred_vllm.dp_openai_server_base_urls}", @@ -1303,23 +1306,21 @@ def init_vllm_then_policy(): results = {k: f.result() for k, f in submitted.items()} if colocated_inference: - policy_generation, vllm_time, policy, policy_time = results[ + policy_generation, vllm_load_time, policy, policy_time = results[ "vllm_policy" ] else: - policy_generation, vllm_time = results["vllm"] + policy_generation, vllm_load_time = results["vllm"] policy, policy_time = results["policy"] nemo_gym_actor, nemo_gym_time = results["nemo_gym"] - worker_init_timing_metrics["vllm_init_time_s"] = vllm_time - worker_init_timing_metrics["policy_init_time_s"] = policy_time - worker_init_timing_metrics["nemo_gym_init_time_s"] = nemo_gym_time + setup_timing_metrics.vllm_init_time_s = vllm_reserve_time + vllm_load_time + setup_timing_metrics.policy_init_time_s = policy_time + setup_timing_metrics.nemo_gym_init_time_s = nemo_gym_time else: policy_generation, policy = initialize_generation_with_policy( init_generation_fn=init_vllm, - generation_name="vLLM", - init_time_key="vllm_init_time_s", colocated_inference=colocated_inference, - worker_init_timing_metrics=worker_init_timing_metrics, + setup_timing_metrics=setup_timing_metrics, ) print( @@ -1336,10 +1337,8 @@ def init_vllm_then_policy(): policy_generation, policy = initialize_generation_with_policy( init_generation_fn=init_sglang, - generation_name="SGLang", - init_time_key="sglang_init_time_s", colocated_inference=colocated_inference, - worker_init_timing_metrics=worker_init_timing_metrics, + setup_timing_metrics=setup_timing_metrics, ) # Capture rollout TP size on the policy once; refit calls no longer need it. @@ -1362,10 +1361,8 @@ def init_trtllm(): policy_generation, policy = initialize_generation_with_policy( init_generation_fn=init_trtllm, - generation_name="TRT-LLM", - init_time_key="trtllm_init_time_s", colocated_inference=colocated_inference, - worker_init_timing_metrics=worker_init_timing_metrics, + setup_timing_metrics=setup_timing_metrics, ) print( @@ -1378,7 +1375,7 @@ def init_trtllm(): policy_generation.dp_openai_server_base_urls, generation_config["model_name"], ) - worker_init_timing_metrics["nemo_gym_init_time_s"] = nemo_gym_time + setup_timing_metrics.nemo_gym_init_time_s = nemo_gym_time # Record when worker initialization completes (for calculating other setup time) worker_init_complete_time = time.perf_counter() - setup_start_time @@ -1457,7 +1454,8 @@ def init_trtllm(): ip, port, world_size, train_world_size=train_world_size ) # type: ignore ray.get(futures_train + futures_inference) - worker_init_timing_metrics["collective_init_time_s"] = time.perf_counter() - t0 + setup_timing_metrics.collective_init_time_s = time.perf_counter() - t0 + if remote_transport is not None: t0 = time.perf_counter() assert isinstance(policy_generation, VllmGeneration) @@ -1475,7 +1473,7 @@ def init_trtllm(): baseline_init_refs=remote_baseline_init_refs, ) policy_generation.weight_synchronizer.init_communicator() - worker_init_timing_metrics[f"vllm_{remote_transport}_sparse_init_time_s"] = ( + setup_timing_metrics.extras[f"vllm_{remote_transport}_sparse_init_time_s"] = ( time.perf_counter() - t0 ) elif checkpoint_engine_config is not None: @@ -1490,7 +1488,7 @@ def init_trtllm(): inference_cluster=inference_cluster, ) policy_generation.weight_synchronizer.init_communicator() - worker_init_timing_metrics["vllm_checkpoint_engine_init_time_s"] = ( + setup_timing_metrics.vllm_checkpoint_engine_init_time_s = ( time.perf_counter() - t0 ) print( @@ -1520,46 +1518,25 @@ def init_trtllm(): ) ) teacher_model_init_time = time.perf_counter() - t0 - worker_init_timing_metrics["teacher_model_init_time_s"] = ( - teacher_model_init_time - ) + setup_timing_metrics.teacher_model_init_time_s = teacher_model_init_time # Preserve the existing metric's end-to-end meaning while exposing the # newly separated reservation and model-initialization phases. - worker_init_timing_metrics["teacher_init_time_s"] = ( + setup_timing_metrics.teacher_init_time_s = ( teacher_reservation_time + teacher_model_init_time ) # Calculate total setup time total_setup_time = time.perf_counter() - setup_start_time - worker_init_timing_metrics["total_setup_time_s"] = total_setup_time + setup_timing_metrics.total_setup_time_s = total_setup_time + setup_timing_metrics.other_setup_time_s = ( + total_setup_time - worker_init_complete_time + ) # Log worker initialization timing metrics to logger - if worker_init_timing_metrics: - print("\n▶ Worker Initialization Timing:") - - vllm_time = worker_init_timing_metrics.get("vllm_init_time_s", 0) - policy_time = worker_init_timing_metrics.get("policy_init_time_s", 0) - total_setup = worker_init_timing_metrics.get("total_setup_time_s", 0) - - if vllm_time: - print(f" vLLM init: {vllm_time:.1f}s") - - if policy_time: - print(f" Policy init: {policy_time:.1f}s") - - teacher_time = worker_init_timing_metrics.get("teacher_init_time_s", 0) - if teacher_time: - print(f" Teacher init: {teacher_time:.1f}s") - - # Calculate "other" time (time after worker init completes) - other_time = total_setup - worker_init_complete_time - worker_init_timing_metrics["other_setup_time_s"] = other_time - print(f" Other setup: {other_time:.1f}s") - - print(f" Total setup: {total_setup:.1f}s") - - # Log all metrics to the logger for analysis - logger.log_metrics(worker_init_timing_metrics, step=0, prefix="timing/setup") + print_setup_timing_summary(setup_timing_metrics, gen_init_time_key) + logger.log_metrics( + setup_timing_metrics.to_metrics_dict(), step=0, prefix="timing/setup" + ) print("\n" + "=" * 60) print(" " * 18 + "SETUP COMPLETE") diff --git a/nemo_rl/algorithms/metric_utils.py b/nemo_rl/algorithms/metric_utils.py new file mode 100644 index 00000000000..8c7cd7c0ffe --- /dev/null +++ b/nemo_rl/algorithms/metric_utils.py @@ -0,0 +1,105 @@ +# 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. + +from __future__ import annotations + +from dataclasses import asdict, dataclass, field +from typing import Any, Optional + + +@dataclass +class SetupTimingMetrics: + """Driver-side per-phase timings collected during setup.""" + + # grpo.py only: generation-backend init (exactly one populated per run). + vllm_init_time_s: Optional[float] = None + sglang_init_time_s: Optional[float] = None + trtllm_init_time_s: Optional[float] = None + megatron_generation_init_time_s: Optional[float] = None + + # SC only: generation init (reserve + load). + generation_init_time_s: Optional[float] = None + generation_init_reserve_time_s: Optional[float] = None + generation_init_load_time_s: Optional[float] = None + + policy_init_time_s: Optional[float] = None + nemo_gym_init_time_s: Optional[float] = None + collective_init_time_s: Optional[float] = None + + # Non-colocated only. (grpo.py only) + parallel_wall_time_s: Optional[float] = None + parallel_init_enabled: Optional[float] = None + + # grpo-only phases (non-colocated OPD teachers, sparse refit, checkpoint-engine). + teacher_reservation_time_s: Optional[float] = None + teacher_model_init_time_s: Optional[float] = None + teacher_init_time_s: Optional[float] = None + vllm_checkpoint_engine_init_time_s: Optional[float] = None + + total_setup_time_s: Optional[float] = None + worker_setup_time_s: Optional[float] = None + other_setup_time_s: Optional[float] = None + + # Overflow bucket for dynamic-keyed metrics (e.g. one entry per active + # sparse refit transport: vllm__sparse_init_time_s). + extras: dict[str, float] = field(default_factory=dict) + + def to_metrics_dict(self) -> dict[str, Any]: + """Serialize for Logger.log_metrics; drops unset (None) fields.""" + base = { + k: v for k, v in asdict(self).items() if k != "extras" and v is not None + } + base.update(self.extras) + return base + + +def print_setup_timing_summary( + metrics: SetupTimingMetrics, gen_init_time_key: Optional[str] = None +) -> None: + """Print the setup-phase summary block. + + Args: + metrics: Populated timing metrics. + gen_init_time_key: grpo.py passes the backend-specific field name + (e.g. "vllm_init_time_s"); SC leaves it None and the summary + reads generation_init_time_s (+ optional reserve/load split). + """ + print("\n▶ Worker Initialization Timing:") + + if metrics.generation_init_reserve_time_s: + # SC + gym-on path + print( + f" Generation init: {metrics.generation_init_time_s:.1f}s" + f" (reserve {metrics.generation_init_reserve_time_s:.1f}s" + f" + load {metrics.generation_init_load_time_s:.1f}s)" + ) + elif gen_init_time_key is None: + # SC + gym-off path + assert metrics.generation_init_time_s is not None + print(f" Generation init: {metrics.generation_init_time_s:.1f}s") + else: + # grpo.py path + assert metrics.generation_init_time_s is None + print(f" Generation init: {getattr(metrics, gen_init_time_key):.1f}s") + + print(f" Policy init: {metrics.policy_init_time_s:.1f}s") + + if metrics.nemo_gym_init_time_s: + print(f" NeMo-Gym init: {metrics.nemo_gym_init_time_s:.1f}s") + + if metrics.teacher_init_time_s: + print(f" Teacher init: {metrics.teacher_init_time_s:.1f}s") + + print(f" Other setup: {metrics.other_setup_time_s:.1f}s") + print(f" Total setup: {metrics.total_setup_time_s:.1f}s", flush=True) diff --git a/nemo_rl/algorithms/single_controller.py b/nemo_rl/algorithms/single_controller.py index 045e3372d90..77325564db8 100644 --- a/nemo_rl/algorithms/single_controller.py +++ b/nemo_rl/algorithms/single_controller.py @@ -43,6 +43,7 @@ import torch from nemo_rl.algorithms.async_utils.staleness_sampler import create_sampler +from nemo_rl.algorithms.metric_utils import SetupTimingMetrics from nemo_rl.algorithms.single_controller_utils.config import ( AdvantageConfig, MasterConfig, @@ -87,12 +88,14 @@ def __init__( self, master_config: MasterConfig, actor_args: SingleControllerActorArgs, + setup_timing_metrics: SetupTimingMetrics, ) -> None: """Initialize the SingleController actor. Args: master_config: SC MasterConfig. actor_args: Pre-built actor args from setup_single_controller. + setup_timing_metrics: Driver-side setup timings; logged here (Logger isn't cloudpickleable). """ validate_single_controller_config(master_config) @@ -125,6 +128,9 @@ def __init__( # _thread.lock that Ray can't cloudpickle into the actor. self._logger = Logger(master_config.logger) # type: ignore self._logger.log_hyperparams(master_config.model_dump()) + self._logger.log_metrics( + setup_timing_metrics.to_metrics_dict(), step=0, prefix="timing/setup" + ) self._timer = Timer() # Pin clusters so RayVirtualCluster.__del__ doesn't remove the PGs. diff --git a/nemo_rl/algorithms/single_controller_utils/setup.py b/nemo_rl/algorithms/single_controller_utils/setup.py index 4a771ceedbe..96eff61231d 100644 --- a/nemo_rl/algorithms/single_controller_utils/setup.py +++ b/nemo_rl/algorithms/single_controller_utils/setup.py @@ -21,9 +21,11 @@ from __future__ import annotations +import time from concurrent.futures import ThreadPoolExecutor from dataclasses import dataclass -from typing import Any, Optional, cast +from functools import partial +from typing import Any, Callable, Optional, cast from torchdata.stateful_dataloader import StatefulDataLoader from transformers import AutoProcessor @@ -37,6 +39,10 @@ ) from nemo_rl.algorithms.loss import ClippedPGLossFn from nemo_rl.algorithms.loss.interfaces import LossFunction +from nemo_rl.algorithms.metric_utils import ( + SetupTimingMetrics, + print_setup_timing_summary, +) from nemo_rl.algorithms.single_controller_utils.config import ( MasterConfig, validate_single_controller_config, @@ -164,19 +170,41 @@ def _build_clusters( def _build_generation( inference_cluster: RayVirtualCluster, master_config: MasterConfig, -): - """Spin up the generation backend (vLLM or SGLang).""" + *, + defer_model_load: bool = False, +) -> tuple[Any, float]: + """Spin up the generation backend (vLLM or SGLang). + + Args: + inference_cluster: Ray virtual cluster the generation workers run on. + master_config: SC MasterConfig. + defer_model_load: If True (for the NeMo-Gym flow), reserve OpenAI server URLs without loading weights; caller runs gen.load_and_start() later. + + Returns: + A tuple of (generation object, wall time spent in this call). The + generation object is a VllmGeneration or SGLangGeneration. + """ + t0 = time.perf_counter() generation_config = master_config.policy["generation"] generation_config["model_name"] = master_config.policy["model_name"] backend = generation_config["backend"] + if backend == "vllm": vllm_config = cast(VllmConfig, generation_config) vllm_config.setdefault("vllm_kwargs", {})["hf_overrides"] = ( master_config.policy.get("hf_config_overrides", {}) ) configure_vllm_for_router_replay(master_config.policy) - gen = VllmGeneration(cluster=inference_cluster, config=vllm_config) + gen = VllmGeneration( + cluster=inference_cluster, + config=vllm_config, + defer_model_load=defer_model_load, + ) + elif backend == "sglang": + assert not defer_model_load, ( + "defer_model_load is only supported for the vllm backend" + ) sglang_config = cast(SGLangConfig, generation_config) sglang_config["sglang_cfg"].setdefault( "model_path", master_config.policy["model_name"] @@ -185,12 +213,31 @@ def _build_generation( cluster=inference_cluster, sglang_cfg=sglang_config, ) + else: raise ValueError( f"single_controller_utils.setup only supports vllm or sglang generation; got {backend!r}" ) - gen.finish_generation() - return gen + + if not defer_model_load: + gen.finish_generation() + + return gen, time.perf_counter() - t0 + + +def _finish_deferred_generation(generation: Any) -> tuple[Any, float]: + """Finish loading and starting the deferred generation. + + Args: + generation: The deferred generation object. + + Returns: + A tuple of (finished generation object, wall time spent in this call). + """ + t0 = time.perf_counter() + generation.load_and_start() + generation.finish_generation() + return generation, time.perf_counter() - t0 def _build_trainer( @@ -198,17 +245,22 @@ def _build_trainer( master_config: MasterConfig, tokenizer, processor, -): +) -> tuple[Any, float]: """Build the TQ-mediated trainer (driver-side TQPolicy). - Driver-side on purpose: instantiating TQPolicy inside another Ray - actor nests runtime_envs and triggers Ray's - get_accelerator_ids_for_accelerator_resource IndexError. Keep this - here until PolicyTrainerActor (PR #2692) lands. + Args: + train_cluster: Ray virtual cluster the trainer workers run on. + master_config: SC MasterConfig. + tokenizer: Tokenizer used by the policy. + processor: Optional AutoProcessor for VLM paths. + + Returns: + A tuple of (TQPolicy trainer, wall time spent in this call). """ + t0 = time.perf_counter() loss_config = master_config.loss_fn init_reference_model = loss_config.reference_policy_kl_penalty > 0 - return TQPolicy( + trainer = TQPolicy( cluster=train_cluster, config=master_config.policy, tokenizer=tokenizer, @@ -219,6 +271,37 @@ def _build_trainer( init_reference_model=init_reference_model, dp_cfg=master_config.data_plane, ) + return trainer, time.perf_counter() - t0 + + +def _spinup_gym(master_config: MasterConfig, base_urls: list[str]) -> tuple[Any, float]: + """Spin up the NeMo-Gym actor against the reserved vLLM URLs. + + Args: + master_config: SC MasterConfig. + base_urls: Reserved vLLM OpenAI server URLs. + + Returns: + A tuple of (NeMo-Gym actor, wall time spent in this call). + """ + t0 = time.perf_counter() + policy_config = master_config.policy + generation_config = policy_config["generation"] + enable_router_replay = router_replay_enabled(policy_config) + routed_experts_dtype = ( + resolve_routed_experts_dtype_name_for_model(generation_config["model_name"]) + if enable_router_replay + else "int16" + ) + actor = spinup_nemo_gym_actor( + env_configs=master_config.env, + base_urls=base_urls, + model_name=generation_config["model_name"], + enable_router_replay=enable_router_replay, + routed_experts_dtype=routed_experts_dtype, + use_fastokens=bool(policy_config["tokenizer"].get("use_fastokens")), + ) + return actor, time.perf_counter() - t0 def _generation_max_seq_len(generation_config) -> int: @@ -267,7 +350,7 @@ def setup_single_controller( *, processor: Optional[AutoProcessor] = None, partition_id: str = "rollout_data", -) -> SingleControllerActorArgs: +) -> tuple[SingleControllerActorArgs, SetupTimingMetrics]: """Build the full SC actor args driver-side. Args: @@ -277,7 +360,8 @@ def setup_single_controller( partition_id: TQ partition the rollout writer + sampler share. Returns: - SingleControllerActorArgs ready to be passed to SingleControllerActor. + A tuple of (pre-built SC actor args, driver-side per-phase timings + logged by the SC actor). """ validate_single_controller_config(master_config) @@ -356,64 +440,138 @@ def setup_single_controller( # ========================== # Setup Clusters & Workers # ========================== + setup_start_time = time.perf_counter() + setup_timing_metrics = SetupTimingMetrics() + + # Create clusters train_cluster, inference_cluster = _build_clusters(master_config) colocated = generation_config["colocated"]["enabled"] - if colocated: - # Colocated: vLLM prefers a clean GPU at load time, so generation - # comes up before the policy. - generation = _build_generation(inference_cluster, master_config) - policy = _build_trainer(train_cluster, master_config, tokenizer, processor) - else: - # Non-colocated: generation + policy run on disjoint GPUs, so - # bring them up in parallel. - with ThreadPoolExecutor(max_workers=2) as executor: - gen_future = executor.submit( - _build_generation, inference_cluster, master_config + + # Create build tasks for generation / trainer / (nemo-gym) workers + build_tasks: dict[str, Callable[[], Any]] = {} + generation = None + defer_generation_model_load = False + gen_reserve_time = 0.0 + + def _build_generation_then_trainer( + defer_generation_model_load: bool, generation=None + ) -> tuple[Any, Any, dict[str, float]]: + """Build generation then trainer serially. + + Args: + defer_generation_model_load: If True, generation is a pre-reserved handle and this call + finishes its model load; if False, builds generation from scratch. + generation: Pre-reserved generation handle when defer_generation_model_load=True; None otherwise. + + Returns: + A tuple of (finalized generation object, TQPolicy trainer, + per-phase wall times keyed as "gen_time" and "trainer_time"). + """ + time_metrics = {} + + # generation + if defer_generation_model_load: + generation, time_metrics["gen_time"] = _finish_deferred_generation( + generation ) - policy_future = executor.submit( - _build_trainer, train_cluster, master_config, tokenizer, processor + else: + generation, time_metrics["gen_time"] = _build_generation( + inference_cluster, master_config ) - generation = gen_future.result() - policy = policy_future.result() - # ========================== - # NeMo-Gym actor (after generation is up so OpenAI URLs are available) - # ========================== + # trainer + trainer, time_metrics["trainer_time"] = _build_trainer( + train_cluster, master_config, tokenizer, processor + ) + + return generation, trainer, time_metrics + if use_nemo_gym: - # TODO(#2625): Mirror GRPO's deferred vLLM load so NeMo-Gym spinup - # overlaps model loading instead of running serially afterward. - enable_router_replay = router_replay_enabled(policy_config) - routed_experts_dtype = ( - resolve_routed_experts_dtype_name_for_model(generation_config["model_name"]) - if enable_router_replay - else "int16" + # defer generation, only get base_urls for nemo_gym spinup + generation, gen_reserve_time = _build_generation( + inference_cluster, + master_config=master_config, + defer_model_load=True, ) - env_handles["nemo_gym"] = spinup_nemo_gym_actor( - env_configs=master_config.env, + defer_generation_model_load = True + # add nemo_gym spinup task + build_tasks["nemo_gym"] = partial( + _spinup_gym, + master_config=master_config, base_urls=generation.dp_openai_server_base_urls, - model_name=generation_config["model_name"], - enable_router_replay=enable_router_replay, - routed_experts_dtype=routed_experts_dtype, - use_fastokens=bool(policy_config["tokenizer"].get("use_fastokens")), ) + if colocated: + # Colocated: vLLM prefers a clean GPU at load time, so generation comes up before the trainer. + build_tasks["generation_trainer"] = partial( + _build_generation_then_trainer, + defer_generation_model_load=defer_generation_model_load, + generation=generation, + ) + else: + # Non-colocated: generation + trainer run on disjoint GPUs, so bring them up in parallel. + if defer_generation_model_load: + build_tasks["generation"] = partial( + _finish_deferred_generation, + generation=generation, + ) + else: + build_tasks["generation"] = partial( + _build_generation, + inference_cluster=inference_cluster, + master_config=master_config, + ) + build_tasks["trainer"] = partial( + _build_trainer, + train_cluster=train_cluster, + master_config=master_config, + tokenizer=tokenizer, + processor=processor, + ) + + # Submit build tasks and get results + with ThreadPoolExecutor(max_workers=len(build_tasks)) as executor: + submitted = {k: executor.submit(fn) for k, fn in build_tasks.items()} + results = {k: f.result() for k, f in submitted.items()} + + if colocated: + generation, trainer, time_metrics = results["generation_trainer"] + gen_load_time = time_metrics["gen_time"] + setup_timing_metrics.policy_init_time_s = time_metrics["trainer_time"] + else: + generation, gen_load_time = results["generation"] + trainer, trainer_time = results["trainer"] + setup_timing_metrics.policy_init_time_s = trainer_time + setup_timing_metrics.generation_init_time_s = gen_reserve_time + gen_load_time + + if use_nemo_gym: + env_handles["nemo_gym"], gym_time = results["nemo_gym"] + setup_timing_metrics.nemo_gym_init_time_s = gym_time + # the two fields are only meaningful when use_nemo_gym enabled + setup_timing_metrics.generation_init_reserve_time_s = gen_reserve_time + setup_timing_metrics.generation_init_load_time_s = gen_load_time + + worker_setup_time = time.perf_counter() - setup_start_time + setup_timing_metrics.worker_setup_time_s = worker_setup_time + # ========================== # Setup Data Plane Client & Weight Sync # ========================== # Connect-only DP client; TQPolicy already bootstrapped the controller. dp_client = build_data_plane_client(dp_config, bootstrap=False) - backend = generation_config["backend"] + t0 = time.perf_counter() weight_synchronizer = create_weight_synchronizer( - policy=policy, + policy=trainer, generation=generation, - generation_backend=backend, + generation_backend=generation_config["backend"], colocated=colocated, train_cluster=train_cluster, inference_cluster=inference_cluster, refit_buffer_size_gb=policy_config.get("refit_buffer_size_gb"), ) weight_synchronizer.init_communicator() + setup_timing_metrics.collective_init_time_s = time.perf_counter() - t0 # ========================== # Setup Algorithm + Rollout Wiring @@ -443,9 +601,16 @@ def setup_single_controller( tq_buffer=tq_buffer, ) - return SingleControllerActorArgs( + # Print setup timing metrics + total_setup_time = time.perf_counter() - setup_start_time + setup_timing_metrics.total_setup_time_s = total_setup_time + setup_timing_metrics.other_setup_time_s = total_setup_time - worker_setup_time + print_setup_timing_summary(setup_timing_metrics) + + # Build actor args and return + actor_args = SingleControllerActorArgs( gen_handle=generation, - trainer_handle=policy, + trainer_handle=trainer, env_handles=env_handles, train_cluster=train_cluster, inference_cluster=inference_cluster, @@ -458,3 +623,4 @@ def setup_single_controller( tq_buffer=tq_buffer, partition_id=partition_id, ) + return actor_args, setup_timing_metrics diff --git a/nemo_rl/environments/nemo_gym.py b/nemo_rl/environments/nemo_gym.py index 375d3b2d41d..00e2361bb61 100644 --- a/nemo_rl/environments/nemo_gym.py +++ b/nemo_rl/environments/nemo_gym.py @@ -850,7 +850,7 @@ def setup_nemo_gym_config(config, tokenizer) -> None: def spinup_nemo_gym_actor( env_configs: dict[str, Any], - base_urls: list[Optional[str]], + base_urls: list[str], model_name: str, *, enable_router_replay: bool, diff --git a/pyrefly.toml b/pyrefly.toml index 5f55643e09b..038beccf6a6 100644 --- a/pyrefly.toml +++ b/pyrefly.toml @@ -55,6 +55,7 @@ project-includes = [ "nemo_rl/algorithms/loss/__init__.py", "nemo_rl/algorithms/loss/interfaces.py", "nemo_rl/algorithms/loss/utils.py", + "nemo_rl/algorithms/metric_utils.py", "nemo_rl/algorithms/opd.py", "nemo_rl/algorithms/reward_functions.py", "nemo_rl/algorithms/single_controller.py", diff --git a/tests/test_suites/llm/grpo-llama3.1-8b-instruct-2n8g-async-1off-single-controller-streaming2.sh b/tests/test_suites/llm/grpo-llama3.1-8b-instruct-2n8g-async-1off-single-controller-streaming2.sh index 0b559f5902f..56a639f74fb 100755 --- a/tests/test_suites/llm/grpo-llama3.1-8b-instruct-2n8g-async-1off-single-controller-streaming2.sh +++ b/tests/test_suites/llm/grpo-llama3.1-8b-instruct-2n8g-async-1off-single-controller-streaming2.sh @@ -36,7 +36,7 @@ if [[ $(jq 'to_entries | .[] | select(.key == "train/loss") | .value | keys | ma uv run tests/check_metrics.py $JSON_METRICS \ 'median(data["train/token_mult_prob_error"]) < 1.1' \ 'data["train/token_mult_prob_error"]["10"] < 1.1' \ - 'mean(data["train/grad_norm"], 2, 0) > 0.10' + 'mean(data["train/grad_norm"], 2, 0) > 0.06' # Clean up checkpoint directory after successful run to save space. rm -rf "$CKPT_DIR" diff --git a/tests/unit/algorithms/test_metric_utils.py b/tests/unit/algorithms/test_metric_utils.py new file mode 100644 index 00000000000..42a55619b42 --- /dev/null +++ b/tests/unit/algorithms/test_metric_utils.py @@ -0,0 +1,144 @@ +# 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. + +"""Unit tests for SetupTimingMetrics + print_setup_timing_summary.""" + +from __future__ import annotations + +import pytest + +from nemo_rl.algorithms.metric_utils import ( + SetupTimingMetrics, + print_setup_timing_summary, +) + + +class TestPrintSetupTimingSummary: + """print_setup_timing_summary has three code paths with assertions.""" + + @staticmethod + def _common_setup(**overrides) -> SetupTimingMetrics: + base = { + "policy_init_time_s": 20.0, + "other_setup_time_s": 1.0, + "total_setup_time_s": 25.0, + } + base.update(overrides) + return SetupTimingMetrics(**base) + + def test_sc_gym_on_prints_reserve_load_split(self, capsys): + """SC + gym-on renders the '(reserve X.Xs + load Y.Ys)' suffix.""" + metrics = self._common_setup( + generation_init_time_s=15.0, + generation_init_reserve_time_s=3.0, + generation_init_load_time_s=12.0, + ) + print_setup_timing_summary(metrics) + out = capsys.readouterr().out + assert "Generation init: 15.0s (reserve 3.0s + load 12.0s)" in out + + def test_sc_gym_off_prints_plain_generation_init(self, capsys): + """SC + gym-off renders only the top-level generation_init_time_s.""" + metrics = self._common_setup(generation_init_time_s=15.0) + print_setup_timing_summary(metrics) + out = capsys.readouterr().out + assert "Generation init: 15.0s\n" in out + # no reserve/load suffix on this path. + assert "reserve" not in out + assert "load" not in out + + def test_sc_gym_off_asserts_generation_init_time_populated(self): + """SC path with gen_init_time_key=None must have generation_init_time_s set.""" + metrics = self._common_setup() + with pytest.raises(AssertionError): + print_setup_timing_summary(metrics) + + def test_grpo_uses_backend_specific_key(self, capsys): + """grpo.py path reads the field named by gen_init_time_key.""" + metrics = self._common_setup(vllm_init_time_s=15.0) + print_setup_timing_summary(metrics, gen_init_time_key="vllm_init_time_s") + out = capsys.readouterr().out + assert "Generation init: 15.0s" in out + assert "reserve" not in out + + def test_grpo_asserts_generation_init_time_unset(self): + """grpo.py path forbids generation_init_time_s from being populated.""" + metrics = self._common_setup( + generation_init_time_s=15.0, + vllm_init_time_s=15.0, + ) + with pytest.raises(AssertionError): + print_setup_timing_summary(metrics, gen_init_time_key="vllm_init_time_s") + + def test_reserve_load_split_takes_precedence_over_gen_key(self, capsys): + """If reserve_time_s is set, the SC+gym-on branch wins even if a key is passed.""" + metrics = self._common_setup( + generation_init_time_s=15.0, + generation_init_reserve_time_s=3.0, + generation_init_load_time_s=12.0, + ) + print_setup_timing_summary(metrics, gen_init_time_key="vllm_init_time_s") + out = capsys.readouterr().out + assert "Generation init: 15.0s (reserve 3.0s + load 12.0s)" in out + + def test_optional_nemo_gym_and_teacher_lines(self, capsys): + """nemo_gym_init_time_s and teacher_init_time_s only print when populated.""" + metrics = self._common_setup( + generation_init_time_s=15.0, + nemo_gym_init_time_s=8.0, + teacher_init_time_s=6.0, + ) + print_setup_timing_summary(metrics) + out = capsys.readouterr().out + assert "NeMo-Gym init: 8.0s" in out + assert "Teacher init: 6.0s" in out + + +class TestSetupTimingMetricsToDict: + """to_metrics_dict serializes into a dict for Logger.log_metrics.""" + + def test_drops_none_fields(self): + """Unset (None) fields are dropped.""" + metrics = SetupTimingMetrics(generation_init_time_s=1.5) + d = metrics.to_metrics_dict() + assert d == {"generation_init_time_s": 1.5} + + def test_zero_is_kept(self): + """Zero survives the None-drop (the filter is 'is not None', not 'truthy').""" + metrics = SetupTimingMetrics(generation_init_time_s=0.0, policy_init_time_s=0.0) + d = metrics.to_metrics_dict() + assert d == {"generation_init_time_s": 0.0, "policy_init_time_s": 0.0} + + def test_extras_merged_into_top_level(self): + """extras dict entries appear as top-level keys, not nested.""" + metrics = SetupTimingMetrics(generation_init_time_s=1.0) + metrics.extras["vllm_nccl_sparse_init_time_s"] = 2.5 + d = metrics.to_metrics_dict() + assert d == { + "generation_init_time_s": 1.0, + "vllm_nccl_sparse_init_time_s": 2.5, + } + # extras itself is not exposed as a nested key. + assert "extras" not in d + + def test_reserve_load_split_serialized(self): + """Reserve/load split fields are included when populated.""" + metrics = SetupTimingMetrics( + generation_init_time_s=15.0, + generation_init_reserve_time_s=3.0, + generation_init_load_time_s=12.0, + ) + d = metrics.to_metrics_dict() + assert d["generation_init_reserve_time_s"] == 3.0 + assert d["generation_init_load_time_s"] == 12.0 diff --git a/tests/unit/single_controller/test_rollout_pump.py b/tests/unit/single_controller/test_rollout_pump.py index dfd4f6c7ca9..440601f16e6 100644 --- a/tests/unit/single_controller/test_rollout_pump.py +++ b/tests/unit/single_controller/test_rollout_pump.py @@ -33,6 +33,7 @@ ) from nemo_rl.algorithms.grpo import GRPOConfig from nemo_rl.algorithms.loss import ClippedPGLossConfig +from nemo_rl.algorithms.metric_utils import SetupTimingMetrics from nemo_rl.algorithms.single_controller import SingleControllerActor from nemo_rl.algorithms.single_controller_utils.config import ( AsyncRLConfig, @@ -353,7 +354,9 @@ def test_rollout_pump_writes_expected_tq_data( partition_id=_PARTITION_ID, ) ctrl = SingleControllerActor.remote( - master_config=master_config, actor_args=actor_args + master_config=master_config, + actor_args=actor_args, + setup_timing_metrics=SetupTimingMetrics(), ) vllm_generation.prepare_for_generation() diff --git a/tests/unit/single_controller/test_run_grpo_single_controller.py b/tests/unit/single_controller/test_run_grpo_single_controller.py index 3a22bd44c82..4433394c8db 100644 --- a/tests/unit/single_controller/test_run_grpo_single_controller.py +++ b/tests/unit/single_controller/test_run_grpo_single_controller.py @@ -19,6 +19,7 @@ import pytest from examples import run_grpo_single_controller +from nemo_rl.algorithms.metric_utils import SetupTimingMetrics @pytest.fixture @@ -77,7 +78,7 @@ def main_context(monkeypatch: pytest.MonkeyPatch) -> SimpleNamespace: monkeypatch.setattr( run_grpo_single_controller, "setup_single_controller", - lambda *_args: actor_args, + lambda *_args: (actor_args, SetupTimingMetrics()), ) monkeypatch.setattr( run_grpo_single_controller.SingleControllerActor, diff --git a/tests/unit/single_controller/test_single_controller.py b/tests/unit/single_controller/test_single_controller.py index d0b09c9011b..952bf18fa79 100644 --- a/tests/unit/single_controller/test_single_controller.py +++ b/tests/unit/single_controller/test_single_controller.py @@ -24,6 +24,7 @@ import nemo_rl.algorithms.single_controller as single_controller from nemo_rl.algorithms.grpo import GRPOConfig from nemo_rl.algorithms.loss import ClippedPGLossConfig +from nemo_rl.algorithms.metric_utils import SetupTimingMetrics from nemo_rl.algorithms.single_controller import SingleControllerActor from nemo_rl.algorithms.single_controller_utils.config import ( AdvantageConfig, @@ -76,6 +77,7 @@ def test_rejects_multiple_optimizer_steps_per_rl_step(monkeypatch) -> None: controller_cls( master_config=master_config, actor_args=actor_args, + setup_timing_metrics=SetupTimingMetrics(), ) @@ -117,6 +119,7 @@ def test_logs_hyperparameters_and_concrete_weight_synchronizer( controller_cls( master_config=master_config, actor_args=actor_args, + setup_timing_metrics=SetupTimingMetrics(), ) logger.log_hyperparams.assert_called_once_with(master_config.model_dump()) @@ -125,6 +128,53 @@ def test_logs_hyperparameters_and_concrete_weight_synchronizer( assert "transport=stub" not in output +def test_logs_setup_timing_metrics(monkeypatch) -> None: + """setup_timing_metrics is forwarded to Logger.log_metrics under timing/setup.""" + logger = MagicMock() + monkeypatch.setattr(single_controller, "Logger", lambda _: logger) + master_config = MasterConfig.model_construct( + policy={"train_global_batch_size": 8}, + grpo=GRPOConfig.model_construct( + num_prompts_per_step=2, + num_generations_per_prompt=4, + ), + loss_fn=ClippedPGLossConfig(force_on_policy_ratio=False), + async_rl=AsyncRLConfig( + min_groups_for_streaming_train=1, + max_buffered_rollouts=4, + ), + logger={}, + ) + setup_metrics = SetupTimingMetrics( + generation_init_time_s=1.5, policy_init_time_s=2.5 + ) + actor_args = SimpleNamespace( + partition_id="rollout_data", + dp_client=None, + gen_handle=None, + trainer_handle=None, + dataloader=None, + weight_synchronizer=FakeWeightSynchronizer(), + advantage_estimator=None, + loss_fn=None, + tq_buffer=None, + rollout_manager=SimpleNamespace(_tq_buffer=None), + train_cluster=None, + inference_cluster=None, + ) + controller_cls = SingleControllerActor.__ray_metadata__.modified_class + + controller_cls( + master_config=master_config, + actor_args=actor_args, + setup_timing_metrics=setup_metrics, + ) + + logger.log_metrics.assert_called_once_with( + setup_metrics.to_metrics_dict(), step=0, prefix="timing/setup" + ) + + @pytest.mark.parametrize( ("recompute_kv_cache", "expected_invalidation_calls"), [(False, 0), (True, 1)], diff --git a/tests/unit/single_controller/test_single_controller_setup.py b/tests/unit/single_controller/test_single_controller_setup.py index 5af07b94500..1fb7bef41f6 100644 --- a/tests/unit/single_controller/test_single_controller_setup.py +++ b/tests/unit/single_controller/test_single_controller_setup.py @@ -100,6 +100,9 @@ def patched_factories(): # len(dataloader) used by the Megatron train_iters injection. fake_dataloader.__len__ = MagicMock(return_value=4) fake_env_handles = {"math": MagicMock(name="math_env")} + # Real return objects; _build_generation and _build_trainer return (obj, elapsed_s) tuples. + fake_gen = MagicMock(name="gen") + fake_policy = MagicMock(name="policy") with ( patch.object( @@ -121,10 +124,10 @@ def patched_factories(): ), ) as mock_clusters, patch.object( - sc_setup_mod, "_build_generation", return_value=MagicMock(name="gen") + sc_setup_mod, "_build_generation", return_value=(fake_gen, 0.0) ) as mock_gen, patch.object( - sc_setup_mod, "_build_trainer", return_value=MagicMock(name="policy") + sc_setup_mod, "_build_trainer", return_value=(fake_policy, 0.0) ) as mock_trainer, patch.object( sc_setup_mod, @@ -162,6 +165,8 @@ def patched_factories(): "ClippedPGLossFn": mock_loss, "dataloader": fake_dataloader, "env_handles": fake_env_handles, + "fake_gen": fake_gen, + "fake_policy": fake_policy, } @@ -173,7 +178,7 @@ def test_build_generation_passes_sglang_config(): inference_cluster = MagicMock(name="inference_cluster") with patch.object(sc_setup_mod, "SGLangGeneration") as mock_sglang: - generation = sc_setup_mod._build_generation( + generation, _ = sc_setup_mod._build_generation( inference_cluster, master_config, ) @@ -249,16 +254,11 @@ def test_returns_actor_args(self, patched_factories): mc = _make_master_config(colocated=True) tokenizer = MagicMock(pad_token_id=0) - actor_args = setup_single_controller(mc, tokenizer) + actor_args, _ = setup_single_controller(mc, tokenizer) assert isinstance(actor_args, SingleControllerActorArgs) - assert ( - actor_args.gen_handle is patched_factories["_build_generation"].return_value - ) - assert ( - actor_args.trainer_handle - is patched_factories["_build_trainer"].return_value - ) + assert actor_args.gen_handle is patched_factories["fake_gen"] + assert actor_args.trainer_handle is patched_factories["fake_policy"] assert actor_args.env_handles is patched_factories["env_handles"] assert ( actor_args.dp_client @@ -289,7 +289,7 @@ def test_router_replay_requires_routes_in_tq_buffer(self, patched_factories): mc = _make_master_config(colocated=True) mc.policy["router_replay"] = {"enabled": True} - actor_args = setup_single_controller(mc, MagicMock(pad_token_id=0)) + actor_args, _ = setup_single_controller(mc, MagicMock(pad_token_id=0)) assert actor_args.tq_buffer._require_routed_experts is True @@ -298,7 +298,7 @@ def test_env_handles_sourced_from_setup_response_data(self, patched_factories): math_env_cfg = {"some": "value"} mc = _make_master_config(env={"math": math_env_cfg}) - actor_args = setup_single_controller(mc, MagicMock(pad_token_id=0)) + actor_args, _ = setup_single_controller(mc, MagicMock(pad_token_id=0)) _, call_kwargs = patched_factories["setup_response_data"].call_args assert call_kwargs["env_configs"] == {"math": math_env_cfg} @@ -312,13 +312,8 @@ def test_weight_sync_factory_args(self, patched_factories): setup_single_controller(mc, tokenizer) _, factory_kwargs = patched_factories["create_weight_synchronizer"].call_args - assert ( - factory_kwargs["policy"] is patched_factories["_build_trainer"].return_value - ) - assert ( - factory_kwargs["generation"] - is patched_factories["_build_generation"].return_value - ) + assert factory_kwargs["policy"] is patched_factories["fake_policy"] + assert factory_kwargs["generation"] is patched_factories["fake_gen"] assert factory_kwargs["generation_backend"] == "vllm" assert factory_kwargs["colocated"] is False @@ -326,7 +321,7 @@ def test_custom_partition_id(self, patched_factories): mc = _make_master_config() tokenizer = MagicMock(pad_token_id=7) - actor_args = setup_single_controller( + actor_args, _ = setup_single_controller( mc, tokenizer, partition_id="custom_partition" ) @@ -423,13 +418,11 @@ def test_nemo_gym_wires_env_handle(self, patched_factories): ) as mock_spinup, patch.object(sc_setup_mod, "router_replay_enabled", return_value=False), ): - actor_args = setup_single_controller(mc, MagicMock(pad_token_id=0)) + actor_args, _ = setup_single_controller(mc, MagicMock(pad_token_id=0)) mock_spinup.assert_called_once_with( env_configs=mc.env, - base_urls=patched_factories[ - "_build_generation" - ].return_value.dp_openai_server_base_urls, + base_urls=patched_factories["fake_gen"].dp_openai_server_base_urls, model_name="test-model", enable_router_replay=False, routed_experts_dtype="int16", @@ -437,6 +430,179 @@ def test_nemo_gym_wires_env_handle(self, patched_factories): ) assert actor_args.env_handles["nemo_gym"] is fake_gym_actor + def test_setup_timing_populated_for_colocated_vllm(self, patched_factories): + """Colocated vLLM records gen+policy+collective+total+worker fields.""" + mc = _make_master_config(colocated=True, backend="vllm") + + _, metrics = setup_single_controller(mc, MagicMock(pad_token_id=0)) + + for field in ( + "generation_init_time_s", + "policy_init_time_s", + "collective_init_time_s", + "worker_setup_time_s", + "total_setup_time_s", + "other_setup_time_s", + ): + value = getattr(metrics, field) + assert value is not None, f"missing {field} on {metrics}" + assert value >= 0 + # parallel_wall_time_s / parallel_init_enabled are grpo.py-only in the + # shared SetupTimingMetrics — SC does not emit them. + assert metrics.parallel_wall_time_s is None + assert metrics.parallel_init_enabled is None + # Reserve/load split is populated on the gym-on path only. + assert metrics.generation_init_reserve_time_s is None + assert metrics.generation_init_load_time_s is None + + def test_setup_timing_populated_for_noncolocated_vllm(self, patched_factories): + """Non-colocated vLLM records the same per-phase fields as colocated.""" + mc = _make_master_config(colocated=False, backend="vllm") + + _, metrics = setup_single_controller(mc, MagicMock(pad_token_id=0)) + + assert metrics.generation_init_time_s is not None + assert metrics.policy_init_time_s is not None + assert metrics.worker_setup_time_s is not None + # parallel_wall_time_s / parallel_init_enabled are grpo.py-only. + assert metrics.parallel_wall_time_s is None + assert metrics.parallel_init_enabled is None + # Reserve/load split is populated on the gym-on path only. + assert metrics.generation_init_reserve_time_s is None + assert metrics.generation_init_load_time_s is None + + def test_setup_timing_backend_agnostic_for_sglang(self, patched_factories): + """SC uses the backend-agnostic generation_init_time_s regardless of backend.""" + mc = _make_master_config(colocated=True, backend="sglang") + + _, metrics = setup_single_controller(mc, MagicMock(pad_token_id=0)) + + assert metrics.generation_init_time_s is not None + # Backend-specific fields are grpo.py-only; SC does not populate them. + assert metrics.vllm_init_time_s is None + assert metrics.sglang_init_time_s is None + + def test_nemo_gym_uses_deferred_vllm_load(self, patched_factories): + """NeMo-Gym path reserves vLLM ports up-front and finishes the load afterwards.""" + mc = _make_master_config(colocated=True, backend="vllm") + mc.policy["generation"]["model_name"] = "test-model" + mc.policy["generation"]["stop_strings"] = None + mc.policy["generation"]["stop_token_ids"] = None + mc.policy["generation"]["top_k"] = None + patched_factories["setup_response_data"].return_value = (list(range(8)), None) + + with ( + patch.object(sc_setup_mod, "_should_use_nemo_gym", return_value=True), + patch.object( + sc_setup_mod, "spinup_nemo_gym_actor", return_value=MagicMock() + ), + patch.object(sc_setup_mod, "router_replay_enabled", return_value=False), + ): + setup_single_controller(mc, MagicMock(pad_token_id=0)) + + # _build_generation must be called with defer_model_load=True so the workers + # only reserve URLs; load_and_start()+finish_generation() run afterwards. + _, gen_kwargs = patched_factories["_build_generation"].call_args + assert gen_kwargs.get("defer_model_load") is True + deferred_vllm = patched_factories["fake_gen"] + deferred_vllm.load_and_start.assert_called_once_with() + deferred_vllm.finish_generation.assert_called_once_with() + + def test_nemo_gym_records_timing_metrics(self, patched_factories): + """NeMo-Gym path records per-phase timings (vllm/policy/gym/worker).""" + mc = _make_master_config(colocated=True, backend="vllm") + mc.policy["generation"]["model_name"] = "test-model" + mc.policy["generation"]["stop_strings"] = None + mc.policy["generation"]["stop_token_ids"] = None + mc.policy["generation"]["top_k"] = None + patched_factories["setup_response_data"].return_value = (list(range(8)), None) + + with ( + patch.object(sc_setup_mod, "_should_use_nemo_gym", return_value=True), + patch.object( + sc_setup_mod, "spinup_nemo_gym_actor", return_value=MagicMock() + ), + patch.object(sc_setup_mod, "router_replay_enabled", return_value=False), + ): + _, metrics = setup_single_controller(mc, MagicMock(pad_token_id=0)) + + assert metrics.nemo_gym_init_time_s is not None + assert metrics.generation_init_time_s is not None + assert metrics.policy_init_time_s is not None + assert metrics.worker_setup_time_s is not None + # parallel_wall_time_s / parallel_init_enabled are grpo.py-only. + assert metrics.parallel_wall_time_s is None + assert metrics.parallel_init_enabled is None + + def test_nemo_gym_noncolocated_finishes_deferred_load(self, patched_factories): + """Non-colocated + gym fans out gym / deferred-load / trainer together.""" + mc = _make_master_config(colocated=False, backend="vllm") + mc.policy["generation"]["model_name"] = "test-model" + mc.policy["generation"]["stop_strings"] = None + mc.policy["generation"]["stop_token_ids"] = None + mc.policy["generation"]["top_k"] = None + patched_factories["setup_response_data"].return_value = (list(range(8)), None) + + with ( + patch.object(sc_setup_mod, "_should_use_nemo_gym", return_value=True), + patch.object( + sc_setup_mod, "spinup_nemo_gym_actor", return_value=MagicMock() + ), + patch.object(sc_setup_mod, "router_replay_enabled", return_value=False), + ): + actor_args, metrics = setup_single_controller(mc, MagicMock(pad_token_id=0)) + + # _build_generation runs once (URL reservation only); the load is finished + # by _finish_deferred_generation inside the executor. + patched_factories["_build_generation"].assert_called_once() + _, gen_kwargs = patched_factories["_build_generation"].call_args + assert gen_kwargs.get("defer_model_load") is True + patched_factories["fake_gen"].load_and_start.assert_called_once_with() + assert actor_args.gen_handle is patched_factories["fake_gen"] + assert metrics.nemo_gym_init_time_s is not None + assert metrics.generation_init_time_s is not None + assert metrics.policy_init_time_s is not None + + @pytest.mark.parametrize("colocated", [True, False]) + def test_nemo_gym_generation_init_time_includes_reserve_time( + self, patched_factories, colocated + ): + """generation_init_time_s folds in the deferred-VllmGeneration reserve time. + + With gym on, _build_generation(defer_model_load=True) does worker-group + spawn + port bind (no weight load). That elapsed time has to end up in + generation_init_time_s alongside the deferred-load elapsed; otherwise + gym-on runs undercount generation setup by the worker-group span. The + reserve/load split is also exposed for overlap analysis. + """ + mc = _make_master_config(colocated=colocated, backend="vllm") + mc.policy["generation"]["model_name"] = "test-model" + mc.policy["generation"]["stop_strings"] = None + mc.policy["generation"]["stop_token_ids"] = None + mc.policy["generation"]["top_k"] = None + patched_factories["setup_response_data"].return_value = (list(range(8)), None) + # Deferred _build_generation returns 3.0s of reserve time; _build_generation + # is only called once (for reservation), so this is the reserve span. + patched_factories["_build_generation"].return_value = ( + patched_factories["fake_gen"], + 3.0, + ) + + with ( + patch.object(sc_setup_mod, "_should_use_nemo_gym", return_value=True), + patch.object( + sc_setup_mod, "spinup_nemo_gym_actor", return_value=MagicMock() + ), + patch.object(sc_setup_mod, "router_replay_enabled", return_value=False), + ): + _, metrics = setup_single_controller(mc, MagicMock(pad_token_id=0)) + + # gen_load_time (from _finish_deferred_generation, unpatched) is ~0 in + # the test — the reserve time dominates and must be present. + assert metrics.generation_init_time_s >= 3.0 + assert metrics.generation_init_reserve_time_s == 3.0 + assert metrics.generation_init_load_time_s is not None + @pytest.mark.parametrize("backend", ["sglang", "megatron"]) def test_nemo_gym_rejects_non_vllm_backend(self, patched_factories, backend): """SC nemo-gym wiring only supports vLLM; every other backend must raise.""" diff --git a/tests/unit/single_controller/test_train_pump.py b/tests/unit/single_controller/test_train_pump.py index 9f19cb0f674..947d0215e49 100644 --- a/tests/unit/single_controller/test_train_pump.py +++ b/tests/unit/single_controller/test_train_pump.py @@ -30,6 +30,7 @@ from nemo_rl.algorithms.async_utils.staleness_sampler import WindowedSamplerConfig from nemo_rl.algorithms.grpo import GRPOConfig from nemo_rl.algorithms.loss import ClippedPGLossConfig +from nemo_rl.algorithms.metric_utils import SetupTimingMetrics from nemo_rl.algorithms.single_controller import SingleControllerActor from nemo_rl.algorithms.single_controller_utils.config import ( AsyncRLConfig, @@ -343,6 +344,7 @@ def test_train_pump_drives_mcore_training_step( metric_log_handle=log, master_config=master_config, actor_args=actor_args, + setup_timing_metrics=SetupTimingMetrics(), ) # train_steps outer steps, each: sampler.select → advantage stage → begin/microbatches/finish → sync. diff --git a/tools/refit_verifier.py b/tools/refit_verifier.py index 44858283341..4a0d4cb0841 100644 --- a/tools/refit_verifier.py +++ b/tools/refit_verifier.py @@ -651,7 +651,7 @@ def initialize_generation_with_policy( worker_init_timing_metrics[init_time_key] = generation_time worker_init_timing_metrics["policy_init_time_s"] = policy_time worker_init_timing_metrics["parallel_wall_time_s"] = parallel_wall_time - worker_init_timing_metrics["parallel_init_enabled"] = True + worker_init_timing_metrics["parallel_init_enabled"] = 1.0 else: print( " ⚙️ Using sequential worker initialization (colocated mode)", From 037fb36bdad40f6336edcdfd39f7c5150b152324 Mon Sep 17 00:00:00 2001 From: jthomson04 Date: Mon, 10 Aug 2026 08:07:39 -0700 Subject: [PATCH 10/10] perf(vllm): return token metadata with chat completions (#3390) Signed-off-by: jthomson04 --- nemo_rl/models/generation/vllm/utils.py | 117 ++++++++++- .../generation/vllm/vllm_worker_async.py | 49 ++++- .../models/generation/test_vllm_generation.py | 52 +++-- .../unit/models/generation/test_vllm_utils.py | 190 +++++++++++++++++- 4 files changed, 369 insertions(+), 39 deletions(-) diff --git a/nemo_rl/models/generation/vllm/utils.py b/nemo_rl/models/generation/vllm/utils.py index b8ec3dee050..689ba95b22b 100644 --- a/nemo_rl/models/generation/vllm/utils.py +++ b/nemo_rl/models/generation/vllm/utils.py @@ -26,6 +26,7 @@ from nemo_rl.utils.routed_experts_codec import encode_routed_experts R3_MISSING_ROUTE_SENTINEL = ROUTED_EXPERTS_MISSING_ROUTE_SENTINEL +VLLM_LOGPROB_FLOOR = -9999.0 # The expert-id range vs carry dtype is model-constant, so it is verified on the # first non-empty routed-experts tensor per process and skipped afterwards. @@ -318,17 +319,119 @@ def attach_routed_experts_to_chat_response_choices( 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.""" +def attach_token_information_to_chat_response_choices( + response: Any, + final_request_output: Any, +) -> Any: + """Attach engine-native token information to OpenAI chat response choices.""" + prompt_token_ids = getattr(final_request_output, "prompt_token_ids", None) + if prompt_token_ids is None: + raise RuntimeError( + "vLLM was asked to return token information for the " + "OpenAI-compatible chat endpoint but the final request output did " + "not include prompt_token_ids." + ) + + generation_outputs = list(getattr(final_request_output, "outputs", [])) + generation_output_indices = [output.index for output in generation_outputs] + outputs_by_index = {output.index: output for output in generation_outputs} + if len(outputs_by_index) != len(generation_outputs): + raise RuntimeError( + "vLLM returned duplicate generation output indices while attaching " + "token information to the OpenAI-compatible chat response." + ) + + choices = list(getattr(response, "choices", [])) + choice_indices = [choice.index for choice in choices] + if len(set(choice_indices)) != len(choice_indices): + raise RuntimeError( + "vLLM returned duplicate response choice indices while attaching " + "token information to the OpenAI-compatible chat response." + ) + + choice_index_set = set(choice_indices) + output_index_set = set(generation_output_indices) + missing_choice_indices = sorted(choice_index_set - output_index_set) + unexpected_output_indices = sorted(output_index_set - choice_index_set) + if missing_choice_indices or unexpected_output_indices: + raise RuntimeError( + "vLLM was asked to return token information for the " + "OpenAI-compatible chat endpoint but response choices could not be " + "matched to generation outputs: " + f"missing_choice_indices={missing_choice_indices}, " + f"unexpected_output_indices={unexpected_output_indices}." + ) + + for choice in choices: + generation_details = outputs_by_index[choice.index] + output_token_ids = getattr(generation_details, "token_ids", None) + if output_token_ids is None: + raise RuntimeError( + "vLLM was asked to return token information for the " + "OpenAI-compatible chat endpoint but generation output " + f"choice_idx={choice.index} did not include token_ids." + ) + generation_token_ids = list(output_token_ids) + + generation_logprob_details = getattr(generation_details, "logprobs", None) + if generation_logprob_details is None: + if generation_token_ids: + raise RuntimeError( + "vLLM was asked to return token information for the " + "OpenAI-compatible chat endpoint but generation output " + f"choice_idx={choice.index} did not include logprobs." + ) + generation_log_probs = [] + else: + if len(generation_token_ids) != len(generation_logprob_details): + raise RuntimeError( + "vLLM returned mismatched generation token IDs and log " + "probabilities for the OpenAI-compatible chat endpoint: " + f"choice_idx={choice.index}, " + f"token_count={len(generation_token_ids)}, " + f"logprob_count={len(generation_logprob_details)}." + ) + generation_log_probs = [] + for token_id, position_logprobs in zip( + generation_token_ids, generation_logprob_details + ): + selected_token_logprob = position_logprobs.get(token_id) + if selected_token_logprob is None: + raise RuntimeError( + "vLLM generation log probabilities did not include the " + "selected token while attaching token information to " + "the OpenAI-compatible chat response: " + f"choice_idx={choice.index}, token_id={token_id}." + ) + generation_log_probs.append( + max(float(selected_token_logprob.logprob), VLLM_LOGPROB_FLOOR) + ) + + choice.message.prompt_token_ids = list(prompt_token_ids) + choice.message.generation_token_ids = generation_token_ids + choice.message.generation_log_probs = generation_log_probs + + return response + + +def model_dump_chat_response_with_dynamic_message_fields( + response: Any, +) -> dict[str, Any]: + """Dump a vLLM OpenAI chat response while preserving dynamic message 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 + message = getattr(choice, "message", None) + for field_name in ( + "routed_experts", + "prompt_token_ids", + "generation_token_ids", + "generation_log_probs", + ): + field_value = getattr(message, field_name, None) + if field_value is not None: + choice_dict.setdefault("message", {})[field_name] = field_value return response_dict diff --git a/nemo_rl/models/generation/vllm/vllm_worker_async.py b/nemo_rl/models/generation/vllm/vllm_worker_async.py index d1da8a3ec1a..86e73932f39 100644 --- a/nemo_rl/models/generation/vllm/vllm_worker_async.py +++ b/nemo_rl/models/generation/vllm/vllm_worker_async.py @@ -45,8 +45,9 @@ ) from nemo_rl.models.generation.vllm.utils import ( attach_routed_experts_to_chat_response_choices, + attach_token_information_to_chat_response_choices, format_prompt_for_vllm_generation, - model_dump_chat_response_with_routed_experts, + model_dump_chat_response_with_dynamic_message_fields, pad_and_align_routed_expert_indices, ) from nemo_rl.models.generation.vllm.vllm_worker import BaseVllmGenerationWorker @@ -594,6 +595,22 @@ async def chat_completion_full_generator( *args, **kwargs, ): + return_as_token_id = ( + request.return_tokens_as_token_ids + if request.return_tokens_as_token_ids is not None + else self.return_tokens_as_token_ids + ) + if ( + request.logprobs + and return_as_token_id + and request.top_logprobs is None + ): + raise VLLMValidationError( + "`top_logprobs` must be set when requesting token " + "information from the NeMo-RL chat endpoint.", + parameter="top_logprobs", + ) + final_res = None async def capture_result_generator(): @@ -609,19 +626,27 @@ async def capture_result_generator(): **kwargs, ) if ( - not worker_self._return_routed_experts_enabled() - or not isinstance(response, ChatCompletionResponse) + not isinstance(response, ChatCompletionResponse) or final_res is None ): return response - return attach_routed_experts_to_chat_response_choices( - response, - final_res, - device=torch.device("cpu"), - logger=LOGGER, - routed_experts_dtype=worker_self.routed_experts_dtype, - ) + if request.logprobs and return_as_token_id: + response = attach_token_information_to_chat_response_choices( + response, + final_res, + ) + + if worker_self._return_routed_experts_enabled(): + response = attach_routed_experts_to_chat_response_choices( + response, + final_res, + device=torch.device("cpu"), + logger=LOGGER, + routed_experts_dtype=worker_self.routed_experts_dtype, + ) + + return response class NeMoRLOpenAIServingChat(NeMoRLOpenAIServingChatMixin, OpenAIServingChat): pass @@ -738,7 +763,9 @@ async def create_chat_completion( elif isinstance(generator, ChatCompletionResponse): return JSONResponse( - content=model_dump_chat_response_with_routed_experts(generator) + content=model_dump_chat_response_with_dynamic_message_fields( + generator + ) ) return StreamingResponse(content=generator, media_type="text/event-stream") diff --git a/tests/unit/models/generation/test_vllm_generation.py b/tests/unit/models/generation/test_vllm_generation.py index 93429ccfb5e..33f0507750e 100644 --- a/tests/unit/models/generation/test_vllm_generation.py +++ b/tests/unit/models/generation/test_vllm_generation.py @@ -1696,7 +1696,6 @@ def test_vllm_http_server(cluster, tokenizer): top_p=generation_config["top_p"], # We want to test the actual train flow and how this is used. So we need to get logprobs here. logprobs=True, - return_tokens_as_token_ids=True, max_tokens=1, ) @@ -1706,6 +1705,21 @@ def test_vllm_http_server(cluster, tokenizer): response = requests.post(url=f"{base_urls[0]}/chat/completions", json=body) actual_result = response.json() + expected_prompt_token_ids = [ + 151644, + 872, + 198, + 1830, + 311, + 220, + 20, + 151645, + 198, + 151644, + 77091, + 198, + ] + # This result assumes this exact model. The expected result here is what the full result looks like before we standardize. expected_result = { "id": "chatcmpl-7b8c0cdeeab34fd58ad260cf44b1a408", @@ -1725,6 +1739,8 @@ def test_vllm_http_server(cluster, tokenizer): # vLLM 0.25 omits tool_calls when empty and dropped # reasoning_content in favor of reasoning. "reasoning": None, + "prompt_token_ids": expected_prompt_token_ids, + "generation_token_ids": [151667], }, "logprobs": { "content": [ @@ -1771,31 +1787,33 @@ def _standardize(d: dict) -> dict: message = d["choices"][0]["message"] for key in ("reasoning", "reasoning_content"): message.pop(key, None) + message.pop("generation_log_probs", None) return d + assert actual_result["choices"][0]["message"]["generation_log_probs"] == [ + actual_result["choices"][0]["logprobs"]["content"][0]["logprob"] + ] assert _standardize(expected_result) == _standardize(actual_result) + # The server default requests token IDs, so top_logprobs=None cannot provide + # the log probabilities required by the training response contract. + response = requests.post( + url=f"{base_urls[0]}/chat/completions", + json=body | {"top_logprobs": None}, + ) + assert response.status_code == 400 + error = response.json()["error"] + assert error["code"] == 400 + assert "top_logprobs" in error["message"] + # Check that tokenization route works response = requests.post(url=f"{base_urls[0]}/../tokenize", json=body) actual_result = response.json() expected_result = { "count": 12, "max_model_len": 1024, - "tokens": [ - 151644, - 872, - 198, - 1830, - 311, - 220, - 20, - 151645, - 198, - 151644, - 77091, - 198, - ], + "tokens": expected_prompt_token_ids, "token_strs": None, } assert expected_result == actual_result @@ -2101,6 +2119,10 @@ async def test_vllm_http_server_correct_merged_tokens_matches_baseline( url=f"{base_urls[0]}/chat/completions", json=body_with_reference_token_ids ) vllm_http_server_result = response.json() + assert ( + vllm_http_server_result["choices"][0]["message"]["prompt_token_ids"] + == initial_tokenized_query_ids + ) vllm_http_server_generated_token = vllm_http_server_result["choices"][0][ "logprobs" ]["content"][0] diff --git a/tests/unit/models/generation/test_vllm_utils.py b/tests/unit/models/generation/test_vllm_utils.py index 835cfc12e2e..372ad098061 100644 --- a/tests/unit/models/generation/test_vllm_utils.py +++ b/tests/unit/models/generation/test_vllm_utils.py @@ -30,9 +30,10 @@ R3_MISSING_ROUTE_SENTINEL, aggregate_spec_decode_counters, attach_routed_experts_to_chat_response_choices, + attach_token_information_to_chat_response_choices, compute_spec_decode_metrics, format_prompt_for_vllm_generation, - model_dump_chat_response_with_routed_experts, + model_dump_chat_response_with_dynamic_message_fields, pad_and_align_routed_expert_indices, ) from nemo_rl.utils.routed_experts_codec import decode_routed_experts @@ -558,13 +559,179 @@ def test_attach_routed_experts_to_chat_response_choices_raises_for_unmatched_cho ) -def test_model_dump_chat_response_with_routed_experts_preserves_dynamic_field(): - routed_experts = [[[1]], [[2]]] +def test_attach_token_information_to_chat_response_choices(): + final_res = SimpleNamespace( + prompt_token_ids=[101, 102, 103], + outputs=[ + SimpleNamespace(index=1, token_ids=[], logprobs=None), + SimpleNamespace( + index=0, + token_ids=[201, 202], + logprobs=[ + {201: SimpleNamespace(logprob=-0.1)}, + {202: SimpleNamespace(logprob=-10000.0)}, + ], + ), + ], + ) + response = SimpleNamespace( + choices=[ + SimpleNamespace( + index=0, + message=SimpleNamespace(), + logprobs=SimpleNamespace( + content=[ + SimpleNamespace(token="decoded token", logprob=-10.0), + SimpleNamespace(token="format is ignored", logprob=-20.0), + ] + ), + ), + SimpleNamespace( + index=1, + message=SimpleNamespace(), + logprobs=SimpleNamespace(content=None), + ), + ], + model_dump=lambda: { + "choices": [ + {"message": {"role": "assistant", "content": "first"}}, + {"message": {"role": "assistant", "content": "second"}}, + ] + }, + ) + + attach_token_information_to_chat_response_choices(response, final_res) + response_dict = model_dump_chat_response_with_dynamic_message_fields(response) + + assert response_dict["choices"][0]["message"]["prompt_token_ids"] == [ + 101, + 102, + 103, + ] + assert response_dict["choices"][0]["message"]["generation_token_ids"] == [201, 202] + assert response_dict["choices"][0]["message"]["generation_log_probs"] == [ + -0.1, + -9999.0, + ] + assert response_dict["choices"][1]["message"]["prompt_token_ids"] == [ + 101, + 102, + 103, + ] + assert response_dict["choices"][1]["message"]["generation_token_ids"] == [] + assert response_dict["choices"][1]["message"]["generation_log_probs"] == [] + + +@pytest.mark.parametrize( + ("token_ids", "logprobs", "error_match"), + [ + ( + [201, 202], + [{201: SimpleNamespace(logprob=-0.1)}], + "mismatched generation token IDs and log probabilities", + ), + ( + [201], + [{999: SimpleNamespace(logprob=-0.1)}], + "did not include the selected token", + ), + ], +) +def test_attach_token_information_to_chat_response_choices_rejects_invalid_logprobs( + token_ids, logprobs, error_match +): + final_res = SimpleNamespace( + prompt_token_ids=[101, 102, 103], + outputs=[SimpleNamespace(index=0, token_ids=token_ids, logprobs=logprobs)], + ) + response = SimpleNamespace( + choices=[SimpleNamespace(index=0, message=SimpleNamespace())] + ) + + with pytest.raises(RuntimeError, match=error_match): + attach_token_information_to_chat_response_choices(response, final_res) + + +@pytest.mark.parametrize( + ("prompt_token_ids", "outputs", "choice_indices", "error_match"), + [ + (None, [], [], "did not include prompt_token_ids"), + ( + [101], + [ + SimpleNamespace(index=0, token_ids=[], logprobs=[]), + SimpleNamespace(index=0, token_ids=[], logprobs=[]), + ], + [0], + "duplicate generation output indices", + ), + ( + [101], + [SimpleNamespace(index=0, token_ids=[], logprobs=[])], + [0, 0], + "duplicate response choice indices", + ), + ( + [101], + [SimpleNamespace(index=1, token_ids=[], logprobs=[])], + [0], + "could not be matched to generation outputs", + ), + ( + [101], + [SimpleNamespace(index=0, logprobs=[])], + [0], + "did not include token_ids", + ), + ( + [101], + [SimpleNamespace(index=0, token_ids=[201], logprobs=None)], + [0], + "did not include logprobs", + ), + ], +) +def test_attach_token_information_to_chat_response_choices_rejects_invalid_structure( + prompt_token_ids, outputs, choice_indices, error_match +): + final_res = SimpleNamespace( + prompt_token_ids=prompt_token_ids, + outputs=outputs, + ) + response = SimpleNamespace( + choices=[ + SimpleNamespace(index=choice_index, message=SimpleNamespace()) + for choice_index in choice_indices + ] + ) + + with pytest.raises(RuntimeError, match=error_match): + attach_token_information_to_chat_response_choices(response, final_res) + + +def test_model_dump_chat_response_with_dynamic_message_fields_preserves_all_fields(): + final_res = SimpleNamespace( + prompt_token_ids=[101, 102], + prompt_routed_experts=torch.tensor( + [[[10]]], dtype=ROUTED_EXPERTS_FALLBACK_DTYPE + ), + outputs=[ + SimpleNamespace( + index=0, + token_ids=[201], + logprobs=[{201: SimpleNamespace(logprob=-0.1)}], + routed_experts=torch.tensor( + [[[20]]], dtype=ROUTED_EXPERTS_FALLBACK_DTYPE + ), + ) + ], + ) class Response: choices = [ SimpleNamespace( - message=SimpleNamespace(routed_experts=routed_experts), + index=0, + message=SimpleNamespace(), ) ] @@ -580,9 +747,20 @@ def model_dump(self): ] } - response_dict = model_dump_chat_response_with_routed_experts(Response()) + response = Response() + attach_token_information_to_chat_response_choices(response, final_res) + attach_routed_experts_to_chat_response_choices( + response, + final_res, + device=torch.device("cpu"), + ) + response_dict = model_dump_chat_response_with_dynamic_message_fields(response) - assert response_dict["choices"][0]["message"]["routed_experts"] == routed_experts + message = response_dict["choices"][0]["message"] + assert message["routed_experts"] == response.choices[0].message.routed_experts + assert message["prompt_token_ids"] == [101, 102] + assert message["generation_token_ids"] == [201] + assert message["generation_log_probs"] == [-0.1] @pytest.mark.vllm