feat(vlm): integrate packed CP and vision sharding for Qwen3.5-MoE by yuhezhang-ai · Pull Request #3186 · NVIDIA-NeMo/Automodel

yuhezhang-ai · 2026-07-22T14:57:47Z

< />!NOTE< />
Final integration PR. #2763, #2937, #2989, and #2990 are merged. This branch includes latest main at d3f3b2b6d, uses distributed.multimodal.vision.frame_sharding with mesh_dims: [cp], and is mergeable at final head 9b24aa970. GitHub shows 25 integration-owned changed files.

What does this PR do?

Integrate packed context parallelism and frame-level CP vision sharding into Qwen3.5-MoE VLM training, with Qwen3.5-122B-A10B at 128K context as the large-scale target.

The implementation follows the merged unified CP contract from #2937: prepare_model_inputs_for_cp is metadata/sharder-only, while embedding, trainable vision encoding, multimodal splicing, and primary sequence sharding happen inside the model forward per microbatch.

Changelog

Select a contiguous block-diagonal ContextParallelSharder for packed Qwen3.5-MoE SDPA batches; ordinary batches retain the merged round-robin CP path.
Run full-sequence multimodal embedding inside forward, then shard the primary embeddings contiguously for packed CP while auxiliary tensors use the identical layout.
Route Qwen3.5-MoE full attention through block-diagonal CP SDPA and pass typed packed-document state to the GDN/linear-attention path, composing packing + GDN + CP.
Consume feat(distributed): frame-level context-parallel vision-tower sharding #2990's frame-level vision sharding for the Qwen3.5-MoE trainable vision tower.
Preserve temporary CP-SDPA suspension during activation-checkpoint recompute.
Handle Qwen3.5-122B split MTP expert checkpoint layout and rebuild loaded expert tensors on the EP mesh.
Add typed, recipe-configurable video processing so Qwen video sampling and total T × H × W resize volume can be bounded explicitly.
Add the Qwen3.5-122B-A10B EP8/CP32 128K example with trainable vision, full activation checkpointing, FP32 master weights, and FP32 optimizer moments.
Add focused model, distributed, activation-checkpoint, state-dict, processor, and recipe-contract coverage.

Landing status

All prerequisite PRs are merged. This branch includes latest main (d3f3b2b6d) at final head 9b24aa970; the remaining 25-file diff is the Qwen3.5-MoE packed VLM integration described below.

Validation

Rebase/refactor validation

After rebuilding on the merged #2937/#2989 design:

129 focused block-diagonal CP, activation-checkpoint, Qwen3.5-MoE, state-dict, and 122B recipe tests passed; 51 CUDA-only/unavailable-dependency tests skipped.
156 unified-sharder, CP utility, and VLM recipe regression tests passed; 7 CUDA-only tests skipped.
The final video-processor and CP32 recipe changes passed 125 targeted tests; 3 CUDA-only tests skipped.
Ruff formatting/checks, shell syntax checks, and git diff --check passed.
Final post-NVIDIA-NeMo/Automodel#2990 merge validation at ef2225f2a: 456 focused CPU tests passed and 51 CUDA-only tests skipped; Ruff formatting/checks and git diff --check passed. The local MDX parser validated all 452 pages; the Fern CLI is not installed locally, so CI supplies the final fern check. GitHub reports this PR mergeable with 25 changed files.
Latest-main merge validation at 9b24aa970: 172 directly affected AutoPipeline, Qwen3.5-MoE, and VLM recipe tests passed; 8 CUDA-only cases skipped. Ruff formatting/checks and git diff --check passed. The resolution preserves upstream mRoPE PP chunking alongside the CP/packing integration.

The previously validated two-GPU packed full-attention + GDN parity run also passed (Slurm job 14214890).

Qwen3.5-122B-A10B, 128K integration runs

100-step mixed-workload run

A 100-step end-to-end CP16 run completed successfully on the pre-rebase integration source revision 7ce8dc955; the history rebase preserves that integration logic while adapting it to the merged sharder API:

16 nodes / 128 H100 GPUs
EP8, CP16, inferred DP8
trainable vision tower with full activation checkpointing
frame-level CP vision sharding
neat-packed 128K sequences
mixed text, image, and bounded-video data, including genuine examples of approximately 120K tokens
TE FusedAdam with FP32 master weights and FP32 first/second moments
100/100 steps completed; no OOM, collective hang, activation-checkpoint error, non-finite gradient, or traceback
peak allocated memory: approximately 64.2 GiB/GPU
final aggregate throughput: 6,340 tokens/s (49.53 tokens/s/GPU)
loss decreased from 1.7345 to effectively zero with finite gradients

