feat(mopd): add Super Omni multimodal distillation

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

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=<account> \
+SLURM_PARTITION=<partition> \
+WANDB_API_KEY=<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.

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

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: {}

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:
+    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
+    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}
+  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

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

examples/nemo_gym/nemotron-3-super-omni/prepare_circle_count_mopd_data.py

@@ -0,0 +1,122 @@
+#!/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.
+
+"""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()

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"

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