feat(vlm): carry a media-token validity mask through sequence packing - #3631
Merged
Conversation
github-actions
Bot
added
the
Documentation
5 tasks
DanialTaheri
force-pushed
the
smohsenitahe/media-token-validity-mask
branch
2 times, most recently
from
August 14, 2026 20:22
ae769e0 to
724412b
Compare
Contributor
Author
|
/ok to test c3559f0 |
@DanialTaheri, there was an error processing your request: See the following link for more information: https://docs.gha-runners.nvidia.com/cpr/e/2/ |
yfw
added
the
CI:Lfast
Contributor
|
/ok to test 724412b |
Contributor
Author
|
/ok to test 87a4c05 |
DanialTaheri
force-pushed
the
smohsenitahe/media-token-validity-mask
branch
from
August 17, 2026 18:19
87a4c05 to
087a985
Compare
1 similar comment
Contributor
Author
|
/ok to test 087a985 |
DanialTaheri
force-pushed
the
smohsenitahe/media-token-validity-mask
branch
from
August 17, 2026 20:19
087a985 to
93f16af
Compare
Contributor
Author
|
/ok to test 93f16af |
yfw
reviewed
Aug 18, 2026
Contributor
|
I think we don't need the Megatron-Bridge bump in this PR anymore because of 04a687d. Also I don't think anything needs to change in uv.lock anymore? |
A media placeholder is an ordinary vocabulary entry, so text can legitimately contain it -- competitive-programming statements, for instance, where <image> 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 <smohsenitahe@nvidia.com>
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 <smohsenitahe@nvidia.com>
Co-authored-by: Yi-Fu Wu <yifu.wu@gmail.com> Signed-off-by: Danial Mohseni Taheri <49656670+DanialTaheri@users.noreply.github.com>
DanialTaheri
force-pushed
the
smohsenitahe/media-token-validity-mask
branch
from
August 18, 2026 06:23
9298ca1 to
525873d
Compare
Contributor
Author
|
/ok to test 525873d |
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 <smohsenitahe@nvidia.com>
Contributor
Author
|
/ok to test c4a7d97 |
Contributor
|
/ok to test 3cd5cbc |
yfw
approved these changes
Aug 18, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What
Lets a row whose text legitimately contains the media placeholder train, instead
of being dropped at data-prep time.
Stacked on #3494 — based on
smohsenitahe/super-omni-mtp-cp, so this diff isonly the two commits. GitHub will retarget the base to
mainonce #3494 merges.Why this is separate
#3494 removes placeholder sanitization, and the VLM blend trains fine without
this: every row there either carries no literal placeholder, or carries exactly
as many as it has images. A 50-step Super Omni run with sanitization off and the
mask entirely unwired completed 50/50 with zero media-alignment failures.
The mask matters only for rows with
placeholders > 0, images == 0— text thatspells the placeholder with no image attached. There are zero such rows in
the filtered VLM blend, so #3494 does not need this. Text blends do: in a 100k
sample of a Nemotron post-training blend, 2,539 rows (2.5%) contain a literal
<image>with no image — competitive-programming statements where it replacedinline math.
Why the model cannot answer this itself
_merge_projected_mediaenforces one projected feature per valid placeholder.The mask it derives comes from the padding or attention mask, and 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.
The caller knows how many media items each row carries, so it can answer. This
carries that answer to the model.
The packing constraint
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_idsrather than being derived downstream, exactlyas
mtp_loss_maskalready does.A mask that is merely misaligned does not raise; it attaches image features to
the wrong positions. The packing tests are the load-bearing ones.
Dependency
Needs NVIDIA-NeMo/Megatron-Bridge#5573, which adds the
media_token_validity_maskargument to
NemotronOmniModel.forward. Until that merges and the pin is bumped,_model_accepts_media_token_validity_maskfinds no such parameter and no mask isever attached — inert, not broken.
Tests
real
<image>statements) reached Step 1/1, rollouts 100%, zero alignment failures