diff --git a/docs/guides/models/nemotron/index.md b/docs/guides/models/nemotron/index.md index a92af13c981..a68bfdfd718 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 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 00000000000..a214bba533b --- /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 5d066eb95db..8f673f02f24 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/grpo_math_1B.yaml b/examples/configs/grpo_math_1B.yaml index d7cd6f528ef..f17e906f12c 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/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 00000000000..bc25445022c --- /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: + 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 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 00000000000..40b340e497c --- /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 00000000000..a17502aeced --- /dev/null +++ b/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() 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 00000000000..e7e94d7868a --- /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/examples/nemo_gym/prefetch_omni_envs.yaml b/examples/nemo_gym/prefetch_omni_envs.yaml index a67fe856883..cdec8c804dc 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 diff --git a/nemo_rl/algorithms/async_utils/trajectory_collector.py b/nemo_rl/algorithms/async_utils/trajectory_collector.py index d56a901375c..b6f6ecbc224 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 6281bb2af18..f6786d1e8e1 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. @@ -182,6 +185,20 @@ 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. 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 " + "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/nemo_rl/models/generation/vllm/vllm_worker_async.py b/nemo_rl/models/generation/vllm/vllm_worker_async.py index 4a0a2c02401..e791fa4a023 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/pyrefly.toml b/pyrefly.toml index bd92f2c37df..49962360b4e 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", diff --git a/tests/test_suites/disabled.txt b/tests/test_suites/disabled.txt index d2dd7da8b47..7c7290f1f1a 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 00000000000..666f7c0dfd6 --- /dev/null +++ b/tests/test_suites/vlm/mopd-nemotron-super-omni-120ba12b-10n8g-megatron-tp8ep16cp2-async-gym.v1.sh @@ -0,0 +1,56 @@ +#!/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. +# ===== 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 + +: "${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 00000000000..90c6c9f2b80 --- /dev/null +++ b/tests/test_suites/vlm/mopd-nemotron-super-omni-120ba12b-4n8g-smoke.v1.sh @@ -0,0 +1,16 @@ +#!/bin/bash +SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd) + +export EXP_NAME="$(basename "$0" .sh)" + +# ===== 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" "$@" diff --git a/tests/unit/algorithms/test_grpo_checkpoint_engine.py b/tests/unit/algorithms/test_grpo_checkpoint_engine.py index c8ac1630177..c766d5bfb88 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/algorithms/test_opd.py b/tests/unit/algorithms/test_opd.py index 46aeb58e8f2..73077687e00 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,131 @@ 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_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) diff --git a/tests/unit/models/generation/test_vllm_backend.py b/tests/unit/models/generation/test_vllm_backend.py index c7e7969ddc9..d81528e07ad 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 diff --git a/tests/unit/reference_configs/grpo_math_1B.yaml b/tests/unit/reference_configs/grpo_math_1B.yaml index e7f64db325c..3e1c44b0ff7 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: {}