feat(model): accept pre-sharded packed Qwen3-VL inputs

[yaoyu-33]

What does this PR do ?

Accept Qwen3-VL packed THD inputs that a caller has already partitioned into Megatron's load-balanced zigzag context-parallel layout.

The current packed-CP path assumes input_ids still contains the full packed stream and always computes a local partition index. Some training frameworks pass rank-local tokens together with global packed-sequence metadata. Partitioning those tokens a second time either selects the wrong rows or fails before vision embeddings can be scattered.

This PR recognizes that contract when the physical global token count equals cp_size * local_token_count, validates that each packed segment can be divided into 2 * cp_size chunks, and skips the second partition. Because local token chunks no longer contain enough information to reconstruct multimodal positions, callers on this path must provide explicit rank-local 3D MRoPE position_ids.

Vision embeddings are still produced in full-sequence order. A count-only CP all-gather reconstructs the vision-token offset of each zigzag chunk and selects the rows owned by the local rank. The same indices are applied to deepstack features; ranks with no local visual tokens receive an empty selection.

Relationship to #4784

#4784 has landed, and this branch is now cleanly rebased onto current main.

This PR deliberately does not duplicate #4784's positional autograd.Function.apply fix, reduce-scatter backward correction, placeholder dtype handling, or zero-image-rank hang fix. The combined CP path still needs to be revalidated on two GPUs before final approval.

Attribution

The algorithm and current-main adaptation come from the radixark fork. This PR preserves those contributions directly:

• Original CP-local vision-row selection by @Zhichenzzz: radixark@eeb84d7
• Current-main double-slice guard by @yushengsu-thu: radixark@390f983
• Original validation PR: fix(qwen3-vl): CP-local vision-embed hook + AllGatherVisionEmbeddings.apply kwarg radixark/Megatron-Bridge#9
• Follow-up sync/validation PR: [model] feat: Sync qwen_vl tree with upstream main, keep CP pre-shard fix radixark/Megatron-Bridge#19

The original downstream validation reported Qwen3-VL-2B CP2 with TP2/TP4, THD packing, healthy logprob differences of 0.0127-0.0131, and no NCCL timeout. Those GPU results belong to the original contributors and were not re-run in this adaptation.

Changelog

• Detect and validate pre-sharded single-row THD packed inputs.
• Skip native packed-CP indexing when the input is already rank-local.
• Require explicit local MRoPE positions for pre-sharded inputs.
• Select local vision and deepstack rows using per-chunk visual-token counts.
• Add focused tests for full-input detection, CP2 multi-segment zigzag ordering, a zero-visual-token rank, deepstack selection, and preservation of the existing full-input partition path.

GitHub Actions CI

Local checks:

• uv run --no-project --with pre-commit pre-commit run --all-files — passed
• Focused CPU tests — 6 passed, 12 deselected
• git diff --check and Python compilation — passed

The focused tests used a temporary local import stub for CUDA/Transformer Engine modules unavailable on macOS; no stub is present in this branch. A normal full environment cannot be resolved locally because nvidia-resiliency-ext is Linux-only.

Before marking ready, run in the standard Linux/GPU environment:

• uv run python -m pytest tests/unit_tests/models/qwen_vl/modelling_qwen3_vl/test_model.py -k 'packed_input_pre_sharded or cp_local_vision_embed_indices or pre_sharded_packed_cp or applies_one_partition_index' -vv
• A CP2 two-GPU forward/backward case after rebasing onto fix(model): correct Qwen3-VL vision_dp_when_cp CP gradient and avoid 0-image-rank hang #4784
• The existing Qwen3-VL functional conversion/generation coverage selected by CI

Before your PR is "Ready for review"

Pre checks:

• Make sure you read and followed Contributor guidelines
• Did you write any new necessary tests?
• Did you add or update any necessary documentation? (The model docstring documents the new input contract; no user guide currently exposes this internal forward path.)
• Does the PR affect components that are optional to install? (No new optional component or dependency.)

Additional Information

• fix(model): correct Qwen3-VL vision_dp_when_cp CP gradient and avoid 0-image-rank hang #4784 is merged and the branch is rebased; combined two-GPU validation remains outstanding.
• No public conversion API, checkpoint mapping, dependency, recipe, workflow, or Megatron-LM submodule changes.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[yaoyu-33]

/ok to test 449ced7

[yaoyu-33]

/ok to test edbd870

[yaoyu-33]

/ok to test edbd870

[yaoyu-33]

/ok to test e7cdad0

[yaoyu-33]

/ok to test 49a3fe5