feat(capture): capture vLLM VLM media alongside token deltas

docs/guides/single-controller.md

@@ -394,6 +394,67 @@ The SC path is still under active development. Feature gaps are tracked in [issu
 - Multimodal/VLM GRPO is supported with Megatron generation. Set
   `policy.is_vlm: true`; see the
   [CLEVR Single-Controller recipe](../../examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-8n4g-megatron-single-controller-async.v1.yaml).
+- NeMo-Gym token capture also supports Omni dynamic-resolution images and native video
+  rollouts with async vLLM generation and a Megatron learner. With
+  `token_capture.enabled: true` and the VLM processor configured, workers capture
+  the processed media used for inference together with each call's token delta.
+  RL hands the owned tensors (`imgs`, `imgs_sizes`, and optional `num_frames`)
+  to Gym's `complete_call_from_response` as opaque attachments, and the TQ sink
+  writes them in the same `put` as the token columns, so `staged` coordinates
+  acknowledge tokens and pixels together and a failed write is `capture_failed`
+  at call time. Tensors keep their native shapes and dtypes on the wire; two
+  per-row flags (`media_present`, `media_has_frames`) mark which rows carry
+  pixels and whether they are video. Media-enabled staging partitions carry
+  these columns on every row; text-only runs register and read none of them.
+  vLLM pixels are rearranged losslessly into packed patches; the finalizer
+  reads the presence flags with the base columns, then issues one batched read
+  of the tensor columns for the terminal-chain calls that carry media, and
+  publishes `pixel_values`, `imgs_sizes`, and `num_frames` for the existing
+  Megatron learner without resampling or normalizing the media again.
+  Only newly introduced occurrences are staged. vLLM-specific `media_spans`
+  extras retain placeholder positions and token hashes for multi-turn prefix
+  replacement, including video's timestamp-separated visual-token spans.
+  Capture requests use vLLM's `skip_mm_cache=True` path to obtain concrete
+  processor tensors. vLLM can still reuse its processor-only cache; this does
+  not guarantee fresh preprocessing. vLLM tiles images from a dummy prompt, so
+  the request text never changes geometry, but the images of one request share
+  the token budget: adding an image under a tight budget can re-tile a retained
+  one, and a warm processor-only cache can keep a geometry a fresh processor
+  would not reproduce. If that changes retained geometry or placeholder tokens, the worker
+  rejects the continuation before inference with HTTP 400 and error code
+  `retained_media_changed` (other capture-time validation failures use
+  `media_capture_rejected`). Gym's current exception middleware wraps the
+  upstream error in HTTP 500, retaining the code in the response body. RL then
+  classifies it as an infrastructure `GymTransportError`, eligible for the
+  configured prompt retry policy; it is not a dedicated terminal rejection.
+  Text-call rows carry sentinels in each column's own dtype, because
+  TransferQueue keeps one dtype per field across live rows.
+  Media-enabled rows also store `media_metadata_digest`, a SHA-256 checksum of
+  the small route-less extras JSON. Sources check it before exposing media
+  descriptors, without reading pixels or routed-expert tensors. Missing or
+  changed checksums reject the row. This detects accidental metadata corruption;
+  it is independent of Gym's receipt and combined `extras_digest` commitment.
+  Media bundles are structurally validated before writing and after reading
+  (required tensors, patch geometry, frame grouping); malformed or missing
+  columns reject the rollout as `invalid_media_columns`, incompatible parts
+  along a chain as `media_chain_incompatible`. Tensor contents are not hashed;
+  retained occurrences are checked by geometry and placeholder tokens. Media
+  must remain immutable for the rollout's lifetime and preprocessing must be
+  deterministic. Same-shape pixel changes and corruption of stored pixel values
+  are outside this check's coverage. Call
+  rows share the existing checkpoint and cleanup lifecycle. TQ has no
+  transactional rollback: a failed combined write is discarded best-effort by
+  the sink, and a failed discard is logged at ERROR.
+  Upgrade the paired Gym and RL changes together. The GB200 functional shard
+  `L1_Functional_Tests_GB200_Vllm_Omni_Single_Controller.sh` smokes this path
+  end to end (CLEVR-style images through Gym `string_match`, native video
+  through Gym `mcqa`) and gates on `train/finalize/media_row_rate == 1`, the
+  metric that reports the fraction of learner rows built from captured media.
+  Media capture requires `policy.generation.backend: vllm`; Megatron inference
+  token capture is text-only. No new Megatron-LM pin is needed. Compaction,
+  mixed image/video conversations, native audio,
+  video token pruning, static tiling (`image_num_patches`), other processor families,
+  and `token_capture.defer_routed_experts_to_policy: true` are not supported.
 - Multi-Teacher On-Policy Distillation (MOPD) is supported for text-only NeMo
   Gym rollouts; multimodal/VLM MOPD is not yet supported. See
   [Multi-Teacher On-Policy Distillation](../about/algorithms/mopd.md#running-mopd).

examples/nemo_gym/vlm_grpo_nemotron_omni_clevr_vllm_token_capture_single_controller.yaml

@@ -0,0 +1,67 @@
+# Nemotron-Omni CLEVR through NeMo-Gym (string_match verifier) with async vLLM
+# generation, a Megatron learner, and vLLM media token capture enabled.
+#
+# Functional-test config (not a nightly recipe, so it lives beside the other
+# Gym smoke configs rather than under examples/configs/recipes/):
+# tests/functional/nemotron_omni_clevr_vllm_token_capture_single_controller_1n2g.sh
+# supplies the data paths, the 1 trainer + 1 generator split, and the step
+# budget. Token capture requires the Gym rollout path, so the native CLEVR
+# dataset/environment of the parent recipe are swapped for Gym rows here.
+# Prefix caching stays off (hybrid-Mamba align cache mode asserts chunked
+# prefill) and the parent's mm_processor_cache_gb=0 keeps concrete pixels
+# available to the capture host.
+defaults: ../configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-8n4g-megatron-single-controller-async.v1.yaml
+data:
+  _override_: true
+  max_input_seq_length: null
+  shuffle: true
+  num_workers: 0
+  use_multiple_dataloader: false
+  train:
+    data_path: /path/to/clevr_gym/train.jsonl
+  validation:
+    data_path: /path/to/clevr_gym/val.jsonl
+  default:
+    dataset_name: NemoGymDataset
+    env_name: nemo_gym
+    prompt_file: null
+    system_prompt_file: null
+    processor: nemo_gym_data_processor
+env:
+  _override_: true
+  should_use_nemo_gym: true
+  should_log_nemo_gym_responses: false
+  should_mask_flagged_samples: true
+  nemo_gym:
+    is_trajectory_collection: false
+    skip_venv_if_present: true
+    port_range_low: 5000
+    port_range_high: 5999
+    config_paths:
+    - responses_api_models/vllm_model/configs/vllm_model_for_training.yaml
+    - resources_servers/string_match/configs/string_match.yaml
+    string_match_simple_agent:
+      responses_api_agents:
+        simple_agent:
+          max_steps: 1
+policy:
+  generation:
+    backend: vllm
+    refit_transport: null
+    vllm_cfg:
+      async_engine: true
+      expose_http_server: true
+      logprobs_mode: raw_logprobs
+      enable_prefix_caching: false
+      enforce_eager: true
+    vllm_kwargs:
+      limit_mm_per_prompt:
+        image: 2
+      skip_mm_profiling: true
+      mamba_ssm_cache_dtype: float32
+      kernel_config:
+        enable_flashinfer_autotune: false
+        moe_backend: triton
+token_capture:
+  enabled: true
+  staging_partition: rollout_staging

nemo_rl/algorithms/single_controller_utils/setup.py

@@ -78,7 +78,7 @@
 )
 from nemo_rl.algorithms.utils import set_seed
 from nemo_rl.data.collate_fn import rl_collate_fn
-from nemo_rl.data.multimodal_utils import WIRE_MULTIMODAL_FIELDS
+from nemo_rl.data.multimodal_utils import WIRE_MULTIMODAL_FIELDS, uses_image_placeholder
 from nemo_rl.data.utils import load_dataloader_state, setup_response_data
 from nemo_rl.data_plane import (
     DATA_PLANE_CHECKPOINT_SCHEMA_VERSION,
@@ -303,11 +303,17 @@ def _register_single_controller_partitions(
     partition_id: str,
     include_multimodal_fields: bool,
 ) -> None:
-    """Warm all SingleController partitions before concurrent data-plane use."""
+    """Warm all SingleController partitions before concurrent data-plane use.
+
+    VLM token capture (``include_multimodal_fields`` with capture enabled) adds
+    the media columns the vLLM worker stages beside each captured call to the
+    staging partition.
+    """
     algo_cfg = algo_config(master_config)
     policy_config = master_config.policy
     token_capture_cfg = master_config.token_capture
     r3_enabled = router_replay_enabled(policy_config)
+    capture_media = token_capture_cfg.enabled and include_multimodal_fields
     group_size = algo_cfg.num_generations_per_prompt
     num_rollout_samples = master_config.async_rl.max_buffered_rollouts * group_size
 
@@ -341,12 +347,16 @@ def _register_single_controller_partitions(
         from nemo_rl.data_plane.schema import (
             ROUTED_EXPERTS_FIELD as STAGING_ROUTED_EXPERTS_FIELD,
         )
-        from nemo_rl.data_plane.tq_token_sink import STAGING_FIELDS
+        from nemo_rl.data_plane.tq_token_sink import (
+            MEDIA_STAGING_FIELDS,
+            STAGING_FIELDS,
+        )
 
         dp_client.register_partition(
             partition_id=token_capture_cfg.staging_partition,
             fields=list(STAGING_FIELDS)
-            + ([STAGING_ROUTED_EXPERTS_FIELD] if r3_enabled else []),
+            + ([STAGING_ROUTED_EXPERTS_FIELD] if r3_enabled else [])
+            + (list(MEDIA_STAGING_FIELDS) if capture_media else []),
             num_samples=num_rollout_samples,
             consumer_tasks=["finalize", "prev_lp", "train"],
         )
@@ -1192,6 +1202,23 @@ def setup_single_controller(
     # nemo_rl/distributed/actor_environments.py), so nothing here needs to
     # change the worker's environment.
     token_capture_cfg = master_config.token_capture
+    capture_media = token_capture_cfg.enabled and processor is not None
+    if capture_media:
+        if generation_config["backend"] != "vllm":
+            raise NotImplementedError(
+                "VLM media token capture is only implemented for the vLLM "
+                f"generation backend; got {generation_config['backend']!r}"
+            )
+        if not uses_image_placeholder(processor):
+            raise ValueError(
+                "VLM token capture currently supports Omni dynamic images and native video"
+            )
+        if not policy_config["megatron_cfg"]["enabled"]:
+            raise ValueError(
+                "Omni media token capture currently requires the Megatron learner"
+            )
+        if token_capture_cfg.defer_routed_experts_to_policy:
+            raise ValueError("VLM token capture requires direct router replay assembly")
     if rollout_checkpoint_cfg.snapshot_attempt_interval_s is not None:
         if not master_config.checkpointing["enabled"]:
             raise ValueError(
@@ -1905,8 +1932,13 @@ def _build_trainer_then_megatron_generation() -> tuple[
             include_multimodal_fields=processor is not None,
         )
     if token_capture_cfg.enabled:
-        # Both active backends stage canonical Gym rows in serving workers.
-        generation.setup_token_capture(dp_config, token_capture_cfg.staging_partition)
+        # Both active backends stage canonical Gym rows in serving workers;
+        # only vLLM workers stage captured media beside them (capture_media).
+        generation.setup_token_capture(
+            dp_config,
+            token_capture_cfg.staging_partition,
+            capture_media=capture_media,
+        )
         generation.set_rollout_weight_version(0)
 
     if weight_synchronizer is None:
@@ -1969,6 +2001,7 @@ def _build_trainer_then_megatron_generation() -> tuple[
                 router_replay_enabled=router_replay_enabled(policy_config),
                 defer_routed_experts_to_policy=token_capture_cfg.defer_routed_experts_to_policy,
                 max_seq_len=_generation_max_seq_len(generation_config),
+                capture_media=capture_media,
             ),
             num_workers=token_capture_cfg.num_reassembler_workers,
         )