diff --git a/docs/guides/models/nemotron/nemotron-3-nano-omni.md b/docs/guides/models/nemotron/nemotron-3-nano-omni.md index d14a958b948..e2418dc413b 100644 --- a/docs/guides/models/nemotron/nemotron-3-nano-omni.md +++ b/docs/guides/models/nemotron/nemotron-3-nano-omni.md @@ -8,8 +8,11 @@ The maintained Nemotron recipes enable `grpo.deduplicate_multimodal_data` to share immutable model-ready media segments across logical GRPO generations and re-intern them after batching, replay, and sharding. The representation supports image, video, and audio payload keys, although the maintained recipes currently -qualify image inputs only. Deduplication currently requires the vLLM generation -backend and `data_plane.enabled=false`. +qualify image inputs only. Deduplication requires the vLLM generation backend. +It works with `data_plane.enabled=true` except on NeMo-Gym runs, where the +TransferQueue trainer does not attach the initial Gym image payloads. On the +data plane it saves driver RAM only: `PackedTensor.to_wire` emits one row per +*logical* row, so the wire payload is `O(G x images)` either way. `grpo.debug_payload_metrics` emits logical, physical, and protocol-5 serialized payload sizes for the exact Ray boundaries used by generation, replay, logprobs, diff --git a/docs/guides/nemotron-3-nano-omni.md b/docs/guides/nemotron-3-nano-omni.md index 1dba0ddb44c..f3da665266c 100644 --- a/docs/guides/nemotron-3-nano-omni.md +++ b/docs/guides/nemotron-3-nano-omni.md @@ -14,8 +14,11 @@ The maintained Nemotron recipes enable `grpo.deduplicate_multimodal_data` to share immutable model-ready media segments across logical GRPO generations and re-intern them after batching, replay, and sharding. The representation supports image, video, and audio payload keys, although the maintained recipes currently -qualify image inputs only. Deduplication currently requires the vLLM generation -backend and `data_plane.enabled=false`. +qualify image inputs only. Deduplication requires the vLLM generation backend. +It works with `data_plane.enabled=true` except on NeMo-Gym runs, where the +TransferQueue trainer does not attach the initial Gym image payloads. On the +data plane it saves driver RAM only: `PackedTensor.to_wire` emits one row per +*logical* row, so the wire payload is `O(G x images)` either way. `grpo.debug_payload_metrics` emits logical, physical, and protocol-5 serialized payload sizes for the exact Ray boundaries used by generation, replay, logprobs, diff --git a/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-1n8g-megatron-tp8ep8.v1-tq_mooncake.yaml b/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-1n8g-megatron-tp8ep8.v1-tq_mooncake.yaml new file mode 100644 index 00000000000..cf4b52f06e5 --- /dev/null +++ b/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-1n8g-megatron-tp8ep8.v1-tq_mooncake.yaml @@ -0,0 +1,9 @@ +defaults: vlm_grpo-nemotron-omni-30ba3b-clevr-1n8g-megatron-tp8ep8.v1.yaml +data_plane: + enabled: true + backend: mooncake_cpu +cluster: + # tp8 on one node leaves dp=1, so the distributed optimizer has nothing to + # shard across and the run lands ~108 MiB short in the MoE forward. dp=2 fits. + # Name stays 1n8g: common-tq.env derives the base recipe from it. + num_nodes: 2 diff --git a/examples/configs/recipes/vlm/vlm_grpo-qwen3.5-35ba3b-geo3k-2n8g-automodel-ep16-tq_simple.yaml b/examples/configs/recipes/vlm/vlm_grpo-qwen3.5-35ba3b-geo3k-2n8g-automodel-ep16-tq_simple.yaml new file mode 100644 index 00000000000..e56d2a92341 --- /dev/null +++ b/examples/configs/recipes/vlm/vlm_grpo-qwen3.5-35ba3b-geo3k-2n8g-automodel-ep16-tq_simple.yaml @@ -0,0 +1,3 @@ +defaults: vlm_grpo-qwen3.5-35ba3b-geo3k-2n8g-automodel-ep16.yaml +data_plane: + enabled: true diff --git a/examples/configs/vlm_grpo_3B.yaml b/examples/configs/vlm_grpo_3B.yaml index 8e0d521582a..e1d994c8f70 100644 --- a/examples/configs/vlm_grpo_3B.yaml +++ b/examples/configs/vlm_grpo_3B.yaml @@ -3,8 +3,10 @@ defaults: "grpo_math_1B.yaml" grpo: - # Share immutable multimodal payloads across logical GRPO rows. Currently - # qualified with the vLLM generation backend and data_plane.enabled=false. + # Share immutable multimodal payloads across logical GRPO rows. Requires the + # vLLM generation backend. Works with data_plane.enabled=true except on + # NeMo-Gym runs; note it saves driver RAM only -- the wire payload is one row + # per logical row either way. deduplicate_multimodal_data: false # Debug-only exact Ray payload telemetry; serialization adds runtime overhead. debug_payload_metrics: false diff --git a/examples/run_grpo.py b/examples/run_grpo.py index b417abe795c..ba7d9ee2baa 100644 --- a/examples/run_grpo.py +++ b/examples/run_grpo.py @@ -21,13 +21,16 @@ from nemo_rl.algorithms.grpo import ( MasterConfig, - grpo_train, setup, shutdown_environments, ) from nemo_rl.algorithms.utils import get_tokenizer from nemo_rl.data.utils import setup_response_data -from nemo_rl.data_plane.factory import maybe_configure_data_plane_env +from nemo_rl.data_plane.factory import ( + make_policy_factory, + maybe_configure_data_plane_env, + select_sync_trainer, +) from nemo_rl.distributed.virtual_cluster import init_ray from nemo_rl.models.generation import configure_generation_config from nemo_rl.telemetry.setup import init_telemetry_driver, shutdown_telemetry @@ -40,22 +43,6 @@ from nemo_rl.utils.timer import Timer -def _select_trainer(master_config: MasterConfig): - """Pick the synchronous trainer based on ``data_plane.enabled``. - - Factored out so test_architecture_invariants can verify dispatch - without the full setup() path. - """ - dp_cfg = master_config.data_plane or {} - if dp_cfg.get("enabled", False): - from nemo_rl.algorithms.grpo_sync import grpo_train_sync - - print("🚀 Running synchronous GRPO training (TransferQueue)") - return grpo_train_sync - print("🚀 Running synchronous GRPO training (legacy)") - return grpo_train - - def parse_args() -> tuple[argparse.Namespace, list[str]]: """Parse command line arguments.""" parser = argparse.ArgumentParser(description="Run GRPO training with configuration") @@ -142,20 +129,6 @@ def main() -> None: tokenizer, config.data, config.env ) - # Pick the policy factory at the launcher level so the legacy trainer - # stays data-plane-agnostic (architectural invariant — see - # tests/data_plane/unit/test_architecture_invariants.py). - _dp_cfg = config.data_plane or {} - if _dp_cfg.get("enabled", False): - from nemo_rl.models.policy.tq_policy import TQPolicy - - def _make_policy(**kwargs): - return TQPolicy(**kwargs, dp_cfg=_dp_cfg) - - _policy_factory = _make_policy - else: - _policy_factory = None # setup() defaults to plain Policy - with rl_init_timer.time("setup"): ( policy, @@ -176,7 +149,7 @@ def _make_policy(**kwargs): tokenizer, dataset, val_dataset, - policy_factory=_policy_factory, + policy_factory=make_policy_factory(config.data_plane), ) rl_init_timer.record("total", time.perf_counter() - main_start) @@ -237,7 +210,7 @@ def _make_policy(**kwargs): else: # Two parallel synchronous trainers (verl-style — main_ppo.py vs # main_ppo_sync.py). data_plane.enabled selects which one runs. - trainer = _select_trainer(master_config) + trainer = select_sync_trainer(master_config) # grpo_train_sync defers checkpoint finalization to the checkpointer's # background threads; the context manager guarantees they are flushed on # exit. (grpo_train also flushes internally; shutdown() is idempotent.) diff --git a/examples/run_vlm_grpo.py b/examples/run_vlm_grpo.py index da363ea4080..f005e1069c7 100644 --- a/examples/run_vlm_grpo.py +++ b/examples/run_vlm_grpo.py @@ -19,10 +19,14 @@ from omegaconf import OmegaConf -from nemo_rl.algorithms.grpo import MasterConfig, async_grpo_train, grpo_train, setup +from nemo_rl.algorithms.grpo import MasterConfig, async_grpo_train, setup from nemo_rl.algorithms.utils import get_tokenizer from nemo_rl.data.utils import setup_response_data -from nemo_rl.data_plane.factory import maybe_configure_data_plane_env +from nemo_rl.data_plane.factory import ( + make_policy_factory, + maybe_configure_data_plane_env, + select_sync_trainer, +) from nemo_rl.distributed.virtual_cluster import init_ray from nemo_rl.models.generation import configure_generation_config from nemo_rl.utils.config import ( @@ -125,7 +129,14 @@ def main() -> None: master_config, teacher_worker_groups, alias_to_group_alias, - ) = setup(config, tokenizer, dataset, val_dataset, processor=processor) + ) = setup( + config, + tokenizer, + dataset, + val_dataset, + processor=processor, + policy_factory=make_policy_factory(config.data_plane), + ) rl_init_timer.record("total", time.perf_counter() - main_start) rl_init_metrics = rl_init_timer.get_timing_metrics(reduction_op="sum") @@ -170,22 +181,27 @@ def main() -> None: processor=processor, ) else: - print("🚀 Running synchronous GRPO training") - grpo_train( - policy, - policy_generation, - dataloader, - val_dataloader, - tokenizer, - loss_fn, - task_to_env, - val_task_to_env, - logger, - checkpointer, - grpo_state, - master_config, - processor=processor, - ) + # ``select_sync_trainer`` prints which sync trainer it picked. + trainer = select_sync_trainer(master_config, label="VLM GRPO") + # grpo_train_sync defers checkpoint finalization to the checkpointer's + # background threads; the context manager guarantees they are flushed on + # exit. (grpo_train also flushes internally; shutdown() is idempotent.) + with checkpointer: + trainer( + policy, + policy_generation, + dataloader, + val_dataloader, + tokenizer, + loss_fn, + task_to_env, + val_task_to_env, + logger, + checkpointer, + grpo_state, + master_config, + processor=processor, + ) if __name__ == "__main__": diff --git a/nemo_rl/algorithms/grpo.py b/nemo_rl/algorithms/grpo.py index 68c4c903c87..82b0b4fd86b 100644 --- a/nemo_rl/algorithms/grpo.py +++ b/nemo_rl/algorithms/grpo.py @@ -461,10 +461,24 @@ def _validate_multimodal_dedup_capability(master_config: MasterConfig) -> None: "grpo.deduplicate_multimodal_data=true is currently qualified " "only with policy.generation.backend=vllm." ) - if (master_config.data_plane or {}).get("enabled", False): + # The data plane accepts deduplicated payloads, so the wire format is not + # the constraint -- but note what dedup buys there. ``to_wire`` emits one + # row per *logical* row, so a shared segment is concatenated once per + # generation: the saving is in driver RAM (the deepcopy memo), not in wire + # or TQ-storage bytes, which stay O(G x images). The one gap is NeMo-Gym: + # ``grpo_train_sync`` does not call + # ``attach_initial_nemo_gym_image_payloads``, which supplies the initial + # image tensors a Gym dataset omits from ``extra_env_info``. That helper is + # itself gated on ``should_use_nemo_gym``, so non-Gym recipes never needed + # it and are unaffected. + if (master_config.data_plane or {}).get("enabled", False) and ( + should_use_nemo_gym(master_config) + ): raise NotImplementedError( - "grpo.deduplicate_multimodal_data=true is currently supported " - "only when data_plane.enabled=false." + "grpo.deduplicate_multimodal_data=true with data_plane.enabled=true " + "is not supported for NeMo-Gym runs: the TransferQueue trainer does " + "not attach the initial Gym image payloads. Non-Gym recipes are " + "supported." ) diff --git a/nemo_rl/algorithms/grpo_sync.py b/nemo_rl/algorithms/grpo_sync.py index e3f1fe46a81..a602b569796 100644 --- a/nemo_rl/algorithms/grpo_sync.py +++ b/nemo_rl/algorithms/grpo_sync.py @@ -37,6 +37,8 @@ from typing import TYPE_CHECKING, Any, Optional if TYPE_CHECKING: + from transformers import AutoProcessor + from nemo_rl.models.policy.tq_policy import TQPolicy import numpy as np @@ -50,6 +52,7 @@ MasterConfig, _clip_grpo_advantages, _create_advantage_estimator, + _initial_policy_generation_stale, _log_mixed_rewards_and_advantages_information, _placeholder_seq_logprob_error_metrics, _policy_dtype, @@ -73,6 +76,7 @@ ) from nemo_rl.data.interfaces import DatumSpec from nemo_rl.data.llm_message_utils import batched_message_log_to_flat_message +from nemo_rl.data.multimodal_utils import present_multimodal_fields from nemo_rl.data_plane.interfaces import KVBatchMeta from nemo_rl.data_plane.schema import DP_CALIB_INPUT_FIELDS, DP_TRAIN_FIELDS from nemo_rl.distributed.batched_data_dict import BatchedDataDict @@ -390,6 +394,19 @@ def grpo_train_sync( checkpointer: CheckpointManager, grpo_save_state: GRPOSaveState, master_config: MasterConfig, + # Unused here, and present only so the shared VLM launcher can pass one + # fixed kwarg set to whichever trainer ``select_sync_trainer`` returns. + # ``grpo_train``'s sole use of it is + # ``attach_initial_nemo_gym_image_payloads``, gated on + # ``grpo.deduplicate_multimodal_data`` *and* ``should_use_nemo_gym`` — the + # combination ``setup()`` rejects via + # ``_validate_multimodal_dedup_capability``. Non-Gym dedup runs never call + # that helper, so they need no processor here either. + # + # TODO: replace this parity kwarg with a ``ProcessorInterface`` both + # trainers consume, rather than threading ``Optional[AutoProcessor]`` + # through every signature — ``grpo.py`` repeats it at seven sites. + processor: Optional["AutoProcessor"] = None, ) -> None: """Run GRPO training algorithm — TransferQueue-mediated. @@ -414,7 +431,13 @@ def grpo_train_sync( kv_scales_cache = None # Cache reused for computed kv scales - POLICY_GENERATION_STALE = True + # Skip a redundant iter-1 refit when setup() already synced weights + # (synchronizer not stale, fresh run). The redundant refit resets + # vLLM CUDA-graph / KV-cache state and yields a step-1 + # token_mult_prob_error spike that converges by step 3. + POLICY_GENERATION_STALE = _initial_policy_generation_stale( + policy_generation, grpo_save_state.total_steps + ) assert policy_generation is not None if master_config.grpo.skip_reference_policy_logprobs_calculation: @@ -577,9 +600,18 @@ def grpo_train_sync( with timer.time("total_step_time"): print("▶ Preparing batch...", flush=True) with timer.time("data_processing"): + # ``share_immutable_media`` must be passed here exactly as + # ``grpo_train`` passes it (grpo.py). Without it + # ``_prepare_multimodal_sharing`` never runs, so + # ``deduplicate_multimodal_data`` becomes a silent no-op on + # this trainer and the deepcopy below makes G independent + # copies of every image in driver RAM. repeated_batch: BatchedDataDict[DatumSpec] = ( batch.repeat_interleave( - master_config.grpo.num_generations_per_prompt + master_config.grpo.num_generations_per_prompt, + share_immutable_media=( + master_config.grpo.deduplicate_multimodal_data + ), ) ) @@ -950,9 +982,14 @@ def grpo_train_sync( # (logprobs/advantages/masks) and wire-only message # log bulk fields are skipped by virtue of not being # in DP_CALIB_INPUT_FIELDS. + # VLM extras cannot be named in + # ``DP_CALIB_INPUT_FIELDS``: the rollout writes + # pixel_values / image_grid_thw / … individually and + # which of them exist is per-processor, so the static + # list alone would calibrate image-blind. _calib_fields = [ f for f in (meta.fields or []) if f in DP_CALIB_INPUT_FIELDS - ] + ] + present_multimodal_fields(meta) calibration_data = policy.read_from_dataplane( meta, select_fields=_calib_fields, diff --git a/nemo_rl/data/multimodal_utils.py b/nemo_rl/data/multimodal_utils.py index 67878e5278a..9932e79ef6b 100644 --- a/nemo_rl/data/multimodal_utils.py +++ b/nemo_rl/data/multimodal_utils.py @@ -22,7 +22,7 @@ from copy import deepcopy from io import BytesIO from pathlib import Path -from typing import Any, Optional, Union +from typing import TYPE_CHECKING, Any, Optional, Union import requests import torch @@ -32,6 +32,11 @@ from transformers.audio_utils import load_audio from transformers.video_utils import load_video +if TYPE_CHECKING: + # Type-only: importing the data plane at module scope here would make + # ``nemo_rl.data_plane`` and ``nemo_rl.data`` mutually importing. + from nemo_rl.data_plane.interfaces import KVBatchMeta + VLLM_MULTIMODAL_DATA_KEYS = frozenset({"vllm_images", "vllm_videos", "vllm_audios"}) NATIVE_MULTIMODAL_KEYS = frozenset({"vllm_content", *VLLM_MULTIMODAL_DATA_KEYS}) IMAGE_CONTENT_TYPES = frozenset({"input_image", "image", "image_url"}) @@ -102,6 +107,185 @@ def uses_image_placeholder(processor: Any) -> bool: return type(processor).__name__ in _PLACEHOLDER_STYLE_PROCESSOR_NAMES +# Wire-transport registries for multimodal fields. These are NOT the origin +# of the packed-vs-per-token classification — ``data/processors.py`` is, where +# every key from ``get_multimodal_keys_from_processor`` is wrapped in a +# ``PackedTensor`` and the type maps are kept as plain tensors. These sets +# mirror that decision by name, because once a value crosses the wire it is a +# plain tensor and the type distinction is gone; the read-side consumers +# (materialize's pad skip, get_multimodal_dict, truncate_tensors, both +# seq-dim validators, the TQ fetch list) can only dispatch on the name. +# Adding a modality therefore means updating processors.py AND the matching +# set here; ``encode_multimodal_for_wire`` raises on anything unregistered. + +# Per-token: rectangular ``[B, S]`` type maps for text/image/video tokens. +PER_TOKEN_MULTIMODAL_FIELDS = frozenset( + { + "token_type_ids", # gemma3: which tokens are image + "mm_token_type_ids", # qwen2.5-vl (transformers>=5.3): text(0)/image(1)/video(2) for 3D RoPE + } +) + +# Packed per-sample: jagged. ``PackedTensor`` in-memory; a single +# ``torch.nested`` value on the wire, whose rows carry their own shapes. +PACKED_MULTIMODAL_FIELDS = frozenset( + { + "pixel_values", + "pixel_values_videos", + "image_grid_thw", + "video_grid_thw", + "second_per_grid_ts", # transformers 4.x spelling; kept for old pins + "input_features", + # qwen2.5/3-omni: ``Qwen2_5OmniProcessor.model_input_names`` appends + # both of these unconditionally, so any omni recipe emits them. + "feature_attention_mask", + "video_second_per_grid", + # nemotron-omni: per-image [H, W] and the RADIO temporal-patching + # frame count. Coupled with pixel_values (see + # ``batched_data_dict._COUPLED_MULTIMODAL_KEYS``); both pack on dim 0. + "imgs_sizes", + "num_frames", + } +) + + +# Suffix for the per-sample tag key carrying what ``to_wire``'s flattening +# removes from the payload. The shapes ride inside ``KVBatchMeta.tags`` -- +# per-sample dicts the data plane transports and projects without knowing what +# a pad shape is. Reassembly lives here, in the layer that owns +# ``PackedTensor``. +ROW_SHAPES_SUFFIX = "__row_shapes" + + +def row_shapes_key(field: str) -> str: + """Companion tag key carrying per-row shapes for ``field``.""" + return field + ROW_SHAPES_SUFFIX + + +# Keys inside the :func:`row_shapes_key` tag value. Kept a plain dict rather +# than a record because ``tags`` is TQ's own per-sample channel and rides its +# serializer; both halves are read back by name, never defaulted. +ROW_GEOMETRY_SHAPES = "shapes" +ROW_GEOMETRY_PAD = "pad" + + +# Include-list of multimodal fields every forward-running dispatch (logprob +# *and* train) must ship so the trainer's forward matches the rollout. One wire +# field per logical field: ``PACKED`` fields travel as a single nested tensor +# whose rows carry their own shapes, and per-token fields are rectangular and +# travel as plain tensors. +WIRE_MULTIMODAL_FIELDS = PER_TOKEN_MULTIMODAL_FIELDS | PACKED_MULTIMODAL_FIELDS + + +def present_multimodal_fields(meta: "KVBatchMeta") -> list[str]: + """Multimodal wire fields the rollout actually wrote for this batch. + + Intersecting with ``meta.fields`` is required, not defensive: the noop + adapter and the TQ contract both raise on a fetch for a field that was + never written, and text-only runs write none of these. Sorted for a + deterministic field list across ranks. + """ + return sorted(WIRE_MULTIMODAL_FIELDS & set(meta.fields or ())) + + +def multimodal_row_tags( + multimodal: dict, sample_count: int +) -> "Optional[list[dict[str, Any]]]": + """Per-sample tag rows carrying what ``to_wire``'s flattening removes. + + ``KVBatchMeta.tags`` is the transport's channel for per-sample primitives: + it is aligned 1:1 with ``sample_ids`` and projected automatically by + ``subset``/``slice``/``concat``, so a worker holding a shard gets its own + rows without anyone re-keying them. The data plane never interprets the + contents. + + Carries ``shapes`` (per-row, and unrecoverable once ``to_wire`` flattens) + and ``pad`` (the field's policy flag). Deliberately *not* a pad target: the + width padding lands at is scratch that the model discards -- mcore crops it + via ``imgs_sizes`` before patchification, and the AutoModel path rejects + mixed-resolution batches outright -- so each consumer pads to its own view + and nothing batch-wide has to be agreed across shards. + """ + tags: list[dict[str, Any]] = [{} for _ in range(sample_count)] + for key, value in multimodal.items(): + if key not in PACKED_MULTIMODAL_FIELDS or not isinstance(value, PackedTensor): + continue + # ``row_shapes()`` rather than ``to_wire()``: the payload is encoded + # separately by ``encode_multimodal_for_wire``, and ``to_wire``'s + # ``torch.cat`` would copy the whole column a second time only for it + # to be discarded here. + shapes = value.row_shapes() + if all(not row_shapes for row_shapes in shapes): + continue # every logical row empty -- to_wire skips the field too + if len(shapes) != sample_count: + raise ValueError( + f"{key!r}: PackedTensor holds {len(shapes)} logical rows but the " + f"batch has {sample_count} samples. Tags are aligned 1:1 with " + "sample_ids, so a disagreement here would pair one sample's " + "pixels with another's shapes." + ) + pad = value.pad_to_max_shape + for row, row_shapes in enumerate(shapes): + tags[row][row_shapes_key(key)] = { + ROW_GEOMETRY_SHAPES: row_shapes, + ROW_GEOMETRY_PAD: pad, + } + # ``None`` rather than ``B`` empty dicts: a text-only run has no packed + # field, and an all-empty tags list would still be pickled on every + # dispatch and re-sliced per DP rank in ``shard_meta_for_dp``. + return tags if any(tags) else None + + +def reassemble_packed_multimodal( + fields: dict, tags: "Optional[list[dict[str, Any]]]" = None +) -> None: + """In place: rebuild ``PackedTensor`` for packed fields, consuming companions. + + Called by the read path once its columns are materialized. Leaves anything + that is not a wire-form packed field untouched, so it is safe to call on any + column dict. + + Raises: + ValueError: A packed field arrived without its shapes companion. Every + producer mints the two together (``multimodal_row_tags`` beside + ``encode_multimodal_for_wire``) and every ``KVBatchMeta`` transform + projects ``tags`` alongside ``sample_ids``, so a missing companion + means the transport dropped it. Reconstructing anyway would hand + the model 1-D pixels and train image-blind without an error. + """ + for key in list(fields): + if key not in PACKED_MULTIMODAL_FIELDS: + continue + value = fields[key] + if not (isinstance(value, torch.Tensor) and value.is_nested): + continue + rows = [t.get(row_shapes_key(key)) for t in tags] if tags is not None else None + present = [r for r in rows if r] if rows else [] + if not present: + raise ValueError( + f"{key!r} arrived as a nested wire value but no sample carries a " + f"{row_shapes_key(key)!r} tag" + + ( + " (tags=None)" + if tags is None + else f" (checked {len(tags)} tag rows)" + ) + + ". to_wire flattens each row, so without the companion the " + "true per-segment shapes are unrecoverable." + ) + # ``from_wire`` returns ``None`` only for a zero-row value, which the + # guard above already excludes: a non-empty ``present`` means at least + # one tag row, and a row-count disagreement raises inside ``from_wire``. + # Indexed, not ``.get``-with-default: a producer-side rename of either + # key must fail here rather than silently restore ``pad=False``, which + # changes what ``as_tensor`` hands the vision encoder. + fields[key] = PackedTensor.from_wire( + value, + [[] if r is None else r[ROW_GEOMETRY_SHAPES] for r in rows], # type: ignore[union-attr] + pad_to_max_shape=bool(present[0][ROW_GEOMETRY_PAD]), + ) + + class PackedTensor: """A logical batch of rows backed by packable tensor segments. @@ -329,6 +513,14 @@ def as_tensor( raise IndexError( f"dim_to_pack={self.dim_to_pack} is invalid for tensors with rank {rank}" ) + # Computed locally, never transported. Two shards padding to + # different widths is harmless because the width is scratch the + # model discards: mcore crops it via ``imgs_sizes`` before + # patchification (see + # ``test_dynamic_resolution_padding_is_cropped_before_radio_patchification``), + # and the AutoModel path rejects mixed-resolution batches outright, + # so every image there is already the same size. Keeping this local + # is what lets the data plane stay ignorant of padding. max_shape = [ max(tensor.shape[dim] for tensor in non_none_tensors) for dim in range(rank) @@ -669,6 +861,244 @@ def flattened_concat( pad_to_max_shape=pad_to_max_shapes[0], ) + # ── Wire encoding (data-plane roundtrip) ───────────────────────── + # PackedTensor is a domain wrapper; the wire only handles + # ``torch.Tensor`` (incl. ``torch.nested``) and ``np.ndarray[object]``. + # ``to_wire`` / ``from_wire`` are the single boundary between the two + # representations. + # + # The geometry rides on ``KVBatchMeta.tags``, not as a companion column: + # ``to_wire`` flattens each row, so the shapes TQ records on the + # controller (``metadata.py::extract_field_schema`` -> ``per_sample_shapes``) + # are the flat lengths, not the true per-segment shapes. See + # :func:`multimodal_row_tags`. + + def _row_segments(self) -> list[list[torch.Tensor]]: + """Per logical row, its non-empty physical segments. + + One entry per *logical* row. Under deduplication a logical row maps to + several shared physical segments (``_row_offsets`` / + ``_segment_indices``), so iterating ``self.tensors`` directly would + yield the physical segment count instead of the batch size. + ``_row_segment_indices`` degrades to ``[row]`` for the legacy + one-tensor-per-row layout. + """ + # Built with an explicit loop rather than a comprehension: the + # ``is not None`` filter does not narrow ``Optional[Tensor]`` inside a + # comprehension, so the result would type as ``list[list[Tensor | None]]``. + row_segments: list[list[torch.Tensor]] = [] + for row in range(len(self)): + segments: list[torch.Tensor] = [] + for i in self._row_segment_indices(row): + segment = self.tensors[i] + if segment is not None: + segments.append(segment) + row_segments.append(segments) + return row_segments + + @staticmethod + def _shapes_of(row_segments: list[list[torch.Tensor]]) -> list[list[list[int]]]: + """``shapes[row][segment]`` for an already-walked row/segment list. + + Shared by :meth:`row_shapes` and :meth:`to_wire` so the shape encoding + cannot drift between the tag minter and the payload encoder. + """ + return [[list(t.shape) for t in segs] for segs in row_segments] + + def row_shapes(self) -> list[list[list[int]]]: + """The ``shapes`` half of :meth:`to_wire`, without encoding the payload. + + Callers that only need the geometry -- :func:`multimodal_row_tags` -- + use this so they do not pay ``to_wire``'s ``torch.cat`` a second time. + """ + return self._shapes_of(self._row_segments()) + + def to_wire( + self, + ) -> tuple[Optional[torch.Tensor], list[list[list[int]]]]: + """Encode as a flattened ``torch.jagged`` tensor plus its row shapes. + + Returns ``(None, [])`` when every logical row is empty so the caller + can skip the field entirely. Otherwise: + + * ``nested`` — one row per *logical* row, each row the 1-D concat of + that row's segments. Flattening is what makes ``torch.jagged`` + total: rows vary only in dim 0, so ragged trailing dims and mixed + rank both encode, and nothing is padded to make a container accept + them. Empty rows become zero-length placeholders. + * ``shapes`` — ``shapes[row][segment]`` is that segment's true shape. + Required because TQ derives ``per_sample_shapes`` from the value it + is handed, so flat rows make it record flat lengths. The caller + ships this on ``KVBatchMeta.tags``; see + :func:`nemo_rl.data.multimodal_utils.multimodal_row_tags`. + + Only ``dim_to_pack=0`` is supported today; other values would + need ``ragged_idx`` on the nested tensor and a matching + transpose in :meth:`from_wire`. + """ + if self.dim_to_pack != 0: + raise NotImplementedError( + f"to_wire only supports dim_to_pack=0, got " + f"{self.dim_to_pack}. Non-zero requires ragged_idx " + "threading in torch.nested and a matching transpose " + "on the read side." + ) + row_segments = self._row_segments() + + # Each segment is flattened to 1-D and the row is the 1-D concat of its + # segments. Two consequences, both deliberate: + # + # * Rows then differ only in dim 0, so ``torch.jagged`` accepts every + # shape -- ragged trailing dims and mixed rank included. No padding + # is materialized into the bytes that cross the wire or land in TQ + # storage, and TQ never falls back to the deprecated strided layout. + # * The per-row concat is 1-D, so it cannot raise on segments whose + # trailing dims differ -- which is what previously forced + # ``pad_to_max_shape`` to pad *before* the concat. + # + # The true shapes travel beside the payload (see the returned + # ``shapes``) because TQ derives ``per_sample_shapes`` from what it is + # handed: give it flat rows and it records flat lengths. Padding still + # happens for ``pad_to_max_shape`` values, but in worker memory at use + # time via :meth:`as_tensor`, not on the wire. + shapes = self._shapes_of(row_segments) + # ``reshape(-1)`` on contiguous processor output is a view, so the + # single-segment row -- one image per sample, the overwhelmingly common + # case -- costs nothing here. That is per-row only: the + # ``as_nested_tensor`` below routes to ``jagged_from_list``, which + # ``torch.cat``s every row into one values buffer, so no row is + # zero-copy end to end. One copy of the column is the floor. + rows: list[Optional[torch.Tensor]] = [ + None + if not segs + else ( + segs[0].reshape(-1) + if len(segs) == 1 + else torch.cat([t.reshape(-1) for t in segs]) + ) + for segs in row_segments + ] + + ref = next((t for t in rows if t is not None), None) + if ref is None: + return None, [] + + if any(t is None for t in rows): + placeholder = torch.zeros(0, dtype=ref.dtype, device=ref.device) + rows = [placeholder if t is None else t for t in rows] + + nested = torch.nested.as_nested_tensor(rows, layout=torch.jagged) # type: ignore[arg-type] + return nested, shapes + + @classmethod + def from_wire( + cls, + nested: torch.Tensor, + shapes: list[list[list[int]]], + *, + pad_to_max_shape: bool = False, + ) -> Optional["PackedTensor"]: + """Reconstruct from the value produced by :meth:`to_wire`. + + Returns ``None`` for an empty batch (no rows); the caller should + skip the field entirely rather than instantiate an empty + ``PackedTensor``. + + ``shapes`` is the companion returned by :meth:`to_wire`, carried on + ``KVBatchMeta.tags`` (see :func:`multimodal_row_tags`). Each flat row is + split by segment ``numel`` and reshaped back to its true shape. It is + required, not optional: reconstructing without it yields 1-D segments, + which train image-blind without erroring. + + A zero-length row means the sample had no media -- it becomes + ``None`` (not a ``(0, ...)`` tensor) so an image-free shard + reconstructs as legacy does: ``as_tensor`` returns ``None`` and + ``logical_segment_counts_by_row`` reports 0 rather than 1. + + ``pad_to_max_shape`` is restored onto the rebuilt value as a flag, not + materialized. Segments come back at their true shapes and stay separate + via the CSR row map, so nothing is padded or concatenated here; + :meth:`as_tensor` pads at use time. + + Mirrors :meth:`to_wire`; both assume ``dim_to_pack=0``. + """ + if not nested.is_nested: + raise TypeError( + "from_wire expects the nested value produced by to_wire, got a " + f"dense tensor of shape {tuple(nested.shape)}. A dense value here " + "means codec.materialize padded the field -- check that its name " + "is in PACKED_MULTIMODAL_FIELDS." + ) + unbound = list(nested.unbind()) + if not unbound: + return None + + if len(shapes) != len(unbound): + raise ValueError( + f"from_wire got {len(unbound)} wire rows but {len(shapes)} shape " + "entries; they are minted together by to_wire and must agree." + ) + + # Segments are rebuilt at their true shapes and kept separate via the + # CSR row map. No padding happens here: the data plane transports, and + # ``as_tensor`` pads at use time because a rectangle is a *model input* + # requirement (the vision encoder consumes one dense tensor), not a + # transport one. An empty row contributes no segments, so + # ``logical_segment_counts_by_row`` reports 0 as it did for ``None``. + segments_flat: list[torch.Tensor] = [] + row_offsets: list[int] = [0] + for flat, row_shapes in zip(unbound, shapes): + # ``torch.split`` cuts the whole row in one dispatch; slicing each + # segment individually costs O(segments) Python-level ops per row, + # and this runs per packed field on every fetch. + numels = [torch.Size(shape).numel() for shape in row_shapes] + segments_flat.extend( + view.reshape(shape) + for view, shape in zip(torch.split(flat, numels), row_shapes) + ) + row_offsets.append(len(segments_flat)) + return cls( + segments_flat, # type: ignore[arg-type] + dim_to_pack=0, + pad_to_max_shape=pad_to_max_shape, + _row_offsets=row_offsets, + _segment_indices=list(range(len(segments_flat))), + ) + + +def encode_multimodal_for_wire( + k: str, v: Union["PackedTensor", torch.Tensor] +) -> Optional[torch.Tensor]: + """The wire value for one multimodal field. Dispatched by registry membership. + + Returns ``None`` when the field has nothing to ship (an all-empty packed + batch), in which case the caller omits the column entirely. The wire key is + always ``k``: one wire field per logical field. + + Payload only. Per-token fields ride rectangular; packed fields ride as one + flattened ``torch.jagged`` value. The geometry :meth:`PackedTensor.from_wire` + needs to undo that flattening -- per-row segment shapes plus the + ``pad_to_max_shape`` flag -- is minted separately by + :func:`multimodal_row_tags` and shipped on ``KVBatchMeta.tags``. TQ cannot + derive it, because it reads ``per_sample_shapes`` off the flattened rows it + is handed. + """ + if k in PACKED_MULTIMODAL_FIELDS: + assert isinstance(v, PackedTensor), ( + f"{k!r}: expected PackedTensor, got {type(v).__name__}" + ) + nested, _shapes = v.to_wire() + return nested # None for an all-empty batch + elif k in PER_TOKEN_MULTIMODAL_FIELDS: + assert isinstance(v, torch.Tensor), ( + f"{k!r}: expected Tensor, got {type(v).__name__}" + ) + return v + else: + raise KeyError( + f"unregistered multimodal field {k!r} — add to PACKED_/PER_TOKEN_MULTIMODAL_FIELDS" + ) + def get_multimodal_keys_from_processor(processor) -> list[str]: """Get keys of the multimodal data that can be used as model inputs. diff --git a/nemo_rl/data_plane/adapters/transfer_queue.py b/nemo_rl/data_plane/adapters/transfer_queue.py index 0ad7b6b3dfc..9201c97c7ae 100644 --- a/nemo_rl/data_plane/adapters/transfer_queue.py +++ b/nemo_rl/data_plane/adapters/transfer_queue.py @@ -24,6 +24,7 @@ import contextlib import glob +import importlib import ipaddress import json import os @@ -53,7 +54,6 @@ backend_config, data_plane_supports_checkpointing, ) -from nemo_rl.data_plane.schema import PROMOTE_1D_FIELDS # ────────────────────────────────────────────────────────────────────────── # Backend init — lifted from rl-arena/arena/backends.py. @@ -343,6 +343,137 @@ def _register_all_buffers(self, ptrs, sizes): # type: ignore[no-untyped-def] cls._nrl_register_checked = True +def _assert_tq_stores_scalar_rows_0d() -> None: + """Confirm a dense 1-D field really is stored as 0-d rows. + + :func:`_patch_scalar_field_schema` rewrites the reported sample shape to + ``()`` on that premise, and nothing reshapes the payload to compensate + any more. If a TQ revision started storing 1-D fields as ``(1,)`` rows + instead — fixing the same bug from the other side — the rewrite would + turn a correct schema into a wrong one, and the symptom would be + corrupt reads rather than an import error. + + So ask TQ directly rather than trusting the pin. + + Raises rather than skipping when the storage module is gone: the caller + reached here only after importing ``transfer_queue.metadata``, so "TQ isn't + installed" is no longer a live explanation — a missing module means the + layout moved, which is exactly what this guard exists to catch. + """ + try: + from transfer_queue.storage.managers.base import KVStorageManager + except ImportError as e: + raise _tq_shape_drift_error( + "storage.managers.base is no longer importable", + "the dense-1-D storage layout the scalar schema patch assumes " + "cannot be verified, and a wrong assumption corrupts reads", + "probe", + ) from e + + generate = getattr(KVStorageManager, "_generate_values", None) + if generate is None: + raise _tq_shape_drift_error( + "KVStorageManager no longer has _generate_values", + "the dense-1-D storage layout the scalar schema patch assumes " + "cannot be verified, and a wrong assumption corrupts reads", + "probe", + ) + + probe = TensorDict({"_nrl_probe": torch.zeros(2)}, batch_size=[2]) + rows = generate(probe) + if len(rows) != 2 or any(getattr(r, "ndim", None) != 0 for r in rows): + shapes = [tuple(getattr(r, "shape", ())) for r in rows] + raise _tq_shape_drift_error( + "a dense 1-D field no longer stores as 0-d rows " + f"(probe yielded {len(rows)} rows with shapes {shapes})", + "rewriting the reported sample shape to () would now disagree " + "with the stored rows and corrupt scalar columns", + "patch (it may simply be unnecessary — check whether upstream " + "fixed extract_field_schema)", + ) + + +def _patch_scalar_field_schema() -> None: + """Report the true ``()`` sample shape for dense 1-D fields. + + Upstream ``transfer_queue.metadata.extract_field_schema`` rebinds a + *local* for 1-D inputs:: + + if len(value.shape) == 1: + value = value.unsqueeze(-1) # local only + first_item = value[0] # -> shape (1,) + + but the value that reaches storage is the original ``(N,)`` tensor, + which ``KVStorageManager._generate_values`` iterates into ``N`` **0-d** + rows. So the schema claims a per-sample shape of ``(1,)`` while the + stored rows are ``()``. + + Only the KV path notices. ``BatchMeta.get_shapes`` repeats the uniform + ``shape`` per sample for non-nested fields, and ``KVStorageManager`` + hands that list to the client, which reshapes raw bytes with it — so + a scalar column reconstructs as ``(1,)`` rows and + ``_merge_tensors_to_tensordict`` then re-nests it instead of taking + its ``all(dim() == 0) -> torch.stack`` branch. ``SimpleStorage`` + fetches stored objects by ``(index, field)`` and never consults the + schema, which is why the symptom is ``mooncake_cpu``-only. + + Byte counts are unaffected either way (``prod(()) == prod((1,)) == 1``); + this is a reshape/dtype-of-container bug, not a sizing one. + + Applied on every backend so one partition's schema cannot disagree with + itself across processes. There is no payload-side fallback, so the + premise is verified against TQ itself before the patch is installed — + see :func:`_assert_tq_stores_scalar_rows_0d`. + """ + try: + from transfer_queue import metadata as _md + except ImportError: + return + if getattr(_md, "_nrl_scalar_schema_patched", False): + return + + orig = getattr(_md, "extract_field_schema", None) + if orig is None: + raise _tq_shape_drift_error( + "metadata module no longer exposes extract_field_schema", + "dense 1-D fields would keep reporting a (1,) sample shape and " + "reconstruct as nested (1,) rows on the KV path", + "function", + ) + + _assert_tq_stores_scalar_rows_0d() + + def extract_field_schema(data): # type: ignore[no-untyped-def] + schema = orig(data) + for name in data.keys(): + value = data.get(name) + if ( + isinstance(value, torch.Tensor) + and not value.is_nested + and value.dim() == 1 + and str(name) in schema + ): + # ``_generate_values`` iterates this into 0-d rows; say so. + schema[str(name)]["shape"] = torch.Size([]) + return schema + + # Both storage managers bound the name at import time + # (``from transfer_queue.metadata import extract_field_schema``), so + # rebinding only the defining module would leave them on the original. + _md.extract_field_schema = extract_field_schema + for mod_path in ( + "transfer_queue.storage.managers.base", + "transfer_queue.storage.managers.simple_storage_manager", + ): + try: + mod = importlib.import_module(mod_path) + except ImportError: + continue + if hasattr(mod, "extract_field_schema"): + mod.extract_field_schema = extract_field_schema + _md._nrl_scalar_schema_patched = True + + def _patch_mooncake_staging_buffers(max_bytes: int) -> None: """Reuse RDMA-registered host buffers for mooncake tensor GETs and PUTs. @@ -574,98 +705,51 @@ def _assert_no_key_loss(src_dict: dict, new_td: TensorDict, fn: str) -> None: ) -def _promote_1d_leaves(td: TensorDict) -> TensorDict: - """Promote declared scalar leaves to ``(N, 1)`` for Mooncake. - - The authoritative field list lives in - :data:`nemo_rl.data_plane.schema.PROMOTE_1D_FIELDS`. Declared fields must - arrive as dense ``(N,)`` tensors. Any other dense 1D tensor is rejected so - it cannot silently encounter TQ v0.1.9's schema/data mismatch. - ``NonTensorStack`` and ``NonTensorData`` leaves pass through. - - Args: - td: TensorDict to validate and encode for the Mooncake wire format. - - Returns: - TensorDict with declared scalar leaves promoted to ``(N, 1)``. - - Raises: - ValueError: If a declared field is not a dense 1D tensor, or an - undeclared field is a dense 1D tensor. - """ - # td.keys() (top-level) includes NonTensorData / NonTensorStack leaves. - # keys(include_nested=True, leaves_only=True) enumerates tensor leaves - # only — non-tensor leaves would silently fall out of the rebuilt dict. - new_dict: dict[str, Any] = {} - changed = False - for k in td.keys(): - v = td.get(k) - field_name = str(k) - if field_name in PROMOTE_1D_FIELDS: - if not isinstance(v, torch.Tensor) or v.is_nested or v.dim() != 1: - shape = tuple(v.shape) if isinstance(v, torch.Tensor) else None - raise ValueError( - f"Mooncake scalar field {field_name!r} must be a dense " - f"1D tensor with shape (N,), got {type(v).__name__} " - f"with shape {shape}." - ) - new_dict[str(k)] = v.unsqueeze(-1).contiguous() - changed = True - elif isinstance(v, torch.Tensor) and not v.is_nested and v.dim() == 1: - raise ValueError( - f"Mooncake field {field_name!r} is a dense 1D tensor but is " - "not declared in data_plane.schema.PROMOTE_1D_FIELDS. Add " - "the field to the schema if it is a per-sample scalar." - ) - else: - new_dict[str(k)] = v - if not changed: - return td - new_td = TensorDict(new_dict, batch_size=td.batch_size) - _assert_no_key_loss(new_dict, new_td, "_promote_1d_leaves") - return new_td - - def _from_wire(td: TensorDict) -> TensorDict: - """Normalize TQ reads and invert :func:`_promote_1d_leaves` when needed. - - Both TQ v0.1.9 storage managers reconstruct every non-scalar field as a - nested tensor, including fields whose rows all have the same shape. - Densify those uniform nested tensors first so regular batched inputs retain - their dense representation. Truly ragged fields remain nested. Finally, - squeeze only singleton dimensions declared in - :data:`nemo_rl.data_plane.schema.PROMOTE_1D_FIELDS`. + """Densify uniform nested tensors coming back from TQ. + + Both storage managers reconstruct every non-scalar field as a nested + tensor, including fields whose rows all share a shape. Densify those so + regular batched inputs retain their dense representation; truly ragged + fields stay nested. + + Per-sample scalar columns need no handling here: with + :func:`_patch_scalar_field_schema` applied they are stored and reported + as 0-d rows, which ``_merge_tensors_to_tensordict`` stacks into a dense + ``(N,)`` column before it ever reaches this function. + + Packed multimodal fields are excluded: their rows are per-sample media, + not a padded sequence, and "all rows share a shape" is a data-dependent + accident (every sample happening to carry one image). Stacking them + discards the row boundaries that ``PackedTensor.from_wire`` needs, and + the dense value then fails the ``is_nested`` check in + ``codec.materialize`` and reaches ``get_multimodal_dict`` unreassembled. + ``codec.materialize`` applies the same exclusion. """ - # Same top-level iteration as `_promote_1d_leaves`: NonTensorData / - # NonTensorStack leaves are only visible via td.keys(), not leaves_only. + # NonTensorData / NonTensorStack leaves are only visible via td.keys(), + # not keys(leaves_only=True) — iterating the latter would silently drop + # them from the rebuilt dict. + # Deferred: ``multimodal_utils`` pulls PIL, requests and a few hundred + # transformers submodules, and this adapter is imported by every process + # that constructs a TQ client. ``codec.materialize`` defers the same import + # for the same reason. + from nemo_rl.data.multimodal_utils import PACKED_MULTIMODAL_FIELDS + new_dict: dict[str, Any] = {} changed = False for k in td.keys(): v = td.get(k) field_name = str(k) - if isinstance(v, torch.Tensor) and v.is_nested: + if ( + isinstance(v, torch.Tensor) + and v.is_nested + and field_name not in PACKED_MULTIMODAL_FIELDS + ): rows = list(v.unbind()) if rows and all(row.shape == rows[0].shape for row in rows[1:]): v = torch.stack(rows) changed = True - if field_name in PROMOTE_1D_FIELDS: - if not isinstance(v, torch.Tensor) or v.is_nested: - raise ValueError( - f"Mooncake scalar field {field_name!r} could not be " - "restored as a dense tensor." - ) - if v.dim() == 1: - new_dict[field_name] = v - elif v.dim() == 2 and v.shape[-1] == 1: - new_dict[field_name] = v.squeeze(-1).contiguous() - changed = True - else: - raise ValueError( - f"Mooncake scalar field {field_name!r} must decode as " - f"(N,) or (N, 1), got shape {tuple(v.shape)}." - ) - else: - new_dict[field_name] = v + new_dict[field_name] = v if not changed: return td new_td = TensorDict(new_dict, batch_size=td.batch_size) @@ -723,14 +807,15 @@ def __init__(self, cfg: DataPlaneConfig, *, bootstrap: bool = True) -> None: if mooncake_cfg.reuse_registered_buffers: _patch_mooncake_staging_buffers(mooncake_cfg.staging_buffer_size) - # Workaround for TQ KVStorageManager's 1D-field schema/data - # mismatch (only `mooncake_cpu` goes through that path; `simple` - # is unaffected). Writer unsqueezes 1D → (N, 1) on put; reader - # squeezes the trailing 1 back on get. Drop when upstream TQ - # unifies the schema/data shapes for 1D fields. self._backend = cfg["backend"] self._supports_checkpointing = data_plane_supports_checkpointing(cfg) - self._promote_1d = cfg["backend"] == "mooncake_cpu" + # Fix TQ's 1-D field schema at the source rather than reshaping the + # payload around it: the schema now reports the ``()`` sample shape + # the stored rows actually have. Applied on every backend and in + # every process that builds a client, before ``_init_tq`` / + # ``_connect_existing``, so no put can land under the old schema. + # Self-verifying — see :func:`_assert_tq_stores_scalar_rows_0d`. + _patch_scalar_field_schema() if bootstrap: _init_tq(cfg) @@ -943,8 +1028,6 @@ def put_samples( TensorDict, fields.detach(), # type: ignore[missing-argument] ) - if self._promote_1d: - detached_fields = _promote_1d_leaves(detached_fields) wire_fields = detached_fields field_names = [str(key) for key in detached_fields.keys()] diff --git a/nemo_rl/data_plane/codec.py b/nemo_rl/data_plane/codec.py index 919bf96a13d..1001869fedc 100644 --- a/nemo_rl/data_plane/codec.py +++ b/nemo_rl/data_plane/codec.py @@ -177,7 +177,10 @@ def pack_jagged_fields( else: raise TypeError( f"pack_jagged_fields: unsupported value type for {k!r}: {type(v)}. " - "Use torch.Tensor or np.ndarray(dtype=object)." + "Use torch.Tensor or np.ndarray(dtype=object). PackedTensor " + "must be converted to torch.nested at the wire boundary " + "(see sync_rollout_actor.py) so the codec's dispatch stays " + "binary." ) return TensorDict(packed, batch_size=[n]) @@ -252,16 +255,17 @@ def materialize( layout: Layout = "padded", pad_value_dict: dict[str, int | float] | None = None, pad_to_seqlen: int = 0, + tags: list[dict[str, Any]] | None = None, ) -> "BatchedDataDict[Any]": """Convert a wire TensorDict to a BatchedDataDict. Trainer/worker code expects rectangular tensors — this is the bridge from the on-wire nested format. - The lazy ``BatchedDataDict`` import keeps + The lazy ``BatchedDataDict`` / ``multimodal_utils`` imports keep ``import nemo_rl.data_plane`` cheap for unit tests that don't - actually call this function (``BatchedDataDict`` transitively - pulls multimodal deps like torchvision / torchaudio). + actually call this function — both transitively pull PIL, + ``requests`` and a few hundred ``transformers`` submodules. Args: td: Wire TensorDict to materialize. @@ -286,6 +290,10 @@ def materialize( """ from tensordict import NonTensorData, NonTensorStack + from nemo_rl.data.multimodal_utils import ( + PACKED_MULTIMODAL_FIELDS, + reassemble_packed_multimodal, + ) from nemo_rl.distributed.batched_data_dict import BatchedDataDict pads = pad_value_dict or {} @@ -314,7 +322,15 @@ def materialize( f"materialize() received unexpected leaf type for {key!r}: " f"{type(val)}. Expected Tensor or NonTensorStack." ) - if val.is_nested and layout == "padded": + # Multimodal packed fields stay nested. Their rows are per-sample + # media, not a padded sequence: ``to_padded_tensor`` would + # rectangularize to ``[B, max_rows, ...]`` and the row boundaries + # would then have to be recovered from a companion field. TQ + # already stores one entry per row and reassembles them + # (``extract_field_schema`` records ``per_sample_shapes``), so the + # nested value handed back here carries the true per-row shapes — + # ``PackedTensor.from_wire`` just unbinds it. + if val.is_nested and layout == "padded" and key not in PACKED_MULTIMODAL_FIELDS: pad = pads.get(key, 0) padded = torch.nested.to_padded_tensor(val, padding=pad) else: @@ -325,11 +341,21 @@ def materialize( # right-padded output) ride the ``else`` branch above, so without # this they'd skip the cross-DP forward pad target and break the # microbatch iterator (truncate_tensors → narrow length>size). + # + # ``not padded.is_nested`` must be tested BEFORE ``shape[1]``: a + # nested tensor's dim 1 is ragged, and comparing it raises + # ``ValueError: ge: relation is indeterminate``. Anything still + # nested here was deliberately left so (multimodal packed fields, + # which skip ``to_padded_tensor`` above) and must not be padded + # anyway — their dim 1 is patch/image count, not seqlen, so + # extending it to a token seqlen inflates pixel_values ~40x. if ( pad_to_seqlen > 0 and isinstance(padded, torch.Tensor) + and not padded.is_nested and padded.dim() >= 2 and padded.shape[1] < pad_to_seqlen + and key not in PACKED_MULTIMODAL_FIELDS ): pad_spec = [0, 0] * (padded.dim() - 2) + [ 0, @@ -337,4 +363,8 @@ def materialize( ] padded = torch.nn.functional.pad(padded, pad_spec, value=pad) out[key] = padded + # Packed multimodal fields arrive flattened with a companion shape column. + # The multimodal layer owns that encoding; the codec only asks it to fix up + # its own fields so no raw nested value reaches a BatchedDataDict consumer. + reassemble_packed_multimodal(out, tags) return BatchedDataDict(out) diff --git a/nemo_rl/data_plane/column_io.py b/nemo_rl/data_plane/column_io.py index f9c49790fdf..2836977dc95 100644 --- a/nemo_rl/data_plane/column_io.py +++ b/nemo_rl/data_plane/column_io.py @@ -35,6 +35,7 @@ import torch from nemo_rl.data.llm_message_utils import attach_message_log_view +from nemo_rl.data.multimodal_utils import PER_TOKEN_MULTIMODAL_FIELDS from nemo_rl.data_plane.codec import materialize, pack_jagged_fields from nemo_rl.data_plane.interfaces import DataPlaneClient, KVBatchMeta from nemo_rl.data_plane.schema import ( @@ -45,7 +46,11 @@ ) from nemo_rl.distributed.batched_data_dict import BatchedDataDict -TOKEN_ALIGNED_FIELDS = frozenset( +# Fields the codec packs jagged via ``pack_per_token_field``. Rest go +# through ``.detach().contiguous()``. Add per-token fields here — no +# separate structural check is needed because the codec's binary +# dispatch (Tensor | ndarray[object]) errors loudly on unknown types. +_TEXT_TOKEN_ALIGNED_FIELDS = frozenset( { "input_ids", "generation_logprobs", @@ -63,6 +68,11 @@ } ) +# Per-token multimodal type maps are sequence-aligned too, so they pack the +# same way. Unioned from the registry rather than re-listed, so a new +# per-token modality cannot be added there and silently forgotten here. +TOKEN_ALIGNED_FIELDS = _TEXT_TOKEN_ALIGNED_FIELDS | PER_TOKEN_MULTIMODAL_FIELDS + def round_up(value: int, multiple: int) -> int: """Smallest ``multiple``-aligned int ≥ ``value`` (no-op when ``multiple <= 1``).""" @@ -108,6 +118,7 @@ def read_columns( layout=layout, pad_value_dict=pad_value_dict, pad_to_seqlen=pad_to_seqlen, + tags=meta.tags, ) attach_message_log_view(data) return data @@ -197,12 +208,26 @@ def kv_first_write( f"kv_first_write: tags ({len(tags)}) must match batch size ({n})" ) lengths = final_batch_cpu["input_lengths"] - fields: dict[str, torch.Tensor | np.ndarray] = { - k: v - for k, v in final_batch_cpu.items() - if isinstance(v, torch.Tensor) - or (isinstance(v, np.ndarray) and v.dtype == object) - } + # Binary wire dispatch: only ``torch.Tensor`` (including + # ``torch.nested``) and ``np.ndarray[object]`` cross the codec. + # Raise loudly on anything else instead of silently dropping — + # that's the class of silent-drop bug that let PackedTensor + # pixel_values disappear from the TQ pipe pre-Option B. + fields: dict[str, torch.Tensor | np.ndarray] = {} + for k, v in final_batch_cpu.items(): + if isinstance(v, torch.Tensor) or ( + isinstance(v, np.ndarray) and v.dtype == object + ): + fields[k] = v + else: + raise TypeError( + f"Field {k!r}: unexpected wire type {type(v).__name__}. " + "Only torch.Tensor (incl. torch.nested) and " + "np.ndarray[object] cross the wire boundary. Convert " + "domain-layer wrappers (e.g. PackedTensor via " + "torch.nested.as_nested_tensor) in the rollout actor " + "before calling kv_first_write." + ) td = pack_jagged_fields( fields, lengths=lengths, diff --git a/nemo_rl/data_plane/driver_mixin.py b/nemo_rl/data_plane/driver_mixin.py index 2dd6df0feed..9e08f2b0a26 100644 --- a/nemo_rl/data_plane/driver_mixin.py +++ b/nemo_rl/data_plane/driver_mixin.py @@ -18,6 +18,7 @@ from dataclasses import replace from typing import Any, Optional +from nemo_rl.data.multimodal_utils import present_multimodal_fields from nemo_rl.data_plane.column_io import read_columns, round_up, write_columns from nemo_rl.data_plane.interfaces import KVBatchMeta from nemo_rl.data_plane.schema import GLOBAL_FORWARD_PAD_SEQLEN @@ -84,12 +85,21 @@ def _isolated_meta( whichever model dispatches first decide the forward pad for the rest -- and with ppo_epochs > 1 the caller's meta is already stamped when the critic dispatches again. + + ``fields`` is the dispatch's *static* column list; the multimodal + columns the rollout actually wrote are unioned in here rather than at + each call site. Every forward-running dispatch needs them -- omitting + them runs a VLM forward image-blind while the sibling dispatches saw + the images -- and nothing in a static list forces the next caller to + remember. The union is a no-op for text-only runs, since + ``present_multimodal_fields`` intersects with ``meta.fields``. """ extra_info = dict(meta.extra_info) extra_info.pop(GLOBAL_FORWARD_PAD_SEQLEN, None) + multimodal = [f for f in present_multimodal_fields(meta) if f not in fields] isolated = replace( meta, - fields=fields, + fields=[*fields, *multimodal], task_name=task_name, extra_info=extra_info, ) diff --git a/nemo_rl/data_plane/factory.py b/nemo_rl/data_plane/factory.py index 08f454442b1..728a5bb2ee9 100644 --- a/nemo_rl/data_plane/factory.py +++ b/nemo_rl/data_plane/factory.py @@ -15,8 +15,64 @@ from __future__ import annotations +from typing import TYPE_CHECKING, Any, Callable, Optional + from nemo_rl.data_plane.interfaces import DataPlaneClient, DataPlaneConfig +if TYPE_CHECKING: + from nemo_rl.algorithms.grpo import MasterConfig + + +def data_plane_enabled(cfg: DataPlaneConfig | None) -> bool: + """Whether the data plane is on. ``None`` (key absent) means off.""" + return cfg is not None and bool(cfg.get("enabled", False)) + + +def select_sync_trainer( + master_config: "MasterConfig", *, label: str = "GRPO" +) -> Callable[..., Any]: + """Pick the synchronous trainer based on ``data_plane.enabled``. + + Shared by every launcher so trainer choice cannot drift between them. + Pairs with :func:`make_policy_factory`: turning the data plane on means + picking *both* the TQ-mediated trainer and a ``TQPolicy``, and a launcher + that picked only one would fail after a full model load. + + Args: + master_config: The resolved config; only ``data_plane`` is read. + label: Algorithm name for the progress line (e.g. ``"VLM GRPO"``). + """ + if data_plane_enabled(master_config.data_plane): + from nemo_rl.algorithms.grpo_sync import grpo_train_sync + + print(f"🚀 Running synchronous {label} training (TransferQueue)") + return grpo_train_sync + + from nemo_rl.algorithms.grpo import grpo_train + + print(f"🚀 Running synchronous {label} training (legacy)") + return grpo_train + + +def make_policy_factory( + cfg: DataPlaneConfig | None, +) -> Optional[Callable[..., Any]]: + """The ``policy_factory`` for ``setup()``, or ``None`` for a plain ``Policy``. + + Lives at the launcher level so the legacy trainer stays data-plane-agnostic + (architectural invariant — see + ``tests/unit/data_plane/test_architecture_invariants.py``). + """ + if not data_plane_enabled(cfg): + return None + + from nemo_rl.models.policy.tq_policy import TQPolicy + + def _make_policy(**kwargs: Any) -> TQPolicy: + return TQPolicy(**kwargs, dp_cfg=cfg) + + return _make_policy + def maybe_configure_data_plane_env(cfg: DataPlaneConfig | None) -> None: """Set backend env vars that must be identical in every process. diff --git a/nemo_rl/data_plane/preshard.py b/nemo_rl/data_plane/preshard.py index f9ce2fdc6c7..5fd196b51bc 100644 --- a/nemo_rl/data_plane/preshard.py +++ b/nemo_rl/data_plane/preshard.py @@ -23,6 +23,7 @@ from __future__ import annotations +from dataclasses import replace from typing import Any, Optional import torch @@ -140,8 +141,6 @@ def shard_meta_for_dp( # pyrefly: ignore # no-matching-overload idx_list: list[int] = shard[META_IDX].tolist() flat_idx.extend(idx_list) - rank_sample_ids = [meta.sample_ids[i] for i in idx_list] - rank_seqlens = [seq_lens[i] for i in idx_list] rank_extra = dict(base_extra) # Per-shard packing metadata — set by ``shard_by_batch_size`` when # sequence_packing or dynamic_batching is enabled. Workers' @@ -158,16 +157,11 @@ def shard_meta_for_dp( val = getattr(shard, attr, None) if val is not None: rank_extra[attr] = val - out.append( - KVBatchMeta( - partition_id=meta.partition_id, - task_name=meta.task_name, - sample_ids=rank_sample_ids, - fields=meta.fields, - sequence_lengths=rank_seqlens, - extra_info=rank_extra, - ) - ) + # ``subset`` owns per-sample projection: it slices ``sample_ids``, + # ``sequence_lengths`` and ``tags`` together, so a sidecar added to + # ``KVBatchMeta`` later cannot go missing on the presharded path. + # Only ``extra_info`` is per-shard rather than per-sample. + out.append(replace(meta.subset(idx_list), extra_info=rank_extra)) # Build inverse permutation: unsorted[orig_idx] = position_in_aggregated. # When workers' results are concatenated in DP-rank order, row `j` of diff --git a/nemo_rl/data_plane/schema.py b/nemo_rl/data_plane/schema.py index 78480dc7948..4722b9ae10e 100644 --- a/nemo_rl/data_plane/schema.py +++ b/nemo_rl/data_plane/schema.py @@ -74,7 +74,7 @@ MALFORMED_THINKING_MASK, ) -# Subset fetched by logprob / ref-logprob workers. +# Core fields the logprob workers require; multimodal extras added by TQPolicy._logprob_dispatch. LP_SEED_FIELDS = ( "input_ids", "input_lengths", @@ -106,34 +106,13 @@ # calibration only handles seq-dim tensor inputs, so we name them # explicitly. Train-side deltas (logprobs/advantages/masks) and # wire-only message-log bulk fields are skipped by virtue of not being -# in this list. ``multi_modal_inputs`` covers VLM extras (pixel values, -# grid metadata, etc.) when present; it's harmlessly absent for -# text-only models so the filter skips it on those. -DP_CALIB_INPUT_FIELDS = (INPUT_IDS, INPUT_LENGTHS, "multi_modal_inputs") +# in this list. VLM extras are not named here — they are per-batch, so +# callers add the ones actually present via +# ``multimodal_utils.present_multimodal_fields``. +DP_CALIB_INPUT_FIELDS = (INPUT_IDS, INPUT_LENGTHS) ROUTED_EXPERTS_FIELD = "routed_experts" -# Per-sample 1D scalar fields. The TQ adapter promotes these to ``(N, 1)`` -# on write to work around TQ v0.1.9's KVStorageManager schema/data mismatch on -# the Mooncake backend, and squeezes them back to ``(N,)`` on read. This is the -# authoritative user-level schema; no per-row shape metadata is carried. -# -# Fields listed here must be dense ``(N,)`` tensors when written through the -# Mooncake adapter. Dense 1D fields not listed here are rejected on that path so -# a new field cannot silently reintroduce the upstream shape mismatch. -# -# Delete this set and the corresponding adapter transforms when upstream TQ -# fixes 1D field schema extraction. -PROMOTE_1D_FIELDS: frozenset[str] = frozenset( - { - INPUT_LENGTHS, - MASK_SAMPLE, - "total_reward", - SAMPLE_MASK, - TRUNCATED, - } -) - def fields_with_optional_routed_experts( fields: Sequence[str], diff --git a/nemo_rl/data_plane/worker_mixin.py b/nemo_rl/data_plane/worker_mixin.py index e545ae48370..def864a957e 100644 --- a/nemo_rl/data_plane/worker_mixin.py +++ b/nemo_rl/data_plane/worker_mixin.py @@ -30,11 +30,13 @@ from typing import TYPE_CHECKING, Any, Literal, Optional +import numpy as np import torch FetchPolicy = Literal["auto", "independent", "leader_broadcast"] from nemo_rl.data.llm_message_utils import attach_message_log_view +from nemo_rl.data.multimodal_utils import PackedTensor from nemo_rl.data_plane.schema import ( ELEM_COUNTS_PER_GB, GLOBAL_FORWARD_PAD_SEQLEN, @@ -72,16 +74,69 @@ def _broadcast_batched_data_dict( backend = torch.distributed.get_backend(group) bcast_device: Any = torch.cuda.current_device() if backend == "nccl" else "cpu" + # Leader-only: the flat payload of each packed field, kept from the + # descriptor pass so ``to_wire``'s ``torch.cat`` of the whole column runs + # once, not once per pass (multimodal_utils.py:203 warns about exactly this). + leader_flat: dict[str, torch.Tensor] = {} + if is_leader: assert data is not None, "leader must provide non-None data" descriptor: list[Any] = [] + empty_packed: list[str] = [] for k, v in data.items(): if isinstance(v, torch.Tensor): descriptor.append( (k, "tensor", str(v.dtype), tuple(v.shape), str(v.device)) ) - else: + elif isinstance(v, PackedTensor): + # A PackedTensor on the ``raw`` branch would be pickled by + # ``broadcast_object_list`` into one contiguous *device* tensor + # -- 26.25 GiB for a VLM batch (job 17648488). Ship it the way + # the wire already does: payload as one flat tensor, geometry as + # ints. Densifying instead loses the row boundaries the model + # needs to match media to placeholders (job 17652563). + nested, shapes = v.to_wire() + if nested is None: + # Every row empty. ``reassemble_packed_multimodal`` drops + # the key in this case; do the same on both sides so the + # ranks agree. + empty_packed.append(k) + continue + values = nested.values() + leader_flat[k] = values + descriptor.append( + ( + k, + "packed_wire", + str(values.dtype), + str(values.device), + nested.offsets().tolist(), + shapes, + v.pad_to_max_shape, + ) + ) + elif ( + v is None + or isinstance(v, (str, int, float, bool)) + or (isinstance(v, np.ndarray) and v.dtype == object) + ): + # Scalars and object arrays are what the raw branch is for: + # small, and cheap to pickle into the object list. descriptor.append((k, "raw", v)) + else: + # Mirrors the write-side gate in ``column_io.kv_first_write``: + # an unknown wrapper on the ``raw`` branch is pickled into + # device memory by ``broadcast_object_list``, which is how the + # PackedTensor payload became a 26 GiB allocation. Fail at the + # boundary instead of discovering it as an OOM. + raise TypeError( + f"Field {k!r}: unexpected broadcast type {type(v).__name__}. " + "The replica-group broadcast carries torch.Tensor, " + "PackedTensor, np.ndarray[object] and scalars; a bulk " + "wrapper must get its own branch rather than being pickled." + ) + for k in empty_packed: + del data[k] payload: list[Any] = [descriptor] else: payload = [None] @@ -114,6 +169,23 @@ def _broadcast_batched_data_dict( and torch.device(src_device).type != torch.device(bcast_device).type ): out[key] = tensor.to(src_device) + elif kind == "packed_wire": + dtype_str, src_device, offsets, shapes, pad_to_max_shape = entry[2:] + if is_leader: + flat = leader_flat[key].to(bcast_device) + else: + dtype = getattr(torch, dtype_str.split(".")[-1]) + flat = torch.empty(offsets[-1], dtype=dtype, device=bcast_device) + torch.distributed.broadcast(flat, src=src, group=group) + if not is_leader: + nested = torch.nested.nested_tensor_from_jagged( + flat, torch.tensor(offsets, dtype=torch.int64, device=flat.device) + ) + if torch.device(src_device).type != torch.device(bcast_device).type: + nested = nested.to(src_device) + out[key] = PackedTensor.from_wire( + nested, shapes, pad_to_max_shape=pad_to_max_shape + ) else: if not is_leader: out[key] = entry[2] @@ -135,12 +207,15 @@ def setup_data_plane(self, cfg: DataPlaneConfig) -> None: Called once by the driver after worker construction. Idempotent. """ - if getattr(self, "model_slices_context_parallel_inputs", False): - raise NotImplementedError( - "TransferQueue/SingleController does not yet support models that " - "insert media before context-parallel input selection. Use the " - "synchronous NeMo-RL policy path for Nemotron Omni." - ) + # Models that insert media before CP input selection + # (``model_slices_context_parallel_inputs``) need the caller to hand + # them full, unsliced THD rows. That is what they get: ``_fetch`` + # leader-fetches one DP slice and NCCL-broadcasts it across the + # replica group, which is TP x CP x PP siblings of a single DP rank, + # so every CP sibling sees identical full rows and the model applies + # its own post-embedding slice. The presharded entrypoints then + # delegate to the same ``train`` / ``get_logprobs`` that carry the flag + # into ``models/megatron/data.py``. if self._dp_client is not None: return from nemo_rl.data_plane import build_data_plane_client @@ -247,6 +322,7 @@ def _fetch( layout=layout, pad_value_dict=pad_value_dict, pad_to_seqlen=pad_to_seqlen, + tags=meta.tags, ) else: data = None @@ -278,6 +354,7 @@ def _fetch( layout=layout, pad_value_dict=pad_value_dict, pad_to_seqlen=pad_to_seqlen, + tags=meta.tags, ) attach_message_log_view(data) trace_tq_fetch_payload( diff --git a/nemo_rl/distributed/batched_data_dict.py b/nemo_rl/distributed/batched_data_dict.py index 57ef895f33c..8269ee18c35 100644 --- a/nemo_rl/distributed/batched_data_dict.py +++ b/nemo_rl/distributed/batched_data_dict.py @@ -34,6 +34,8 @@ from nemo_rl.data.multimodal_utils import ( MULTIMODAL_CONTENT_TYPES, NATIVE_MULTIMODAL_KEYS, + PACKED_MULTIMODAL_FIELDS, + PER_TOKEN_MULTIMODAL_FIELDS, PackedTensor, ) from nemo_rl.data.packing import get_packer @@ -125,12 +127,6 @@ class DynamicBatchingArgs(TypedDict): class BatchedDataDict(UserDict, Generic[DictT]): _PIXEL_DTYPE_CAST_KEYS = frozenset({"pixel_values", "pixel_values_videos"}) - # keys that are model specific, but not part of the PackedTensor - ADDITIONAL_OPTIONAL_KEY_TENSORS = [ - "token_type_ids", # specific to gemma3 that tells where the image tokens are in the sequence, not required for llm-only inference/training - "mm_token_type_ids", # specific to qwen2.5-vl (transformers>=5.3): tells model which tokens are text(0)/image(1)/video(2) for 3D RoPE position encoding - ] - def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) @@ -144,7 +140,15 @@ def get_multimodal_dict( device: Optional[torch.device] = None, pixel_dtype: Optional[torch.dtype] = None, ) -> dict[str, Any]: - """Return a regular dict of tensors or packed multimodal data items. + """Return the multimodal fields as a dict. + + Four cases per (k, v): + * ``PackedTensor`` — in-memory form, keep as-is. + * ``k`` in ``PACKED_MULTIMODAL_FIELDS`` — data-plane wire form + (a nested tensor), reassemble via ``PackedTensor.from_wire``. + * ``k`` in ``PER_TOKEN_MULTIMODAL_FIELDS`` — plain per-token + tensor, keep as-is. + * anything else — not multimodal, skip. ``pixel_dtype`` converts pixel tensors without materializing repeated logical segments. This is used to reduce policy-bound Ray payloads. @@ -176,16 +180,36 @@ def get_multimodal_dict( f"{metadata_counts}." ) - multimodal_dict = {} + result: dict[str, Any] = {} for k, v in self.data.items(): if isinstance(v, PackedTensor): + # In-memory PackedTensor (or a per-token field a caller + # happened to wrap; matches the pre-refactor behavior of + # unwrapping via as_tensor). if pixel_dtype is not None and k in self._PIXEL_DTYPE_CAST_KEYS: v = v.to_dtype(pixel_dtype) - multimodal_dict[k] = v.as_tensor(device=device) if as_tensors else v - elif k in self.ADDITIONAL_OPTIONAL_KEY_TENSORS: - multimodal_dict[k] = v - - return multimodal_dict + result[k] = v.as_tensor(device=device) if as_tensors else v + elif k in PER_TOKEN_MULTIMODAL_FIELDS: + # Plain per-token tensor: emit as-is. + result[k] = v + elif k in PACKED_MULTIMODAL_FIELDS: + # Data-plane wire form: a value that reached here without + # ``codec.materialize`` reassembling it. Purely a fail-loud + # guard -- never a reconstruction path, because neither case is + # recoverable from here. A *nested* value still needs the + # per-segment shapes, which live only on ``KVBatchMeta.tags``; + # taking its flat rows as-is would emit 1-D pixels and train + # image-blind. A *dense* value means the field was padded and + # the row boundaries are already gone. + raise ValueError( + f"{k!r} is still in data-plane wire form " + f"({'nested' if getattr(v, 'is_nested', False) else 'dense'}). " + "Packed multimodal fields must be rebuilt by " + "multimodal_utils.reassemble_packed_multimodal (which " + "codec.materialize calls) before get_multimodal_dict." + ) + # else: not a multimodal field, silently skip. + return result @classmethod def from_batches( @@ -951,6 +975,15 @@ def repeat_interleave( def truncate_tensors(self, dim: int, truncated_len: int): """Truncates tensors in this dict of a given dim to a given length.""" for k, v in self.items(): + # Packed multimodal fields are not sequence-aligned — their + # dim 1 is patch/image count — so narrowing them to the token + # seqlen silently corrupts images (or raises when the patch + # count is smaller than the seqlen). The in-memory + # ``PackedTensor`` form is skipped by ``torch.is_tensor`` + # below, but the data-plane wire form is a nested tensor, so + # name it here. + if k in PACKED_MULTIMODAL_FIELDS: + continue if torch.is_tensor(v) and len(v.shape) >= dim + 1: self.data[k] = torch.narrow(v, dim=dim, start=0, length=truncated_len) diff --git a/nemo_rl/experience/sync_rollout_actor.py b/nemo_rl/experience/sync_rollout_actor.py index 9f7f7e9ea03..48bf096d530 100644 --- a/nemo_rl/experience/sync_rollout_actor.py +++ b/nemo_rl/experience/sync_rollout_actor.py @@ -43,6 +43,10 @@ import ray import torch +from nemo_rl.data.multimodal_utils import ( + encode_multimodal_for_wire, + multimodal_row_tags, +) from nemo_rl.data_plane.column_io import kv_first_write from nemo_rl.data_plane.interfaces import KVBatchMeta from nemo_rl.data_plane.schema import ROUTED_EXPERTS_FIELD @@ -211,6 +215,8 @@ def rollout_to_tq( …) — stays on the driver, never crosses an actor boundary. """ # Lazy imports keep rollout-specific dependencies off the actor startup path. + # ``_policy_dtype`` sizes the VLM pixel tensors below. + from nemo_rl.algorithms.grpo import _policy_dtype from nemo_rl.algorithms.utils import get_gdpo_reward_component_keys from nemo_rl.data.llm_message_utils import ( MESSAGE_LOG_BULK_FIELDS, @@ -315,9 +321,19 @@ def rollout_to_tq( ) if ROUTED_EXPERTS_FIELD in flat: bulk_batch[ROUTED_EXPERTS_FIELD] = flat[ROUTED_EXPERTS_FIELD] - for k, v in flat.get_multimodal_dict(as_tensors=False).items(): - if isinstance(v, torch.Tensor): - bulk_batch[k] = v + # ``pixel_dtype`` mirrors the legacy analogs (``grpo._build_async_grpo_train_data`` + # and the sync train-data builders): cast pixels to the policy precision + # once here, at the same point they'd be cast in-memory. No worker + # re-applies it, so without this the largest column crosses the wire in + # fp32 where legacy shipped bf16. ``PackedTensor.to_dtype`` leaves + # integer segments (grid_thw / imgs_sizes / num_frames) untouched. + multimodal = flat.get_multimodal_dict( + as_tensors=False, pixel_dtype=_policy_dtype(cfg.policy) + ) + for k, v in multimodal.items(): + wire_value = encode_multimodal_for_wire(k, v) + if wire_value is not None: + bulk_batch[k] = wire_value # ``content`` (raw assistant text per sample) — rides TQ as a # NonTensorStack so the driver can fetch it back at jsonl time # (kv_first_write wraps it via NonTensorStack). @@ -404,6 +420,10 @@ def rollout_to_tq( dp_client=self._dp_client, partition_id=partition_id, extra_info={"rollout_metrics": rollout_metrics}, + # Per-row shapes the flattening removes from the payload. ``tags`` + # is the transport's per-sample channel and is projected with the + # rows, so no consumer re-keys them. + tags=multimodal_row_tags(multimodal, len(sample_ids)), task_name=partition_id, pad_to_multiple=int( cfg.policy.get("make_sequence_length_divisible_by") or 1 diff --git a/nemo_rl/models/automodel/data.py b/nemo_rl/models/automodel/data.py index 9f2170d61a0..dfd5422c222 100644 --- a/nemo_rl/models/automodel/data.py +++ b/nemo_rl/models/automodel/data.py @@ -25,6 +25,7 @@ from transformers import AutoTokenizer from nemo_rl.algorithms.loss.interfaces import LossFunction, LossType +from nemo_rl.data.multimodal_utils import PACKED_MULTIMODAL_FIELDS from nemo_rl.distributed.batched_data_dict import BatchedDataDict from nemo_rl.models.huggingface.common import ( get_flash_attention_kwargs, @@ -400,6 +401,15 @@ def check_sequence_dim( for k, v in data.items(): if k in skip_set: continue + # Multimodal fields are never sequence-aligned: dim 1 is + # num_images / num_patches. In-memory these ride as + # ``PackedTensor`` and are skipped by ``torch.is_tensor`` below, but + # the data-plane wire form is a nested tensor, so name it here. + # Mirrors ``megatron/data.py::get_and_validate_seqlen``; kept inside + # this helper rather than pushed onto ``skip_keys`` because all seven + # call sites need it and none of them should know the wire format. + if k in PACKED_MULTIMODAL_FIELDS: + continue if torch.is_tensor(v) and len(v.shape) > 1: assert v.shape[sequence_dim] == seq_dim_size, ( f"Dim 1 must be the sequence dim, expected dim 1={seq_dim_size} but got shape {v.shape}" diff --git a/nemo_rl/models/megatron/data.py b/nemo_rl/models/megatron/data.py index ff303b96a18..acf1173abd8 100644 --- a/nemo_rl/models/megatron/data.py +++ b/nemo_rl/models/megatron/data.py @@ -29,6 +29,7 @@ from megatron.core.utils import StragglerDetector from nemo_rl.algorithms.loss.interfaces import LossFunction, LossType +from nemo_rl.data.multimodal_utils import PACKED_MULTIMODAL_FIELDS from nemo_rl.distributed.batched_data_dict import BatchedDataDict from nemo_rl.distributed.model_utils import _get_tokens_on_this_cp_rank from nemo_rl.models.megatron.common import _round_up_to_multiple @@ -1471,10 +1472,14 @@ def _unpack_sequences_from_megatron( def get_and_validate_seqlen(data: BatchedDataDict[Any]): - # dim 1 is always assumed to be the sequence dim, sanity check this here + # dim 1 is always assumed to be the sequence dim, sanity check this here. + # Skip multimodal fields: their dim 1 is num_images / num_patches, not + # seqlen. sequence_dim = 1 seq_dim_size = data["input_ids"].shape[sequence_dim] for k, v in data.items(): + if k in PACKED_MULTIMODAL_FIELDS: + continue if torch.is_tensor(v) and len(v.shape) > 1: assert v.shape[sequence_dim] == seq_dim_size, ( f"Dim 1 must be the sequence dim, expected dim 1={seq_dim_size} but got shape {v.shape} for key {k}" diff --git a/nemo_rl/models/policy/tq_policy.py b/nemo_rl/models/policy/tq_policy.py index bd332e54195..f80ca61c5e5 100644 --- a/nemo_rl/models/policy/tq_policy.py +++ b/nemo_rl/models/policy/tq_policy.py @@ -216,14 +216,22 @@ def _logprob_dispatch( ) -> None: """Shared body of get_logprobs_from_meta / get_reference_policy_logprobs_from_meta. - Logprob workers need only LP_SEED_FIELDS — narrow the meta's - field list so ``_fetch`` doesn't pull rollout-only payload (e.g. - multimodal). The same shape is used for both prev_lp and ref_lp. - Workers compute the per-token tensor and commit it to TQ via the - leader-rank ``_write_back_result_field``; the Ray return is - always None, so this dispatcher just waits for completion. + Logprob workers fetch ``LP_SEED_FIELDS`` plus the multimodal + columns ``_isolated_meta`` unions in, so prev/ref logprobs see the + same model inputs as the training forward, which is narrowed through + the same helper. Narrowing the + meta's field list still keeps rollout-only payload (message-log + bulk, ``content``) in TQ. The same shape is used for both prev_lp + and ref_lp. Workers compute the per-token tensor and commit it to + TQ via the leader-rank ``_write_back_result_field``; the Ray + return is always None, so this dispatcher just waits for + completion. """ spa, dba = self._packing_args("logprob_mb_tokens") + # Narrow the fetch to LP_SEED_FIELDS + optional routed_experts under + # R3 replay. ``_isolated_meta`` unions in the multimodal columns the + # rollout wrote, for this dispatch and the training one alike, so the + # prev/ref logprobs and the training forward see identical model inputs. lp_meta = self._isolated_meta( meta, fields=fields_with_optional_routed_experts( @@ -331,10 +339,14 @@ def train_from_meta( # default ``DP_TRAIN_FIELDS``) must be in TQ before this call — written # by workers + driver delta-writes. Caller may narrow to drop columns # skipped this step (e.g. ``prev_logprobs`` under force_on_policy_ratio). + # The multimodal columns are per-batch, not part of the static schema, + # so ``_isolated_meta`` unions them in — without them a VLM training + # forward would run image-blind while the logprob forwards saw images. train_meta = self._isolated_meta( meta, fields=fields_with_optional_routed_experts( - train_fields, enabled=self._router_replay_enabled + train_fields, + enabled=self._router_replay_enabled, ), task_name="train", ) @@ -463,7 +475,8 @@ def train_microbatches_from_meta( train_meta = self._isolated_meta( meta, fields=fields_with_optional_routed_experts( - train_fields, enabled=self._router_replay_enabled + train_fields, + enabled=self._router_replay_enabled, ), task_name="train", ) diff --git a/nemo_rl/models/policy/workers/dtensor_policy_worker.py b/nemo_rl/models/policy/workers/dtensor_policy_worker.py index d460b2c6506..55d62ab88ea 100644 --- a/nemo_rl/models/policy/workers/dtensor_policy_worker.py +++ b/nemo_rl/models/policy/workers/dtensor_policy_worker.py @@ -65,7 +65,10 @@ distributed_vocab_topk, get_logprobs_from_vocab_parallel_logits, ) -from nemo_rl.models.automodel.data import filter_multimodal_kwargs_for_model +from nemo_rl.models.automodel.data import ( + check_sequence_dim, + filter_multimodal_kwargs_for_model, +) from nemo_rl.models.dtensor.parallelize import ( _parallelize_model, clip_grad_by_total_norm_, @@ -636,13 +639,9 @@ def train( "cross-tokenizer distillation requires dtensor_cfg._v2=True." ) # dim 1 is always assumed to be the sequence dim, sanity check this here. - sequence_dim = 1 - seq_dim_size = data.get("input_ids").shape[sequence_dim] - for k, v in data.items(): - if torch.is_tensor(v) and len(v.shape) > 1: - assert v.shape[sequence_dim] == seq_dim_size, ( - f"Dim 1 must be the sequence dim, expected dim 1={seq_dim_size} but got shape {v.shape}" - ) + # Shared with the v2 worker so the multimodal skip (packed wire + # fields are not sequence-aligned) lives in exactly one place. + sequence_dim, seq_dim_size = check_sequence_dim(data) if eval_mode: ctx: AbstractContextManager[Any] = torch.no_grad() @@ -1049,13 +1048,9 @@ def get_logprobs( logprob_chunk_size = self.cfg.get("logprob_chunk_size", None) # dim 1 is always assumed to be the sequence dim, sanity check this here - sequence_dim = 1 - seq_dim_size = data.get("input_ids").shape[sequence_dim] - for k, v in data.items(): - if torch.is_tensor(v) and len(v.shape) > 1: - assert v.shape[sequence_dim] == seq_dim_size, ( - f"Dim 1 must be the sequence dim, expected dim 1={seq_dim_size} but got shape {v.shape}" - ) + # Shared with the v2 worker so the multimodal skip (packed wire + # fields are not sequence-aligned) lives in exactly one place. + sequence_dim, seq_dim_size = check_sequence_dim(data) all_log_probs = [] self.model.eval() @@ -1352,13 +1347,9 @@ def get_logprobs( def score(self, data: BatchedDataDict) -> BatchedDataDict[ScoreOutputSpec]: global_batch_size = min(self.cfg["batch_size"], data.size) - sequence_dim = 1 - seq_dim_size = data.get("input_ids").shape[sequence_dim] - for k, v in data.items(): - if torch.is_tensor(v) and len(v.shape) > 1: - assert v.shape[sequence_dim] == seq_dim_size, ( - f"Dim 1 must be the sequence dim, expected dim 1={seq_dim_size} but got shape {v.shape}" - ) + # Shared with the v2 worker so the multimodal skip (packed wire + # fields are not sequence-aligned) lives in exactly one place. + sequence_dim, seq_dim_size = check_sequence_dim(data) self.model.eval() with unshard_fsdp2_model(self.model), torch.no_grad(): diff --git a/tests/test_suites/llm/common-tq.env b/tests/test_suites/llm/common-tq.env index 80d187fcf6e..fe4f654ab04 100644 --- a/tests/test_suites/llm/common-tq.env +++ b/tests/test_suites/llm/common-tq.env @@ -2,15 +2,17 @@ # Helper sourced by TQ wrapper scripts. Computes: # TQ_EXP_NAME — this wrapper's basename (used for log/ckpt dirs + wandb) # BASE_RECIPE — the underlying recipe name (this wrapper delegates to its .sh) -# Validates that the wrapper targets a supported algorithm (grpo/dapo/prorlv2). +# Validates that the wrapper targets a supported algorithm — grpo/dapo/prorlv2, +# with or without the VLM suite's ``vlm_`` prefix (tests/test_suites/vlm/common-tq.env +# is a symlink to this file, as common.env already is). # data_plane.* overrides live in the wrapper's YAML (defaults inheritance), # not here. set -euo pipefail TQ_EXP_NAME=$(basename "$0" .sh) BASE_RECIPE=$(echo "$TQ_EXP_NAME" | sed -E 's/-tq_(simple|mooncake)$//') -if [[ ! "$BASE_RECIPE" =~ ^(grpo|dapo|prorlv2)- ]]; then - echo "[ERROR] TQ coverage is only for grpo/dapo/prorlv2 recipes; got: $BASE_RECIPE" >&2 +if [[ ! "$BASE_RECIPE" =~ ^(vlm_)?(grpo|dapo|prorlv2)- ]]; then + echo "[ERROR] TQ coverage is only for (vlm_)grpo/dapo/prorlv2 recipes; got: $BASE_RECIPE" >&2 exit 1 fi diff --git a/tests/test_suites/nightly.txt b/tests/test_suites/nightly.txt index 42f19e86b56..32d690530af 100644 --- a/tests/test_suites/nightly.txt +++ b/tests/test_suites/nightly.txt @@ -63,6 +63,14 @@ tests/test_suites/vlm/vlm_grpo-nemotron-omni-30ba3b-mmpr-4n8g-megatron-tp8ep16.v # hangs in sample_tokens and stays disabled -- see disabled.txt. tests/test_suites/vlm/vlm_grpo-qwen3.5-35ba3b-geo3k-2n8g-automodel-ep16.sh +# TQ data-plane coverage for VLM: AutoModel on backend=simple, Megatron on +# backend=mooncake_cpu. Each wrapper delegates to its base recipe with +# data_plane.enabled=true and adds a gate the base cannot use: Qwen3.5 trains one +# inner step per rollout so probs_ratio is an exact identity; Nemotron-Omni +# trains several, so it gates token_mult_prob_error instead. +tests/test_suites/vlm/vlm_grpo-qwen3.5-35ba3b-geo3k-2n8g-automodel-ep16-tq_simple.sh +tests/test_suites/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-1n8g-megatron-tp8ep8.v1-tq_mooncake.sh + # Removing this until this issue is resolved: https://github.com/huggingface/transformers/issues/41190 # tests/test_suites/vlm/vlm_grpo-smolvlm2-2.2b-instruct-clevr-1n2g-dtensor2tp1.v2.sh diff --git a/tests/test_suites/vlm/common-tq.env b/tests/test_suites/vlm/common-tq.env new file mode 120000 index 00000000000..5b123ebdd3e --- /dev/null +++ b/tests/test_suites/vlm/common-tq.env @@ -0,0 +1 @@ +../llm/common-tq.env \ No newline at end of file diff --git a/tests/test_suites/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-1n8g-megatron-tp8ep8.v1-tq_mooncake.sh b/tests/test_suites/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-1n8g-megatron-tp8ep8.v1-tq_mooncake.sh new file mode 100755 index 00000000000..8b8754e62be --- /dev/null +++ b/tests/test_suites/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-1n8g-megatron-tp8ep8.v1-tq_mooncake.sh @@ -0,0 +1,30 @@ +#!/bin/bash +SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd) + +# ===== BEGIN CONFIG ===== +# Mirrors vlm_grpo-nemotron-omni-30ba3b-clevr-1n8g-megatron-tp8ep8.v1.sh (delegated +# base) except NUM_NODES -- see cluster.num_nodes in the matching yaml. +NUM_NODES=2 +GPUS_PER_NODE=8 +STEPS_PER_RUN=10 +MAX_STEPS=10 +NUM_RUNS=$(( (MAX_STEPS + STEPS_PER_RUN - 1) / STEPS_PER_RUN )) # Round up +NUM_MINUTES=120 +# ===== END CONFIG ===== + +source "$SCRIPT_DIR/common-tq.env" +# Run base script under this wrapper's identity (own log/ckpt dirs, wandb name). +# The matching TQ YAML inherits from .yaml and turns on data_plane. +export EXP_NAME="$TQ_EXP_NAME" +bash "$SCRIPT_DIR/$BASE_RECIPE.sh" "$@" + +# No TQ-specific gate: turning the data plane on must not change what the recipe +# is held to, so this wrapper passes exactly when the base recipe's own +# max(train/reward) > 0.5 passes. +# +# An earlier version added 'max(train/token_mult_prob_error) < 1.02', a bound +# taken from the automodel sibling (measured 1.0129-1.0155 there). It does not +# transfer to this recipe: the legacy control -- same node count, same sequence +# length, data_plane.enabled=false -- measured 1.035-1.063 across nine steps +# (job 17686235), so the gate would fail the no-data-plane path too. A check +# that the control cannot pass is testing the backend, not the data plane. diff --git a/tests/test_suites/vlm/vlm_grpo-qwen3.5-35ba3b-geo3k-2n8g-automodel-ep16-tq_simple.sh b/tests/test_suites/vlm/vlm_grpo-qwen3.5-35ba3b-geo3k-2n8g-automodel-ep16-tq_simple.sh new file mode 100755 index 00000000000..ed7c4f7da61 --- /dev/null +++ b/tests/test_suites/vlm/vlm_grpo-qwen3.5-35ba3b-geo3k-2n8g-automodel-ep16-tq_simple.sh @@ -0,0 +1,34 @@ +#!/bin/bash +SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd) + +# ===== BEGIN CONFIG ===== +# Mirrors vlm_grpo-qwen3.5-35ba3b-geo3k-2n8g-automodel-ep16.sh (delegated base). +NUM_NODES=2 +STEPS_PER_RUN=20 +MAX_STEPS=20 +NUM_RUNS=$(( (MAX_STEPS + STEPS_PER_RUN - 1) / STEPS_PER_RUN )) # Round up +NUM_MINUTES=240 +# ===== END CONFIG ===== + +source "$SCRIPT_DIR/common-tq.env" +# Run base script under this wrapper's identity (own log/ckpt dirs, wandb name). +# The matching TQ YAML inherits from .yaml and turns on data_plane. +export EXP_NAME="$TQ_EXP_NAME" +bash "$SCRIPT_DIR/$BASE_RECIPE.sh" "$@" + +# TQ-specific gate, on top of the base recipe's own reward check. +# +# This recipe trains one inner step per rollout (train_global_batch_size == +# num_prompts_per_step * num_generations_per_prompt), so the training forward +# runs on the very weights that produced prev_logprobs and the importance ratio +# is an identity. Any deviation means the two passes saw different data -- +# exactly what a data-plane defect looks like. Measured 1.000000 on all 20 +# steps across three wire formats. +# +# Deliberately NOT applied to recipes with >1 inner step: there the ratio +# measures policy drift, and its max swings 5.85-29.21 run to run on identical +# code (it does so on the non-data-plane path too). +source "$SCRIPT_DIR/common.env" +uv run tests/check_metrics.py "$JSON_METRICS" \ + 'max(data["train/probs_ratio_max"]) < 1.0001' \ + 'min(data["train/probs_ratio_min"]) > 0.9999' diff --git a/tests/unit/algorithms/test_grpo.py b/tests/unit/algorithms/test_grpo.py index 3fca2187900..f91beef3ddc 100644 --- a/tests/unit/algorithms/test_grpo.py +++ b/tests/unit/algorithms/test_grpo.py @@ -841,13 +841,35 @@ def test_multimodal_dedup_rejects_unqualified_transfer_paths( _validate_multimodal_dedup_capability(master_config) master_config.policy["generation"]["backend"] = "vllm" + + # Data plane + NeMo-Gym stays rejected: ``grpo_train_sync`` never calls + # ``attach_initial_nemo_gym_image_payloads``, so the run would silently + # train on the media a Gym dataset omits from ``extra_env_info``. master_config.data_plane = {"enabled": True} - with pytest.raises(NotImplementedError, match="data_plane.enabled=false"): + with patch("nemo_rl.algorithms.grpo.should_use_nemo_gym", return_value=True): + with pytest.raises(NotImplementedError, match="NeMo-Gym"): + _validate_multimodal_dedup_capability(master_config) + + # Data plane without Gym is supported: the wire format carries dedup + # (``PackedTensor.to_wire`` emits one row per *logical* row), and the Gym + # attach helper is itself gated on ``should_use_nemo_gym``. Rejecting this + # blocked every Nemotron-Omni recipe, since all of them set + # ``deduplicate_multimodal_data: true``. + with patch("nemo_rl.algorithms.grpo.should_use_nemo_gym", return_value=False): _validate_multimodal_dedup_capability(master_config) master_config.data_plane = {"enabled": False} _validate_multimodal_dedup_capability(master_config) + # And with dedup off, nothing is gated — the guard returns before it looks + # at the backend or at NeMo-Gym, so a text-only sync GRPO + Gym run is not + # blocked by a multimodal validator. + master_config.grpo.deduplicate_multimodal_data = False + master_config.policy["generation"]["backend"] = "sglang" + master_config.data_plane = {"enabled": True} + with patch("nemo_rl.algorithms.grpo.should_use_nemo_gym", return_value=True): + _validate_multimodal_dedup_capability(master_config) + def test_grpo_sync_seq_logprob_error_helper_accepts_dict_result(monkeypatch): from nemo_rl.algorithms import grpo_sync as grpo_sync_mod diff --git a/tests/unit/data/test_multimodal_dict.py b/tests/unit/data/test_multimodal_dict.py index 2544a4218c1..c88883316a2 100644 --- a/tests/unit/data/test_multimodal_dict.py +++ b/tests/unit/data/test_multimodal_dict.py @@ -19,7 +19,12 @@ from nemo_rl.data.llm_message_utils import batched_message_log_to_flat_message from nemo_rl.data.multimodal_utils import ( + PACKED_MULTIMODAL_FIELDS, + PER_TOKEN_MULTIMODAL_FIELDS, PackedTensor, + encode_multimodal_for_wire, + multimodal_row_tags, + reassemble_packed_multimodal, ) from nemo_rl.distributed.batched_data_dict import ( BatchedDataDict, @@ -111,6 +116,28 @@ def test_truncate_tensors_with_packed_data(): assert batch["image_features"].as_tensor().shape == (5, 6, 128, 4, 2, 2) +def test_truncate_tensors_skips_wire_form_multimodal(): + """Dynamic batching narrows dim 1 to the microbatch seqlen. The + data-plane wire form of a packed multimodal field has patch count on + dim 1, not seqlen — narrowing it would corrupt the images (or raise + when patches < seqlen).""" + batch = BatchedDataDict( + { + "input_ids": torch.arange(8).reshape(2, 4), + # [B, max_patches, feat] — 3 patches, fewer than seqlen=4. + "pixel_values": torch.randn(2, 3, 16), + # Per-token multimodal IS sequence-aligned and must truncate. + "mm_token_type_ids": torch.ones((2, 4), dtype=torch.long), + } + ) + + batch.truncate_tensors(dim=1, truncated_len=2) + + assert torch.equal(batch["input_ids"], torch.tensor([[0, 1], [4, 5]])) + assert batch["mm_token_type_ids"].shape == (2, 2) + assert batch["pixel_values"].shape == (2, 3, 16) + + def test_multiturn_rollout_with_packed_data(): """Test multiturn conversations with packed multimodal data.""" message_log_1 = [ @@ -640,3 +667,378 @@ def test_packedtensor_empty_legacy_rows_survive_copy_pickle_and_slice(): assert sum(value.logical_segment_counts_by_row()) == 0 assert not value.deduplication_enabled assert value.as_tensor() is None + + +def test_to_wire_emits_one_row_per_logical_row_under_dedup(): + """Deduplicated values map one logical row to several shared physical + segments, so the wire encoder must walk logical rows — iterating + ``tensors`` would emit the physical segment count as the batch size + and desync every downstream column.""" + seg_a = torch.ones(2, 3) + seg_b = 2 * torch.ones(4, 3) + # 3 logical rows over 2 physical segments: [a], [a, b], [b]. + packed = PackedTensor( + [seg_a, seg_b], + dim_to_pack=0, + _row_offsets=[0, 1, 3, 4], + _segment_indices=[0, 0, 1, 1], + ).enable_deduplication() + assert len(packed) == 3 + + nested, shapes = packed.to_wire() + # Rows are flattened, so the logical row count shows up as one shape entry + # per row; the per-row element counts follow the 2/6/4 row heights. + assert len(shapes) == 3 + assert [t.numel() for t in nested.unbind()] == [2 * 3, 6 * 3, 4 * 3] + assert torch.equal( + PackedTensor.from_wire(nested, shapes).as_tensor(), + packed.as_tensor(), + ) + + +def test_to_wire_does_not_pad_segments_before_concat_under_dedup(): + """A dedup row spanning segments of differing trailing dims. + + ``to_wire`` flattens each segment, so the per-row concat is 1-D and cannot + hit ``RuntimeError: Sizes of tensors must match except in dimension 0``. + The padding ``as_tensor`` needs is applied on the read side by + ``from_wire`` instead, so no padded bytes cross the wire. + """ + # One logical row referencing two segments: 2x4 and 4x2 spatial dims. + packed = PackedTensor( + [torch.ones(1, 3, 2, 4), 2 * torch.ones(1, 3, 4, 2)], + dim_to_pack=0, + pad_to_max_shape=True, + _row_offsets=[0, 2], + _segment_indices=[0, 1], + ) + assert len(packed) == 1 + expected = packed.as_tensor() + assert expected.shape == (2, 3, 4, 4) # padded to the batch max + + nested, shapes = packed.to_wire() + # Natural size: 1*3*2*4 + 1*3*4*2 = 48 elements, no padding materialized. + assert [t.numel() for t in nested.unbind()] == [48] + assert shapes == [[[1, 3, 2, 4], [1, 3, 4, 2]]] + + restored = PackedTensor.from_wire(nested, shapes, pad_to_max_shape=True).as_tensor() + assert torch.equal(restored, expected) + + +def test_from_wire_rejects_dense_input(): + """A dense value here means ``materialize`` padded the field, which + silently loses the row boundaries. Fail loud instead.""" + with pytest.raises(TypeError, match="expects the nested value"): + PackedTensor.from_wire(torch.zeros(3, 2, 4), []) + + +def test_from_wire_empty_rows_match_legacy_none_semantics(): + """An image-free shard must reconstruct as legacy does: ``as_tensor`` + returns ``None`` and the per-row counts are 0, not a ``(0, ...)`` + tensor with counts of 1. A zero-length row is how absence travels.""" + nested = torch.nested.as_nested_tensor( + [torch.zeros(0, 3, 4), torch.zeros(0, 3, 4)], layout=torch.jagged + ) + + # An empty row contributes no segments, so its shapes entry is empty too + # — what ``to_wire`` mints for an all-``None`` row. + restored = PackedTensor.from_wire(nested, [[], []]) + assert restored.as_tensor() is None + assert restored.logical_segment_counts_by_row() == [0, 0] + + +def test_to_wire_does_not_materialize_pad_to_max_shape(): + """Dynamic-resolution values travel at natural size. + + Flattening each row to 1-D makes ``torch.jagged`` total, so differing + trailing dims no longer force the write side to pad up to the batch max. + The padding ``as_tensor`` returns is reapplied on read from the carried + shapes, keeping the padded bytes out of the wire and out of TQ storage. + """ + # Same rank, different trailing dims — nemotron-omni style tiles. + first = torch.ones(1, 3, 2, 4) + second = 2 * torch.ones(2, 3, 4, 2) + packed = PackedTensor([first, second], dim_to_pack=0, pad_to_max_shape=True) + + nested, shapes = packed.to_wire() + rows = list(nested.unbind()) + # Natural sizes: 1*3*2*4=24 and 2*3*4*2=48. Padding to the batch max + # (3, 4, 4) would have cost 48 and 96 -- 3x the bytes for this batch. + assert [t.numel() for t in rows] == [24, 48] + assert shapes == [[[1, 3, 2, 4]], [[2, 3, 4, 2]]] + + # Padding is reapplied on read, reproducing the pre-wire as_tensor(). + assert torch.equal( + PackedTensor.from_wire(nested, shapes, pad_to_max_shape=True).as_tensor(), + packed.as_tensor(), + ) + + +def test_get_multimodal_dict_rejects_wire_form_field(): + """Either wire form reaching here is unrecoverable, so both fail loud. + + Dense means ``codec.materialize`` padded the field and the row boundaries + are gone; nested means the shapes companion on ``KVBatchMeta.tags`` was + never applied, and taking the flat rows as-is would train image-blind. + Neither is reconstructible from inside ``get_multimodal_dict``. + """ + dense = BatchedDataDict({"pixel_values": torch.zeros(2, 3, 4, 4)}) + with pytest.raises(ValueError, match=r"wire form \(dense\)"): + dense.get_multimodal_dict(as_tensors=False) + + nested_value, _ = PackedTensor( + [torch.ones(3, 4), torch.ones(1, 4)], dim_to_pack=0 + ).to_wire() + nested = BatchedDataDict({"pixel_values": nested_value}) + with pytest.raises(ValueError, match=r"wire form \(nested\)"): + nested.get_multimodal_dict(as_tensors=False) + + +def test_image_free_shard_emits_no_wire_field_and_reads_back_empty(): + """An all-empty packed field never reaches the wire, end to end. + + Replaces an earlier "0-row wire field" guard: that state is now + unconstructible. ``to_wire`` returns ``None`` for an all-``None`` + value so the field is never emitted, and ``kv_first_write`` rejects a + zero-row batch outright — so the read side has nothing to skip rather + than an empty column to tolerate. + """ + packed = PackedTensor([None, None], dim_to_pack=0) + + assert encode_multimodal_for_wire("pixel_values", packed) is None + + # What the trainer actually receives for an image-free shard. + data = BatchedDataDict({"input_ids": torch.zeros(2, 4, dtype=torch.long)}) + assert "pixel_values" not in data.get_multimodal_dict(as_tensors=False) + + +def test_encode_multimodal_for_wire_packed_emits_single_nested_entry(): + """A ``PackedTensor`` arrives on the wire as exactly one entry. Row + boundaries live in the nested tensor itself and are preserved by TQ + (one stored entry per row), so no companion field is emitted.""" + packed = PackedTensor( + [torch.ones(3, 4), torch.ones(1, 4)], + dim_to_pack=0, + ) + + # One wire value under the field's own key: the payload. Shapes ride on + # ``KVBatchMeta.tags``, not as a companion column, so the field count on + # the wire is unchanged. + value = encode_multimodal_for_wire("pixel_values", packed) + + assert value is not None + assert value.is_nested + # Flattened rows: 3*4 and 1*4 elements. + assert [t.numel() for t in value.unbind()] == [12, 4] + + tags = multimodal_row_tags({"pixel_values": packed}, len(packed)) + assert [t["pixel_values__row_shapes"]["shapes"] for t in tags] == [ + [[3, 4]], + [[1, 4]], + ] + assert tags[0]["pixel_values__row_shapes"]["pad"] is False + + +def test_multimodal_row_tags_does_not_encode_the_payload(): + """``multimodal_row_tags`` needs geometry only, and must not pay for bytes. + + It used to call ``to_wire``, whose ``torch.cat`` copies the whole column, + and then throw the nested value away — a full copy of the largest field in + the batch, discarded, once per rollout step. + """ + packed = PackedTensor([torch.ones(3, 4), torch.ones(1, 4)], dim_to_pack=0) + calls = [] + packed.to_wire = lambda: calls.append(1) # type: ignore[method-assign] + + tags = multimodal_row_tags({"pixel_values": packed}, len(packed)) + + assert calls == [], "multimodal_row_tags must not call to_wire()" + assert [t["pixel_values__row_shapes"]["shapes"] for t in tags] == [ + [[3, 4]], + [[1, 4]], + ] + + +def test_multimodal_row_tags_rejects_row_count_disagreement(): + """A tags list shorter than the batch would leave a trailing sample with no + companion, which the read side then cannot distinguish from a lost one.""" + packed = PackedTensor([torch.ones(3, 4), torch.ones(1, 4)], dim_to_pack=0) + + with pytest.raises(ValueError, match="logical rows but the batch has"): + multimodal_row_tags({"pixel_values": packed}, 3) + + +def test_reassemble_packed_multimodal_raises_without_companion(): + """No companion means the true shapes are gone; reconstructing anyway + yields 1-D pixels and trains image-blind with no error.""" + packed = PackedTensor([torch.ones(3, 4), torch.ones(1, 4)], dim_to_pack=0) + nested, _ = packed.to_wire() + + with pytest.raises(ValueError, match="tags=None"): + reassemble_packed_multimodal({"pixel_values": nested}, None) + + with pytest.raises(ValueError, match="checked 2 tag rows"): + reassemble_packed_multimodal({"pixel_values": nested}, [{}, {}]) + + +def test_reassemble_packed_multimodal_round_trips_with_companion(): + packed = PackedTensor([torch.ones(3, 4), torch.ones(1, 4)], dim_to_pack=0) + nested, _ = packed.to_wire() + tags = multimodal_row_tags({"pixel_values": packed}, len(packed)) + + fields = {"pixel_values": nested} + reassemble_packed_multimodal(fields, tags) + + assert torch.equal(fields["pixel_values"].as_tensor(), packed.as_tensor()) + + +def test_encode_multimodal_for_wire_per_token_passes_through(): + """Per-token fields are rectangular ``[B, S]`` — they ride as plain + tensors, not nested ones.""" + ids = torch.ones((2, 6), dtype=torch.long) + + assert encode_multimodal_for_wire("mm_token_type_ids", ids) is ids + + +def test_pixel_dtype_cast_survives_to_the_wire_and_spares_integers(): + """The rollout casts pixels once, at ``get_multimodal_dict``. + + No worker re-applies it, so if the cast did not survive + ``encode_multimodal_for_wire`` the largest column would cross the wire in + fp32 where the legacy path shipped bf16 — a silent 2x on the dominant + field. Integer geometry must not be dragged along with it. + """ + batch = BatchedDataDict( + { + "pixel_values": PackedTensor( + [torch.ones(3, 4), torch.ones(1, 4)], dim_to_pack=0 + ), + "image_grid_thw": PackedTensor( + [torch.ones(1, 3, dtype=torch.long)] * 2, dim_to_pack=0 + ), + } + ) + + multimodal = batch.get_multimodal_dict(as_tensors=False, pixel_dtype=torch.bfloat16) + wire = {k: encode_multimodal_for_wire(k, v) for k, v in multimodal.items()} + + assert wire["pixel_values"].dtype == torch.bfloat16 + assert wire["image_grid_thw"].dtype == torch.int64 + + +def test_encode_multimodal_for_wire_skips_all_empty_packed(): + """An all-``None`` packed field has nothing to ship at all.""" + packed = PackedTensor([None, None], dim_to_pack=0) + + assert encode_multimodal_for_wire("pixel_values", packed) is None + + +def test_encode_multimodal_for_wire_rejects_unregistered_field(): + """A new modality added to ``get_multimodal_dict`` without a registry + entry must fail loud here — the silent-drop class this PR fixes.""" + with pytest.raises(KeyError, match="unregistered multimodal field"): + encode_multimodal_for_wire("audio_values", torch.zeros(2, 4)) + + +def test_encode_multimodal_for_wire_rejects_unregistered_packed_field(): + """The realistic drift case, and the one the registries exist to catch. + + ``get_multimodal_dict`` emits *any* ``PackedTensor`` value regardless of + key, so a processor field that is in neither registry reaches + ``encode_multimodal_for_wire`` as a ``PackedTensor``. It must raise + rather than be dropped: a dropped image column is a silently + image-blind forward, not a crash. + + The key below is a stand-in — the behavior is name-independent. No + specific model's ``model_input_names`` is asserted here, because that + could not be verified against the pinned transformers version. + """ + packed = PackedTensor([torch.ones(2, 4), torch.ones(3, 4)], dim_to_pack=0) + + with pytest.raises(KeyError, match="unregistered multimodal field"): + encode_multimodal_for_wire("pixel_attention_mask", packed) + + +def test_unregistered_packed_field_survives_get_multimodal_dict_then_raises(): + """Pin the whole leak path, not just the encoder in isolation. + + ``get_multimodal_dict``'s first branch is ``isinstance(v, PackedTensor)`` + — no registry check — so an unregistered packed field passes straight + through it. The registry gate is therefore load-bearing only at the wire + boundary; this test proves the two halves compose into a loud failure. + """ + data = BatchedDataDict( + { + "input_ids": torch.arange(8).reshape(2, 4), + "pixel_attention_mask": PackedTensor( + [torch.ones(2, 4), torch.ones(3, 4)], dim_to_pack=0 + ), + } + ) + + mm = data.get_multimodal_dict(as_tensors=False) + assert "pixel_attention_mask" in mm # slipped past the read-side dispatch + + with pytest.raises(KeyError, match="unregistered multimodal field"): + for k, v in mm.items(): + encode_multimodal_for_wire(k, v) + + +def test_encode_multimodal_for_wire_rejects_wrong_type_per_registry(): + """Registry membership decides the branch, so a mismatched value type + is a contract break, not something to coerce.""" + with pytest.raises(AssertionError, match="expected PackedTensor"): + encode_multimodal_for_wire("pixel_values", torch.zeros(2, 4)) + + with pytest.raises(AssertionError, match="expected Tensor"): + encode_multimodal_for_wire( + "mm_token_type_ids", PackedTensor([torch.ones(2, 3)], dim_to_pack=0) + ) + + +def test_multimodal_registries_are_disjoint(): + """A field in both registries would make ``encode_multimodal_for_wire`` + dispatch order-dependent and silently pick the packed branch.""" + assert not (PACKED_MULTIMODAL_FIELDS & PER_TOKEN_MULTIMODAL_FIELDS) + + +# ── to_wire guard rails ────────────────────────────────────────── + + +def test_to_wire_rejects_nonzero_dim_to_pack(): + """Only ``dim_to_pack=0`` round-trips; anything else needs + ``ragged_idx`` threading, so it must raise instead of silently + encoding along the wrong axis.""" + packed = PackedTensor([torch.ones(2, 3), torch.ones(2, 5)], dim_to_pack=1) + + with pytest.raises(NotImplementedError, match="only supports dim_to_pack=0"): + packed.to_wire() + + +def test_to_wire_all_none_returns_none(): + """All-empty batch signals 'skip this field' with ``None`` rather than + an empty nested tensor the read side cannot interpret.""" + packed = PackedTensor([None, None, None], dim_to_pack=0) + + nested, shapes = packed.to_wire() + assert nested is None + assert shapes == [] + + +def test_to_wire_carries_mixed_rank_rows(): + """Rows of differing rank now encode. + + The old encoder rejected these: padding to a batch max is undefined across + ranks, and no ``torch.nested`` layout holds them. Flattening sidesteps both + -- every row becomes rank-1 -- so the rank check was removable rather than + load-bearing. Reshaping on read restores the original ranks. + """ + rows = [torch.ones(1, 3, 2), torch.ones(2, 3)] + packed = PackedTensor(list(rows), dim_to_pack=0, pad_to_max_shape=True) + + nested, shapes = packed.to_wire() + assert [t.numel() for t in nested.unbind()] == [6, 6] + assert shapes == [[[1, 3, 2]], [[2, 3]]] + + restored = PackedTensor.from_wire(nested, shapes, pad_to_max_shape=True) + assert [tuple(t.shape) for t in restored.tensors] == [(1, 3, 2), (2, 3)] diff --git a/tests/unit/data_plane/README.md b/tests/unit/data_plane/README.md index 9ef9f60d417..a2efd6f3388 100644 --- a/tests/unit/data_plane/README.md +++ b/tests/unit/data_plane/README.md @@ -32,11 +32,9 @@ Generated audit of every test function under `tests/unit/data_plane/` with a one ## `test_codec_mooncake.py` -- `test_promote_1d_leaves_unsqueezes_1d` — `_promote_1d_leaves` turns schema-declared scalar fields from `(N,)` into `(N, 1)` for the Mooncake wire. -- `test_promote_1d_roundtrip_via_from_wire` — `_promote_1d_leaves` + `_from_wire` restores the schema-declared field's original `(N,)` shape and values. +- `test_from_wire_densifies_uniform_nested_rows` — Uniform nested rows come back dense on the read side. - `test_from_wire_rejects_invalid_declared_field_shape` — Corrupt or incompatible scalar wire shapes fail at the data-plane boundary. -- `test_promote_1d_leaves_rejects_undeclared_1d_field` — Unknown dense 1D Mooncake fields fail loudly instead of silently hitting TQ's shape mismatch. -- `test_put_samples_uses_schema_without_private_shape_tags` — Promotion does not add per-row adapter metadata to user tags. +- `test_put_samples_uses_schema_without_private_shape_tags` — The scalar-field schema patch does not add per-row adapter metadata to user tags. - `test_get_samples_uses_static_shape_schema` — Reads restore scalar fields by the shared schema while preserving genuine `(N, 1)` columns. - `test_pack_per_token_field_truncates_sp_padding` — pack_per_token_field slices each row to its own length, dropping SP padding. - `test_pack_per_token_field_exact_fit_matches_to_nested_by_length` — At exact fit, `pack_per_token_field` matches `to_nested_by_length`. diff --git a/tests/unit/data_plane/_rollout_shapes.py b/tests/unit/data_plane/_rollout_shapes.py index 147f0a70611..daa129af3e8 100644 --- a/tests/unit/data_plane/_rollout_shapes.py +++ b/tests/unit/data_plane/_rollout_shapes.py @@ -108,9 +108,13 @@ def _lp() -> torch.Tensor: } if multimodal: - # VLM extras as flat top-level fields (the codec wire format — - # nested dicts aren't valid leaves). Real production writes these - # with similar shapes; we keep them small for fast tests. + # VLM extras as flat *dense* top-level fields. This is the codec's + # plain-tensor path, not the packed wire form a real VLM rollout + # writes: production sends these as one flattened ``torch.jagged`` + # value per field with the geometry on ``KVBatchMeta.tags`` (see + # ``multimodal_utils.encode_multimodal_for_wire``). Tests that need + # the packed form build it locally; this stays dense so the codec's + # dtype/shape handling is covered without TQ tag plumbing. T, H, W = 1, 8, 8 n_image_tokens = T * H * W out["pixel_values"] = torch.randn(n, n_image_tokens, 3, generator=g).to( diff --git a/tests/unit/data_plane/test_arbitrary_shape_roundtrip.py b/tests/unit/data_plane/test_arbitrary_shape_roundtrip.py new file mode 100644 index 00000000000..5ca636f35a0 --- /dev/null +++ b/tests/unit/data_plane/test_arbitrary_shape_roundtrip.py @@ -0,0 +1,253 @@ +# Copyright (c) 2025, 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. +"""Can a row of arbitrary shape survive put -> get unchanged? + +The producer hands the data plane **one entry per row** and gets the same +rows back. No ``torch.nested`` container, so no layout can reject a shape +and nothing has to be padded to make a container accept it. + +Deliberately does NOT go through ``kv_first_write`` / ``pack_jagged_fields``: +those apply the application layer's padded-rectangle conventions, and +whether those are convenient is a separate argument from whether the +transport can carry the shape at all. + +Three escalating cases, matching what real VLM fields look like: + + A. ragged dim 0, uniform trailing dims -- ``pixel_values [n_patches, D]`` + B. non-uniform trailing dims, same rank -- dynamic-resolution ``[n, 3, H_i, W_i]`` + C. differing rank per row -- image ``[p, D]`` beside video ``[T, p, D]`` + +All three work today. ``PackedTensor.to_wire`` flattens each row to 1-D +before handing it to ``torch.jagged``, so rows differ only in dim 0 and B and +C encode as easily as A -- see ``test_to_wire_carries_mixed_rank_rows``. The +true shapes ride beside the payload on ``KVBatchMeta.tags``. What +``pad_to_max_shape`` still materializes is applied in worker memory by +``as_tensor``, never on the wire, and mixed rank is rejected only there. + +The ``test_jagged_rejects_*`` cases need no backend and never skip. They +pin *why* the padding exists, so the workaround can be deleted with +evidence rather than by assertion. + +Open assumption this exists to settle: TQ's msgpack encoder must +round-trip a ``torch.Tensor`` inside a ``NonTensorStack``. +``codec.unwrap_wire_stripped_payload`` exists because that path has bitten +before, and nothing currently verifies it for tensor payloads. +""" + +from __future__ import annotations + +import numpy as np +import pytest +import torch +from tensordict import TensorDict + +# -------------------------------------------------------------------------- +# Row fixtures. Small, but structurally faithful to the real fields. +# -------------------------------------------------------------------------- + +# A: only dim 0 varies. Trailing dim D is fixed by the vision tower. +ROWS_RAGGED_DIM0 = [ + torch.arange(4 * 8, dtype=torch.float32).reshape(4, 8), + torch.arange(7 * 8, dtype=torch.float32).reshape(7, 8), + torch.arange(1 * 8, dtype=torch.float32).reshape(1, 8), +] + +# B: same rank, H/W differ per row (dynamic-resolution processors). +ROWS_RAGGED_TRAILING = [ + torch.zeros(2, 3, 4, 4), + torch.ones(1, 3, 6, 6), + torch.full((3, 3, 2, 8), 7.0), +] + +# C: rank differs per row (a still image beside a video clip). +ROWS_MIXED_RANK = [ + torch.zeros(5, 8), # image: [patches, D] + torch.ones(2, 5, 8), # video: [frames, patches, D] + torch.full((3, 8), 2.0), # image: [patches, D] +] + +ALL_CASES = [ + (ROWS_RAGGED_DIM0, "ragged_dim0"), + (ROWS_RAGGED_TRAILING, "ragged_trailing"), + (ROWS_MIXED_RANK, "mixed_rank"), +] +CASE_IDS = ["ragged_dim0", "ragged_trailing", "mixed_rank"] + + +def _to_wire(rows: list[torch.Tensor]) -> np.ndarray: + """One object cell per row -- the whole encoder. + + ``pack_jagged_fields`` already forwards ``np.ndarray(dtype=object)`` + untouched and ``materialize`` already returns it without padding, so + this needs no new field registry, no ``__lengths`` companion, and no + shape bookkeeping in the application layer. TQ stores one entry per + (sample, field) either way -- see ``base.py::_generate_values`` -- so + this is not a fragmentation regression against the nested form. + """ + arr = np.empty(len(rows), dtype=object) + for i, row in enumerate(rows): + arr[i] = row + return arr + + +def _rows_back(value) -> list: + """Rows out of whatever container the adapter handed back. + + Deliberately permissive: ``_from_wire`` may densify or re-nest + depending on what the storage manager reassembled. The container is + not the contract -- the rows are. + """ + if isinstance(value, torch.Tensor): + return list(value.unbind()) if value.is_nested else list(value) + return list(value) + + +def _assert_rows_equal(got: list, want: list[torch.Tensor]) -> None: + assert len(got) == len(want), f"row count {len(got)} != {len(want)}" + for i, (have, expect) in enumerate(zip(got, want, strict=True)): + assert isinstance(have, torch.Tensor), ( + f"row {i} came back as {type(have).__name__}, not a Tensor -- " + "the wire path did not preserve the payload" + ) + assert have.shape == expect.shape, ( + f"row {i}: got {tuple(have.shape)}, want {tuple(expect.shape)}" + ) + assert have.dtype == expect.dtype, f"row {i}: dtype changed" + assert torch.equal(have, expect), f"row {i}: values differ" + + +# -------------------------------------------------------------------------- +# No backend needed: pin the torch.nested constraint that makes ``to_wire`` +# flatten each segment to 1-D before building the jagged value. +# -------------------------------------------------------------------------- + + +def test_jagged_rejects_nonuniform_trailing_dims() -> None: + """Why ``to_wire`` flattens rather than handing rows over as-is. + + ``torch.jagged`` allows exactly one ragged dim; every other dim must agree + across rows, so these rows cannot be a nested value in their natural shape. + Flattening each segment to 1-D moves all the variation onto dim 0, which is + what makes the encoding total -- and it is why nothing has to be padded to + satisfy the container. ``test_arbitrary_shape_rows_roundtrip[ragged_trailing]`` + is the same rows going through the real encoder. + """ + with pytest.raises((RuntimeError, TypeError)): + torch.nested.as_nested_tensor(ROWS_RAGGED_TRAILING, layout=torch.jagged) + + +def test_jagged_rejects_mixed_rank() -> None: + """Same constraint, the case padding could never have solved. + + No amount of padding reconciles rank 2 with rank 3, so image-beside-video in + one field is unrepresentable in a nested value built from natural shapes. + ``to_wire`` carries it anyway because 1-D rows have no rank to disagree on + -- see ``test_to_wire_carries_mixed_rank_rows``. + """ + with pytest.raises((RuntimeError, TypeError)): + torch.nested.as_nested_tensor(ROWS_MIXED_RANK, layout=torch.jagged) + + +@pytest.mark.parametrize("rows,case", ALL_CASES, ids=CASE_IDS) +def test_wire_form_preserves_rows_locally(rows, case) -> None: + """``_to_wire`` round-trips in-process, before any transport. + + Separated from the backend tests so a failure here is unambiguously + an encoding problem rather than a storage one. + """ + _assert_rows_equal(_rows_back(_to_wire(rows)), rows) + + +# -------------------------------------------------------------------------- +# Through the real adapter, on both backends. +# -------------------------------------------------------------------------- + + +@pytest.mark.parametrize("rows,case", ALL_CASES, ids=CASE_IDS) +def test_arbitrary_shape_rows_roundtrip(tq_client_backends, rows, case) -> None: + """Put arbitrary-shaped rows, get identical rows back. + + ``ragged_dim0`` should pass today. ``ragged_trailing`` and + ``mixed_rank`` are the open question: TQ's controller can represent + them (``extract_field_schema`` records ``per_sample_shapes`` as full + per-row tuples), but nothing verifies the storage managers accept + rows that disagree on trailing dims or rank. + """ + client = tq_client_backends + # conftest requires a partition id unique to each test. + partition_id = f"arbshape-{case}" + sample_ids = [f"s{i}" for i in range(len(rows))] + + client.register_partition( + partition_id=partition_id, + fields=["pixel_values"], + num_samples=len(rows), + consumer_tasks=["train"], + ) + try: + client.put_samples( + sample_ids=sample_ids, + partition_id=partition_id, + fields=TensorDict({"pixel_values": _to_wire(rows)}, batch_size=[len(rows)]), + ) + out = client.get_samples( + sample_ids=sample_ids, + partition_id=partition_id, + select_fields=["pixel_values"], + ) + _assert_rows_equal(_rows_back(out["pixel_values"]), rows) + finally: + client.clear_samples(sample_ids=sample_ids, partition_id=partition_id) + + +def test_no_padding_is_introduced(tq_client_backends) -> None: + """Element count must not grow -- the regression this change is for. + + ``pad_to_max_shape`` inflates every row to the batch max on each + non-ragged dim; for these rows that is 3*3*6*8 = 432 elements each + (1296 total) against 336 actual. Asserting on the total catches a + silent reintroduction of padding anywhere in the path. + """ + client = tq_client_backends + partition_id = "arbshape-nopad" + sample_ids = [f"p{i}" for i in range(len(ROWS_RAGGED_TRAILING))] + want_elems = sum(r.numel() for r in ROWS_RAGGED_TRAILING) + + client.register_partition( + partition_id=partition_id, + fields=["pixel_values"], + num_samples=len(ROWS_RAGGED_TRAILING), + consumer_tasks=["train"], + ) + try: + client.put_samples( + sample_ids=sample_ids, + partition_id=partition_id, + fields=TensorDict( + {"pixel_values": _to_wire(ROWS_RAGGED_TRAILING)}, + batch_size=[len(ROWS_RAGGED_TRAILING)], + ), + ) + out = client.get_samples( + sample_ids=sample_ids, + partition_id=partition_id, + select_fields=["pixel_values"], + ) + got_elems = sum(r.numel() for r in _rows_back(out["pixel_values"])) + assert got_elems == want_elems, ( + f"element count changed: {got_elems} != {want_elems} -- " + "padding was reintroduced somewhere in the path" + ) + finally: + client.clear_samples(sample_ids=sample_ids, partition_id=partition_id) diff --git a/tests/unit/data_plane/test_architecture_invariants.py b/tests/unit/data_plane/test_architecture_invariants.py index ab961960098..9dbb967e812 100644 --- a/tests/unit/data_plane/test_architecture_invariants.py +++ b/tests/unit/data_plane/test_architecture_invariants.py @@ -13,8 +13,10 @@ # limitations under the License. """Minimal behavioral invariants for the data-plane wiring. -* ``examples/run_grpo._select_trainer`` dispatches the legacy trainer - when ``data_plane`` is absent and the sync trainer when enabled. +* ``factory.select_sync_trainer`` dispatches the legacy trainer when + ``data_plane`` is absent and the sync trainer when enabled. +* Every launcher picks trainer *and* policy through those shared helpers, + so the two choices cannot drift apart. * The ``DataPlaneClient`` ABC carries every method adapters depend on. """ @@ -26,26 +28,134 @@ REPO = pathlib.Path(__file__).resolve().parents[3] +LAUNCHERS = ["run_grpo.py", "run_vlm_grpo.py"] -def test_run_grpo_dispatches_both_trainers(): - """``examples/run_grpo._select_trainer`` returns the TQ-mediated - ``grpo_train_sync`` iff ``data_plane.enabled`` is true, and the - legacy ``grpo_train`` otherwise.""" - import sys - sys.path.insert(0, str(REPO / "examples")) - try: - from run_grpo import _select_trainer - finally: - sys.path.pop(0) +def test_select_sync_trainer_dispatches_both_trainers(): + """Returns the TQ-mediated ``grpo_train_sync`` iff ``data_plane.enabled`` + is true, and the legacy ``grpo_train`` otherwise.""" from nemo_rl.algorithms.grpo import MasterConfig, grpo_train from nemo_rl.algorithms.grpo_sync import grpo_train_sync + from nemo_rl.data_plane.factory import select_sync_trainer cfg_legacy = MasterConfig.model_construct(data_plane=None) - assert _select_trainer(cfg_legacy) is grpo_train + assert select_sync_trainer(cfg_legacy) is grpo_train cfg_sync = MasterConfig.model_construct(data_plane={"enabled": True}) - assert _select_trainer(cfg_sync) is grpo_train_sync + assert select_sync_trainer(cfg_sync) is grpo_train_sync + + +def test_make_policy_factory_pairs_tq_policy_with_the_sync_trainer(): + """The other half of the dispatch. + + Turning the data plane on means picking *two* things that have to agree: + ``grpo_train_sync`` (covered above) and a ``TQPolicy`` factory. A launcher + that picked only one would run the sync trainer against a plain ``Policy`` + — which fails only after a full model load. + """ + from nemo_rl.data_plane.factory import make_policy_factory + + assert make_policy_factory(None) is None + assert make_policy_factory({"enabled": False}) is None + + dp_cfg = {"enabled": True} + factory = make_policy_factory(dp_cfg) + assert factory is not None + + captured = {} + + class _FakeTQPolicy: + def __init__(self, **kwargs): + captured.update(kwargs) + + import nemo_rl.models.policy.tq_policy as tq_policy_module + + real = tq_policy_module.TQPolicy + tq_policy_module.TQPolicy = _FakeTQPolicy + try: + make_policy_factory(dp_cfg)(cluster="c") + finally: + tq_policy_module.TQPolicy = real + + assert captured == {"cluster": "c", "dp_cfg": dp_cfg} + + +@pytest.mark.parametrize("launcher", LAUNCHERS) +def test_launchers_dispatch_through_the_shared_helpers(launcher: str): + """A launcher that re-implements the dispatch inline can drift from the + helpers the two tests above pin. Cheaper to require the call than to + diff two copies of the source.""" + source = (REPO / "examples" / launcher).read_text() + + assert "select_sync_trainer(" in source, ( + f"examples/{launcher} does not route its sync-trainer choice through " + "nemo_rl.data_plane.factory.select_sync_trainer." + ) + assert "make_policy_factory(" in source, ( + f"examples/{launcher} does not route its policy choice through " + "nemo_rl.data_plane.factory.make_policy_factory, so it can run the " + "sync trainer against a plain Policy." + ) + + +def test_sync_trainer_is_call_compatible_with_legacy_trainer(): + """Both trainers must accept the same call, because the VLM launcher + picks one at runtime and passes a single fixed kwarg set. + + Caught a real break: ``run_vlm_grpo`` passes ``processor=`` (VLM-only), + which ``grpo_train_sync`` did not accept — so every + ``data_plane.enabled=true`` VLM run died with ``TypeError: + grpo_train_sync() got an unexpected keyword argument 'processor'`` + after full model load. A signature check is cheap; the e2e that + surfaces it costs two nodes and ~12 minutes of setup. + """ + import inspect + + from nemo_rl.algorithms.grpo import grpo_train + from nemo_rl.algorithms.grpo_sync import grpo_train_sync + + # Mirror of the call in examples/run_vlm_grpo.py::main — 12 positional + # args (policy .. master_config) plus the VLM-only ``processor`` kwarg. + # Asserted via ``bind`` rather than as full signature parity: parity + # would force every future grpo_train parameter to be mirrored into + # grpo_train_sync as dead weight, which is a cost the dispatch does not + # actually impose. Only the shape the launchers really pass matters. + launcher_args = (None,) * 12 + launcher_kwargs = {"processor": None} + + for fn in (grpo_train, grpo_train_sync): + try: + inspect.signature(fn).bind(*launcher_args, **launcher_kwargs) + except TypeError as e: + raise AssertionError( + f"{fn.__module__}.{fn.__name__} cannot accept the call made by " + f"examples/run_vlm_grpo.py: {e}. Both trainers must bind the " + f"same launcher call, or the data_plane dispatch fails at " + f"runtime after a full model load." + ) from e + + +def test_both_trainers_wire_deduplicate_multimodal_data_into_repeat_interleave(): + """``deduplicate_multimodal_data`` must not become a silent no-op. + + ``enable_deduplication`` is reached only through + ``BatchedDataDict.repeat_interleave(..., share_immutable_media=True)``. A + trainer that omits the kwarg makes the flag do nothing: provenance is never + assigned, the deepcopy runs with an empty memo, and the user gets G + independent copies of every image in driver RAM with no warning. + """ + import inspect + + from nemo_rl.algorithms import grpo, grpo_sync + + for module in (grpo, grpo_sync): + source = inspect.getsource(module) + assert "share_immutable_media=" in source, ( + f"{module.__name__} calls repeat_interleave without " + "share_immutable_media, so grpo.deduplicate_multimodal_data is a " + "silent no-op on that trainer." + ) + assert "deduplicate_multimodal_data" in source def test_sync_trainer_rejects_message_level_advantage_penalties(): diff --git a/tests/unit/data_plane/test_codec_mooncake.py b/tests/unit/data_plane/test_codec_mooncake.py index bced5792717..f5705d1ac63 100644 --- a/tests/unit/data_plane/test_codec_mooncake.py +++ b/tests/unit/data_plane/test_codec_mooncake.py @@ -29,70 +29,26 @@ from ._rollout_shapes import make_rollout_batch -# ── P1: promote_1d — writer unsqueezes, reader squeezes ────────────────────── - - -def test_promote_1d_leaves_unsqueezes_1d() -> None: - """`_promote_1d_leaves` turns 1D ``(N,)`` leaves into ``(N, 1)``. - - Guards the mooncake_cpu path where TQ's extract_field_schema silently - unsqueezes 1D fields in metadata; the wire layer pre-unsqueezes so the - per-row data shape matches the metadata-recorded shape. - """ - from tensordict import TensorDict - - from nemo_rl.data_plane.adapters.transfer_queue import _promote_1d_leaves - - n = 8 - t = torch.arange(n, dtype=torch.float32) - td = TensorDict({"input_lengths": t}, batch_size=[n]) - - out = _promote_1d_leaves(td) - assert out["input_lengths"].shape == (n, 1), ( - "Expected input_lengths to use the Mooncake wire shape " - f"({n}, 1), got {tuple(out['input_lengths'].shape)}." - ) - - -def test_promote_1d_roundtrip_via_from_wire() -> None: - """`_promote_1d_leaves` then `_from_wire` restores the original ``(N,)`` shape and values.""" - from tensordict import TensorDict - - from nemo_rl.data_plane.adapters.transfer_queue import ( - _from_wire, - _promote_1d_leaves, - ) - - n = 6 - original = torch.arange(n, dtype=torch.float32) - td = TensorDict({"input_lengths": original}, batch_size=[n]) - - wire = _promote_1d_leaves(td) - assert wire["input_lengths"].shape == (n, 1) - - back = _from_wire(wire) - assert back["input_lengths"].shape == (n,) - assert torch.equal(back["input_lengths"], original) - @pytest.mark.parametrize("field_name", ["mask_sample", "truncated"]) def test_raw_sample_filter_fields_roundtrip_as_dense_1d(field_name: str) -> None: - """Raw loss-filter fields use the Mooncake scalar wire workaround.""" + """Raw loss-filter fields survive the Mooncake scalar wire as dense ``(N,)``. + + Ported from the ``_promote_1d_leaves`` pair this file used to carry. That + writer/reader pair was replaced by ``_patch_scalar_field_schema``, which + fixes the same TQ bug at the schema layer and covers *every* dense 1-D + field instead of a hand-kept allowlist — so there is no promote step to + assert on any more, and the property to pin is just the round trip. + """ from tensordict import TensorDict - from nemo_rl.data_plane.adapters.transfer_queue import ( - _from_wire, - _promote_1d_leaves, - ) + from nemo_rl.data_plane.adapters.transfer_queue import _from_wire n = 4 original = torch.tensor([False, True, False, True]) td = TensorDict({field_name: original}, batch_size=[n]) - wire = _promote_1d_leaves(td) - assert wire[field_name].shape == (n, 1) - - back = _from_wire(wire) + back = _from_wire(td) assert back[field_name].shape == (n,) assert torch.equal(back[field_name], original) @@ -116,8 +72,15 @@ def test_from_wire_densifies_uniform_nested_rows() -> None: assert torch.equal(back["input_ids"], torch.stack(rows)) -def test_from_wire_preserves_genuine_length_one_token_column() -> None: - """Only fields promoted from ``(N,)`` are squeezed after a TQ read.""" +def test_from_wire_squeezes_nothing_even_for_scalar_field_names() -> None: + """``_from_wire`` densifies; it never reinterprets a row's rank. + + A ``(1,)`` row is genuine data now. Per-sample scalar columns are + stored as 0-d rows (``_patch_scalar_field_schema``) and TQ stacks them + into a dense ``(N,)`` before this function runs, so a nested ``(1,)`` + row arriving here means the producer really wrote length-1 rows — + squeezing it would corrupt them. Field name must not change that. + """ from tensordict import TensorDict from nemo_rl.data_plane.adapters.transfer_queue import _from_wire @@ -137,49 +100,40 @@ def test_from_wire_preserves_genuine_length_one_token_column() -> None: back = _from_wire(wire) - assert back["total_reward"].shape == (n,) + # ``total_reward`` is a per-sample scalar by name, but these rows are + # length-1 vectors — both columns densify identically. + assert back["total_reward"].shape == (n, 1) assert back["input_ids"].shape == (n, 1) assert torch.equal(back["input_ids"], torch.arange(n).unsqueeze(-1)) -def test_from_wire_rejects_invalid_declared_field_shape() -> None: - """A corrupted scalar wire shape fails at the data-plane boundary.""" +def test_from_wire_passes_dense_fields_through_untouched() -> None: + """Dense inputs are returned as-is — no rank policing by field name. + + Replaces a guard that rejected a declared scalar arriving as ``(3, 2)``. + That check belonged to the writer-unsqueeze/reader-squeeze pair, which + no longer exists: the schema now reports the shape the rows actually + have, so there is no promoted encoding for a malformed value to + violate. + """ from tensordict import TensorDict from nemo_rl.data_plane.adapters.transfer_queue import _from_wire wire = TensorDict({"input_lengths": torch.ones(3, 2)}, batch_size=[3]) - with pytest.raises(ValueError, match=r"input_lengths.*\(N, 1\)"): - _from_wire(wire) - - -def test_promote_1d_leaves_rejects_undeclared_1d_field() -> None: - """New scalar fields must be added to the authoritative schema.""" - from tensordict import TensorDict - - from nemo_rl.data_plane.adapters.transfer_queue import _promote_1d_leaves - - fields = TensorDict({"new_scalar": torch.arange(3)}, batch_size=[3]) - - with pytest.raises(ValueError, match="not declared.*PROMOTE_1D_FIELDS"): - _promote_1d_leaves(fields) - - -def test_promote_1d_leaves_rejects_invalid_declared_field_shape() -> None: - """A schema-declared scalar cannot silently change its user-level rank.""" - from tensordict import TensorDict - - from nemo_rl.data_plane.adapters.transfer_queue import _promote_1d_leaves - - fields = TensorDict({"input_lengths": torch.ones(3, 2)}, batch_size=[3]) + back = _from_wire(wire) + assert back["input_lengths"].shape == (3, 2) - with pytest.raises(ValueError, match=r"input_lengths.*shape \(N,\)"): - _promote_1d_leaves(fields) +def test_put_samples_passes_fields_and_tags_through_unchanged(monkeypatch) -> None: + """``put_samples`` reshapes nothing and does not touch user tags. -def test_put_samples_uses_schema_without_private_shape_tags(monkeypatch) -> None: - """Mooncake promotion changes tensors but not user-provided TQ tags.""" + The writer-unsqueeze half of the old 1-D workaround is gone — the + schema now reports the ``()`` rows TQ actually stores — so a ``(N,)`` + column reaches the wire as ``(N,)``. Tag passthrough was the other + half of this test's intent and is unchanged. + """ from tensordict import TensorDict import nemo_rl.data_plane.adapters.transfer_queue as tq_adapter @@ -203,13 +157,12 @@ def fake_kv_batch_put( ) -> None: assert keys == ["a", "b", "c"] assert partition_id == "train" - assert fields["input_lengths"].shape == (n, 1) + assert fields["input_lengths"].shape == (n,) # not promoted any more assert fields["input_ids"].shape == (n, 1) assert tags == user_tags monkeypatch.setattr(tq_adapter.tq, "kv_batch_put", fake_kv_batch_put) client = object.__new__(tq_adapter.TQDataPlaneClient) - client._promote_1d = True meta = client.put_samples( ["a", "b", "c"], "train", fields=original_fields, tags=user_tags @@ -218,8 +171,14 @@ def fake_kv_batch_put( assert meta.tags == user_tags -def test_get_samples_uses_static_shape_schema(monkeypatch) -> None: - """The Mooncake adapter restores scalar ranks without row metadata.""" +def test_get_samples_returns_scalar_columns_dense(monkeypatch) -> None: + """A per-sample scalar column arrives dense and is passed through. + + TQ stores these as 0-d rows and ``_merge_tensors_to_tensordict`` stacks + them into ``(N,)`` before the adapter sees them, so ``get_samples`` has + no rank to restore — it just must not disturb the column. Previously + this arrived nested with ``(1,)`` rows and was squeezed back. + """ from tensordict import TensorDict import nemo_rl.data_plane.adapters.transfer_queue as tq_adapter @@ -234,10 +193,8 @@ def test_get_samples_uses_static_shape_schema(monkeypatch) -> None: ) wire_data = TensorDict( { - "total_reward": torch.nested.as_nested_tensor( - [row for row in original["total_reward"].unsqueeze(-1)], - layout=torch.jagged, - ), + # Dense: what TQ hands back for a 0-d-row scalar column. + "total_reward": original["total_reward"], "input_ids": torch.nested.as_nested_tensor( [row for row in original["input_ids"]], layout=torch.jagged ), @@ -255,7 +212,6 @@ def fake_kv_batch_get( monkeypatch.setattr(tq_adapter.tq, "kv_batch_get", fake_kv_batch_get) client = object.__new__(tq_adapter.TQDataPlaneClient) - client._promote_1d = True client._data_operations_started = False restored = client.get_samples( @@ -290,7 +246,6 @@ def fake_kv_batch_get( monkeypatch.setattr(tq_adapter.tq, "kv_batch_get", fake_kv_batch_get, raising=False) client = object.__new__(tq_adapter.TQDataPlaneClient) - client._promote_1d = False client._data_operations_started = False restored = client.get_samples(["a", "b"], "train", ["input_ids"]) diff --git a/tests/unit/data_plane/test_leader_broadcast.py b/tests/unit/data_plane/test_leader_broadcast.py index 5a74f438c42..f44a5711685 100644 --- a/tests/unit/data_plane/test_leader_broadcast.py +++ b/tests/unit/data_plane/test_leader_broadcast.py @@ -25,6 +25,7 @@ import torch.distributed as dist import torch.multiprocessing as mp +from nemo_rl.data.multimodal_utils import PackedTensor from nemo_rl.data_plane.worker_mixin import _broadcast_batched_data_dict from nemo_rl.distributed.batched_data_dict import BatchedDataDict @@ -40,12 +41,27 @@ def _worker(rank: int, world_size: int, tmp_init_file: str, q): world_size=world_size, ) try: + # ``pixel_values`` is the case that mattered: a PackedTensor is not a + # torch.Tensor, so before the ``packed_wire`` branch it rode the object + # list and ``broadcast_object_list`` pickled the pixels into device + # memory (26.25 GiB on a real VLM batch). Rows differ in their trailing + # dims and one sample has no media, which is what the format exists for. + pixel_rows = [ + torch.arange(2 * 3 * 4, dtype=torch.float32).reshape(2, 3, 4), + torch.arange(1 * 5 * 4, dtype=torch.float32).reshape(1, 5, 4) + 100, + None, + ] if rank == 0: data = BatchedDataDict( { "input_ids": torch.arange(12, dtype=torch.long).reshape(3, 4), "input_lengths": torch.tensor([4, 3, 2], dtype=torch.int32), "scalar_meta": "step_42", + "pixel_values": PackedTensor( + [r.clone() if r is not None else None for r in pixel_rows], + dim_to_pack=0, + pad_to_max_shape=True, + ), } ) else: @@ -62,6 +78,25 @@ def _worker(rank: int, world_size: int, tmp_init_file: str, q): out["input_lengths"], torch.tensor([4, 3, 2], dtype=torch.int32) ) assert out["scalar_meta"] == "step_42" + + packed = out["pixel_values"] + assert isinstance(packed, PackedTensor), type(packed).__name__ + # Compare on logical rows, not ``.tensors``: ``from_wire`` returns + # segments flat with a CSR row map, so an empty row contributes no + # entry there. Per-row segment counts are what the model uses to match + # media against placeholder tokens -- densifying before the wire + # collapses them and misaligns the forward. + expected = PackedTensor( + [r.clone() if r is not None else None for r in pixel_rows], + dim_to_pack=0, + pad_to_max_shape=True, + ) + assert ( + packed.logical_segment_counts_by_row() + == expected.logical_segment_counts_by_row() + == [1, 1, 0] + ) + assert torch.equal(packed.as_tensor(), expected.as_tensor()) q.put((rank, "ok")) except Exception as e: # pragma: no cover — surface failures to parent q.put((rank, f"err: {type(e).__name__}: {e}")) diff --git a/tests/unit/data_plane/test_multimodal_wire_roundtrip.py b/tests/unit/data_plane/test_multimodal_wire_roundtrip.py new file mode 100644 index 00000000000..f165105c673 --- /dev/null +++ b/tests/unit/data_plane/test_multimodal_wire_roundtrip.py @@ -0,0 +1,539 @@ +# Copyright (c) 2025, 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. +"""End-to-end VLM multimodal wire roundtrip through the data plane. + +Runs a realistic VLM rollout write → TQ (NoOp adapter) → materialize → +trainer read, exercising every layer the multimodal fix touches: + + * ``PackedTensor.to_wire`` on the write side + * ``kv_first_write`` wire-type filter + * ``codec.pack_jagged_fields`` (write-time layout transform) + * ``codec.materialize`` (read-time padded conversion), including the + ``PACKED_MULTIMODAL_FIELDS`` pad_to_seqlen exclusion + * ``PackedTensor.from_wire`` on the read side + * ``BatchedDataDict.get_multimodal_dict`` dispatch + +Guards the silent-drop regression class that motivated the PR. +""" + +from __future__ import annotations + +from dataclasses import replace +from types import SimpleNamespace + +import torch + +from nemo_rl.data.multimodal_utils import ( + PACKED_MULTIMODAL_FIELDS, + PER_TOKEN_MULTIMODAL_FIELDS, + PackedTensor, + encode_multimodal_for_wire, + multimodal_row_tags, +) +from nemo_rl.data_plane.adapters.noop import NoOpDataPlaneClient +from nemo_rl.data_plane.column_io import kv_first_write, read_columns +from nemo_rl.data_plane.interfaces import KVBatchMeta +from nemo_rl.data_plane.schema import GLOBAL_FORWARD_PAD_SEQLEN +from nemo_rl.distributed.batched_data_dict import BatchedDataDict + +from ._rollout_shapes import keys_from_uids, register_train_partition + + +def _make_vlm_rollout_output( + n: int = 4, seqlen: int = 32 +) -> tuple[BatchedDataDict, dict[str, PackedTensor], dict[str, torch.Tensor]]: + """Build the ``final_batch_cpu`` a VLM rollout would hand to + ``kv_first_write``, plus the per-key ground truth for round-trip + assertions. + + Ground truth is captured as ``as_tensor()``-form so the read side + can compare after concat, regardless of jagged reshuffle. + """ + torch.manual_seed(0) + input_lengths = torch.tensor([seqlen, seqlen, seqlen, seqlen], dtype=torch.long) + + # PackedTensor multimodal — sample 2 is text-only (None) to + # exercise the placeholder / lengths[i]==0 path. + pixel_values = PackedTensor( + [ + torch.randn(3, 4, 4), # sample 0: 3 patches + torch.randn(5, 4, 4), # sample 1: 5 patches + None, # sample 2: no image + torch.randn(2, 4, 4), # sample 3: 2 patches + ], + dim_to_pack=0, + ) + image_grid_thw = PackedTensor( + [ + torch.tensor([[1, 2, 2]], dtype=torch.int64), + torch.tensor([[1, 3, 3], [1, 2, 2]], dtype=torch.int64), + None, + torch.tensor([[1, 2, 1]], dtype=torch.int64), + ], + dim_to_pack=0, + ) + # Per-token multimodal (rectangular) — model-specific type map. + mm_token_type_ids = torch.randint(0, 3, (n, seqlen), dtype=torch.int64) + + fb = BatchedDataDict() + fb["input_ids"] = torch.arange(n * seqlen, dtype=torch.long).reshape(n, seqlen) + fb["input_lengths"] = input_lengths + fb["token_mask"] = torch.ones((n, seqlen), dtype=torch.long) + fb["sample_mask"] = torch.ones((n,), dtype=torch.long) + fb["generation_logprobs"] = torch.zeros((n, seqlen), dtype=torch.float32) + fb["pixel_values"] = pixel_values + fb["image_grid_thw"] = image_grid_thw + fb["mm_token_type_ids"] = mm_token_type_ids + + packed_truth = { + "pixel_values": pixel_values, + "image_grid_thw": image_grid_thw, + } + plain_truth = {"mm_token_type_ids": mm_token_type_ids} + return fb, packed_truth, plain_truth + + +def _sync_rollout_write_loop(fb: BatchedDataDict, bulk_batch: BatchedDataDict) -> None: + """Replay the write loop in ``sync_rollout_actor.rollout_to_tq``. + + Copies non-multimodal fields verbatim, then hands every multimodal + field to the *production* encoder ``encode_multimodal_for_wire`` — + the same call the actor makes. Reimplementing the encode branches + here would make this test blind to exactly the drift (new key added, + one branch forgotten) that the silent-drop regression was. + """ + for k, v in fb.items(): + if k in PACKED_MULTIMODAL_FIELDS or k in PER_TOKEN_MULTIMODAL_FIELDS: + continue # handled below via get_multimodal_dict + if isinstance(v, PackedTensor): + continue # multimodal PackedTensor; also handled below + bulk_batch[k] = v + + multimodal = fb.get_multimodal_dict(as_tensors=False) + for k, v in multimodal.items(): + wire_value = encode_multimodal_for_wire(k, v) + if wire_value is not None: + bulk_batch[k] = wire_value + + +def test_vlm_wire_roundtrip_through_noop_data_plane(): + """Full trip: build a VLM rollout batch → write via kv_first_write + → materialize on read → verify ``get_multimodal_dict`` on the + reassembled batch matches the pre-wire ``.as_tensor()`` output. + + Uses the NoOp data-plane adapter so the test runs without TQ + installed but still exercises the real ABC contract. + """ + # Registry sanity — if a rename is landed without updating the + # module-level sets, the failure surfaces here, not in production. + assert "pixel_values" in PACKED_MULTIMODAL_FIELDS + assert "image_grid_thw" in PACKED_MULTIMODAL_FIELDS + assert "mm_token_type_ids" in PER_TOKEN_MULTIMODAL_FIELDS + + n = 4 + seqlen = 32 + fb, packed_truth, plain_truth = _make_vlm_rollout_output(n=n, seqlen=seqlen) + + # ── Write path ────────────────────────────────────────────────── + client = NoOpDataPlaneClient() + fields = list(fb.keys()) + register_train_partition(client, num_samples=n, fields=fields) + + bulk_batch: BatchedDataDict = BatchedDataDict() + _sync_rollout_write_loop(fb, bulk_batch) + meta = kv_first_write( + bulk_batch, + sample_ids=keys_from_uids(["a", "b", "c", "d"]), + dp_client=client, + partition_id="train", + tags=multimodal_row_tags(fb.get_multimodal_dict(as_tensors=False), n), + ) + # kv_first_write ships one wire field per logical field — the nested + # parent carries its own row boundaries, so there is no companion. + for k in packed_truth: + assert k in meta.fields + assert f"{k}__lengths" not in meta.fields + assert "mm_token_type_ids" in meta.fields + + # ── Read path ─────────────────────────────────────────────────── + fetched = read_columns( + client, + meta, + select_fields=meta.fields, + layout="padded", + ) + + # ── Contract checks on the materialized batch ─────────────────── + # Packed multimodal fields are never rectangularized to + # ``[B, max_per_sample, ...]`` — no zero padding on the wire or in + # worker memory, and the row boundaries never have to be recovered + # from a companion column. + # + # ``materialize`` hands them back as ``PackedTensor``, not as the raw + # nested wire value: generic ``BatchedDataDict`` consumers dispatch on + # ``PackedTensor`` and otherwise fall through to ``tensor[start:end]`` + # (``slice``) or ``tuple(v.shape)`` (``_broadcast_batched_data_dict``), + # neither of which torch.nested supports on the ragged dim 0. + assert isinstance(fetched["pixel_values"], PackedTensor) + assert isinstance(fetched["image_grid_thw"], PackedTensor) + assert [t.numel() for t in fetched["pixel_values"].to_wire()[0].unbind()] == [ + t.numel() for t in packed_truth["pixel_values"].to_wire()[0].unbind() + ] + assert "pixel_values__lengths" not in fetched + assert "image_grid_thw__lengths" not in fetched + + # ── Reassembly via BatchedDataDict.get_multimodal_dict ────────── + mm = fetched.get_multimodal_dict(as_tensors=True) + assert set(mm.keys()) == {"pixel_values", "image_grid_thw", "mm_token_type_ids"} + # Packed fields concatenated back to their pre-wire ``.as_tensor()``. + assert torch.equal(mm["pixel_values"], packed_truth["pixel_values"].as_tensor()) + assert torch.equal(mm["image_grid_thw"], packed_truth["image_grid_thw"].as_tensor()) + # Per-token field passes through untouched. + assert torch.equal(mm["mm_token_type_ids"], plain_truth["mm_token_type_ids"]) + + # ``as_tensors=False`` returns ``PackedTensor`` wrappers for the + # packed keys so callers can slice per-sample if needed. + mm_wrapped = fetched.get_multimodal_dict(as_tensors=False) + assert isinstance(mm_wrapped["pixel_values"], PackedTensor) + assert isinstance(mm_wrapped["image_grid_thw"], PackedTensor) + # Non-multimodal keys (input_ids, token_mask, ...) must NOT leak + # into the multimodal dict. + assert "input_ids" not in mm_wrapped + assert "token_mask" not in mm_wrapped + + +def test_materialize_forward_pad_skips_packed_multimodal_fields(): + """``codec.materialize`` right-pads dim 1 up to the cross-DP forward pad + target, which is a *token* seqlen. For a packed multimodal field dim 1 is + patch/image count, so padding it there would inflate ``pixel_values`` by + the ratio seqlen/patches — the ~40x blow-up the exclusion averts. + + Those fields are now excluded one step earlier: they skip + ``to_padded_tensor`` entirely and stay nested, so they are never + rectangular for ``pad_to_seqlen`` to extend. This pins that the + token-aligned fields still pad while the multimodal ones are untouched. + + The pad only fires when ``GLOBAL_FORWARD_PAD_SEQLEN`` is stamped on the + meta (``TQPolicy._stamp_pad_seqlen`` does this in production). The plain + roundtrip above never stamps it, so without this test the exclusion + branch in ``codec.materialize`` is never entered by any unit test. + """ + n, seqlen = 4, 32 + pad_to = 64 # cross-DP forward pad target, > seqlen + fb, _, _ = _make_vlm_rollout_output(n=n, seqlen=seqlen) + + client = NoOpDataPlaneClient() + fields = list(fb.keys()) + register_train_partition(client, num_samples=n, fields=fields) + + bulk_batch: BatchedDataDict = BatchedDataDict() + _sync_rollout_write_loop(fb, bulk_batch) + meta = kv_first_write( + bulk_batch, + sample_ids=keys_from_uids(["a", "b", "c", "d"]), + dp_client=client, + partition_id="train", + tags=multimodal_row_tags(fb.get_multimodal_dict(as_tensors=False), n), + ) + meta = replace( + meta, extra_info={**(meta.extra_info or {}), GLOBAL_FORWARD_PAD_SEQLEN: pad_to} + ) + + fetched = read_columns(client, meta, select_fields=meta.fields, layout="padded") + + # Token-aligned fields DO get padded up to the forward target. + assert fetched["input_ids"].shape == (n, pad_to) + assert fetched["mm_token_type_ids"].shape == (n, pad_to) + # Packed multimodal fields are never rectangularized, so the forward pad + # target cannot reach them. ``materialize`` returns them as PackedTensor; + # ``to_wire()`` recovers the per-row view to check the row counts. + assert isinstance(fetched["pixel_values"], PackedTensor) + assert isinstance(fetched["image_grid_thw"], PackedTensor) + assert ( + max(t.shape[0] for t in fetched["pixel_values"].tensors if t is not None) == 5 + ) + assert ( + max(t.shape[0] for t in fetched["image_grid_thw"].tensors if t is not None) == 2 + ) + + # Reassembly still reproduces the pre-wire values under forward padding. + mm = fetched.get_multimodal_dict(as_tensors=True) + assert torch.equal(mm["pixel_values"], fb["pixel_values"].as_tensor()) + assert torch.equal(mm["image_grid_thw"], fb["image_grid_thw"].as_tensor()) + + +# ── Dispatch field selection (logprob vs train parity) ────────────────── + + +def _stub_tq_policy(monkeypatch, captured: dict[str, KVBatchMeta]): + """A ``TQPolicy`` with only the surface the dispatch bodies touch. + + ``shard_meta_for_dp`` is stubbed to record the meta each dispatch + builds, so the assertions are on the real field-selection code + rather than a reimplementation of it. + """ + from nemo_rl.models.policy import tq_policy as tq_policy_mod + + def fake_shard(meta, **kwargs): + captured[meta.task_name] = meta + return [meta], None + + monkeypatch.setattr(tq_policy_mod, "shard_meta_for_dp", fake_shard) + + # ``Policy.shutdown`` short-circuits on a policy with no ``worker_group``, + # but the dispatch bodies need one, so that escape hatch is unavailable + # here. Override the destructor instead: this object never owned Ray + # workers or a TQ client, so running live teardown on it at GC is simply + # wrong. Faking ``shutdown``/``dp_client`` would work too, but it would + # report a successful shutdown that never happened and would silently + # absorb any future change to the real teardown contract. + # + # It must be a real no-op method, not ``__del__ = None``: CPython + # installs ``tp_finalize`` whenever ``__del__`` is present in the class + # dict, then calls it — ``None()`` raises ``TypeError`` at GC and pytest + # reports it as a ``PytestUnraisableExceptionWarning`` charged to + # whichever test happens to be running. + class _StubTQPolicy(tq_policy_mod.TQPolicy): + def __del__(self) -> None: + pass + + pol = object.__new__(_StubTQPolicy) + pol.cfg = {} + pol._router_replay_enabled = False + pol.flops_tracker = None + pol.sharding_annotations = SimpleNamespace(get_axis_size=lambda _axis: 1) + pol.worker_group = SimpleNamespace( + run_all_workers_sharded_data=lambda *a, **k: [], + get_all_worker_results=lambda _futures: [ + {"global_loss": 0.0, "grad_norm": 0.0, "all_mb_metrics": {}} + ], + ) + return pol + + +def test_train_dispatch_ships_the_same_multimodal_fields_as_logprob(monkeypatch): + """The training forward must see the images the logprob forwards saw. + + ``DP_TRAIN_FIELDS`` is a static text-only schema, so without the + per-batch multimodal add-on the GRPO update would run image-blind + against prev/ref logprobs that were computed *with* images — a + silent objective mismatch, not a crash. + """ + mm_fields = [ + "pixel_values", + "image_grid_thw", + "mm_token_type_ids", + ] + meta = KVBatchMeta( + partition_id="train", + task_name="rollout", + sample_ids=["a", "b"], + fields=["input_ids", "input_lengths", "token_mask", *mm_fields], + sequence_lengths=[8, 8], + ) + + captured: dict[str, KVBatchMeta] = {} + pol = _stub_tq_policy(monkeypatch, captured) + + pol.get_logprobs_from_meta(meta) + pol.train_from_meta(meta, loss_fn=None, gbs=2, mbs=1) + + lp_fields = set(captured["prev_lp"].fields) + train_fields = set(captured["train"].fields) + # Parity across the FULL registry, not just the hand-listed mm_fields: + # the two supersets below would be satisfied by a dispatch that shipped + # extra multimodal columns to one side only. + registry = PACKED_MULTIMODAL_FIELDS | PER_TOKEN_MULTIMODAL_FIELDS + assert set(mm_fields) <= lp_fields + assert set(mm_fields) <= train_fields + assert lp_fields & registry == train_fields & registry + + +def test_text_only_dispatch_requests_no_multimodal_fields(monkeypatch): + """Text-only runs never write the multimodal columns; requesting + them would raise at the adapter, so the add-on must stay empty.""" + meta = KVBatchMeta( + partition_id="train", + task_name="rollout", + sample_ids=["a", "b"], + fields=["input_ids", "input_lengths", "token_mask"], + sequence_lengths=[8, 8], + ) + + captured: dict[str, KVBatchMeta] = {} + pol = _stub_tq_policy(monkeypatch, captured) + + pol.get_logprobs_from_meta(meta) + pol.train_from_meta(meta, loss_fn=None, gbs=2, mbs=1) + + all_mm = PACKED_MULTIMODAL_FIELDS | PER_TOKEN_MULTIMODAL_FIELDS + assert not (set(captured["prev_lp"].fields) & all_mm) + assert not (set(captured["train"].fields) & all_mm) + + +def test_ref_logprob_dispatch_ships_multimodal_fields(monkeypatch): + """``get_reference_policy_logprobs_from_meta`` shares ``_logprob_dispatch`` + with the prev-logprob path, but it is the ref forward whose output ends up + in the KL term — an image-blind ref logprob is a silent objective skew, so + pin the ref task explicitly rather than trusting the shared body.""" + mm_fields = [ + "pixel_values", + "mm_token_type_ids", + ] + meta = KVBatchMeta( + partition_id="train", + task_name="rollout", + sample_ids=["a", "b"], + fields=["input_ids", "input_lengths", "token_mask", *mm_fields], + sequence_lengths=[8, 8], + ) + + captured: dict[str, KVBatchMeta] = {} + pol = _stub_tq_policy(monkeypatch, captured) + + pol.get_reference_policy_logprobs_from_meta(meta) + + assert set(mm_fields) <= set(captured["ref_lp"].fields) + + +def test_sc_microbatch_dispatch_ships_multimodal_fields(monkeypatch): + """The single-controller split-API path (``train_microbatches_from_meta``) + carries its own copy of the multimodal add-on, separate from + ``train_from_meta``. Without a test here the SC path could regress to an + image-blind forward while the sync path stays green.""" + mm_fields = [ + "pixel_values", + "image_grid_thw", + "mm_token_type_ids", + ] + meta = KVBatchMeta( + partition_id="train", + task_name="rollout", + sample_ids=["a", "b"], + fields=["input_ids", "input_lengths", "token_mask", *mm_fields], + sequence_lengths=[8, 8], + ) + + # One policy, one patch: ``_stub_tq_policy`` patches the module-level + # ``shard_meta_for_dp``, so two stubs would share the last patch and the + # first capture dict would stay empty. + captured: dict[str, KVBatchMeta] = {} + pol = _stub_tq_policy(monkeypatch, captured) + + pol.train_from_meta(meta, loss_fn=None, gbs=2, mbs=1) + sync_fields = set(captured["train"].fields) + pol.train_microbatches_from_meta(meta) + sc_fields = set(captured["train"].fields) + + assert set(mm_fields) <= sc_fields + # Both train entrypoints must request the identical multimodal set. + all_mm = PACKED_MULTIMODAL_FIELDS | PER_TOKEN_MULTIMODAL_FIELDS + assert sc_fields & all_mm == sync_fields & all_mm + + +def test_materialized_multimodal_batch_survives_microbatch_slicing(): + """A materialized batch must be sliceable before any multimodal reassembly. + + The training path slices the fetched batch long before anything calls + ``get_multimodal_dict``:: + + get_logprobs -> make_processed_microbatch_iterator + -> make_microbatch_iterator_with_dynamic_shapes + -> BatchedDataDict.slice -> self.data[k][start:end] + + ``slice`` dispatches on ``PackedTensor`` and otherwise falls through to + plain indexing. A raw nested value satisfies ``isinstance(v, Tensor)`` and + has a valid ``shape[0]``, so it clears both guards and only fails at the + indexing itself with ``slice(): not supported for NestedTensor on dim=0``. + That took down a Qwen3.5 VLM GRPO run at the first logprob step. + """ + n = 4 + fb, packed_truth, _ = _make_vlm_rollout_output(n=n, seqlen=32) + + client = NoOpDataPlaneClient() + register_train_partition(client, num_samples=n, fields=list(fb.keys())) + bulk_batch: BatchedDataDict = BatchedDataDict() + _sync_rollout_write_loop(fb, bulk_batch) + meta = kv_first_write( + bulk_batch, + sample_ids=keys_from_uids(["a", "b", "c", "d"]), + dp_client=client, + partition_id="train", + tags=multimodal_row_tags(fb.get_multimodal_dict(as_tensors=False), n), + ) + fetched = read_columns(client, meta, select_fields=meta.fields, layout="padded") + + # The regression: this raised RuntimeError when the packed fields were + # left as raw nested tensors. + sliced = fetched.slice(1, 3) + + for key in ("pixel_values", "image_grid_thw"): + assert isinstance(sliced[key], PackedTensor) + expected = packed_truth[key].slice([1, 2]).as_tensor() + assert torch.equal(sliced.get_multimodal_dict(as_tensors=True)[key], expected) + + +def test_from_wire_keeps_packed_multimodal_fields_nested(): + """``_from_wire`` must not densify packed multimodal fields. + + It stacks any nested field whose rows all share a shape, to restore the + dense representation of ordinary batched inputs. For a packed multimodal + field, uniform rows are a *data-dependent accident* -- every sample simply + happened to carry one image -- and stacking discards the row boundaries. + + The dense value then fails the ``is_nested`` check in + ``codec.materialize``, so it is never reassembled, and reaches + ``get_multimodal_dict`` as a dense tensor:: + + ValueError: 'image_grid_thw' is still in data-plane wire form (dense). + Packed multimodal fields must be rebuilt by + multimodal_utils.reassemble_packed_multimodal ... + + That killed a Qwen3.5 VLM GRPO run on a batch where every sample had + exactly one image. + """ + from tensordict import TensorDict + + from nemo_rl.data_plane.adapters.transfer_queue import _from_wire + + # Rows that flatten to equal lengths -- one image per sample. + uniform_grid = PackedTensor( + [ + torch.tensor([[1, 2, 2]], dtype=torch.int64), + torch.tensor([[1, 3, 3]], dtype=torch.int64), + torch.tensor([[1, 2, 1]], dtype=torch.int64), + ], + dim_to_pack=0, + ) + # A non-multimodal field with equally uniform rows, to pin that the + # densification still happens for everything else. + ordinary = torch.nested.as_nested_tensor( + [torch.ones(4), torch.ones(4), torch.ones(4)], layout=torch.jagged + ) + + td = TensorDict( + {"image_grid_thw": uniform_grid.to_wire()[0], "input_ids": ordinary}, + batch_size=(3,), + ) + out = _from_wire(td) + + assert out["image_grid_thw"].is_nested, ( + "packed multimodal field was densified; row boundaries are lost" + ) + assert not out["input_ids"].is_nested, ( + "ordinary uniform field should still be densified" + ) + # Still reassembles into the original per-row values. + rebuilt = PackedTensor.from_wire(out["image_grid_thw"], uniform_grid.to_wire()[1]) + assert torch.equal(rebuilt.as_tensor(), uniform_grid.as_tensor()) diff --git a/tests/unit/data_plane/test_preshard_extras.py b/tests/unit/data_plane/test_preshard_extras.py index 0c5b9e0d62f..752f635b7c1 100644 --- a/tests/unit/data_plane/test_preshard_extras.py +++ b/tests/unit/data_plane/test_preshard_extras.py @@ -133,6 +133,34 @@ def test_shard_meta_for_dp_preserves_partition_id(): assert all(m.partition_id == "train" for m in metas) +def test_shard_meta_for_dp_permutes_tags_with_sample_ids(): + """Tags must ride each sample, not each position. + + ``tags`` carries the per-row multimodal shapes. Sharding sorts by sequence + length, so a tags list left un-permuted would pair one sample's pixel bytes + with another sample's shapes — and ``from_wire`` would reshape into it + silently instead of raising. + """ + n, dp = 8, 4 + meta = _meta(n) + # Tag each sample with its own id so a mispairing is directly visible. + meta.tags = [{"owner": sid} for sid in meta.sample_ids] + + metas, _ = shard_meta_for_dp(meta, dp_world=dp, batch_size=n) + + for shard in metas: + assert shard.tags is not None + assert len(shard.tags) == len(shard.sample_ids) + for sid, tag in zip(shard.sample_ids, shard.tags): + assert tag["owner"] == sid + + +def test_shard_meta_for_dp_leaves_tags_none_when_absent(): + """Text-only runs write no tags; sharding must not invent an empty list.""" + metas, _ = shard_meta_for_dp(_meta(4), dp_world=2, batch_size=4) + assert all(m.tags is None for m in metas) + + def test_shard_meta_for_dp_unsorted_round_trip(): """unsorted_indices must reconstruct the input order from DP-rank concat.""" n, dp = 8, 4 diff --git a/tests/unit/data_plane/test_scalar_field_schema_patch.py b/tests/unit/data_plane/test_scalar_field_schema_patch.py new file mode 100644 index 00000000000..3c106cc598d --- /dev/null +++ b/tests/unit/data_plane/test_scalar_field_schema_patch.py @@ -0,0 +1,155 @@ +# Copyright (c) 2025, 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. +"""The dense 1-D sample shape TQ reports must match the rows it stores. + +Upstream ``extract_field_schema`` rebinds a *local* for 1-D inputs +(``value = value.unsqueeze(-1)``) and derives the sample shape from it, +while ``_generate_values`` iterates the *original* ``(N,)`` tensor into +``N`` 0-d rows. Schema says ``(1,)``, storage holds ``()``. + +Only the KV path notices, because it is the only one that reconstructs +from the schema: ``BatchMeta.get_shapes`` repeats the uniform ``shape`` +per sample and the mooncake client reshapes raw bytes with it. That turns +a scalar column into ``(1,)`` rows, which then re-nest instead of taking +``_merge_tensors_to_tensordict``'s ``all(dim() == 0) -> torch.stack`` +branch. ``SimpleStorage`` fetches stored objects by ``(index, field)`` and +never consults the schema, so it is unaffected. + +These tests pin the invariant on TQ's own function rather than on a +recorded shape, so they fail if a future TQ revision changes the +derivation in either direction — the fix landing upstream shows up here as +a still-passing test, and a different 1-D convention shows up as a failure +rather than as silent nested scalars in production. +""" + +from __future__ import annotations + +import pytest +import torch + +from nemo_rl.data_plane.adapters import transfer_queue as tq_adapter + +tq_metadata = pytest.importorskip( + "transfer_queue.metadata", + reason="transfer_queue not installed", +) +TensorDict = pytest.importorskip("tensordict").TensorDict + + +@pytest.fixture(autouse=True) +def _patched(): + """Apply the patch for each test; it is idempotent and process-global.""" + tq_adapter._patch_scalar_field_schema() + yield + + +def _schema(fields: dict, n: int) -> dict: + return tq_metadata.extract_field_schema(TensorDict(fields, batch_size=[n])) + + +def test_dense_1d_field_reports_scalar_sample_shape() -> None: + """``(N,)`` in means ``()`` per sample — what ``_generate_values`` stores. + + A ``(1,)`` here is the upstream bug: the KV read would reshape each + row to ``(1,)`` and the column would come back nested. + """ + schema = _schema({"input_lengths": torch.arange(4, dtype=torch.int64)}, 4) + assert tuple(schema["input_lengths"]["shape"]) == () + + +def test_byte_count_is_unchanged_by_the_fix() -> None: + """``prod(()) == prod((1,)) == 1`` — this is a reshape fix, not a sizing + one. Pinned because a sizing change would corrupt reads rather than + mis-shape them, and that is a much worse failure to discover late.""" + from transfer_queue.utils.tensor_utils import get_nbytes + + n = 4 + shape = _schema({"input_lengths": torch.arange(n, dtype=torch.int64)}, n)[ + "input_lengths" + ]["shape"] + assert get_nbytes([torch.int64] * n, [shape] * n) == [8] * n + + +def test_2d_and_nested_fields_are_left_alone() -> None: + """The patch must touch dense 1-D only. + + ``(N, S)`` already agrees with its stored rows, and nested fields carry + exact ``per_sample_shapes`` — rewriting either would break the fields + that were never broken. + """ + rows = [torch.ones(3), torch.ones(5)] + schema = _schema( + { + "input_ids": torch.zeros(2, 7, dtype=torch.long), + "logprobs": torch.nested.as_nested_tensor(rows, layout=torch.jagged), + }, + 2, + ) + assert tuple(schema["input_ids"]["shape"]) == (7,) + assert schema["logprobs"]["is_nested"] + assert [tuple(s) for s in schema["logprobs"]["per_sample_shapes"]] == [(3,), (5,)] + + +def test_probe_confirms_tq_stores_scalar_rows_as_0d() -> None: + """The premise the whole patch rests on, checked against TQ itself. + + There is no payload-side fallback any more, so if a TQ revision started + storing dense 1-D fields as ``(1,)`` rows the rewrite would turn a + correct schema into a wrong one — and the symptom would be corrupt + reads, not an import error. ``_patch_scalar_field_schema`` runs this + probe before installing itself; here it is asserted directly so the + failure names the cause. + """ + from transfer_queue.storage.managers.base import KVStorageManager + + rows = KVStorageManager._generate_values( + TensorDict({"probe": torch.zeros(2)}, batch_size=[2]) + ) + assert len(rows) == 2 + assert all(r.ndim == 0 for r in rows), ( + f"TQ now stores dense 1-D rows as {[tuple(r.shape) for r in rows]}; " + "the scalar schema patch must be re-checked" + ) + + +def test_patch_is_idempotent() -> None: + """Every process that builds a client calls it; double-application must + not stack wrappers (which would still be correct but unboundedly deep).""" + first = tq_metadata.extract_field_schema + tq_adapter._patch_scalar_field_schema() + assert tq_metadata.extract_field_schema is first + + +def test_storage_managers_see_the_patched_function() -> None: + """Both managers bind the name at import time + (``from transfer_queue.metadata import extract_field_schema``), so + rebinding only the defining module would leave them on the original and + the fix would silently not apply where it is actually called.""" + from transfer_queue.storage.managers import base as _base + + assert _base.extract_field_schema is tq_metadata.extract_field_schema + + +def test_scalar_rows_round_trip_as_a_dense_column() -> None: + """End state the fix exists for: reconstructing with the reported shape + yields 0-d rows, which ``_merge_tensors_to_tensordict`` stacks back into + a dense ``(N,)`` column instead of re-nesting it.""" + n = 4 + src = torch.arange(n, dtype=torch.int64) + shape = _schema({"input_lengths": src}, n)["input_lengths"]["shape"] + + # What the KV client does: one stored row -> reshape(reported shape). + rebuilt = [row.reshape(tuple(shape)) for row in src] + assert all(r.dim() == 0 for r in rebuilt) + assert torch.equal(torch.stack(rebuilt), src) diff --git a/tests/unit/data_plane/test_tq_lifecycle.py b/tests/unit/data_plane/test_tq_lifecycle.py index e51fe398741..77642c1dee3 100644 --- a/tests/unit/data_plane/test_tq_lifecycle.py +++ b/tests/unit/data_plane/test_tq_lifecycle.py @@ -215,7 +215,6 @@ def test_each_public_data_operation_marks_the_client_dirty( client._data_operations_started = False client._warmed_fields = {} client._poll_interval_s = 0 - client._promote_1d = False _DATA_OPERATION_INVOKERS[operation_name](client) @@ -297,7 +296,6 @@ def test_checkpoint_load_rejects_client_after_data_operation( client = object.__new__(tq_adapter.TQDataPlaneClient) client._backend = "simple" client._supports_checkpointing = True - client._promote_1d = False client._data_operations_started = False client.put_samples( sample_ids=["sample-0"], @@ -454,9 +452,13 @@ def test_smoke_round_trip_backends(tq_client_backends) -> None: def test_smoke_round_trip_1d_fields(tq_client_backends) -> None: """A 1D (N,) tensor put into TQ must come back as (N,), not (N,1). - Regression guard for R-C2: TQ's KVStorageManager path silently unsqueezes - 1D fields. The adapter's `_promote_1d_leaves` + `_from_wire` pair fixes - this for mooncake_cpu; simple passes the tensor through unchanged. + Regression guard for R-C2, and the end-to-end proof of + ``_patch_scalar_field_schema``: upstream ``extract_field_schema`` + reports a ``(1,)`` sample shape for dense 1-D fields while storage + holds 0-d rows, so the KV path would rebuild this column as nested + ``(1,)`` rows. The patch makes the reported shape match what is + stored; ``simple`` never consulted the schema and is unaffected either + way, so running both backends here pins that they agree. """ n = 6 total_reward = torch.arange(n, dtype=torch.float32) @@ -650,34 +652,34 @@ def test_object_and_tensor_mixed_round_trip_backends(tq_client_backends) -> None client.clear_samples(sample_ids=None, partition_id=partition_id) -def test_promote_1d_leaves_object_array_roundtrip() -> None: - """``_promote_1d_leaves`` + ``_from_wire`` preserves non-tensor leaves. +def test_from_wire_preserves_object_arrays_through_densify() -> None: + """``_from_wire`` must not drop non-tensor leaves when it rebuilds. - Pins the production TD shape (1D tensor + object array + 2D tensor) - against tensordict 0.12.2 reconstruction bugs that could silently - strip ``NonTensorStack`` / ``NonTensorData`` leaves. Symmetric to - the documented ``.contiguous()`` bug in - ``adapters/transfer_queue.py`` lines 558–562. + It only constructs a new ``TensorDict`` when it densifies a uniform + nested field, and that rebuild is where tensordict 0.12.2 can silently + strip ``NonTensorStack`` / ``NonTensorData`` leaves. Symmetric to the + documented ``.contiguous()`` bug in ``adapters/transfer_queue.py``. """ - from nemo_rl.data_plane.adapters.transfer_queue import ( - _from_wire, - _promote_1d_leaves, - ) + from nemo_rl.data_plane.adapters.transfer_queue import _from_wire arr = np.empty(4, dtype=object) arr[:] = [["a", "b"], ["c"], ["d", "e"], ["f"]] + # Uniform rows, so _from_wire densifies and takes the rebuild path. + nested = torch.nested.as_nested_tensor( + [torch.zeros(8, dtype=torch.long) for _ in range(4)], layout=torch.jagged + ) td = TensorDict( { - "input_ids": torch.zeros(4, 8, dtype=torch.long), - "input_lengths": torch.tensor([4, 3, 2, 1]), # 1D → promoted + "input_ids": nested, + "input_lengths": torch.tensor([4, 3, 2, 1]), "content": arr, }, batch_size=[4], ) - promoted = _promote_1d_leaves(td) - assert promoted["input_lengths"].shape == (4, 1) - np.testing.assert_array_equal(promoted["content"], arr) - restored = _from_wire(promoted) + restored = _from_wire(td) + assert not restored["input_ids"].is_nested + assert restored["input_ids"].shape == (4, 8) + # Scalar column passes through untouched — no squeeze step any more. assert restored["input_lengths"].shape == (4,) np.testing.assert_array_equal(restored["content"], arr) diff --git a/tests/unit/models/automodel/test_automodel_train.py b/tests/unit/models/automodel/test_automodel_train.py index 9e6f423dfa7..92737b121e1 100644 --- a/tests/unit/models/automodel/test_automodel_train.py +++ b/tests/unit/models/automodel/test_automodel_train.py @@ -870,6 +870,38 @@ def test_ignores_1d_tensors(self): assert seq_dim == 1 assert seq_dim_size == 64 + def test_ignores_packed_multimodal_wire_fields(self): + """A packed multimodal field has patch/image count on dim 1, not + seqlen. Without the skip every VLM step through the data plane trips + this assert.""" + data = BatchedDataDict( + { + "input_ids": torch.randint(0, 1000, (4, 64)), + # [B, max_patches, C, H, W] — dim 1 is 7 patches, not 64 tokens. + "pixel_values": torch.randn(4, 7, 3, 2, 2), + "image_grid_thw": torch.zeros(4, 2, 3, dtype=torch.long), + } + ) + + seq_dim, seq_dim_size = check_sequence_dim(data) + + assert seq_dim == 1 + assert seq_dim_size == 64 + + def test_per_token_multimodal_is_still_checked(self): + """Only the *packed* registry is exempt. Per-token maps like + ``mm_token_type_ids`` are ``[B, S]`` and a seqlen mismatch there is a + real bug that must keep failing.""" + data = BatchedDataDict( + { + "input_ids": torch.randint(0, 1000, (4, 64)), + "mm_token_type_ids": torch.zeros(4, 32, dtype=torch.long), + } + ) + + with pytest.raises(AssertionError, match="Dim 1 must be the sequence dim"): + check_sequence_dim(data) + # ===================== # Test ProcessedInputs properties diff --git a/tests/unit/models/megatron/test_megatron_data.py b/tests/unit/models/megatron/test_megatron_data.py index 9390a112c8d..13fae82dc3c 100644 --- a/tests/unit/models/megatron/test_megatron_data.py +++ b/tests/unit/models/megatron/test_megatron_data.py @@ -146,6 +146,54 @@ def test_get_and_validate_seqlen_skips_1d_tensors(self): sequence_dim, seq_dim_size = get_and_validate_seqlen(data) assert seq_dim_size == 10 + def test_get_and_validate_seqlen_skips_packed_multimodal_fields(self): + """mcore twin of the automodel ``check_sequence_dim`` skip. + + A packed multimodal field has patch/image count on dim 1, not seqlen. + Without the skip every VLM step through the data plane trips the assert. + """ + from nemo_rl.models.megatron.data import get_and_validate_seqlen + + data = MagicMock() + data.__getitem__ = MagicMock( + side_effect=lambda k: torch.zeros(2, 10) if k == "input_ids" else None + ) + data.items = MagicMock( + return_value=[ + ("input_ids", torch.zeros(2, 10)), + # [B, max_patches, C, H, W] — dim 1 is 7 patches, not 10 tokens. + ("pixel_values", torch.zeros(2, 7, 3, 2, 2)), + ("image_grid_thw", torch.zeros(2, 4, 3)), + ] + ) + + sequence_dim, seq_dim_size = get_and_validate_seqlen(data) + + assert sequence_dim == 1 + assert seq_dim_size == 10 + + def test_get_and_validate_seqlen_still_checks_per_token_multimodal(self): + """Only the *packed* registry is exempt. Per-token maps like + ``mm_token_type_ids`` are ``[B, S]``, so a seqlen mismatch there is a + real bug that must keep failing.""" + from nemo_rl.models.megatron.data import get_and_validate_seqlen + + data = MagicMock() + data.__getitem__ = MagicMock( + side_effect=lambda k: torch.zeros(2, 10) if k == "input_ids" else None + ) + data.items = MagicMock( + return_value=[ + ("input_ids", torch.zeros(2, 10)), + ("mm_token_type_ids", torch.zeros(2, 15)), # Mismatched! + ] + ) + + with pytest.raises(AssertionError) as exc_info: + get_and_validate_seqlen(data) + + assert "Dim 1 must be the sequence dim" in str(exc_info.value) + @pytest.mark.mcore class TestProcessMicrobatch: diff --git a/tests/unit/models/policy/test_megatron_worker.py b/tests/unit/models/policy/test_megatron_worker.py index 2e9aea76076..8c4fdae81a5 100644 --- a/tests/unit/models/policy/test_megatron_worker.py +++ b/tests/unit/models/policy/test_megatron_worker.py @@ -149,18 +149,29 @@ class ModelSlicesContextParallelInputs: assert not _model_slices_context_parallel_inputs(object()) -def test_model_cp_slicing_rejects_transfer_queue_setup(): +def test_model_cp_slicing_accepts_transfer_queue_setup(monkeypatch): + """Media-before-CP models are served by the leader-broadcast fetch. + + ``_fetch`` broadcasts one DP slice across the replica group (TP x CP x PP + siblings of a DP rank), so every CP sibling gets identical full THD rows and + the model applies its own post-embedding slice. Setup used to reject these + models outright; the contract is satisfied, so it must not. + """ from nemo_rl.models.policy.workers.megatron_policy_worker import ( MegatronPolicyWorkerImpl, ) worker = object.__new__(MegatronPolicyWorkerImpl) worker.model_slices_context_parallel_inputs = True + worker._dp_client = None + + client = MagicMock() + monkeypatch.setattr( + "nemo_rl.data_plane.build_data_plane_client", lambda cfg, bootstrap: client + ) + worker.setup_data_plane(MagicMock()) - with pytest.raises( - NotImplementedError, match="TransferQueue/SingleController does not yet support" - ): - worker.setup_data_plane(MagicMock()) + assert worker._dp_client is client def test_refit_size_estimate_preserves_integral_buffer_dtype():