|
| 1 | +--- |
| 2 | +name: nemo-mbridge-perf-vision-dp-over-cp |
| 3 | +description: Operational guide for sharding a VLM vision encoder across the language model's context-parallel ranks in Megatron-Bridge, including config knobs, code anchors, load-balance pitfalls, and measured impact. |
| 4 | +license: Apache-2.0 |
| 5 | +when_to_use: The vision tower replicates work and activations on every CP rank while the language model is already sharded, or a VLM run at CP>1 OOMs in the encoder; 'vision_dp_over_cp', 'vision CP', 'vision DP over CP', 'image sharding across CP ranks', 'encoder replicated across CP'. |
| 6 | +--- |
| 7 | + |
| 8 | +# Vision DP Over CP Skill |
| 9 | + |
| 10 | +This skill covers `vision_dp_over_cp`: reusing the language model's |
| 11 | +context-parallel process group to shard the **images** (or temporal tubelets) of |
| 12 | +a microbatch across CP ranks, so each rank encodes `1/CP` of them instead of all |
| 13 | +of them. |
| 14 | + |
| 15 | +It is **not** ring attention and shares nothing with |
| 16 | +`nemo-mbridge-perf-hierarchical-context-parallel` beyond the process group. |
| 17 | +Vision encoders such as RADIO already attend *within* each image via |
| 18 | +`cu_seqlens`, so there is no long sequence to split and no cross-rank attention |
| 19 | +communication. The correct mental model is **data parallelism over images, |
| 20 | +borrowing the CP group** — a rank encodes whole images, never partial ones. |
| 21 | + |
| 22 | +Shipped for Nemotron Omni. The MCore split/gather helpers are model-agnostic, so |
| 23 | +the pattern generalizes to any Bridge VLM with a dynamic-resolution encoder. |
| 24 | + |
| 25 | +## Do Not Confuse With `dist_train.vision_context_parallel_size` |
| 26 | + |
| 27 | +Bridge already has a similarly named and completely different knob: |
| 28 | +`vision_context_parallel_size` on the dist-train path |
| 29 | +(`src/megatron/bridge/models/qwen_vl/qwen3_vl_provider.py`, consumed in |
| 30 | +`src/megatron/bridge/training/initialize.py`). |
| 31 | + |
| 32 | +| | `model.vision_dp_over_cp` (bool) | `dist_train.vision_context_parallel_size` (int) | |
| 33 | +|---|---|---| |
| 34 | +| Ranks | Reuses the language model's CP ranks | Gives the vision tower its **own** ranks | |
| 35 | +| Splits | Whole images across those ranks | The vision encoder's own sequence dimension | |
| 36 | +| Encoder `context_parallel_size` | Stays `1` | Set to this value | |
| 37 | +| Models | Nemotron Omni | Qwen3-VL dist-train | |
| 38 | + |
| 39 | +They are not alternatives and do not compose in any tested configuration. |
| 40 | + |
| 41 | +## Enablement |
| 42 | + |
| 43 | +```python |
| 44 | +cfg.model.context_parallel_size = 2 # the flag is a no-op at CP=1 |
| 45 | +cfg.model.vision_dp_over_cp = False # opt out; on by default |
| 46 | +``` |
| 47 | + |
| 48 | +Defaults to `True`. It self-disables when `context_parallel_size == 1`, so a |
| 49 | +recipe that does not raise CP behaves exactly as it did before the flag existed |
| 50 | +and reports no error. Opting out is only useful for A/B measurement or to |
| 51 | +isolate a suspected split/gather bug. |
| 52 | + |
| 53 | +## Code Anchors |
| 54 | + |
| 55 | +Provider field, `src/megatron/bridge/models/nemotron_omni/nemotron_omni_provider.py`: |
| 56 | + |
| 57 | +```python |
| 58 | +vision_dp_over_cp: bool = True |
| 59 | +``` |
| 60 | + |
| 61 | +Gating, in `NemotronOmniModel.__init__` |
| 62 | +(`src/megatron/bridge/models/nemotron_omni/modeling_nemotron_omni.py`): |
| 63 | + |
| 64 | +```python |
| 65 | +self.context_parallel_lm = language_transformer_config.context_parallel_size |
| 66 | +self.vision_dp_over_cp = vision_dp_over_cp and self.context_parallel_lm > 1 |
| 67 | +``` |
| 68 | + |
| 69 | +The split itself wraps an upstream MCore helper: |
| 70 | + |
| 71 | +```python |
| 72 | +return split_to_context_parallel_ranks_dynamic_res( |
| 73 | + images, |
| 74 | + imgs_sizes, |
| 75 | + vision_packed_seq_params, |
| 76 | + patch_dim=self.patch_dim, |
| 77 | + fp8_enabled=False, |
| 78 | + fp8_recipe=getattr(self.config, "fp8_recipe", None), |
| 79 | + num_frames=num_frames, |
| 80 | + temporal_patch_size=temporal_patch_size, |
| 81 | +) |
| 82 | +``` |
| 83 | + |
| 84 | +Project-then-gather, at the end of `_encode_images`: |
| 85 | + |
| 86 | +```python |
| 87 | +if shard_vision: |
| 88 | + # Project before the gather so the projector stays sharded, then |
| 89 | + # restore the global feature set every CP rank's media merge needs. |
| 90 | + projected = _project_multimodal_embeddings(self.vision_projection, encoded) |
| 91 | + gathered = gather_from_context_parallel_ranks_dynamic_res(projected, num_padded_ranks) |
| 92 | + return gathered.contiguous() |
| 93 | +``` |
| 94 | + |
| 95 | +Upstream helpers (do not modify — changes go through the MCore repo): |
| 96 | + |
| 97 | +- `3rdparty/Megatron-LM/megatron/core/models/multimodal/context_parallel.py` |
| 98 | + — `split_to_context_parallel_ranks_dynamic_res`, |
| 99 | + `gather_from_context_parallel_ranks_dynamic_res` |
| 100 | +- `3rdparty/Megatron-LM/megatron/core/models/multimodal/llava_model.py` |
| 101 | + — performs the identical split/encode/gather on its temporal path |
| 102 | + |
| 103 | +## Implementation Map |
| 104 | + |
| 105 | +1. **Split.** The MCore helper partitions the patched image tensor on image (or |
| 106 | + tubelet) boundaries and rewrites `imgs_sizes` and the vision |
| 107 | + `PackedSeqParams` (`cu_seqlens`, `max_seqlen`) to describe only the rank-local |
| 108 | + share. Sharding is by **image count**, not patch count; the remainder goes to |
| 109 | + the **last** rank. |
| 110 | +2. **Placeholder padding.** When a microbatch has fewer images than CP ranks, |
| 111 | + the helper injects `num_padded_imgs = max(0, cp_size - num_imgs)` dummy `1x1` |
| 112 | + images so no rank runs an empty encoder, and returns `num_padded_ranks` so the |
| 113 | + gather can drop them again. |
| 114 | +3. **Encode.** Each rank runs the full vision tower on its own images only. |
| 115 | + Activations, not just compute, are what drop by `1/CP`. The vision encoder's |
| 116 | + own `context_parallel_size` stays hardcoded to `1` in |
| 117 | + `_build_vision_config` — that is correct and unchanged by this feature, since |
| 118 | + the sharding happens *outside* the encoder, over its inputs. |
| 119 | +4. **Even-grid padding.** `_pad_patch_grid_to_even` zero-pads an odd patch grid |
| 120 | + before pixel shuffle. Only the injected `1x1` placeholders hit this path; |
| 121 | + real frames are already even in both extents. |
| 122 | +5. **Project, then gather.** The vision projector runs on the rank-local shard, |
| 123 | + so it stays sharded too, and only the smaller post-projector features cross |
| 124 | + the wire. The gather is the **only** collective this feature adds. |
| 125 | +6. **Merge.** After the gather, every rank holds the full global feature set, |
| 126 | + which is exactly what the unchanged media-merge path expects. |
| 127 | + |
| 128 | +## Measured Impact |
| 129 | + |
| 130 | +Nemotron-3 Nano Omni 30B-A3B, 2 nodes / 16 GPUs, PP=1 DP=2 EP=8 ETP=1 MBS=2 |
| 131 | +GBS=64 GA=16, with only `vision_dp_over_cp` flipped: |
| 132 | + |
| 133 | +| Dataset | TP | CP | Step time | Peak memory | |
| 134 | +|-------------|----|----|---------------|-------------| |
| 135 | +| InfoVQA | 4 | 2 | 13.2 -> 12.5 s | 81 -> 73 GB | |
| 136 | +| Mantis | 2 | 4 | 12.2 -> 11.6 s | 81 -> 77 GB | |
| 137 | +| llava_video | 4 | 2 | 13.6 -> 13.0 s | 84 -> 75 GB | |
| 138 | + |
| 139 | +It is both a memory win (-5% to -11% peak) and a throughput win (-4% to -5% step |
| 140 | +time). Do not describe it as memory-only; a single early video A/B suggested step |
| 141 | +time was flat, and the three-dataset sweep contradicted that. |
| 142 | + |
| 143 | +The Mantis row is the weakest memory result, and that is the load-balance |
| 144 | +signature described under Pitfalls, not noise. |
| 145 | + |
| 146 | +## Pitfalls |
| 147 | + |
| 148 | +1. **Silent no-op at CP=1.** The flag ANDs with `context_parallel_lm > 1`. Both |
| 149 | + checked-in Nemotron Omni recipes ship CP=1, so the on-by-default value changes |
| 150 | + nothing there. Always confirm CP>1 before attributing a result to this flag, |
| 151 | + and note that raising CP now silently enables image sharding too. |
| 152 | +2. **Count-based split leaves the tail rank hot.** Images are divided by count |
| 153 | + with the remainder on the last rank. Multi-image data with a wide per-sample |
| 154 | + image-count spread (e.g. Mantis, 1-5 images/sample at CP=4) balances poorly: |
| 155 | + peak memory is set by the busiest rank, so the win shrinks. A |
| 156 | + patch-count-balanced greedy split on image boundaries is the fix, and it |
| 157 | + belongs **upstream in MCore**, not reimplemented in Bridge. Since the flag is |
| 158 | + on by default, this imbalance is what a multi-image CP>1 run gets unless the |
| 159 | + user opts out. |
| 160 | +3. **Process-group divergence is checked, not assumed.** The MCore splitter |
| 161 | + resolves the CP group from global parallel state, so a `ProcessGroupCollection` |
| 162 | + whose `cp` size disagrees with `context_parallel_size` would shard against the |
| 163 | + wrong ranks. Bridge raises instead: |
| 164 | + `"Nemotron Omni vision_dp_over_cp does not match its process group"`. |
| 165 | +4. **Odd patch grids break pixel shuffle.** Omni's factor-2 pixel shuffle rejects |
| 166 | + an odd grid, and MCore's `1x1` placeholder images are odd by construction. |
| 167 | + Without `_pad_patch_grid_to_even`, the CP>images case fails only for the |
| 168 | + specific microbatches that trigger placeholder injection — an intermittent |
| 169 | + failure that looks data-dependent. |
| 170 | +5. **Do not gather before projecting.** Gathering encoder output first would |
| 171 | + replicate the projector on every rank and move a larger tensor. The ordering |
| 172 | + is deliberate. |
| 173 | +6. **The placeholder path is unvalidated end-to-end.** A two-rank unit test |
| 174 | + proves the sharded encoder matches the unsharded one when a microbatch holds |
| 175 | + fewer images than ranks, but no multi-node training run has exercised it, and |
| 176 | + nothing covers CP=8. |
| 177 | + |
| 178 | +## Verification |
| 179 | + |
| 180 | +Unit tests. The second file self-launches a two-rank `torch.distributed` job and |
| 181 | +asserts the sharded tower matches the unsharded one for fewer images than ranks, |
| 182 | +an exact split, and a remainder; it skips when fewer than two GPUs are visible: |
| 183 | + |
| 184 | +```bash |
| 185 | +uv run python -m pytest tests/unit_tests/models/nemotron_omni/test_nemotron_omni_model.py -q |
| 186 | +uv run python -m pytest tests/unit_tests/models/nemotron_omni/test_vision_dp_over_cp_distributed.py -q |
| 187 | +``` |
| 188 | + |
| 189 | +Multi-GPU A/B. Run the same recipe twice at CP>1, flipping only the flag, and |
| 190 | +compare reported step time and peak allocated memory: |
| 191 | + |
| 192 | +```bash |
| 193 | +uv run python -m torch.distributed.run --nproc_per_node=8 \ |
| 194 | + scripts/training/run_recipe.py \ |
| 195 | + --recipe nemotron_omni_cord_v2_sft_config \ |
| 196 | + model.context_parallel_size=2 \ |
| 197 | + model.vision_dp_over_cp=false \ |
| 198 | + train.train_iters=20 |
| 199 | +``` |
| 200 | + |
| 201 | +Success criteria: |
| 202 | + |
| 203 | +- Peak memory drops relative to the `false` arm; step time does not regress. |
| 204 | +- Loss curves of the two arms track each other. The gather makes this |
| 205 | + mathematically equivalent, so a divergence means the split or gather is wrong, |
| 206 | + not that the feature is lossy. |
| 207 | +- If nothing changes at all, check that `context_parallel_size > 1` — the flag |
| 208 | + silently self-disables at CP=1. |
| 209 | + |
| 210 | +## Related |
| 211 | + |
| 212 | +- @skills/nemo-mbridge-perf-moe-vlm-training/SKILL.md — VLM training strategy, |
| 213 | + FSDP vs 3D parallel, and why mock image data invalidates VLM benchmarks |
| 214 | +- @skills/nemo-mbridge-perf-memory-tuning/SKILL.md |
| 215 | +- @skills/nemo-mbridge-perf-hierarchical-context-parallel/SKILL.md — the |
| 216 | + *sequence*-splitting use of the same process group; unrelated mechanism |
| 217 | +- @skills/nemo-mbridge-perf-vision-dp-over-cp/card.yaml |
0 commit comments