From 796de752fe6a8e7e7cdf4b22220de453fca59f64 Mon Sep 17 00:00:00 2001 From: DanialTaheri Date: Thu, 13 Aug 2026 12:58:24 -0700 Subject: [PATCH 1/4] feat(vlm): carry a media-token validity mask through sequence packing A media placeholder is an ordinary vocabulary entry, so text can legitimately contain it -- competitive-programming statements, for instance, where replaced inline math. The model derives which placeholder positions anchor a projected feature from the padding or attention mask, but both answer "is this a real token", which is a different question from "is this a media anchor". They coincide only while every media token in a valid position anchors an image, so a text row that spells the placeholder makes the model demand a feature that was never meant to exist. Build the answer where it is known. The caller knows how many media items each row carries, so build_media_token_validity_mask marks the placeholder positions of media-less rows. Rows that do carry media keep every position valid, so a genuine placeholder/feature disagreement is still reported rather than masked away. The mask is built in sample space, where a row is a sample. Sequence packing then concatenates samples into one THD sequence and CP-shards it, after which no per-row question can be asked -- so the mask travels through the same packing transform as input_ids rather than being derived downstream, exactly as mtp_loss_mask already does. It is packed in the token dtype because packing pads with 0, which is a valid token id but not a valid bool, and a model that slices CP itself receives the unsharded row so the mask lines up after media insertion. Only models whose forward declares the argument receive it. 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. Attached at all three forward sites: logprobs run the same forward as training, so computing them without the mask would score a different media alignment than the one trained on. Models that pack internally raise rather than silently dropping the mask: a misaligned media mask attaches features to the wrong positions without erroring. Signed-off-by: DanialTaheri --- nemo_rl/data/multimodal_utils.py | 58 ++++++++++ nemo_rl/models/megatron/data.py | 62 +++++++++++ nemo_rl/models/megatron/train.py | 11 ++ .../policy/workers/megatron_policy_worker.py | 103 ++++++++++++++++++ .../data/test_media_token_validity_mask.py | 71 ++++++++++++ .../megatron/test_media_mask_packing.py | 101 +++++++++++++++++ 6 files changed, 406 insertions(+) create mode 100644 tests/unit/data/test_media_token_validity_mask.py create mode 100644 tests/unit/models/megatron/test_media_mask_packing.py diff --git a/nemo_rl/data/multimodal_utils.py b/nemo_rl/data/multimodal_utils.py index 156465e9c1a..c7c99e3da33 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,60 @@ 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 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..90d79b048b5 100644 --- a/nemo_rl/models/policy/workers/megatron_policy_worker.py +++ b/nemo_rl/models/policy/workers/megatron_policy_worker.py @@ -50,6 +50,10 @@ from nemo_rl.algorithms.logits_sampling_utils import TrainingSamplingParams from nemo_rl.algorithms.loss.interfaces import LossFunction +from nemo_rl.data.multimodal_utils import ( + PackedTensor, + build_media_token_validity_mask, +) 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 +185,86 @@ def _model_slices_context_parallel_inputs(model: Any) -> bool: ) +def _model_media_placeholder_token_id(model: Any) -> Optional[int]: + """The vocabulary id this model treats as a media placeholder, if any.""" + from megatron.core.utils import unwrap_model + + unwrapped = unwrap_model(model) + chunks = unwrapped if isinstance(unwrapped, (list, tuple)) else [unwrapped] + 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 _model_accepts_media_token_validity_mask(model: Any) -> bool: + """Whether the model'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. + """ + import inspect + + from megatron.core.utils import unwrap_model + + unwrapped = unwrap_model(model) + chunks = unwrapped if isinstance(unwrapped, (list, tuple)) else [unwrapped] + 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 _attach_media_token_validity_mask( + batch: BatchedDataDict[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 ``process_microbatch`` carries this through the + same packing transform as ``input_ids`` rather than deriving it downstream. + """ + 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 + + +def _image_counts_by_row( + batch: BatchedDataDict[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 _estimate_refit_tensor_size_in_bytes( param: torch.Tensor, *, @@ -596,6 +680,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 +906,10 @@ 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 +1744,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 +1968,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..43afb870ecc --- /dev/null +++ b/tests/unit/data/test_media_token_validity_mask.py @@ -0,0 +1,71 @@ +# 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 build_media_token_validity_mask + +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]) 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]) From 4e3eed85540fc71babb15f3397f22dc919f130e8 Mon Sep 17 00:00:00 2001 From: DanialTaheri Date: Thu, 13 Aug 2026 12:59:15 -0700 Subject: [PATCH 2/4] docs(vlm): design note for the media-token validity mask Records why placeholder sanitization was removed and what replaced it: the mask the model derives answers "is this a real token", the merge needs "is this a media anchor", and those only coincide while every media token anchors an image. Covers the packing constraint that makes the mask meaningful, and notes that filtering the blend is now required rather than optional. Signed-off-by: DanialTaheri --- docs/design-docs/media-token-validity-mask.md | 214 ++++++++++++++++++ docs/index.md | 1 + 2 files changed, 215 insertions(+) create mode 100644 docs/design-docs/media-token-validity-mask.md 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..440c5d66cd4 --- /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. Text can legitimately +contain it, and until now that text could not be trained on: the model read +every occurrence as an anchor for a projected image feature and failed when no +feature existed. The workaround was a sanitizer that rewrote the offending +tokens before they reached the model, which changed the data to fit the model. + +This design removes the sanitizer and replaces it with 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 keep +their ordinary token embedding. + +## 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} From 525873d71f796baf7d39d928c8d9a6b79216053a Mon Sep 17 00:00:00 2001 From: Danial Mohseni Taheri <49656670+DanialTaheri@users.noreply.github.com> Date: Mon, 17 Aug 2026 23:16:45 -0700 Subject: [PATCH 3/4] Update docs/design-docs/media-token-validity-mask.md Co-authored-by: Yi-Fu Wu Signed-off-by: Danial Mohseni Taheri <49656670+DanialTaheri@users.noreply.github.com> --- docs/design-docs/media-token-validity-mask.md | 20 +++++++++---------- 1 file changed, 10 insertions(+), 10 deletions(-) diff --git a/docs/design-docs/media-token-validity-mask.md b/docs/design-docs/media-token-validity-mask.md index 440c5d66cd4..4bf5caa5a47 100644 --- a/docs/design-docs/media-token-validity-mask.md +++ b/docs/design-docs/media-token-validity-mask.md @@ -2,16 +2,16 @@ ## Summary -A media placeholder is an ordinary vocabulary entry. Text can legitimately -contain it, and until now that text could not be trained on: the model read -every occurrence as an anchor for a projected image feature and failed when no -feature existed. The workaround was a sanitizer that rewrote the offending -tokens before they reached the model, which changed the data to fit the model. - -This design removes the sanitizer and replaces it with 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 keep -their ordinary token embedding. +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 From c4a7d97b4788eacdf2fafc1348ce3fb4a5db856a Mon Sep 17 00:00:00 2001 From: Danial Mohseni Taheri Date: Tue, 18 Aug 2026 09:59:50 -0700 Subject: [PATCH 4/4] refactor(vlm): move the mask's batch logic out of the Megatron worker Building the mask needs only input_ids, the placeholder id and the per-row image counts -- none of which are Megatron concerns. Sitting in the policy worker put that logic behind a module that cannot be imported without megatron, so it could not be unit tested and the plumbing had to be trusted. Move the four pure pieces next to build_media_token_validity_mask, where they have no such dependency, and leave the worker with the unwrapping glue that genuinely does. The batch is duck-typed there because batched_data_dict imports this module, so naming its type would be circular. Adds twelve tests over the moved pieces, including the two that decide whether the feature engages at all: the capability probe reads the forward signature rather than a flag, and a batch with no pixel_values counts as text-only rather than unreadable. Signed-off-by: Danial Mohseni Taheri --- nemo_rl/data/multimodal_utils.py | 70 ++++++++++++ .../policy/workers/megatron_policy_worker.py | 93 +++------------ .../data/test_media_token_validity_mask.py | 107 +++++++++++++++++- 3 files changed, 191 insertions(+), 79 deletions(-) diff --git a/nemo_rl/data/multimodal_utils.py b/nemo_rl/data/multimodal_utils.py index c7c99e3da33..68dd4e9d097 100644 --- a/nemo_rl/data/multimodal_utils.py +++ b/nemo_rl/data/multimodal_utils.py @@ -1209,3 +1209,73 @@ def build_media_token_validity_mask( 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/policy/workers/megatron_policy_worker.py b/nemo_rl/models/policy/workers/megatron_policy_worker.py index 90d79b048b5..7099a181cf2 100644 --- a/nemo_rl/models/policy/workers/megatron_policy_worker.py +++ b/nemo_rl/models/policy/workers/megatron_policy_worker.py @@ -51,8 +51,9 @@ from nemo_rl.algorithms.logits_sampling_utils import TrainingSamplingParams from nemo_rl.algorithms.loss.interfaces import LossFunction from nemo_rl.data.multimodal_utils import ( - PackedTensor, - build_media_token_validity_mask, + 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 @@ -185,84 +186,22 @@ def _model_slices_context_parallel_inputs(model: Any) -> bool: ) -def _model_media_placeholder_token_id(model: Any) -> Optional[int]: - """The vocabulary id this model treats as a media placeholder, if any.""" +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) - chunks = unwrapped if isinstance(unwrapped, (list, tuple)) else [unwrapped] - for chunk in chunks: - token_id = getattr(chunk, "image_token_index", None) - if token_id is not None: - return int(token_id) - return None + return list(unwrapped) if isinstance(unwrapped, (list, tuple)) else [unwrapped] -def _model_accepts_media_token_validity_mask(model: Any) -> bool: - """Whether the model'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. - """ - import inspect - - from megatron.core.utils import unwrap_model - - unwrapped = unwrap_model(model) - chunks = unwrapped if isinstance(unwrapped, (list, tuple)) else [unwrapped] - 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 _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 _attach_media_token_validity_mask( - batch: BatchedDataDict[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 ``process_microbatch`` carries this through the - same packing transform as ``input_ids`` rather than deriving it downstream. - """ - 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 - - -def _image_counts_by_row( - batch: BatchedDataDict[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 _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( @@ -906,9 +845,7 @@ def train( ].unsqueeze(-1) batch["mtp_loss_mask"] = mtp_loss_mask - _attach_media_token_validity_mask( - batch, self.media_placeholder_token_id - ) + attach_media_token_validity_mask(batch, self.media_placeholder_token_id) ( data_iterator, @@ -1747,7 +1684,7 @@ def get_logprobs( # 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) + attach_media_token_validity_mask(data, self.media_placeholder_token_id) ( mb_iterator, @@ -1968,7 +1905,7 @@ def get_topk_logits( self.model.eval() - _attach_media_token_validity_mask(data, self.media_placeholder_token_id) + attach_media_token_validity_mask(data, self.media_placeholder_token_id) ( mb_iterator, diff --git a/tests/unit/data/test_media_token_validity_mask.py b/tests/unit/data/test_media_token_validity_mask.py index 43afb870ecc..2f39c775d5a 100644 --- a/tests/unit/data/test_media_token_validity_mask.py +++ b/tests/unit/data/test_media_token_validity_mask.py @@ -15,7 +15,14 @@ import pytest import torch -from nemo_rl.data.multimodal_utils import build_media_token_validity_mask +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 @@ -69,3 +76,101 @@ def test_rejects_non_2d_input_ids(): 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