Weights & Biases run

The validation set has 48 packed sequences and was intentionally reused for the 100-step mechanics/convergence smoke. The near-zero terminal loss demonstrates that the configuration can train and overfit this dataset; it is not evidence of production-scale convergence or generalization.

Post-refactor CP32 runs

Two exact post-refactor implementation runs used revision df4c51ac9 and the CP32/GBS4 settings now published by ded4f01eb:

Workload	Slurm	Result	Peak allocated	Steady measured step	Evidence
Public MedPix, 128K packs	`14293833`	2/2 steps, finite loss and gradients	59.62 GiB/GPU	~92 s	W&B
Mixed text/image/video, including an approximately 121.5K-token example	`14293834`	2/2 steps, finite loss and gradients	58.49 GiB/GPU	~74 s	W&B

Both runs used 16 nodes / 128 H100 GPUs, EP8, CP32, inferred DP4, a trainable vision tower, full activation checkpointing, and FP32 optimizer state. Video preprocessing was bounded to at most 8 sampled frames and a total sampled volume of 524,288 pixels (approximately 8 × 256 × 256).

The first logged step in both runs took approximately five minutes because it included compilation and autotuning; it is not representative steady-state throughput. The public MedPix workload is also intentionally vision-heavy—its 128K packs combine many independent single-image samples—so its approximately 3K aggregate steady TPS is a topology/correctness result, not a throughput claim. For representative long-text or bounded-video training, use the smallest CP degree that fits; CP16 preserves DP8 on this 128-GPU allocation and was more throughput-efficient than CP32/DP4.

Before ready for review

Rebase onto the merged refactor(distributed): unify CP input prep and dispatch across models #2937/feat(distributed): block-diagonal varlen CP for packed sequences #2989 CP design.
Remove obsolete merge commits and already-landed dependency history.
Resolve the unified-sharder/in-forward embedding conflicts.
Rerun focused CPU regression, formatting, and lint checks.
Confirm DCO sign-off on all integration commits.
Run post-refactor CP32 public-image and mixed-video GPU validation.
Merge feat(distributed): frame-level context-parallel vision-tower sharding #2990.
Merge latest main (d3f3b2b6d), preserve upstream Qwen3.5 mRoPE PP chunking, and confirm the remaining diff is integration-only.
Run final GitHub CI on exact head 9b24aa970.

Additional information

copy-pr-bot · 2026-07-22T14:57:51Z

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.

github-actions · 2026-07-22T14:59:38Z

🌿 Preview your docs: https://nvidia-preview-yuhez-feat-cp-vlm-integration-2937.docs.buildwithfern.com/nemo/automodel

…ding Under context parallelism the VLM pre-embed step runs the entire vision tower on the full, unsharded set of images on every CP rank. A CP size of N therefore means N-fold redundant vision computation and all-media vision activations on each rank. Qwen3-VL-style vision towers attend per image or frame through cu_seqlens built from grid_thw, with spatial position embeddings repeated per frame. maybe_distribute_visual expands each (t, h, w) entry into t frame units, partitions them contiguously across the sharding group, runs the vision tower on each rank's slice, and reassembles the complete embedding set with a differentiable variable-length all-gather. Details: - Frame granularity allows a single large video, or a pack with fewer videos than ranks, to be divided across ranks. Images are one unit. - When there are fewer frames than ranks, minimal dummy frames keep the collective sequence uniform. Their embeddings are removed before the loss and contribute no gradient. - Batches below distributed.cp_vision_sharding.min_tokens and batches without media keep the replicated path. Disabling distributed.cp_vision_sharding provides the A/B baseline. - maybe_distribute_independent_units exposes the same partition, run, gather, and slice scheme for towers with other forward signatures. A frozen vision tower can use a wider group with replicated weights. A trainable tower must use the CP group only because the vision tower is replicated across TP ranks and a wider reduce-scatter SUM would over-count its gradient. Vision sharding is inactive when cp_size is 1. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Yuhe Zhang <yuhez@nvidia.com>

