diff --git a/docs/design-docs/media-token-validity-mask.md b/docs/design-docs/media-token-validity-mask.md new file mode 100644 index 00000000000..4bf5caa5a47 --- /dev/null +++ b/docs/design-docs/media-token-validity-mask.md @@ -0,0 +1,214 @@ +# Training on text that contains `` + +## Summary + +A media placeholder is an ordinary vocabulary entry, so text can legitimately +contain it — a competitive-programming statement that spells `` in its +prose, for instance. The model reads every occurrence as an anchor for a +projected image feature and fails when no feature exists, so those rows cannot +be trained on. + +This design adds a **media-token validity mask**: the caller, which knows how +many media items each row carries, marks which placeholder positions are real +anchors. Positions it does not mark are skipped by the media merge and keep +whatever language embedding the forward gave them. + +## The problem + +`NemotronOmniModel._merge_projected_media` enforces a strict 1:1 contract — one +projected feature per valid placeholder: + +```python +media_mask = input_ids == media_token_id +if attention_mask is not None: + media_mask = media_mask & attention_mask.bool() + +expected_features = int(media_mask.sum().item()) +actual_features = media_embeddings.shape[0] +if expected_features != actual_features: + raise ValueError("Expanded-sequence media alignment failed: ...") +``` + +The contract is right. The question is where `media_mask` comes from. Before +this change the model derived it from whichever mask it had: + +```python +media_token_validity_mask = None +if padding_mask is not None: + media_token_validity_mask = ~padding_mask +elif attention_mask is not None and attention_mask.dim() == input_ids.dim(): + media_token_validity_mask = attention_mask +``` + +Both of those answer **"is this a real token?"**. The merge needs **"is this a +media anchor?"**. Those coincide only while every media token in a non-padding +position anchors an image. A text row that spells the placeholder breaks the +equivalence: the position is a real token, so the derived mask marks it valid, +so the merge demands a feature that was never meant to exist. + +This is not hypothetical. In a 100k-row sample of a Nemotron post-training +blend, 2,539 rows (2.5%) contain a literal `` with **zero** attached +images — competitive-programming statements where `` replaced inline +math, e.g. `"for any character there is exactly one character"`. They +carry between 1 and 12 placeholders each. + +There is a second, subtler source. Chat templates that treat a literal +`` in the prose as the placement anchor suppress their own generated +image block. A row with two literal tokens and one attached image therefore +renders two placeholders for one feature. + +## Why not sanitize + +The previous approach rewrote the data at rollout time: drop the literal token +when images were attached, or replace it with the word `image` when not. + +It worked, but it had three problems: + +1. **It edits the user's prose.** `"for any character there is"` became + `"for any character image there is"`. The model trains on text the author + did not write. +2. **It hides malformed data.** A row with two placeholders and one image is a + defect; sanitizing silently patched it. Filtering the blend surfaced 392 such + rows, which turned out to be unanswerable questions. +3. **It cannot express the legitimate case.** There is no rewrite that both + preserves the prose and tells the model "this one is not an anchor". + +## Design + +Give the caller a way to state the answer directly, since the caller is the only +party that knows it. The model keeps its strict contract; it just stops guessing +the input to that contract. + +### Model change + +`NemotronOmniModel.forward` takes a keyword-only argument that takes precedence +over the derived masks: + +```python +if media_token_validity_mask is None: + if padding_mask is not None: + media_token_validity_mask = ~padding_mask + elif attention_mask is not None and attention_mask.dim() == input_ids.dim(): + media_token_validity_mask = attention_mask +``` + +Behavior is unchanged when the argument is omitted. This is the only change to +the model, and it is additive. + +### Building the mask + +`build_media_token_validity_mask(input_ids, media_token_id, media_counts_by_row)` +marks media-token positions in rows whose media count is zero: + +- Every row carries media → returns `None`; the model derives its own mask. +- Text rows exist but none spell the token → returns `None` for the same reason. +- Otherwise → a `[B, S]` bool mask, `False` at media-token positions of + media-less rows. + +Rows that **do** carry media keep every position valid. A genuine +placeholder/feature disagreement there is still reported rather than masked +away, so the mask cannot be used to silence real misalignment. + +### Carrying it through sequence packing + +This is the part that determines whether the mask means anything. + +The mask is built in **sample space**, where row `i` of `input_ids` pairs with +`media_counts_by_row[i]`. Sequence packing then concatenates many samples into +one THD sequence and context-parallel-shards it. After that, one "row" holds +many samples and each rank holds a slice — the per-row question the mask answers +can no longer be asked. + +So the mask must travel through the *same* transform as `input_ids` rather than +be derived downstream. It is packed alongside them, exactly as `mtp_loss_mask` +already is: + +```python +if "media_token_validity_mask" in data_dict: + packed_media_mask, local_media_mask, _, _, _ = _pack_sequences_for_megatron( + data_dict["media_token_validity_mask"].to(data_dict["input_ids"].dtype), + seq_lengths, + pad_individual_seqs_to_multiple_of, + pad_packed_seq_to_multiple_of, + pad_full_seq_to, + cp_rank=get_context_parallel_rank(), + cp_size=get_context_parallel_world_size(), + ) + media_token_validity_mask = ( + packed_media_mask + if model_slices_context_parallel_inputs + else local_media_mask + ).bool() +``` + +Two details matter: + +- **Packed in token dtype.** Packing pads with `value=0`, which is a valid token + id but not a valid bool. It is converted back to bool after packing. Padding + positions become `False`, which is harmless because padding holds no media + token. +- **Packed vs CP-local.** A model that slices context parallelism itself + receives the full THD row so it can insert media before selecting its + CP-owned embeddings. Its mask must stay unsharded to line up. Every other + model consumes the CP-local shard. This mirrors the `mtp_loss_mask` choice + for the same reason. `NemotronOmniModel` sets + `model_slices_context_parallel_inputs = True`, so it gets the full row. + +### Capability detection + +The mask is only sent to models whose `forward` actually declares it: + +```python +def _model_accepts_media_token_validity_mask(model) -> bool: + ... + if "media_token_validity_mask" in inspect.signature(chunk.forward).parameters: + return True +``` + +The check is on the signature rather than a class flag because a model that does +not know about the mask would absorb it into `**kwargs` and ignore it — which +looks identical to the mask having been applied. Failing to send it is visible; +sending it into a void is not. + +### Where it is attached + +At all three forward sites: training and **both** logprob paths. Logprobs run +the same forward as training, so a batch that needs the mask needs it there too +— otherwise logprobs would be computed against a different media alignment than +the one trained on. + +### Self-packing models + +Models that pack internally (`delegate_pack_to_model`) raise +`NotImplementedError` rather than silently dropping the mask. A mask built +against caller-side rows would reach the merge in a layout that no longer +matches its tokens, and a misaligned media mask attaches features to the wrong +positions without erroring. + +## Consequences + +**Removed.** `sanitize_nemo_gym_example_image_placeholders`, +`_normalize_image_placeholders`, `_count_image_payloads`, the +`sanitize_image_placeholders` config flag, and their tests — 212 lines. + +**Data prep is now required, not optional.** Without the sanitizer there is no +safety net: a blend whose placeholder and image counts disagree fails loudly at +the media merge. That is the intended behavior — the failure is a real defect — +but it means malformed rows must be filtered before training rather than being +absorbed at rollout time. + +**Text rows containing `` can be trained on** instead of dropped, which +is what the mask exists for. + +## Validation + +| Test | Result | +|---|---| +| 50-step Super Omni run, sanitization removed | 50/50 steps, zero media-alignment / device / IMA errors; reward 0.458 (steps 1–8) → 0.520 (steps 26–49), peak 0.5879 vs. a 0.4932 baseline best | +| Mixed batch: 128 rows with real images + 128 text-only rows carrying real `` statements | Reached Step 1/1, rollouts 100%, zero alignment failures | +| Unit: mask construction | 6 tests | +| Unit: mask survives packing, lands on the same tokens | 2 tests (require mcore) | + +The packing tests are the load-bearing ones: they assert the mask still marks +the intended tokens after `_pack_sequences_for_megatron`, because a mask that is +merely misaligned does not raise — it attaches features to the wrong positions. diff --git a/docs/index.md b/docs/index.md index 8f673f02f24..7fe07862184 100644 --- a/docs/index.md +++ b/docs/index.md @@ -371,6 +371,7 @@ design-docs/env-vars.md design-docs/nemo-gym-integration.md design-docs/modelopt-real-quant-architecture.md design-docs/nccl-reshard-refit.md +design-docs/media-token-validity-mask.md ``` ```{toctree} diff --git a/nemo_rl/data/multimodal_utils.py b/nemo_rl/data/multimodal_utils.py index 156465e9c1a..68dd4e9d097 100644 --- a/nemo_rl/data/multimodal_utils.py +++ b/nemo_rl/data/multimodal_utils.py @@ -18,6 +18,7 @@ import re import uuid from collections import defaultdict +from collections.abc import Sequence from copy import deepcopy from io import BytesIO from typing import Any, Optional, Union @@ -1151,3 +1152,130 @@ def load_media_from_message( loaded_media["video"].append(vid) return loaded_media + + +def build_media_token_validity_mask( + input_ids: torch.Tensor, + media_token_id: int, + media_counts_by_row: Sequence[int], + base_mask: Optional[torch.Tensor] = None, +) -> Optional[torch.Tensor]: + """Mark media-token positions in rows that carry no media of that modality. + + A media token is an ordinary vocabulary entry with its own embedding row. + It only means "a projected feature belongs here" when media is attached; + in a text-only row the same id is whatever the author wrote, and counting + it as a placeholder makes the model demand a feature that does not exist. + + Rows that do carry media keep every position valid, so a real + placeholder/feature disagreement there is still reported rather than + silently masked away. + + Args: + input_ids: ``[B, S]`` token ids, one row per sample. + media_token_id: Vocabulary id the model treats as a media placeholder. + media_counts_by_row: Attached media items per row, e.g. from + :meth:`PackedTensor.logical_segment_counts_by_row`. + base_mask: Optional ``[B, S]`` validity mask to refine, so masks for + several modalities can be combined. + + Returns: + A ``[B, S]`` bool mask, or ``None`` when nothing needs masking and the + caller should leave the model's own derivation untouched. + """ + if input_ids.ndim != 2: + raise ValueError(f"input_ids must be [B, S], got {tuple(input_ids.shape)}") + if len(media_counts_by_row) != input_ids.shape[0]: + raise ValueError( + "media_counts_by_row must have one entry per row: got " + f"{len(media_counts_by_row)} for {input_ids.shape[0]} rows" + ) + + empty_rows = [row for row, count in enumerate(media_counts_by_row) if count == 0] + if not empty_rows: + return base_mask + + is_media_token = input_ids == media_token_id + if not bool(is_media_token[empty_rows].any()): + # Text-only rows exist but none of them spell the token, so the + # model's own derivation is already correct. + return base_mask + + mask = ( + torch.ones_like(input_ids, dtype=torch.bool) + if base_mask is None + else base_mask.clone() + ) + for row in empty_rows: + mask[row] &= ~is_media_token[row] + return mask + + +def media_placeholder_token_id_from_chunks(chunks: Sequence[Any]) -> Optional[int]: + """The vocabulary id these model chunks treat as a media placeholder, if any.""" + for chunk in chunks: + token_id = getattr(chunk, "image_token_index", None) + if token_id is not None: + return int(token_id) + return None + + +def chunks_accept_media_token_validity_mask(chunks: Sequence[Any]) -> bool: + """Whether a model chunk's forward takes an explicit media-token validity mask. + + Checked against the signature rather than a class flag so a model that does + not know about the mask never receives it: such a forward would absorb it + into ``**kwargs`` and silently ignore it, which looks identical to the mask + having been applied. Failing to send it is visible; sending it into a void + is not. + """ + for chunk in chunks: + try: + parameters = inspect.signature(chunk.forward).parameters + except (TypeError, ValueError): + continue + if "media_token_validity_mask" in parameters: + return True + return False + + +def image_counts_by_row(batch: Any, num_rows: int) -> Optional[list[int]]: + """How many images each row of the batch actually carries. + + Returns None when the batch describes its images in a way this cannot read, + so the caller leaves the model's own derivation alone rather than guessing a + count and masking against it. + """ + pixel_values = batch.get("pixel_values", None) + if pixel_values is None: + # A text-only batch genuinely has no images anywhere, which is exactly + # the case the mask exists for -- not a missing-data case. + return [0] * num_rows + if not isinstance(pixel_values, PackedTensor): + return None + counts = pixel_values.logical_segment_counts_by_row() + return counts if len(counts) == num_rows else None + + +def attach_media_token_validity_mask(batch: Any, media_token_id: Optional[int]) -> None: + """Mark media tokens that anchor nothing, so the model keeps their embedding. + + Builds the mask while rows still are samples. Sequence packing later + concatenates those rows into one THD sequence, after which no per-row + question can be asked, so the packing step carries this through the same + transform as ``input_ids`` rather than deriving it downstream. + + The batch is duck-typed rather than annotated as ``BatchedDataDict``: + that module imports this one, so naming it here would be circular. + """ + if media_token_id is None: + return + input_ids = batch.get("input_ids", None) + if not isinstance(input_ids, torch.Tensor) or input_ids.ndim != 2: + return + counts = image_counts_by_row(batch, input_ids.shape[0]) + if counts is None: + return + mask = build_media_token_validity_mask(input_ids, media_token_id, counts) + if mask is not None: + batch["media_token_validity_mask"] = mask diff --git a/nemo_rl/models/megatron/data.py b/nemo_rl/models/megatron/data.py index a0124333995..c55740b1b7e 100644 --- a/nemo_rl/models/megatron/data.py +++ b/nemo_rl/models/megatron/data.py @@ -50,6 +50,7 @@ class ProcessedInputs: mtp_loss_mask: Optional[torch.Tensor] = None routed_experts: Optional[torch.Tensor] = None routed_experts_cp_sharded: Optional[torch.Tensor] = None + media_token_validity_mask: Optional[torch.Tensor] = None @dataclass @@ -72,6 +73,9 @@ class ProcessedMicrobatch: None when MTP is disabled or token/sample masks are absent. routed_experts: Optional token-aligned routed expert ids routed_experts_cp_sharded: Context-parallel sharded routed expert ids + media_token_validity_mask: Which media-token positions actually anchor a + projected feature, in the model's own token layout. None when the + batch needs no correction and the model should derive its own. """ data_dict: BatchedDataDict[Any] @@ -84,6 +88,7 @@ class ProcessedMicrobatch: mtp_loss_mask: Optional[torch.Tensor] = None routed_experts: Optional[torch.Tensor] = None routed_experts_cp_sharded: Optional[torch.Tensor] = None + media_token_validity_mask: Optional[torch.Tensor] = None def make_processed_microbatch_iterator( @@ -146,6 +151,7 @@ def make_processed_microbatch_iterator( mtp_loss_mask=processed_inputs.mtp_loss_mask, routed_experts=processed_inputs.routed_experts, routed_experts_cp_sharded=processed_inputs.routed_experts_cp_sharded, + media_token_validity_mask=processed_inputs.media_token_validity_mask, ) @@ -286,6 +292,7 @@ def process_microbatch( cu_seqlens = None cu_seqlens_padded = None mtp_loss_mask = None + media_token_validity_mask = None if pack_sequences: # For packed sequences with padded input, we need sequence lengths @@ -305,6 +312,17 @@ def process_microbatch( "MTP training requires a self-packing VLM that advertises " "model_owns_mtp_loss_mask_packing" ) + if "media_token_validity_mask" in data_dict: + # A self-packing model repacks internally, so a mask built + # against caller-side rows would reach the merge in a layout + # that no longer matches its tokens -- and a media mask that + # is merely misaligned silently attaches features to the + # wrong positions rather than failing. + raise NotImplementedError( + "media_token_validity_mask is not supported for models " + "that pack sequences internally (delegate_pack_to_model); " + "the mask would need to be packed inside the model." + ) # VLM path: model (e.g. mbridge Qwen3VL) does its own # preprocess_packed_seqs; NeMo-RL must NOT pre-pack + CP-shard, # or the double-processing produces shape mismatches downstream @@ -501,6 +519,43 @@ def process_microbatch( else local_mtp_loss_mask ) + # Pack the media-token validity mask the same way as input_ids. + # The mask answers a per-token question, so it only means + # anything while it sits in the same layout as the tokens the + # model will compare it against. Packing is what destroys the + # per-sample rows it was built from, so it has to travel through + # the identical transform rather than be rebuilt afterwards. + if "media_token_validity_mask" in data_dict: + ( + packed_media_mask, + local_media_mask, + _, + _, + _, + ) = _pack_sequences_for_megatron( + # Pack in the token dtype: padding is filled with 0, + # which is a valid token id but not a valid bool. Read + # the dtype off the unpacked ids, since the local + # input_ids is already the packed tensor here. + data_dict["media_token_validity_mask"].to( + data_dict["input_ids"].dtype + ), + seq_lengths, + pad_individual_seqs_to_multiple_of, + pad_packed_seq_to_multiple_of, + pad_full_seq_to, + cp_rank=get_context_parallel_rank(), + cp_size=get_context_parallel_world_size(), + ) + # Mirror the input_ids layout choice above, for the same + # reason the MTP mask does: a model that slices CP itself + # merges media against the full THD row. + media_token_validity_mask = ( + packed_media_mask + if model_slices_context_parallel_inputs + else local_media_mask + ).bool() + # For packed sequences, position_ids and attention_mask are typically None # The PackedSeqParams handles all necessary sequence information position_ids = None @@ -550,6 +605,12 @@ def process_microbatch( ) if "mtp_loss_mask" in data_dict: mtp_loss_mask = data_dict["mtp_loss_mask"] + # Unpacked: rows still are samples, so the mask is already in the + # layout the model will see. + if "media_token_validity_mask" in data_dict: + media_token_validity_mask = data_dict[ + "media_token_validity_mask" + ].bool() return ProcessedInputs( input_ids=input_ids, input_ids_cp_sharded=input_ids_cp_sharded, @@ -560,6 +621,7 @@ def process_microbatch( mtp_loss_mask=mtp_loss_mask, routed_experts=routed_experts, routed_experts_cp_sharded=routed_experts_cp_sharded, + media_token_validity_mask=media_token_validity_mask, ) diff --git a/nemo_rl/models/megatron/train.py b/nemo_rl/models/megatron/train.py index d8ecb4f6590..6e2cfd6e6a1 100644 --- a/nemo_rl/models/megatron/train.py +++ b/nemo_rl/models/megatron/train.py @@ -127,6 +127,7 @@ def model_forward( mtp_loss_mask: Optional[torch.Tensor] = None, straggler_timer: Optional[StragglerDetector] = None, use_fused_linear_logprobs: bool = False, + media_token_validity_mask: Optional[torch.Tensor] = None, ) -> torch.Tensor: """Perform a single forward pass through the model. @@ -143,6 +144,9 @@ def model_forward( straggler_timer: Straggler detector for profiling the forward pass use_fused_linear_logprobs: Whether to compute logprobs with the fused chunked linear cross-entropy kernel (directly from hidden states) + media_token_validity_mask: Which media-token positions actually anchor a + projected feature, already in this model's token layout. Only passed + when the model accepts it; otherwise the model derives its own. Returns: torch.Tensor: Output tensor from the model (logits) @@ -162,6 +166,11 @@ def model_forward( if mtp_loss_mask is not None: additional_kwargs["loss_mask"] = mtp_loss_mask + # Only sent when the model advertises the parameter, so it never reaches a + # forward that would swallow it into **kwargs and quietly ignore it. + if media_token_validity_mask is not None: + additional_kwargs["media_token_validity_mask"] = media_token_validity_mask + if defer_fp32_logits: additional_kwargs["fp32_output"] = False if use_fused_linear_logprobs: @@ -259,6 +268,7 @@ def forward_with_post_processing_fn( cu_seqlens_padded = processed_mb.cu_seqlens_padded mtp_loss_mask = processed_mb.mtp_loss_mask routed_experts_cp_sharded = processed_mb.routed_experts_cp_sharded + media_token_validity_mask = processed_mb.media_token_validity_mask if use_router_replay: if routed_experts_cp_sharded is None: @@ -282,6 +292,7 @@ def forward_with_post_processing_fn( mtp_loss_mask=mtp_loss_mask, straggler_timer=straggler_timer, use_fused_linear_logprobs=use_fused_linear_logprobs, + media_token_validity_mask=media_token_validity_mask, ) except Exception: # The forward above armed the router-replay action (set_router_replay_forward); diff --git a/nemo_rl/models/policy/workers/megatron_policy_worker.py b/nemo_rl/models/policy/workers/megatron_policy_worker.py index b6ae9a23a95..7099a181cf2 100644 --- a/nemo_rl/models/policy/workers/megatron_policy_worker.py +++ b/nemo_rl/models/policy/workers/megatron_policy_worker.py @@ -50,6 +50,11 @@ from nemo_rl.algorithms.logits_sampling_utils import TrainingSamplingParams from nemo_rl.algorithms.loss.interfaces import LossFunction +from nemo_rl.data.multimodal_utils import ( + attach_media_token_validity_mask, + chunks_accept_media_token_validity_mask, + media_placeholder_token_id_from_chunks, +) from nemo_rl.data_plane.worker_mixin import TQWorkerMixin from nemo_rl.distributed.batched_data_dict import BatchedDataDict from nemo_rl.distributed.named_sharding import NamedSharding @@ -181,6 +186,24 @@ def _model_slices_context_parallel_inputs(model: Any) -> bool: ) +def _unwrapped_chunks(model: Any) -> list[Any]: + """Model chunks as a flat list, whatever wrapping the caller handed us.""" + from megatron.core.utils import unwrap_model + + unwrapped = unwrap_model(model) + return list(unwrapped) if isinstance(unwrapped, (list, tuple)) else [unwrapped] + + +def _model_media_placeholder_token_id(model: Any) -> Optional[int]: + """The vocabulary id this model treats as a media placeholder, if any.""" + return media_placeholder_token_id_from_chunks(_unwrapped_chunks(model)) + + +def _model_accepts_media_token_validity_mask(model: Any) -> bool: + """Whether the model's forward takes an explicit media-token validity mask.""" + return chunks_accept_media_token_validity_mask(_unwrapped_chunks(model)) + + def _estimate_refit_tensor_size_in_bytes( param: torch.Tensor, *, @@ -596,6 +619,14 @@ def __init__( self.model_slices_context_parallel_inputs = ( _model_slices_context_parallel_inputs(self.model) ) + # A media placeholder is an ordinary vocabulary entry, so text that + # legitimately contains it must not be read as an anchor demanding a + # projected feature. Only models that accept the mask are sent one. + self.media_placeholder_token_id = ( + _model_media_placeholder_token_id(self.model) + if _model_accepts_media_token_validity_mask(self.model) + else None + ) if self.model_slices_context_parallel_inputs: if self.delegate_pack_to_model: raise RuntimeError( @@ -814,6 +845,8 @@ def train( ].unsqueeze(-1) batch["mtp_loss_mask"] = mtp_loss_mask + attach_media_token_validity_mask(batch, self.media_placeholder_token_id) + ( data_iterator, num_microbatches, @@ -1648,6 +1681,11 @@ def get_logprobs( self.model.eval() + # Logprobs run the same forward as training, so a batch that needs the + # mask needs it here too -- otherwise these logprobs would be taken + # against a different media alignment than the one trained on. + attach_media_token_validity_mask(data, self.media_placeholder_token_id) + ( mb_iterator, num_microbatches, @@ -1867,6 +1905,8 @@ def get_topk_logits( self.model.eval() + attach_media_token_validity_mask(data, self.media_placeholder_token_id) + ( mb_iterator, num_microbatches, diff --git a/tests/unit/data/test_media_token_validity_mask.py b/tests/unit/data/test_media_token_validity_mask.py new file mode 100644 index 00000000000..2f39c775d5a --- /dev/null +++ b/tests/unit/data/test_media_token_validity_mask.py @@ -0,0 +1,176 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import pytest +import torch + +from nemo_rl.data.multimodal_utils import ( + PackedTensor, + attach_media_token_validity_mask, + build_media_token_validity_mask, + chunks_accept_media_token_validity_mask, + image_counts_by_row, + media_placeholder_token_id_from_chunks, +) + +IMG = 7 # stand-in media token id +TXT = 1 + + +def test_returns_base_mask_when_every_row_has_media(): + # Nothing to correct: each row's media token really does mark a feature. + input_ids = torch.tensor([[TXT, IMG, TXT], [IMG, TXT, TXT]]) + assert build_media_token_validity_mask(input_ids, IMG, [1, 1]) is None + + +def test_returns_base_mask_when_text_rows_do_not_spell_the_token(): + # Text-only rows exist, but none contain the token, so the model's own + # derivation is already right and we must not hand it a mask. + input_ids = torch.tensor([[TXT, IMG, TXT], [TXT, TXT, TXT]]) + assert build_media_token_validity_mask(input_ids, IMG, [1, 0]) is None + + +def test_masks_the_token_only_in_rows_without_media(): + # Row 0 has an image so its token is a real placeholder; row 1 has none, so + # its identical token is just prose the author wrote. + input_ids = torch.tensor([[TXT, IMG, TXT], [IMG, TXT, IMG]]) + mask = build_media_token_validity_mask(input_ids, IMG, [1, 0]) + assert mask is not None + torch.testing.assert_close( + mask, + torch.tensor([[True, True, True], [False, True, False]]), + ) + + +def test_refines_rather_than_replaces_a_base_mask(): + # A caller combining modalities must keep the positions the base mask + # already invalidated. + input_ids = torch.tensor([[IMG, TXT], [IMG, TXT]]) + base = torch.tensor([[True, False], [True, True]]) + mask = build_media_token_validity_mask(input_ids, IMG, [1, 0], base_mask=base) + assert mask is not None + torch.testing.assert_close( + mask, + torch.tensor([[True, False], [False, True]]), + ) + # The caller's tensor must not be mutated in place. + torch.testing.assert_close(base, torch.tensor([[True, False], [True, True]])) + + +def test_rejects_non_2d_input_ids(): + with pytest.raises(ValueError, match=r"input_ids must be \[B, S\]"): + build_media_token_validity_mask(torch.tensor([TXT, IMG]), IMG, [0]) + + +def test_rejects_count_length_mismatch(): + with pytest.raises(ValueError, match="one entry per row"): + build_media_token_validity_mask(torch.tensor([[TXT, IMG]]), IMG, [0, 1]) + + +# -------------------------------------------------------------------------- +# capability probes and batch-level attach +# -------------------------------------------------------------------------- + + +class _ChunkWithMask: + image_token_index = IMG + + def forward(self, input_ids, *, media_token_validity_mask=None): + raise AssertionError("not called") + + +class _ChunkWithoutMask: + image_token_index = IMG + + def forward(self, input_ids): + raise AssertionError("not called") + + +def test_placeholder_token_id_is_read_from_the_first_chunk_that_has_one(): + assert media_placeholder_token_id_from_chunks([_ChunkWithMask()]) == IMG + # A chunk without the attribute must not stop the search. + assert media_placeholder_token_id_from_chunks([object(), _ChunkWithMask()]) == IMG + + +def test_placeholder_token_id_is_none_when_no_chunk_declares_one(): + assert media_placeholder_token_id_from_chunks([object(), object()]) is None + assert media_placeholder_token_id_from_chunks([]) is None + + +def test_capability_is_read_from_the_forward_signature(): + """A flag could lie; the signature is what decides whether the kwarg lands.""" + assert chunks_accept_media_token_validity_mask([_ChunkWithMask()]) is True + assert chunks_accept_media_token_validity_mask([_ChunkWithoutMask()]) is False + assert chunks_accept_media_token_validity_mask([]) is False + + +def test_capability_skips_chunks_whose_signature_cannot_be_read(): + """An unintrospectable forward must not abort the search.""" + + class _Opaque: + forward = print # builtins raise ValueError from inspect.signature + + assert ( + chunks_accept_media_token_validity_mask([_Opaque(), _ChunkWithMask()]) is True + ) + + +def test_image_counts_treats_a_batch_without_pixel_values_as_text_only(): + """No images anywhere is the case the mask exists for, not missing data.""" + assert image_counts_by_row({}, 3) == [0, 0, 0] + + +def test_image_counts_returns_none_for_unreadable_media(): + """Rather than guess a count and mask against it.""" + assert image_counts_by_row({"pixel_values": torch.ones(2, 3)}, 2) is None + + +def test_image_counts_reads_logical_segments_per_row(): + packed = PackedTensor([torch.ones(1, 3, 2, 2)], dim_to_pack=0) + assert image_counts_by_row({"pixel_values": packed}, 1) == [1] + + +def test_image_counts_returns_none_on_row_count_mismatch(): + packed = PackedTensor([torch.ones(1, 3, 2, 2)], dim_to_pack=0) + assert image_counts_by_row({"pixel_values": packed}, 2) is None + + +def test_attach_sets_the_mask_for_a_text_row_that_spells_the_token(): + batch = {"input_ids": torch.tensor([[TXT, IMG, TXT]])} + attach_media_token_validity_mask(batch, IMG) + torch.testing.assert_close( + batch["media_token_validity_mask"], + torch.tensor([[True, False, True]]), + ) + + +def test_attach_is_a_noop_without_a_media_token_id(): + """Models that never declare the kwarg must not get a mask.""" + batch = {"input_ids": torch.tensor([[TXT, IMG]])} + attach_media_token_validity_mask(batch, None) + assert "media_token_validity_mask" not in batch + + +def test_attach_is_a_noop_when_nothing_needs_masking(): + """No key at all, so the model keeps deriving its own.""" + packed = PackedTensor([torch.ones(1, 3, 2, 2)], dim_to_pack=0) + batch = {"input_ids": torch.tensor([[TXT, IMG]]), "pixel_values": packed} + attach_media_token_validity_mask(batch, IMG) + assert "media_token_validity_mask" not in batch + + +def test_attach_ignores_a_batch_whose_input_ids_are_not_2d(): + batch = {"input_ids": torch.tensor([TXT, IMG])} + attach_media_token_validity_mask(batch, IMG) + assert "media_token_validity_mask" not in batch diff --git a/tests/unit/models/megatron/test_media_mask_packing.py b/tests/unit/models/megatron/test_media_mask_packing.py new file mode 100644 index 00000000000..c2089f5e18d --- /dev/null +++ b/tests/unit/models/megatron/test_media_mask_packing.py @@ -0,0 +1,101 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""The media-token validity mask must survive sequence packing. + +Packing concatenates per-sample rows into one THD sequence, which is exactly +the step that destroys the per-row structure the mask was built from. These +tests pin the mask to the tokens it is supposed to describe on the other side +of that transform -- a mask that is merely misaligned does not raise, it +attaches image features to the wrong positions. +""" + +import pytest +import torch + +from nemo_rl.data.multimodal_utils import build_media_token_validity_mask + +IMG = 7 +TXT = 1 + + +@pytest.mark.mcore +def test_mask_lands_on_the_same_tokens_after_packing(): + from nemo_rl.models.megatron.data import _pack_sequences_for_megatron + + # Row 0 carries an image, so its IMG is a real anchor. + # Row 1 is text that merely spells IMG, so its IMG must be masked off. + input_ids = torch.tensor([[TXT, IMG, TXT, TXT], [IMG, TXT, IMG, TXT]]) + seq_lengths = torch.tensor([3, 4]) + + mask = build_media_token_validity_mask(input_ids, IMG, [1, 0]) + assert mask is not None + + packed_ids, _, _, _, _ = _pack_sequences_for_megatron( + input_ids, seq_lengths, cp_rank=0, cp_size=1 + ) + packed_mask, _, _, _, _ = _pack_sequences_for_megatron( + mask.to(input_ids.dtype), seq_lengths, cp_rank=0, cp_size=1 + ) + packed_mask = packed_mask.bool() + + assert packed_ids.shape == packed_mask.shape + # Row 0 contributes tokens 0..2 (IMG at packed index 1, a true anchor); + # row 1 contributes 3..6 (IMG at packed indices 3 and 5, both bogus). + torch.testing.assert_close( + packed_ids, torch.tensor([[TXT, IMG, TXT, IMG, TXT, IMG, TXT]]) + ) + torch.testing.assert_close( + packed_mask, + torch.tensor([[True, True, True, False, True, False, True]]), + ) + # Every masked-off position is in fact a media token, and the anchoring row + # keeps its own. + assert bool((packed_ids[~packed_mask] == IMG).all()) + + +@pytest.mark.mcore +def test_padding_introduced_by_packing_is_not_treated_as_an_anchor(): + from nemo_rl.models.megatron.data import _pack_sequences_for_megatron + + # Padding is filled with token id 0; the mask packs to False there. That is + # only safe while 0 is not the media token, which this asserts explicitly. + input_ids = torch.tensor([[TXT, IMG, TXT, TXT]]) + seq_lengths = torch.tensor([3]) + mask = build_media_token_validity_mask(input_ids, IMG, [0]) + assert mask is not None + + packed_ids, _, _, _, _ = _pack_sequences_for_megatron( + input_ids, + seq_lengths, + pad_individual_seqs_to_multiple_of=4, + cp_rank=0, + cp_size=1, + ) + packed_mask, _, _, _, _ = _pack_sequences_for_megatron( + mask.to(input_ids.dtype), + seq_lengths, + pad_individual_seqs_to_multiple_of=4, + cp_rank=0, + cp_size=1, + ) + packed_mask = packed_mask.bool() + + assert IMG != 0, "padding fill value must not collide with the media token" + padded_positions = packed_ids[0, 3:] + assert bool((padded_positions == 0).all()) + # The real IMG at index 1 is masked off (row has no image); padding is + # False too, which is harmless because padding holds no media token. + assert not bool(packed_mask[0, 1]) + assert bool(packed_mask[0, 0]) and bool(packed_mask[0, 2])