fix(model): correct Qwen3-VL vision_dp_when_cp CP gradient and avoid 0-image-rank hang

[going-song]

What does this PR do ?

Fixes three bugs that break Qwen3-VL training with context parallelism when vision_dp_when_cp=True: the vision-embedding all-gather is entered with a TypeError, its backward silently drops cross-rank gradient contributions to the vision encoder, and CP ranks that receive 0 images cause an NCCL hang once the backward collective is fixed.

We have been running this fix in our internal production VLM training (multi-node, context parallelism with vision_dp_when_cp enabled): long runs train stably, and downstream evaluation scores improved after applying the fix — consistent with the vision encoder finally receiving the complete gradient.

Under vision_dp_when_cp the rank that ENCODES an image's tokens is not the rank
that PROCESSES them in the LLM: the fused text+vision sequence is sharded across
CP ranks in the usual zigzag pattern (rank r takes sequence chunks r and
2*cp_size-1-r), so each rank's backward produces gradient only for the rows
sitting on its own sequence shard. Printing the nonzero-row mask of grad_output
in a real CP=2 training step (one image, all of its tokens encoded by rank 0):

  grad_output on rank 0:   111111...1000000...0
  grad_output on rank 1:   000000...0111111...1   ← exactly complementary

BEFORE (no all_reduce — each rank slices its local grad):
  rank 0 keeps only its 1-rows; every 0-row of its image's gradient stays zero
  → vision grad does not flow
  rank 1 (0 images, no-grad placeholder) never runs backward at all

AFTER (all_reduce, then slice = reduce-scatter):
  the masks sum to 111111...1 — every rank holds the full gradient and slices
  out its own tokens → complete gradient ✓
  a 0-image rank joins the collective with a 0-length grad → no hang ✓

Changelog

• modelling_qwen3_vl/model.py: call AllGatherVisionEmbeddings.apply with positional arguments — torch.autograd.Function.apply does not accept keyword arguments, so the previous cp_group=... call raised TypeError: apply() takes no keyword arguments as soon as the vision_dp_when_cp CP path was entered.
• modelling_qwen3_vl/utils.py: implement the correct backward of the all-gather (a reduce-scatter). The forward all-gathers each rank's vision embeddings to every CP rank, and the fused text+vision sequence is then sharded across CP ranks, so the gradient for one rank's vision tokens can be produced on any rank. The backward now all_reduces grad_output over the CP group before slicing out the local range; previously each rank kept only its local slice and the cross-rank contributions were silently dropped (incomplete vision-encoder gradients). forward stashes ctx.cp_group for this.
• modelling_qwen3_vl/model.py: create the 0-image placeholder tensors (vision_embeds and each deepstack feature) with requires_grad=True. When num_images < cp_size a CP rank receives 0 images; with a no-grad placeholder, autograd skips the Function backward on that rank, the all_reduce above is never issued there, and the remaining ranks hang in the collective. The 0-length placeholders contribute nothing to the forward output.
• tests/unit_tests/.../test_utils.py: add two distributed unit tests (2-GPU, follow the file's existing WORLD_SIZE>=2 pattern):
  • test_allgather_vision_embeddings_backward_reduce_scatter — asserts the gradient value is the CP-summed gradient sliced to the rank's range (the pre-existing test only checked grad is not None, which the buggy backward also passes).
  • test_allgather_vision_embeddings_empty_rank — a 0-token rank participates in the backward collective (no hang) and the token-owning rank receives the CP-summed gradient.

GitHub Actions CI

See the CI section in the Contributing doc for how to trigger the CI. A Nvidia developer will need to approve and trigger the CI for external contributors.

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? (bug fix — no doc changes needed)
• Does the PR affect components that are optional to install? (No)

Additional Information

Verified beyond the unit tests: a mock-data Qwen3-VL pretrain (tiny 4-layer model, MockVLMConversationProvider) with context_parallel_size=8, vision_dp_when_cp=True on 8×A100 runs 10/10 iterations with no hang and decreasing loss (9.34 → 7.82); 0-image ranks occur naturally in every batch at CP=8. Without this fix the same setup either raises the apply() TypeError or, with only the reduce-scatter applied, hangs at the first backward.

Signed-off-by: Yoonsik Kim yoonsik.kim90@navercorp.com
Signed-off-by: Kayeon Song kayeon.song@navercorp.com
Signed-off-by: Chanwoo Park chanwoo.park98@navercorp.com

[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]

The reduce-scatter correction and positional apply change look sound. I left two inline comments about dtype consistency on empty ranks and coverage of the model-side autograd safeguard.

[yaoyu-33]

Could we also make the empty-rank placeholder dtype match the real vision output? The empty branch above still hard-codes torch.bfloat16, while a non-empty rank will send the configured vision dtype. With FP16/FP32, CP ranks would enter the same NCCL all_gather with different dtypes, so the zero-image path can still fail. Using the actual model/vision output dtype consistently across ranks, and covering a non-BF16 case, would make this fix robust.

[going-song]

Placeholders now follow self.config.params_dtype, and the empty-rank / frozen-input tests are parametrized over bf16/fp16/fp32. I couldn't get a full non-BF16 run in my environment (TE has no fp32 fused-attention backend, and the fp16 mock NaNs at iter 2 even with CP=1), so the non-BF16 coverage stays at the unit test level.

[yaoyu-33]

Could we add coverage for the model-side requires_grad safeguard itself? Both the empty-rank and frozen-input tests pre-set requires_grad=True before calling apply, so reverting the new block in model.py would leave these tests green. A small model-path test, or extracting the "ensure collective inputs require grad" step into a helper, would protect the deadlock fix rather than only the custom Function behavior.

[going-song]

Extracted the step into ensure_requires_grad_for_cp_collective() in utils.py
as you suggested — the model calls it right before the gather, both tests go
through it now, and there's a direct unit test for the helper (forces in grad
mode, no-op under torch.no_grad).

[yaoyu-33]

/ok to test 194ecb5

[huvunvidia]

Hi @going-song ,
All CICD tests have passed but currently the code coverage is less than 80%. Can you add more tests?

[going-song]

/ok to test d4fec2b

[going-song]

@huvunvidia Thanks for the review. Added a test that checks the reduce-scatter backward
per rank. Not sure how to trigger the CI from my side though, would you mind
kicking it off?

[huvunvidia]

/ok to test 5a67233

[huvunvidia]

/ok to test e34d288

[yaoyu-33]

/ok to test 092a5fe