40 CPU-only tests (CP ranks simulated in-process, collectives mocked; no GPU / no distributed init) covering: - partition: contiguous, complete, >=1 entry per rank, attention-cost balanced (including high-resolution hotspots), None when fewer entries than ranks - the core property: visual(full) == concat_r visual(slice_r) for forward AND vision-parameter gradients at world sizes 2-4, including video (t > 1) entries - variable-token differentiable all-gather forward + backward - maybe_distribute_independent_units: order preservation, single differentiable gather, dummy padding with exact zero gradient - maybe_distribute_visual end-to-end: replicate parity, deepstack gathering, a single video split across ranks at frame granularity, every fallback (no group, env-disabled, below min tokens), and the pad path backward through the production gather/slice code asserting sum_r grad_r == replicate grad with zero gradient on padded ranks Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Yuhe Zhang <yuhez@nvidia.com>

Add an Advanced Training guide for frame-level CP vision-tower sharding: design (frame independence, balanced contiguous partition, differentiable all-gather, dummy-frame pad path), the sharding-scope design constraint (frozen tower may shard across CP x TP; trainable tower shards within the CP group only, so the reduce-scatter backward never double-counts TP-replicated gradients; pure TP is not enabled), usage and environment toggles, and a validation summary. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Yuhe Zhang <yuhez@nvidia.com>

- maybe_distribute_visual: return the plain visual(...) call when pixel_values or grid_thw is None; the no-media fallback previously routed through _grid_for_visual and crashed on grid_thw.device. - Machine-check the frozen-vs-trainable sharding-scope constraint: set_cp_vision_group(group, *, spans_only_cp=True) records the caller's declaration and maybe_distribute_visual raises ValueError when a trainable vision tower meets a group not declared CP-only (gathering across CP x TP would accumulate the vision gradient tp-fold in the reduce-scatter(SUM) backward). - Drop dead surface with no production callers: the generic maybe_distribute_independent_units helper, the CPU-grid metadata side channel (set/reset_cp_vision_cpu_grids, cp_vision_grid_metadata), and the never-set _nemo_cpu_grid_flash_ok escape hatch; add __all__ with the four names the wiring PR calls. - Validate each rank's local visual token count (pooler + deepstack) before the gather instead of silently mis-slicing gathered blocks. - Annotate the public API signatures (Tensor | None, ProcessGroup | None, typed reset token) and downgrade the informational SHARDED / small-workload logs from WARNING to INFO. - Tests: cover the None-input fallback, the trainable-scope guard (raise) and the frozen-tower spans_only_cp=False path; drop tests of the removed helper/side channel. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Yuhe Zhang <yuhez@nvidia.com>

Drive maybe_distribute_visual forward AND backward through real all_gather / reduce_scatter(SUM) collectives on a 2-rank gloo CPU group (mp.spawn, matching the existing 2-proc gloo test conventions): each rank backprops only its sequence shard of the gathered embeds, and the group-summed vision-parameter gradients must equal the single-process replicated reference -- for the balanced-partition path (with deepstack) and for the dummy-pad path, where the padded rank must receive exactly zero gradient. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Yuhe Zhang <yuhez@nvidia.com>

The frozen-vs-trainable sharding-group constraint is now machine-checked via set_cp_vision_group(..., spans_only_cp=...); describe the raise, fix the pure-TP row to match the actual group-size <= 1 gate, drop the removed generic helper, and list the new real-collective gloo tests. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Yuhe Zhang <yuhez@nvidia.com>

Balance the frame-level CP vision-tower partition by p*(p+alpha) instead of pure p**2. The quadratic term captures per-frame attention, while alpha*p represents linear per-patch work such as QKV and MLP projections. This improves balance for packs that mix large image frames with many smaller video frames. By default, alpha is inferred as 3 * vision_hidden_size from the model without class-name checks. Unknown towers fall back to alpha=0. distributed.cp_vision_sharding.cost_alpha accepts an exact nonnegative integer override; 0 restores p**2 and auto selects model discovery. This setting affects partitioning only. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Yuhe Zhang <yuhez@nvidia.com>

_contiguous_balanced_bounds used bisect.bisect_left to place each rank's cut point, which undershoots by one when a frame lands exactly on a cumulative-cost boundary: 4 equal-cost frames at world=2 split [1, 3] instead of [2, 2]. This directly skews packed-sequence vision load balancing (the cost model exists precisely to balance packs mixing big image frames with many small video frames). Switch to bisect.bisect_right so the cut is placed past the frame that closes an exact-target cumulative sum, giving an even split. The >=1 entry-per-rank invariant is preserved by the existing lo/hi clamps. Add unit tests: equal-cost even splits (4/world2, 6/world3, 8/world4), an uneven-cost exact-boundary case (patches [3,4,5] -> loads [25,25]), and cuts landing exactly on a cumulative sum. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: Yuhe Zhang <yuhez@nvidia.com>

