fix: scatter Qwen3-VL text-model embeddings under SP for standalone LM forwards

[kevalmorabia97]

What

Fix a sequence-parallel bug in the Qwen3-VL text model that breaks standalone language-model forwards (e.g. distilling only the LM tower of a VLM) under tensor + sequence parallelism.

Why

Qwen3VLModel builds its text model with scatter_embedding_sequence_parallel=False because the outer VLM model manually scatters the merged vision+text embeddings for sequence parallelism (modelling_qwen3_vl/model.py).

When the language model is run on its own — no outer VLM, e.g. ModelOpt language-model distillation of a VLM — that manual scatter is bypassed. Under SP the embeddings therefore stay unscattered, the decoder runs the full sequence on every TP rank, and the output-side sequence-parallel gather then doubles the sequence length.

This surfaces downstream as a TP_size × seq_length vs seq_length shape mismatch. Concretely, ModelOpt VLM language-model distillation at TP=2 + SP fails in the KD loss-mask step:

RuntimeError: The size of tensor a (32) must match the size of tensor b (16) at non-singleton dimension 0

(32 = TP_size(2) × seq_length(16).) Plain TP (no SP) works; the plain LM path works; only the standalone-VLM-LM + SP combination is affected.

Fix

In Qwen3VLTextModel.forward, when no external decoder_input is provided (i.e. the LM embeds internally) and sequence parallelism is enabled, scatter the internally-embedded sequence — mirroring the existing MTP-path scatter a few lines below. The normal VLM forward (which passes pre-scattered decoder_input from the outer model) is unaffected.

if not _embeddings_provided and self.pre_process and self.config.sequence_parallel:
    decoder_input = tensor_parallel.scatter_to_sequence_parallel_region(
        decoder_input, group=self.pg_collection.tp
    )

Testing

Validated on nvcr.io/nvidia/nemo:26.06: Qwen3.5-VL language-model distillation at TP=2 + SP passes with this change (previously failed at the loss-mask step). Single-GPU and TP-without-SP were already passing and remain so.

🤖 Generated with Claude Code

[claude]

Light review - LGTM.

The fix is minimal and correct:

• _embeddings_provided (line 147) is captured before the _preprocess unpack reassigns decoder_input (line 157), so it faithfully reflects the original caller argument.
• The guard "not _embeddings_provided and self.pre_process and self.config.sequence_parallel" correctly isolates the standalone-LM path: the normal VLM forward passes pre-scattered decoder_input (model.py:554-558), so _embeddings_provided is True and the new scatter is skipped.
• The scatter mirrors the existing MTP-path scatter (lines 196-204) and uses the same self.pg_collection.tp group, consistent with the outer model.

Minor notes (non-blocking):

• Consider a short comment noting this matches the MTP scatter below, so the two SP-scatter sites stay in sync if one is ever refactored.
• No automated regression test accompanies the fix. An SP-only path is hard to cover cheaply, but a 2-GPU (TP=2 + SP) standalone-LM-forward shape assertion would guard against regressions if it fits the 2-GPU budget.

Suggested test cases

• No perf tests impacted. Only models/qwen_vl/modelling_qwen3_vl/text_model.py is changed; no scripts/performance/configs/ entries were touched.
• If a regression test is added: a TP=2 + sequence-parallel standalone Qwen3VLTextModel.forward (no external decoder_input) asserting the output hidden_states sequence length equals seq_length rather than TP_size x seq_length.

[kevalmorabia97]

Superseded by a general fix in Megatron-LM (mcore): NVIDIA/Megatron-LM#5628.

Investigation showed this is not Qwen3-VL-specific — every Megatron-Bridge VLM/omni/audio model builds its language tower with scatter_embedding_sequence_parallel=False (the outer multimodal model scatters the merged embeddings), so a standalone-LM forward under sequence parallelism hits the same TP×seq doubling. Confirmed for both Qwen3.5-VL and Gemma3-VL.

The Qwen3-VL text-model change here only covered Qwen3-VL/3.5-VL and would double-scatter if combined with the mcore fix. The mcore GPTModel._preprocess fix (#5628) covers all such models in one place and is a no-op for standard LMs, so closing this in favor of it.