From 1f58f5d12fbdbff6ca493f45290560393b0cecb6 Mon Sep 17 00:00:00 2001 From: Ilia Karmanov Date: Fri, 14 Aug 2026 02:15:31 -0700 Subject: [PATCH 01/11] fix(mopd): support image-aware async teacher refits Route row-aligned multimodal inputs to non-colocated teachers and optionally clear vLLM encoder outputs after quiesced weight updates. Signed-off-by: Ilia Karmanov --- .../async_utils/trajectory_collector.py | 24 +++++ nemo_rl/models/generation/vllm/config.py | 3 + .../generation/vllm/vllm_worker_async.py | 12 +++ tests/unit/algorithms/test_opd.py | 98 +++++++++++++++++++ .../models/generation/test_vllm_backend.py | 82 +++++++++++++++- 5 files changed, 217 insertions(+), 2 deletions(-) diff --git a/nemo_rl/algorithms/async_utils/trajectory_collector.py b/nemo_rl/algorithms/async_utils/trajectory_collector.py index d56a901375..b6f6ecbc22 100644 --- a/nemo_rl/algorithms/async_utils/trajectory_collector.py +++ b/nemo_rl/algorithms/async_utils/trajectory_collector.py @@ -30,6 +30,7 @@ from nemo_rl.algorithms.grpo import MasterConfig from nemo_rl.algorithms.opd import resolve_reference_aliases, teacher_seq_pad_multiple from nemo_rl.data.interfaces import DatumSpec +from nemo_rl.data.multimodal_utils import PackedTensor from nemo_rl.distributed.batched_data_dict import BatchedDataDict from nemo_rl.environments.interfaces import EnvironmentInterface from nemo_rl.experience.interfaces import ( @@ -686,6 +687,7 @@ def _compute_teacher_logprobs( input_ids: torch.Tensor, agent_refs: list[dict[str, Any]], input_lengths: Optional[torch.Tensor] = None, + multimodal_data: Optional[dict[str, Any]] = None, ) -> tuple[torch.Tensor, float]: """Compute teacher logprobs for non-colocated teachers. @@ -695,6 +697,8 @@ def _compute_teacher_logprobs( input_ids: [B, S] tokenized input tensor agent_refs: list of B agent reference dicts input_lengths: [B] per-sample lengths (required for sequence packing) + multimodal_data: batch-level multimodal inputs, row-aligned with + ``input_ids`` and sliced per teacher Returns: ([B, S] teacher logprobs tensor, total_time_seconds) @@ -734,6 +738,7 @@ def _get_logprobs_for_group(group_key, indices): twg = self.teacher_worker_groups[group_key] sub_input_ids = input_ids[indices] sub_lengths = input_lengths[indices] if input_lengths is not None else None + row_indices = list(indices) # Pad batch to multiple of dp_size (required for DP sharding) dp_size = twg.sharding_annotations.get_axis_size("data_parallel") @@ -749,10 +754,26 @@ def _get_logprobs_for_group(group_key, indices): sub_lengths = torch.cat( [sub_lengths, sub_lengths[-1:].expand(pad_count)], dim=0 ) + row_indices.extend([row_indices[-1]] * pad_count) sub_data = BatchedDataDict({"input_ids": sub_input_ids}) if sub_lengths is not None: sub_data["input_lengths"] = sub_lengths + if multimodal_data: + selected_multimodal = BatchedDataDict(multimodal_data).select_indices( + row_indices + ) + sub_data.update( + { + key: value + for key, value in selected_multimodal.items() + if value is not None + and not ( + isinstance(value, PackedTensor) + and not any(value.logical_segment_counts_by_row()) + ) + } + ) # Serialize calls per teacher to prevent NCCL collective desync t_lock_start = time.time() @@ -1003,6 +1024,9 @@ async def _enqueue_rollout_group( flat_for_teacher["token_ids"], agent_refs, input_lengths=teacher_input_lengths, + multimodal_data=flat_for_teacher.get_multimodal_dict( + as_tensors=False + ), ) # Keep the tensor inside the batch so replay-buffer collation can # pad variable-length prompt groups correctly. diff --git a/nemo_rl/models/generation/vllm/config.py b/nemo_rl/models/generation/vllm/config.py index 6281bb2af1..eb8d3f2f42 100644 --- a/nemo_rl/models/generation/vllm/config.py +++ b/nemo_rl/models/generation/vllm/config.py @@ -56,6 +56,9 @@ class VllmSpecificArgs(TypedDict): expose_http_server: NotRequired[bool] # Environment variable containing the internal refit API key. http_refit_api_key_env_var: NotRequired[str | None] + # Invalidate weight-dependent multimodal encoder outputs after a successful + # async refit. Enable only when generation is quiesced during weight updates. + reset_encoder_cache_after_weight_update: NotRequired[bool] # Fixed internal refit endpoint port for stable Kubernetes targetPorts. http_refit_server_port: NotRequired[int | None] # Fixed ZeroMQ relay port for stable Kubernetes targetPorts. diff --git a/nemo_rl/models/generation/vllm/vllm_worker_async.py b/nemo_rl/models/generation/vllm/vllm_worker_async.py index 4a0a2c0240..e791fa4a02 100644 --- a/nemo_rl/models/generation/vllm/vllm_worker_async.py +++ b/nemo_rl/models/generation/vllm/vllm_worker_async.py @@ -1407,6 +1407,15 @@ async def prepare_refit_info_async(self, state_dict_info: dict[str, Any]) -> Non """Async version of prepare_refit_info.""" await self.llm.collective_rpc("prepare_refit_info", args=(state_dict_info,)) + async def _reset_encoder_cache_after_weight_update(self) -> None: + """Invalidate weight-dependent multimodal encoder outputs when enabled.""" + if not self.cfg["vllm_cfg"].get( + "reset_encoder_cache_after_weight_update", False + ): + return + assert self.llm is not None + await self.llm.reset_encoder_cache() + async def update_weights_via_ipc_zmq_async( self, ) -> bool: @@ -1438,6 +1447,7 @@ async def update_weights_via_ipc_zmq_async( f"Error: Worker failed to update weights. Results: {worker_results}" ) return False + await self._reset_encoder_cache_after_weight_update() return True except Exception as e: print(f"Exception during collective_rpc for weight update: {e}") @@ -1474,6 +1484,7 @@ async def update_weights_from_collective_async(self) -> bool: f"Error: Worker failed to update weights. Results: {worker_results}" ) return False + await self._reset_encoder_cache_after_weight_update() return True except Exception as e: print(f"Exception during collective_rpc for weight update: {e}") @@ -1533,6 +1544,7 @@ async def nccl_reshard_refit_async(self) -> bool: f"Error: Worker failed nccl_reshard_refit. Result: {worker_result}" ) return False + await self._reset_encoder_cache_after_weight_update() return True except Exception as e: print(f"Exception during nccl_reshard_refit: {e}", flush=True) diff --git a/tests/unit/algorithms/test_opd.py b/tests/unit/algorithms/test_opd.py index 46aeb58e8f..18a65c1994 100644 --- a/tests/unit/algorithms/test_opd.py +++ b/tests/unit/algorithms/test_opd.py @@ -15,6 +15,7 @@ import pytest import torch +from nemo_rl.data.multimodal_utils import PackedTensor from nemo_rl.distributed.batched_data_dict import BatchedDataDict # --------------------------------------------------------------------------- @@ -112,6 +113,103 @@ def test_compute_teacher_logprobs_dp_padding(batch_size, dp_size): assert torch.allclose(result, torch.tensor(2.0)) +class _RecordingTeacherWorkerGroup(_MockTeacherWorkerGroup): + """Capture the batch passed to a teacher for row-alignment assertions.""" + + def __init__(self, fill_value=1.0, dp_size=4): + super().__init__(fill_value=fill_value, dp_size=dp_size) + self.received: BatchedDataDict | None = None + + def get_logprobs(self, data): + self.received = data + return super().get_logprobs(data) + + +def _row_marked_packed_tensor(markers): + return PackedTensor( + [ + None + if marker is None + else torch.full((1, 2), float(marker), dtype=torch.float32) + for marker in markers + ], + dim_to_pack=0, + ).enable_deduplication() + + +def _received_row_markers(packed): + return [ + None if tensor is None else float(tensor[0, 0]) + for tensor in packed.iter_logical_segments() + ] + + +def test_compute_teacher_logprobs_selects_multimodal_rows_per_teacher(): + """Each teacher receives media rows aligned with its selected token rows.""" + vision_twg = _RecordingTeacherWorkerGroup(fill_value=1.0, dp_size=1) + text_twg = _RecordingTeacherWorkerGroup(fill_value=2.0, dp_size=1) + collector = _make_collector( + teacher_worker_groups={"vision": vision_twg, "text": text_twg}, + alias_to_group_alias={"vision_agent": "vision", "text_agent": "text"}, + on_policy_distillation_cfg={ + "teacher_model_by_agent_name": { + "vision_agent": "/ckpt/vision", + "text_agent": "/ckpt/text", + }, + }, + _has_distillation_teachers=True, + ) + + collector._compute_teacher_logprobs( + torch.randint(0, 100, (4, 8)), + [ + {"name": "vision_agent"}, + {"name": "text_agent"}, + {"name": "vision_agent"}, + {"name": "text_agent"}, + ], + multimodal_data={ + "pixel_values": _row_marked_packed_tensor([0, None, 2, None]), + "imgs_sizes": _row_marked_packed_tensor([10, None, 12, None]), + }, + ) + + assert vision_twg.received is not None + assert text_twg.received is not None + assert _received_row_markers(vision_twg.received["pixel_values"]) == [0.0, 2.0] + assert _received_row_markers(vision_twg.received["imgs_sizes"]) == [10.0, 12.0] + assert "pixel_values" not in text_twg.received + assert "imgs_sizes" not in text_twg.received + + +def test_compute_teacher_logprobs_dp_padding_repeats_multimodal_row(): + """DP padding repeats the media row paired with the repeated token row.""" + twg = _RecordingTeacherWorkerGroup(fill_value=3.0, dp_size=4) + collector = _make_collector( + teacher_worker_groups={"vision": twg}, + alias_to_group_alias={"vision_agent": "vision"}, + on_policy_distillation_cfg={ + "teacher_model_by_agent_name": {"vision_agent": "/ckpt/vision"}, + }, + _has_distillation_teachers=True, + ) + + result, _ = collector._compute_teacher_logprobs( + torch.randint(0, 100, (1, 8)), + [{"name": "vision_agent"}], + multimodal_data={ + "pixel_values": _row_marked_packed_tensor([7]), + "num_frames": _row_marked_packed_tensor([1]), + }, + ) + + assert twg.received is not None + assert twg.received["input_ids"].shape[0] == 4 + assert _received_row_markers(twg.received["pixel_values"]) == [7.0] * 4 + assert _received_row_markers(twg.received["num_frames"]) == [1.0] * 4 + assert result.shape == (1, 8) + + def test_compute_teacher_logprobs_routes_to_correct_teacher(): """Samples are routed to the right teacher and results stitched back.""" math_twg = _MockTeacherWorkerGroup(fill_value=1.0, dp_size=1) diff --git a/tests/unit/models/generation/test_vllm_backend.py b/tests/unit/models/generation/test_vllm_backend.py index c7e7969ddc..d81528e07a 100644 --- a/tests/unit/models/generation/test_vllm_backend.py +++ b/tests/unit/models/generation/test_vllm_backend.py @@ -241,10 +241,88 @@ async def test_async_weight_updates_check_every_internal_worker( ) worker = VllmAsyncGenerationWorkerImpl.__new__(VllmAsyncGenerationWorkerImpl) - worker.cfg = {"vllm_cfg": {"async_engine": True}} - worker.llm = SimpleNamespace(collective_rpc=AsyncMock(return_value=worker_results)) + worker.cfg = { + "vllm_cfg": { + "async_engine": True, + "reset_encoder_cache_after_weight_update": True, + } + } + worker.llm = SimpleNamespace( + collective_rpc=AsyncMock(return_value=worker_results), + reset_encoder_cache=AsyncMock(), + ) assert await getattr(worker, method_name)() is expected + if expected: + worker.llm.reset_encoder_cache.assert_awaited_once_with() + else: + worker.llm.reset_encoder_cache.assert_not_awaited() + + +@pytest.mark.vllm +@pytest.mark.asyncio +async def test_async_weight_update_skips_encoder_cache_reset_when_disabled(): + """Text-only and in-flight refit users retain the existing cache behavior.""" + from nemo_rl.models.generation.vllm.vllm_worker_async import ( + VllmAsyncGenerationWorkerImpl, + ) + + worker = VllmAsyncGenerationWorkerImpl.__new__(VllmAsyncGenerationWorkerImpl) + worker.cfg = {"vllm_cfg": {"async_engine": True}} + worker.llm = SimpleNamespace( + collective_rpc=AsyncMock(return_value=[True]), + reset_encoder_cache=AsyncMock(), + ) + + assert await worker.update_weights_from_collective_async() is True + worker.llm.reset_encoder_cache.assert_not_awaited() + + +@pytest.mark.vllm +@pytest.mark.asyncio +async def test_async_weight_update_fails_when_encoder_cache_reset_fails(): + """A successful refit must not resume with stale multimodal encoder outputs.""" + from nemo_rl.models.generation.vllm.vllm_worker_async import ( + VllmAsyncGenerationWorkerImpl, + ) + + worker = VllmAsyncGenerationWorkerImpl.__new__(VllmAsyncGenerationWorkerImpl) + worker.cfg = { + "vllm_cfg": { + "async_engine": True, + "reset_encoder_cache_after_weight_update": True, + } + } + worker.llm = SimpleNamespace( + collective_rpc=AsyncMock(return_value=[True]), + reset_encoder_cache=AsyncMock(side_effect=RuntimeError("reset failed")), + ) + + assert await worker.update_weights_from_collective_async() is False + + +@pytest.mark.vllm +@pytest.mark.asyncio +async def test_nccl_reshard_refit_resets_encoder_cache(): + """NCCL-reshard refits invalidate encoder outputs just like other transports.""" + from nemo_rl.models.generation.vllm.vllm_worker_async import ( + VllmAsyncGenerationWorkerImpl, + ) + + worker = VllmAsyncGenerationWorkerImpl.__new__(VllmAsyncGenerationWorkerImpl) + worker.cfg = { + "vllm_cfg": { + "async_engine": True, + "reset_encoder_cache_after_weight_update": True, + } + } + worker.llm = SimpleNamespace( + collective_rpc=AsyncMock(return_value=[True]), + reset_encoder_cache=AsyncMock(), + ) + + assert await worker.nccl_reshard_refit_async() is True + worker.llm.reset_encoder_cache.assert_awaited_once_with() @pytest.mark.vllm From 940f40cf509bc985022fd71a2e52eb88fb2394b7 Mon Sep 17 00:00:00 2001 From: Ilia Karmanov Date: Fri, 14 Aug 2026 02:16:03 -0700 Subject: [PATCH 02/11] feat(mopd): add Super Omni image distillation recipe Add production and smoke overlays, a thin launcher, deterministic circle-count data preparation, and focused usage documentation. Signed-off-by: Ilia Karmanov --- .../nemotron/nemotron-3-super-omni-mopd.md | 68 +++++++++++ docs/index.md | 1 + ...0n8g-megatron-tp8ep16cp2-async-gym.v1.yaml | 78 ++++++++++++ ...ron-super-omni-120ba12b-4n8g-smoke.v1.yaml | 20 ++++ .../prepare_circle_count_mopd_data.py | 112 ++++++++++++++++++ .../run_mopd_circle_count.sh | 22 ++++ tests/test_suites/disabled.txt | 2 + ...-10n8g-megatron-tp8ep16cp2-async-gym.v1.sh | 54 +++++++++ ...otron-super-omni-120ba12b-4n8g-smoke.v1.sh | 8 ++ 9 files changed, 365 insertions(+) create mode 100644 docs/guides/models/nemotron/nemotron-3-super-omni-mopd.md create mode 100644 examples/configs/recipes/vlm/mopd-nemotron-super-omni-120ba12b-10n8g-megatron-tp8ep16cp2-async-gym.v1.yaml create mode 100644 examples/configs/recipes/vlm/mopd-nemotron-super-omni-120ba12b-4n8g-smoke.v1.yaml create mode 100755 examples/nemo_gym/nemotron-3-super-omni/prepare_circle_count_mopd_data.py create mode 100755 examples/nemo_gym/nemotron-3-super-omni/run_mopd_circle_count.sh create mode 100755 tests/test_suites/vlm/mopd-nemotron-super-omni-120ba12b-10n8g-megatron-tp8ep16cp2-async-gym.v1.sh create mode 100755 tests/test_suites/vlm/mopd-nemotron-super-omni-120ba12b-4n8g-smoke.v1.sh diff --git a/docs/guides/models/nemotron/nemotron-3-super-omni-mopd.md b/docs/guides/models/nemotron/nemotron-3-super-omni-mopd.md new file mode 100644 index 0000000000..a214bba533 --- /dev/null +++ b/docs/guides/models/nemotron/nemotron-3-super-omni-mopd.md @@ -0,0 +1,68 @@ +# Nemotron 3 Super Omni Image MOPD + +This recipe distills a non-colocated Nemotron 3 Super Omni teacher into a +Super Omni policy over multimodal NeMo Gym trajectories. It extends the +MTP-disabled Super Omni GRPO recipe with OPD advantages, teacher resources, +and image-aware teacher log-probability computation. + +## Data + +Generate deterministic circle-count examples from the pinned NeMo Gym +submodule: + +```bash +uv run python \ + examples/nemo_gym/nemotron-3-super-omni/prepare_circle_count_mopd_data.py \ + --out /shared/data/circle_count_train.jsonl \ + --num-samples 512 +``` + +Each row contains one structured `input_image` data URL and an `agent_ref` +routing it to `circle_count_simple_agent`. The verifier metadata remains +outside `responses_create_params` and is not included in the model prompt. + +## Launch + +The production recipe uses ten nodes with eight GPUs per node: + +- one vLLM generation node; +- one non-colocated teacher node; +- eight Megatron policy nodes using TP8, EP16, and CP2. + +Set the paths and Slurm values required by the shared Super Omni launcher: + +```bash +MODEL_PATH=/shared/models/super-omni-hf \ +TEACHER_MODEL_PATH=/shared/models/super-omni-teacher-hf \ +TRAIN_PATH=/shared/data/circle_count_train.jsonl \ +CONTAINER=/shared/containers/nemo-rl.sqsh \ +SANDBOX_CONTAINER=/shared/containers/nemo-skills-sandbox.sqsh \ +PERSISTENT_CACHE=/shared/cache/nemo-rl-super-omni \ +EXTRA_MOUNTS=/shared:/shared \ +SLURM_ACCOUNT= \ +SLURM_PARTITION= \ +WANDB_API_KEY= \ +examples/nemo_gym/nemotron-3-super-omni/run_mopd_circle_count.sh +``` + +`TEACHER_MODEL_PATH` is optional. When omitted, the recipe uses +`MODEL_PATH` for self-distillation. A self-distillation run should have a +near-zero mean OPD advantage while retaining non-zero token-level spread. + +The recipe disables in-flight weight updates and enables vLLM encoder-cache +invalidation. This orders each encoder-cache reset after refit and before the +next image request when the vision tower is trainable. + +## Three-step smoke + +Use the four-node smoke before a production run: + +```bash +CONFIG_PATH=examples/configs/recipes/vlm/mopd-nemotron-super-omni-120ba12b-4n8g-smoke.v1.yaml \ +EXP_NAME=mopd-super-omni-circle-count-smoke \ +examples/nemo_gym/nemotron-3-super-omni/run_mopd_circle_count.sh +``` + +The smoke runs three optimizer/refit steps. With one-step asynchronous +trajectory staleness, the third step uses trajectories generated after the +first weight update. diff --git a/docs/index.md b/docs/index.md index 5d066eb95d..8f673f02f2 100644 --- a/docs/index.md +++ b/docs/index.md @@ -313,6 +313,7 @@ guides/models/nemotron/nemotron-3-nano.md guides/models/nemotron/nemotron-3-nano-omni.md guides/models/nemotron/nemotron-3.5-lightning.md guides/models/nemotron/nemotron-3-super.md +guides/models/nemotron/nemotron-3-super-omni-mopd.md guides/models/nemotron/nemotron-3-ultra.md guides/models/qwen/index.md guides/models/qwen/qwen3-5.md diff --git a/examples/configs/recipes/vlm/mopd-nemotron-super-omni-120ba12b-10n8g-megatron-tp8ep16cp2-async-gym.v1.yaml b/examples/configs/recipes/vlm/mopd-nemotron-super-omni-120ba12b-10n8g-megatron-tp8ep16cp2-async-gym.v1.yaml new file mode 100644 index 0000000000..6fe6c1a567 --- /dev/null +++ b/examples/configs/recipes/vlm/mopd-nemotron-super-omni-120ba12b-10n8g-megatron-tp8ep16cp2-async-gym.v1.yaml @@ -0,0 +1,78 @@ +# Super Omni image MOPD: 1 generation, 1 teacher, and 8 policy nodes. +# The default mapping is self-distillation; override it for a stronger teacher. +defaults: vlm_grpo-nemotron-super-omni-120ba12b-16n8g-megatron-tp8ep16cp2-async-gym.v1.yaml +checkpointing: + checkpoint_dir: results/mopd-nemotron-super-omni-circle-count + keep_top_k: 1 +cluster: + num_nodes: 10 +grpo: + num_prompts_per_step: 8 + num_generations_per_prompt: 8 + max_num_epochs: 5 + max_num_steps: 200 + async_grpo: + in_flight_weight_updates: false + adv_estimator: + name: opd +loss_fn: + disable_ppo_ratio: true + truncated_importance_sampling_type: icepop + truncated_importance_sampling_ratio: 5.0 + truncated_importance_sampling_ratio_min: 0.2 + force_on_policy_ratio: false + use_kl_in_reward: false +policy: + train_global_batch_size: 64 + tokenizer: + chat_template_kwargs: + enable_thinking: true + truncate_history_thinking: false + generation: + max_new_tokens: 4096 + colocated: + resources: + num_nodes: 1 + vllm_cfg: + gpu_memory_utilization: 0.5 + enforce_eager: true + reset_encoder_cache_after_weight_update: true +data: + train: + data_path: /path/to/circle_count_train.jsonl + validation: + data_path: /path/to/circle_count_train.jsonl +env: + nemo_gym: + config_paths: + - responses_api_models/vllm_model/configs/vllm_model_for_training.yaml + - resources_servers/circle_count/configs/circle_count.yaml + policy_model: + responses_api_models: + vllm_model: + max_input_tokens: ${policy.max_total_sequence_length} +on_policy_distillation: + enabled: true + teacher_model_by_agent_name: + circle_count_simple_agent: ${policy.model_name} + default_teacher_alias: circle_count_simple_agent + strict_agent_name_match: true + deduplicate_shared_teacher_checkpoints: true + non_colocated_teachers: + enabled: true + default_teacher_cfg: + tensor_model_parallel_size: 8 + pipeline_model_parallel_size: 1 + expert_model_parallel_size: 8 + context_parallel_size: 1 + num_nodes: 1 + gpus_per_node: 8 + precision: bf16 + micro_batch_size: 1 + moe_shared_expert_overlap: false +logger: + log_dir: logs/mopd-nemotron-super-omni-circle-count + wandb: + project: mopd-nemotron-super-omni + name: mopd-nemotron-super-omni-circle-count + log_nemo_gym_full_result_tables: true diff --git a/examples/configs/recipes/vlm/mopd-nemotron-super-omni-120ba12b-4n8g-smoke.v1.yaml b/examples/configs/recipes/vlm/mopd-nemotron-super-omni-120ba12b-4n8g-smoke.v1.yaml new file mode 100644 index 0000000000..40b340e497 --- /dev/null +++ b/examples/configs/recipes/vlm/mopd-nemotron-super-omni-120ba12b-4n8g-smoke.v1.yaml @@ -0,0 +1,20 @@ +# Three steps ensure at least one batch was generated after a weight refit. +defaults: mopd-nemotron-super-omni-120ba12b-10n8g-megatron-tp8ep16cp2-async-gym.v1.yaml +checkpointing: + enabled: false +cluster: + num_nodes: 4 +grpo: + num_prompts_per_step: 2 + num_generations_per_prompt: 2 + max_num_epochs: 1 + max_num_steps: 3 +policy: + train_global_batch_size: 4 + max_total_sequence_length: 4096 + generation: + max_new_tokens: 256 +logger: + log_dir: logs/mopd-nemotron-super-omni-circle-count-smoke + wandb: + name: mopd-nemotron-super-omni-circle-count-smoke diff --git a/examples/nemo_gym/nemotron-3-super-omni/prepare_circle_count_mopd_data.py b/examples/nemo_gym/nemotron-3-super-omni/prepare_circle_count_mopd_data.py new file mode 100755 index 0000000000..9b577ab80b --- /dev/null +++ b/examples/nemo_gym/nemotron-3-super-omni/prepare_circle_count_mopd_data.py @@ -0,0 +1,112 @@ +#!/usr/bin/env python3 +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Generate deterministic NeMo Gym circle-count rows for image MOPD.""" + +from __future__ import annotations + +import argparse +import importlib.util +import json +from pathlib import Path +from types import ModuleType +from typing import Any + + +AGENT_REF = { + "type": "responses_api_agents", + "name": "circle_count_simple_agent", +} + + +def _load_circle_count_generator() -> ModuleType: + repo_root = Path(__file__).resolve().parents[3] + generator_path = ( + repo_root + / "3rdparty" + / "Gym-workspace" + / "Gym" + / "resources_servers" + / "circle_count" + / "generate_data.py" + ) + spec = importlib.util.spec_from_file_location( + "_nemo_gym_circle_count_generate_data", generator_path + ) + if spec is None or spec.loader is None: + raise RuntimeError(f"Could not load circle-count generator: {generator_path}") + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +def _validate_example(example: dict[str, Any]) -> None: + if example.get("agent_ref") != AGENT_REF: + raise ValueError("circle-count MOPD row has an invalid agent_ref") + + responses_create_params = example.get("responses_create_params") + if not isinstance(responses_create_params, dict): + raise ValueError("row is missing responses_create_params") + + image_urls: list[str] = [] + for message in responses_create_params.get("input", []): + content = message.get("content", []) if isinstance(message, dict) else [] + if not isinstance(content, list): + continue + for item in content: + if isinstance(item, dict) and item.get("type") == "input_image": + image_urls.append(str(item.get("image_url", ""))) + + if len(image_urls) != 1 or not image_urls[0].startswith("data:image/"): + raise ValueError( + "each circle-count MOPD row must contain exactly one data-URL input_image" + ) + + request_text = json.dumps(responses_create_params) + if '"circles"' in request_text or '"target_color"' in request_text: + raise ValueError("answer metadata leaked into responses_create_params") + + +def main() -> None: + parser = argparse.ArgumentParser( + description="Generate image MOPD data routed to circle_count_simple_agent." + ) + parser.add_argument("--out", type=Path, required=True) + parser.add_argument("--num-samples", type=int, default=512) + parser.add_argument("--seed-offset", type=int, default=0) + parser.add_argument("--image-size", type=int, default=1000) + parser.add_argument("--radius-min", type=int, default=30) + parser.add_argument("--radius-max", type=int, default=60) + parser.add_argument("--num-circles-min", type=int, default=5) + parser.add_argument("--num-circles-max", type=int, default=20) + parser.add_argument("--num-colors-min", type=int, default=2) + parser.add_argument("--num-colors-max", type=int, default=4) + args = parser.parse_args() + + if args.num_samples <= 0: + raise ValueError("--num-samples must be positive") + + generator = _load_circle_count_generator() + args.out.parent.mkdir(parents=True, exist_ok=True) + with args.out.open("w") as output: + for index in range(args.num_samples): + example = generator.make_example( + args.seed_offset + index, + img_size_range=(args.image_size, args.image_size), + circle_radius_range=(args.radius_min, args.radius_max), + num_circles_range=( + args.num_circles_min, + args.num_circles_max, + ), + num_colors_range=(args.num_colors_min, args.num_colors_max), + ) + example["agent_ref"] = dict(AGENT_REF) + _validate_example(example) + output.write(json.dumps(example) + "\n") + + print(f"Generated {args.num_samples} image-MOPD rows: {args.out}") + + +if __name__ == "__main__": + main() diff --git a/examples/nemo_gym/nemotron-3-super-omni/run_mopd_circle_count.sh b/examples/nemo_gym/nemotron-3-super-omni/run_mopd_circle_count.sh new file mode 100755 index 0000000000..e7e94d7868 --- /dev/null +++ b/examples/nemo_gym/nemotron-3-super-omni/run_mopd_circle_count.sh @@ -0,0 +1,22 @@ +#!/usr/bin/env bash +set -euo pipefail + +# Thin MOPD wrapper around the shared Super Omni launcher. The launcher +# validates MODEL_PATH, TRAIN_PATH, CONTAINER, SANDBOX_CONTAINER, +# PERSISTENT_CACHE, SLURM_ACCOUNT, and WANDB_API_KEY (for online logging). + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" + +export EXP_NAME="${EXP_NAME:-mopd-super-omni-circle-count}" +export CONFIG_PATH="${CONFIG_PATH:-examples/configs/recipes/vlm/mopd-nemotron-super-omni-120ba12b-10n8g-megatron-tp8ep16cp2-async-gym.v1.yaml}" +export WANDB_PROJ="${WANDB_PROJ:-mopd-nemotron-super-omni}" + +if [[ -n "${TEACHER_MODEL_PATH:-}" ]]; then + while [[ "${TEACHER_MODEL_PATH}" == */ && "${TEACHER_MODEL_PATH}" != "/" ]]; do + TEACHER_MODEL_PATH="${TEACHER_MODEL_PATH%/}" + done + teacher_override="on_policy_distillation.teacher_model_by_agent_name.circle_count_simple_agent=${TEACHER_MODEL_PATH}" + export EXTRA_HYDRA_ARGS="${EXTRA_HYDRA_ARGS:+${EXTRA_HYDRA_ARGS} }${teacher_override}" +fi + +exec "${SCRIPT_DIR}/super_omni_launch.sh" diff --git a/tests/test_suites/disabled.txt b/tests/test_suites/disabled.txt index d2dd7da8b4..7c7290f1f1 100644 --- a/tests/test_suites/disabled.txt +++ b/tests/test_suites/disabled.txt @@ -21,3 +21,5 @@ tests/test_suites/vlm/vlm_grpo-nemotron-omni-30ba3b-circle-click-2n8g-megatron-t tests/test_suites/vlm/vlm_grpo-nemotron-super-omni-120ba12b-16n8g-megatron-tp8ep16cp2-async-gym.v1.sh tests/test_suites/vlm/vlm_grpo-nemotron-super-omni-120ba12b-16n8g-megatron-tp8ep16cp2-async-gym-mtp.v1.sh tests/test_suites/vlm/vlm_grpo-nemotron-super-omni-120ba12b-16n8g-megatron-tp8ep16cp2-async-gym-mtp-specdec.v1.sh +tests/test_suites/vlm/mopd-nemotron-super-omni-120ba12b-10n8g-megatron-tp8ep16cp2-async-gym.v1.sh +tests/test_suites/vlm/mopd-nemotron-super-omni-120ba12b-4n8g-smoke.v1.sh diff --git a/tests/test_suites/vlm/mopd-nemotron-super-omni-120ba12b-10n8g-megatron-tp8ep16cp2-async-gym.v1.sh b/tests/test_suites/vlm/mopd-nemotron-super-omni-120ba12b-10n8g-megatron-tp8ep16cp2-async-gym.v1.sh new file mode 100755 index 0000000000..982571d8af --- /dev/null +++ b/tests/test_suites/vlm/mopd-nemotron-super-omni-120ba12b-10n8g-megatron-tp8ep16cp2-async-gym.v1.sh @@ -0,0 +1,54 @@ +#!/bin/bash +SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd) +source "$SCRIPT_DIR/common.env" + +# Parameterized because the Super Omni checkpoint and image data are external. +NUM_NODES="${NUM_NODES:-10}" +GPUS_PER_NODE=8 +STEPS_PER_RUN=3 +MAX_STEPS=3 +NUM_RUNS=1 +NUM_MINUTES="${NUM_MINUTES:-840}" + +exit_if_max_steps_reached + +: "${MODEL_PATH:?MODEL_PATH must point at the Super Omni HF checkpoint}" +: "${TRAIN_PATH:?TRAIN_PATH must point at circle-count Gym JSONL data}" +VAL_PATH="${VAL_PATH:-$TRAIN_PATH}" + +teacher_args=() +if [[ -n "${TEACHER_MODEL_PATH:-}" ]]; then + teacher_args+=( + "on_policy_distillation.teacher_model_by_agent_name.circle_count_simple_agent=${TEACHER_MODEL_PATH}" + ) +fi + +cd "$PROJECT_ROOT" +uv run examples/nemo_gym/run_grpo_nemo_gym.py \ + --config "$CONFIG_PATH" \ + grpo.max_num_steps=$MAX_STEPS \ + policy.model_name="$MODEL_PATH" \ + policy.tokenizer.chat_template="$MODEL_PATH/chat_template.jinja" \ + policy.generation.vllm_cfg.http_server_serving_chat_kwargs.chat_template="$MODEL_PATH/chat_template.jinja" \ + data.train.data_path="$TRAIN_PATH" \ + data.validation.data_path="$VAL_PATH" \ + 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=False \ + "${teacher_args[@]}" \ + "$@" \ + 2>&1 | tee "$RUN_LOG" + +uv run tests/json_dump_tb_logs.py "$LOG_DIR" --output_path "$JSON_METRICS" + +if [[ $(jq 'to_entries | .[] | select(.key == "train/loss") | .value | keys | map(tonumber) | max' "$JSON_METRICS") -ge $MAX_STEPS ]]; then + uv run tests/check_metrics.py "$JSON_METRICS" \ + 'max(data["train/loss"]) < 1000000.0' \ + 'min(data["train/loss"]) > -1000000.0' \ + 'max(data["train/grad_norm"]) < 1000000.0' \ + 'min(data["train/grad_norm"]) >= 0.0' +fi diff --git a/tests/test_suites/vlm/mopd-nemotron-super-omni-120ba12b-4n8g-smoke.v1.sh b/tests/test_suites/vlm/mopd-nemotron-super-omni-120ba12b-4n8g-smoke.v1.sh new file mode 100755 index 0000000000..61ff664b89 --- /dev/null +++ b/tests/test_suites/vlm/mopd-nemotron-super-omni-120ba12b-4n8g-smoke.v1.sh @@ -0,0 +1,8 @@ +#!/bin/bash +SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd) + +export EXP_NAME="$(basename "$0" .sh)" +export NUM_NODES=4 +export NUM_MINUTES=480 + +exec "$SCRIPT_DIR/mopd-nemotron-super-omni-120ba12b-10n8g-megatron-tp8ep16cp2-async-gym.v1.sh" "$@" From 441229eaad6f7ba22fa0f847b65b010159989d7b Mon Sep 17 00:00:00 2001 From: Ilia Karmanov Date: Fri, 14 Aug 2026 03:47:52 -0700 Subject: [PATCH 03/11] fix(mopd): target Gym service virtual environments Install each Gym service into its newly created venv when a managed Python interpreter is selected globally. Signed-off-by: Ilia Karmanov --- ...per-omni-120ba12b-10n8g-megatron-tp8ep16cp2-async-gym.v1.yaml | 1 + 1 file changed, 1 insertion(+) diff --git a/examples/configs/recipes/vlm/mopd-nemotron-super-omni-120ba12b-10n8g-megatron-tp8ep16cp2-async-gym.v1.yaml b/examples/configs/recipes/vlm/mopd-nemotron-super-omni-120ba12b-10n8g-megatron-tp8ep16cp2-async-gym.v1.yaml index 6fe6c1a567..e33a0e0976 100644 --- a/examples/configs/recipes/vlm/mopd-nemotron-super-omni-120ba12b-10n8g-megatron-tp8ep16cp2-async-gym.v1.yaml +++ b/examples/configs/recipes/vlm/mopd-nemotron-super-omni-120ba12b-10n8g-megatron-tp8ep16cp2-async-gym.v1.yaml @@ -44,6 +44,7 @@ data: data_path: /path/to/circle_count_train.jsonl env: nemo_gym: + uv_pip_set_python: true config_paths: - responses_api_models/vllm_model/configs/vllm_model_for_training.yaml - resources_servers/circle_count/configs/circle_count.yaml From 4430840bd03f9edd60997bb2d94a45c4f6cbfd4b Mon Sep 17 00:00:00 2001 From: Yi-Fu Wu Date: Fri, 14 Aug 2026 12:13:01 -0700 Subject: [PATCH 04/11] fix(mopd): add tools/launch config blocks to the mopd test-suite drivers tools/launch extracts NUM_NODES/GPUS_PER_NODE/STEPS_PER_RUN/MAX_STEPS/ NUM_RUNS/NUM_MINUTES from a delimited CONFIG block and exits 1 when the markers are absent, which failed test_dry_run_does_not_fail_and_prints_total_gpu_hours (disabled.txt does not exempt drivers from that glob). Wrap the existing vars in the 10n8g driver and add a mirroring block to the smoke wrapper, keeping NUM_NODES/NUM_MINUTES exported for the exec'd child. Co-Authored-By: Claude Fable 5 Signed-off-by: Yi-Fu Wu --- ...20ba12b-10n8g-megatron-tp8ep16cp2-async-gym.v1.sh | 2 ++ ...opd-nemotron-super-omni-120ba12b-4n8g-smoke.v1.sh | 12 ++++++++++-- 2 files changed, 12 insertions(+), 2 deletions(-) diff --git a/tests/test_suites/vlm/mopd-nemotron-super-omni-120ba12b-10n8g-megatron-tp8ep16cp2-async-gym.v1.sh b/tests/test_suites/vlm/mopd-nemotron-super-omni-120ba12b-10n8g-megatron-tp8ep16cp2-async-gym.v1.sh index 982571d8af..666f7c0dfd 100755 --- a/tests/test_suites/vlm/mopd-nemotron-super-omni-120ba12b-10n8g-megatron-tp8ep16cp2-async-gym.v1.sh +++ b/tests/test_suites/vlm/mopd-nemotron-super-omni-120ba12b-10n8g-megatron-tp8ep16cp2-async-gym.v1.sh @@ -3,12 +3,14 @@ SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd) source "$SCRIPT_DIR/common.env" # Parameterized because the Super Omni checkpoint and image data are external. +# ===== BEGIN CONFIG ===== NUM_NODES="${NUM_NODES:-10}" GPUS_PER_NODE=8 STEPS_PER_RUN=3 MAX_STEPS=3 NUM_RUNS=1 NUM_MINUTES="${NUM_MINUTES:-840}" +# ===== END CONFIG ===== exit_if_max_steps_reached diff --git a/tests/test_suites/vlm/mopd-nemotron-super-omni-120ba12b-4n8g-smoke.v1.sh b/tests/test_suites/vlm/mopd-nemotron-super-omni-120ba12b-4n8g-smoke.v1.sh index 61ff664b89..90c6c9f2b8 100755 --- a/tests/test_suites/vlm/mopd-nemotron-super-omni-120ba12b-4n8g-smoke.v1.sh +++ b/tests/test_suites/vlm/mopd-nemotron-super-omni-120ba12b-4n8g-smoke.v1.sh @@ -2,7 +2,15 @@ SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd) export EXP_NAME="$(basename "$0" .sh)" -export NUM_NODES=4 -export NUM_MINUTES=480 + +# ===== BEGIN CONFIG ===== +NUM_NODES=4 +GPUS_PER_NODE=8 +STEPS_PER_RUN=3 +MAX_STEPS=3 +NUM_RUNS=1 +NUM_MINUTES=480 +# ===== END CONFIG ===== +export NUM_NODES NUM_MINUTES exec "$SCRIPT_DIR/mopd-nemotron-super-omni-120ba12b-10n8g-megatron-tp8ep16cp2-async-gym.v1.sh" "$@" From fba10957c56365437c94cac838d90248513e54ac Mon Sep 17 00:00:00 2001 From: Yi-Fu Wu Date: Fri, 14 Aug 2026 12:13:16 -0700 Subject: [PATCH 05/11] fix(vllm): document and guard reset_encoder_cache_after_weight_update Document the new flag's default in the exemplar YAML and its paired reference config (test_reference_configs_up_to_date requires both), and make normalize_vllm_refit_config raise when the flag is combined with a refit transport whose weight-landing path never resets the encoder cache (nixl, sparse-delta, custom checkpoint engines). Only the collective/IPC and nccl_reshard async paths implement the reset; on any other transport the flag was a silent no-op that kept stale vision embeddings across refits. Co-Authored-By: Claude Fable 5 Signed-off-by: Yi-Fu Wu --- examples/configs/grpo_math_1B.yaml | 1 + nemo_rl/models/generation/vllm/config.py | 12 ++++++++ .../algorithms/test_grpo_checkpoint_engine.py | 28 +++++++++++++++++++ .../unit/reference_configs/grpo_math_1B.yaml | 1 + 4 files changed, 42 insertions(+) diff --git a/examples/configs/grpo_math_1B.yaml b/examples/configs/grpo_math_1B.yaml index d7cd6f528e..f17e906f12 100644 --- a/examples/configs/grpo_math_1B.yaml +++ b/examples/configs/grpo_math_1B.yaml @@ -413,6 +413,7 @@ policy: enable_vllm_metrics_logger: true # Set to true to enable vLLM internal metrics logger, turn off for better performance vllm_metrics_logger_interval: 0.5 # Interval in seconds to collect vLLM logger metrics http_refit_api_key_env_var: null # Optional env var containing the internal refit API key. + reset_encoder_cache_after_weight_update: false # Invalidate cached multimodal encoder outputs after a successful async refit. Only safe when generation is quiesced (grpo.async_grpo.in_flight_weight_updates=false). http_refit_server_port: null # Optional fixed port for Kubernetes targetPorts. zmq_refit_server_port: null # Optional fixed ZeroMQ relay port for Kubernetes targetPorts. vllm_kwargs: {} diff --git a/nemo_rl/models/generation/vllm/config.py b/nemo_rl/models/generation/vllm/config.py index eb8d3f2f42..43bbc1a587 100644 --- a/nemo_rl/models/generation/vllm/config.py +++ b/nemo_rl/models/generation/vllm/config.py @@ -185,6 +185,18 @@ def normalize_vllm_refit_config(config: VllmConfig) -> VllmRefitConfig | None: "'nccl_reshard', 'vllm_s3_sparse', 'vllm_zmq_sparse', 'nixl', or a " "'module:ClassName' checkpoint-engine path." ) + # The encoder-cache reset is implemented only on the collective/IPC and + # nccl_reshard async refit paths (both returned above). Fail loudly rather + # than let other transports silently keep stale multimodal encoder outputs + # across weight updates. + if config["vllm_cfg"].get("reset_encoder_cache_after_weight_update"): + raise ValueError( + "vllm_cfg.reset_encoder_cache_after_weight_update is not supported " + f"with refit_transport={transport!r}: this transport's refit path " + "does not reset the multimodal encoder cache, so stale vision " + "embeddings would silently survive weight updates. Supported " + "transports: null (collective/IPC) and 'nccl_reshard'." + ) refit_config = VllmRefitConfig.model_validate(config.get("refit_cfg") or {}) if ":" in transport: plugin_config = (refit_config.model_extra or {}).get(transport) diff --git a/tests/unit/algorithms/test_grpo_checkpoint_engine.py b/tests/unit/algorithms/test_grpo_checkpoint_engine.py index c8ac163017..c766d5bfb8 100644 --- a/tests/unit/algorithms/test_grpo_checkpoint_engine.py +++ b/tests/unit/algorithms/test_grpo_checkpoint_engine.py @@ -47,6 +47,34 @@ def test_nixl_example_is_an_enabled_non_colocated_overlay(): assert config.cluster["num_nodes"] == 2 +def test_reset_encoder_cache_flag_rejected_on_unsupported_refit_transports(): + """The encoder-cache reset is honored only on collective/IPC and nccl_reshard.""" + import pytest + + from nemo_rl.models.generation.vllm.config import ( + VllmConfig, + normalize_vllm_refit_config, + ) + + def _config(transport): + return cast( + VllmConfig, + { + "vllm_cfg": {"reset_encoder_cache_after_weight_update": True}, + "refit_transport": transport, + }, + ) + + # Supported transports pass through unchanged. + assert normalize_vllm_refit_config(_config(None)) is None + assert normalize_vllm_refit_config(_config("nccl_reshard")) is None + + # Transports whose refit path never resets the encoder cache fail loudly. + for transport in ("nixl", "vllm_s3_sparse", "vllm_zmq_sparse"): + with pytest.raises(ValueError, match="reset_encoder_cache_after_weight_update"): + normalize_vllm_refit_config(_config(transport)) + + def test_refit_policy_generation_uses_attached_checkpoint_engine_synchronizer(): from nemo_rl.algorithms import grpo as grpo_mod from nemo_rl.models.generation.vllm import VllmGeneration diff --git a/tests/unit/reference_configs/grpo_math_1B.yaml b/tests/unit/reference_configs/grpo_math_1B.yaml index e7f64db325..3e1c44b0ff 100644 --- a/tests/unit/reference_configs/grpo_math_1B.yaml +++ b/tests/unit/reference_configs/grpo_math_1B.yaml @@ -409,6 +409,7 @@ policy: enable_vllm_metrics_logger: true # Set to true to enable vLLM internal metrics logger, turn off for better performance vllm_metrics_logger_interval: 0.5 # Interval in seconds to collect vLLM logger metrics http_refit_api_key_env_var: null # Optional env var containing the internal refit API key. + reset_encoder_cache_after_weight_update: false # Invalidate cached multimodal encoder outputs after a successful async refit. Only safe when generation is quiesced (grpo.async_grpo.in_flight_weight_updates=false). http_refit_server_port: null # Optional fixed port for Kubernetes targetPorts. zmq_refit_server_port: null # Optional fixed ZeroMQ relay port for Kubernetes targetPorts. vllm_kwargs: {} From c23e481bb8ca2aaf7eff91c4887e31428cd7914a Mon Sep 17 00:00:00 2001 From: Yi-Fu Wu Date: Fri, 14 Aug 2026 12:13:25 -0700 Subject: [PATCH 06/11] test(opd): cover mixed media and text-only rows within one teacher group The existing tests cover all-image and all-text groups; a single teacher receiving both takes the filter branch that keeps a PackedTensor with a None segment, and nothing asserted the surviving None stays row-aligned with its token row. Co-Authored-By: Claude Fable 5 Signed-off-by: Yi-Fu Wu --- tests/unit/algorithms/test_opd.py | 28 ++++++++++++++++++++++++++++ 1 file changed, 28 insertions(+) diff --git a/tests/unit/algorithms/test_opd.py b/tests/unit/algorithms/test_opd.py index 18a65c1994..73077687e0 100644 --- a/tests/unit/algorithms/test_opd.py +++ b/tests/unit/algorithms/test_opd.py @@ -210,6 +210,34 @@ def test_compute_teacher_logprobs_dp_padding_repeats_multimodal_row(): assert result.shape == (1, 8) +def test_compute_teacher_logprobs_mixed_media_and_text_rows_per_teacher(): + """Mixed image and text-only rows in one group keep the empty rows aligned.""" + twg = _RecordingTeacherWorkerGroup(fill_value=4.0, dp_size=1) + collector = _make_collector( + teacher_worker_groups={"mixed": twg}, + alias_to_group_alias={"mixed_agent": "mixed"}, + on_policy_distillation_cfg={ + "teacher_model_by_agent_name": {"mixed_agent": "/ckpt/mixed"}, + }, + _has_distillation_teachers=True, + ) + + result, _ = collector._compute_teacher_logprobs( + torch.randint(0, 100, (3, 8)), + [{"name": "mixed_agent"}] * 3, + multimodal_data={ + "pixel_values": _row_marked_packed_tensor([5, None, 6]), + "imgs_sizes": _row_marked_packed_tensor([15, None, 16]), + }, + ) + + assert twg.received is not None + # The text-only row keeps its slot so media rows stay paired with token rows. + assert _received_row_markers(twg.received["pixel_values"]) == [5.0, None, 6.0] + assert _received_row_markers(twg.received["imgs_sizes"]) == [15.0, None, 16.0] + assert result.shape == (3, 8) + + def test_compute_teacher_logprobs_routes_to_correct_teacher(): """Samples are routed to the right teacher and results stitched back.""" math_twg = _MockTeacherWorkerGroup(fill_value=1.0, dp_size=1) From 3b9fe6d99a2efcbb69e881551fc8ee3e02909a35 Mon Sep 17 00:00:00 2001 From: Yi-Fu Wu Date: Fri, 14 Aug 2026 12:13:27 -0700 Subject: [PATCH 07/11] docs(mopd): link the Super Omni MOPD guide from the Nemotron model index The guide was added to the docs/index.md toctree but not to the Nemotron landing page that lists every sibling guide. Co-Authored-By: Claude Fable 5 Signed-off-by: Yi-Fu Wu --- docs/guides/models/nemotron/index.md | 3 +++ 1 file changed, 3 insertions(+) diff --git a/docs/guides/models/nemotron/index.md b/docs/guides/models/nemotron/index.md index a92af13c98..a68bfdfd71 100644 --- a/docs/guides/models/nemotron/index.md +++ b/docs/guides/models/nemotron/index.md @@ -17,6 +17,9 @@ For the full list of supported models, see (CLEVR-CoGenT and MMPR-Tiny recipes). - **[Nemotron 3 Super](nemotron-3-super.md)** — the multi-stage Nemotron 3 Super post-training recipe (RLVR, SWE, and RLHF stages). +- **[Nemotron 3 Super Omni Image MOPD](nemotron-3-super-omni-mopd.md)** — image + on-policy distillation for the Super Omni vision-language model with a + non-colocated teacher (10-node production recipe plus a 4-node smoke). - **[Nemotron 3 Ultra](nemotron-3-ultra.md)** — RLVR, teacher training, and MOPD stages on GB200 NVL72 hardware. - **[Nemotron 3.5 Lightning](nemotron-3.5-lightning.md)** — RLVR with NeMo Gym From b56ac06b9804ac63b3a8c81478bbfb457c8b155a Mon Sep 17 00:00:00 2001 From: Yi-Fu Wu Date: Fri, 14 Aug 2026 12:13:38 -0700 Subject: [PATCH 08/11] chore(mopd): use the full copyright header and drop dead default_teacher_alias Replace the SPDX two-liner with the standard Apache block (and fix the ruff isort blank-line failure it sat next to), and remove default_teacher_alias from the recipe: resolve_reference_aliases checks strict_agent_name_match before the fallback branch, so with strict matching enabled the alias could never be used. Co-Authored-By: Claude Fable 5 Signed-off-by: Yi-Fu Wu --- ...b-10n8g-megatron-tp8ep16cp2-async-gym.v1.yaml | 1 - .../prepare_circle_count_mopd_data.py | 16 +++++++++++++--- 2 files changed, 13 insertions(+), 4 deletions(-) diff --git a/examples/configs/recipes/vlm/mopd-nemotron-super-omni-120ba12b-10n8g-megatron-tp8ep16cp2-async-gym.v1.yaml b/examples/configs/recipes/vlm/mopd-nemotron-super-omni-120ba12b-10n8g-megatron-tp8ep16cp2-async-gym.v1.yaml index e33a0e0976..bc25445022 100644 --- a/examples/configs/recipes/vlm/mopd-nemotron-super-omni-120ba12b-10n8g-megatron-tp8ep16cp2-async-gym.v1.yaml +++ b/examples/configs/recipes/vlm/mopd-nemotron-super-omni-120ba12b-10n8g-megatron-tp8ep16cp2-async-gym.v1.yaml @@ -56,7 +56,6 @@ on_policy_distillation: enabled: true teacher_model_by_agent_name: circle_count_simple_agent: ${policy.model_name} - default_teacher_alias: circle_count_simple_agent strict_agent_name_match: true deduplicate_shared_teacher_checkpoints: true non_colocated_teachers: diff --git a/examples/nemo_gym/nemotron-3-super-omni/prepare_circle_count_mopd_data.py b/examples/nemo_gym/nemotron-3-super-omni/prepare_circle_count_mopd_data.py index 9b577ab80b..a17502aece 100755 --- a/examples/nemo_gym/nemotron-3-super-omni/prepare_circle_count_mopd_data.py +++ b/examples/nemo_gym/nemotron-3-super-omni/prepare_circle_count_mopd_data.py @@ -1,6 +1,17 @@ #!/usr/bin/env python3 -# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 +# 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. """Generate deterministic NeMo Gym circle-count rows for image MOPD.""" @@ -13,7 +24,6 @@ from types import ModuleType from typing import Any - AGENT_REF = { "type": "responses_api_agents", "name": "circle_count_simple_agent", From 684e4bc5c9454b64803971b0799522defcf84928 Mon Sep 17 00:00:00 2001 From: Yi-Fu Wu Date: Fri, 14 Aug 2026 13:59:55 -0700 Subject: [PATCH 09/11] build(gym): prefetch circle_click and circle_count venvs for omni images prefetch_omni_envs.yaml mirrored only the Super Omni GRPO recipe's five Gym servers, so images built with NEMO_GYM_PREFETCH_CONFIGS=examples/nemo_gym/prefetch_omni_envs.yaml had no baked venv for the circle-click GRPO recipe or the new circle-count MOPD recipe, forcing an on-node uv venv build (network egress plus a writable /opt/gym_venvs) at first startup. Add both servers; these are the only Gym-using omni recipes not already covered. Co-Authored-By: Claude Fable 5 Signed-off-by: Yi-Fu Wu --- examples/nemo_gym/prefetch_omni_envs.yaml | 2 ++ 1 file changed, 2 insertions(+) diff --git a/examples/nemo_gym/prefetch_omni_envs.yaml b/examples/nemo_gym/prefetch_omni_envs.yaml index a67fe85688..cdec8c804d 100644 --- a/examples/nemo_gym/prefetch_omni_envs.yaml +++ b/examples/nemo_gym/prefetch_omni_envs.yaml @@ -29,3 +29,5 @@ env: - resources_servers/mcqa/configs/mcqa.yaml - resources_servers/gui_coordinate/configs/gui_coordinate.yaml - resources_servers/string_match/configs/string_match.yaml + - resources_servers/circle_click/configs/circle_click.yaml + - resources_servers/circle_count/configs/circle_count.yaml From eec4023ee1447ec399b35b5df1331ea4d8290138 Mon Sep 17 00:00:00 2001 From: Yi-Fu Wu Date: Fri, 14 Aug 2026 14:49:22 -0700 Subject: [PATCH 10/11] ci: add circle-count data prep script to pyrefly project-includes The lint workflow's ratchet requires every file with zero pyrefly errors to be listed in project-includes; the new examples/nemo_gym/nemotron-3-super-omni/prepare_circle_count_mopd_data.py type-checks clean, so CI failed until it was whitelisted. Co-Authored-By: Claude Fable 5 Signed-off-by: Yi-Fu Wu --- pyrefly.toml | 1 + 1 file changed, 1 insertion(+) diff --git a/pyrefly.toml b/pyrefly.toml index bd92f2c37d..49962360b4 100644 --- a/pyrefly.toml +++ b/pyrefly.toml @@ -44,6 +44,7 @@ project-includes = [ "examples/custom_parallel/custom_parallel.py", "examples/custom_parallel/llama_nemotron_super_49b_custom_plan.py", "examples/modelopt/export_quantized_to_hf.py", + "examples/nemo_gym/nemotron-3-super-omni/prepare_circle_count_mopd_data.py", "nemo_rl/algorithms/__init__.py", "nemo_rl/algorithms/advantage_estimator.py", "nemo_rl/algorithms/async_utils/__init__.py", From ec994310cdbfcec784a5872fe5479319d22b6325 Mon Sep 17 00:00:00 2001 From: Yi-Fu Wu Date: Fri, 14 Aug 2026 15:58:31 -0700 Subject: [PATCH 11/11] fix(vllm): tolerate partial configs in the encoder-cache transport guard normalize_vllm_refit_config is also called by worker-side NIXL setup with partial generation configs that carry no vllm_cfg key, so the new guard's direct subscript raised KeyError (test_configure_nixl_worker_ignores_other_configs). Use a presence check instead; the guard still raises when the flag is actually set on an unsupported transport. Co-Authored-By: Claude Fable 5 Signed-off-by: Yi-Fu Wu --- nemo_rl/models/generation/vllm/config.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/nemo_rl/models/generation/vllm/config.py b/nemo_rl/models/generation/vllm/config.py index 43bbc1a587..f6786d1e8e 100644 --- a/nemo_rl/models/generation/vllm/config.py +++ b/nemo_rl/models/generation/vllm/config.py @@ -188,8 +188,10 @@ def normalize_vllm_refit_config(config: VllmConfig) -> VllmRefitConfig | None: # The encoder-cache reset is implemented only on the collective/IPC and # nccl_reshard async refit paths (both returned above). Fail loudly rather # than let other transports silently keep stale multimodal encoder outputs - # across weight updates. - if config["vllm_cfg"].get("reset_encoder_cache_after_weight_update"): + # across weight updates. Some callers re-validate partial generation + # configs (e.g. worker-side NIXL setup), so vllm_cfg may be absent here. + vllm_cfg = config.get("vllm_cfg") + if vllm_cfg and vllm_cfg.get("reset_encoder_cache_after_weight_update"): raise ValueError( "vllm_cfg.reset_encoder_cache_after_weight_update is not supported " f"with refit_transport={transport!r}: this transport's refit path "