…ivalent The vision-shard parity tests assert torch.allclose (with tolerances), not bitwise equality, so "exactly"/"byte-for-byte"/"bit-identical" overclaimed the guarantee. Reword the module docstring/comments, the cp-vision-sharding guide, and the test module docstring to say "numerically equivalent (allclose)". No tolerances or behavior change; the exact-call replicate FALLBACK wording (a literally identical call) is left as-is. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: Yuhe Zhang <yuhez@nvidia.com>

…atch The per-rank token-count validation (and the deepstack-features check) raised a ValueError BEFORE the differentiable all-gather. If one rank's visual output diverged from its planned frame slice, that rank raised while every peer blocked forever in all_gather -> deadlock. Route both checks through _raise_if_any_rank_failed, which all-reduces a boolean flag (MAX) over the sharding group so every rank takes the same path: either all raise or all proceed. The diverging rank raises its own actionable message; peers raise a group-level message pointing at it. The flag tensor is placed on the compute device so the reduce matches the active backend (CUDA under NCCL). Add a 2-rank gloo test (_DivergentVisual drops one token on the last rank) asserting BOTH ranks raise ValueError with no hang, and mock all_reduce as a no-op in the single-process CPU simulation helper. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: Yuhe Zhang <yuhez@nvidia.com>

Qwen3_5ForConditionalGeneration.prepare_model_inputs_for_cp ran the full vision tower redundantly on every CP rank, including the visual forward and all-media activations. Route its image and video calls through cp_vision_shard.maybe_distribute_visual so each CP rank processes a frame slice and gathers the resulting embeddings in their original order. The typed distributed.cp_vision_sharding policy controls the feature. When it is disabled or no CP vision group is published, the helper preserves the replicated get_image_features and get_video_features path. The published group is CP-only, which is gradient-correct for frozen and trainable towers. The integration preserves the upstream prepare_model_inputs_for_cp structure and its per-modality calls. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: Yuhe Zhang <yuhez@nvidia.com>

…ward The VLM fine-tuning recipe drives context-parallel vision pre-embedding through model(_pre_embed_only=True, ...). Wrap the training and evaluation call sites in _run_cp_pre_embed, which publishes the CP submesh process group for the duration of the forward and restores it in a finally block. The group is CP-only and is gradient-correct for frozen and trainable towers. The wrapper is a no-op when cp_size is 1 or the typed distributed.cp_vision_sharding policy is disabled. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: Yuhe Zhang <yuhez@nvidia.com>

Add a single-node dense Qwen3.5-4B VLM configuration that exercises context-parallel vision-tower sharding on eight GPUs with CP2 and DP4. Vision sharding is enabled under distributed.cp_vision_sharding, and the header documents a same-topology baseline using the typed CLI override. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: Yuhe Zhang <yuhez@nvidia.com>

…n-shard Add coverage that exercises the ACTUAL transformers Qwen3VLVisionModel (the tower the dense Qwen3.5 VLM CP pre-embed shards), closing the "no real-tower coverage" gap: - test_cp_vision_shard_qwen3_5.py (CPU): build a tiny real Qwen3-VL tower and assert visual(all) == concat_r visual(slice_r) for pooler_output AND every deepstack feature, partitioned by the module's own _contiguous_balanced_bounds -- image entries at world 2-4, a single 8-frame video split at frame granularity, and the flat-vs-split pooler_output equivalence the Qwen3.5 wiring relies on. - test_cp_vision_shard_gloo.py: add a 2-rank gloo test driving maybe_distribute_visual on the real tiny tower through real all_gather / reduce_scatter collectives, asserting pooler + deepstack forward parity with the replicated reference (skipped when transformers Qwen3-VL is unavailable). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: Yuhe Zhang <yuhez@nvidia.com>

Signed-off-by: Yuhe Zhang <yuhez@nvidia.com>

yuhezhang-ai · 2026-07-28T17:17:40Z

/ok to test dcca00d

akoumpa · 2026-07-29T13:37:12Z

/ok to test ae0ff22

athitten

Thank you @yuhezhang-ai ! Just re-approving based on @akoumpa 's approval

yuhezhang-ai mentioned this pull request Jul 22, 2026

[Tracking] Large-scale MoE-VLM training improvements #2985

Closed

Base automatically changed from huiyingl/refactor/cp-unify to main July 23, 2026 09:28

