feat(mopd): add Super Omni multimodal distillation

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/recipes/vlm/mopd-nemotron-super-omni-120ba12b-10n8g-megatron-tp8ep16cp2-async-gym.v1.yaml

@@ -0,0 +1,79 @@
+# 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}
+  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

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,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()

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"

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.

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.