yuhezhang-ai force-pushed the yuhez/feat/cp-vlm-integration-2937 branch 2 times, most recently from c05725a to fd5faa7 Compare July 23, 2026 14:33

yuhezhang-ai changed the base branch from main to yuhez/upstream-cp-vision-shard July 23, 2026 14:48

yuhezhang-ai and others added 23 commits July 23, 2026 08:00

fix(cp-vision): validate dense qwen3.5 integration

902bf3d

Signed-off-by: Yuhe Zhang <yuhez@nvidia.com>

fix(cp-vision): use typed recipe configuration

ce36287

Signed-off-by: Yuhe Zhang <yuhez@nvidia.com>

feat(qwen3-vl): add context-parallel vision sharding

a3b3ea4

Signed-off-by: Yuhe Zhang <yuhez@nvidia.com>

feat(cp-vision): accept explicit auto cost model

be6b491

Signed-off-by: Yuhe Zhang <yuhez@nvidia.com>

fix(cp-vision): keep execution grid on device

5fe7a89

Signed-off-by: Yuhe Zhang <yuhez@nvidia.com>

refactor(cp-vision): nest recipe policy under distributed

1e886ca

Signed-off-by: Yuhe Zhang <yuhez@nvidia.com>

fix(cp-vision): align production metadata handling

6ef339b

Signed-off-by: Yuhe Zhang <yuhez@nvidia.com>

docs(cp): apply technical publications review

91c24ca

Signed-off-by: Yuhe Zhang <yuhez@nvidia.com>

fix(cp-vision): align dense callers with unified CP sharder

dafb86d

Signed-off-by: Yuhe Zhang <yuhez@nvidia.com>

copy-pr-bot Bot temporarily deployed to public July 28, 2026 15:48 Inactive

fix(vlm): resolve CI regressions after rebase

dcca00d

Signed-off-by: Yuhe Zhang <yuhez@nvidia.com>

copy-pr-bot Bot temporarily deployed to test July 28, 2026 17:18 Inactive

copy-pr-bot Bot temporarily deployed to nemo-ci July 28, 2026 17:18 Inactive

copy-pr-bot Bot temporarily deployed to public July 28, 2026 17:18 Inactive

copy-pr-bot Bot temporarily deployed to public July 28, 2026 17:20 Inactive

copy-pr-bot Bot had a problem deploying to nemo-ci July 28, 2026 20:52 Failure

copy-pr-bot Bot temporarily deployed to nemo-ci July 28, 2026 20:52 Inactive

Merge branch 'main' into yuhez/feat/cp-vlm-integration-2937

ae0ff22

copy-pr-bot Bot temporarily deployed to test July 29, 2026 13:37 Inactive

copy-pr-bot Bot temporarily deployed to nemo-ci July 29, 2026 13:37 Inactive

copy-pr-bot Bot temporarily deployed to public July 29, 2026 13:37 Inactive

copy-pr-bot Bot temporarily deployed to public July 29, 2026 13:40 Inactive

athitten approved these changes Jul 29, 2026

View reviewed changes

Navigation Menu

Uh oh!

feat(vlm): integrate packed CP and vision sharding for Qwen3.5-MoE - #3186

feat(vlm): integrate packed CP and vision sharding for Qwen3.5-MoE#3186
yuhezhang-ai merged 61 commits into
mainfrom
yuhez/feat/cp-vlm-integration-2937

yuhezhang-ai commented Jul 22, 2026 •

edited

Loading

Uh oh!

copy-pr-bot Bot commented Jul 22, 2026

Uh oh!

github-actions Bot commented Jul 22, 2026

Uh oh!

yuhezhang-ai commented Jul 28, 2026

Uh oh!

akoumpa commented Jul 29, 2026

Uh oh!

athitten left a comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

Uh oh!

Conversation

yuhezhang-ai commented Jul 22, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What does this PR do?

Changelog

Landing status

Validation

Rebase/refactor validation

Qwen3.5-122B-A10B, 128K integration runs

100-step mixed-workload run

Post-refactor CP32 runs

Before ready for review

Additional information

Uh oh!

copy-pr-bot Bot commented Jul 22, 2026

Uh oh!

github-actions Bot commented Jul 22, 2026

Uh oh!

yuhezhang-ai commented Jul 28, 2026

Uh oh!

akoumpa commented Jul 29, 2026

Uh oh!

athitten left a comment

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

yuhezhang-ai commented Jul 22, 2026 •

edited

Loading