[Fork review] NMFW-17: ColocatedBridgeCommunicator (PP=1, CP=1) by yashaswikarnati · Pull Request #10 · yashaswikarnati/Megatron-LM

yashaswikarnati · 2026-04-17T22:54:42Z

Fork-internal PR for reviewing the pr-a diff before upstreaming.

Commits

9cc2d37e19 — Initial ColocatedBridgeCommunicator (heterogeneous TP/DP, PP=1, CP=1)
eb8191e6ea — Fix fan-out backward: all-gather across sibling dest ranks
e238bcb184 — Apply review fixes (8 items):
- Reject CP>1 with a clear error (follow-up PR will add CP>1 support)
- Validate 'tp' and 'dp' dims present in both grids
- Guard _get_fan_in_slice_info against non-src ranks (symmetric with fan-out)
- Drop sorted() in group construction; append order is already slot order
- Add destroy() for NCCL process group cleanup
- Remove dead ctx.input_batch_size and _my_all_gather_group_idx
- Strengthen E2E test with cross-TP consistency + iter-over-iter change checks
- Direct unit test for fan_out_gather_group_ranks + slot-order invariant
- Copyright 2026 → 2025

Stack

pr-a (this): bridge, PP=1, CP=1
pr-c (next): adds dest CP>1 on top of pr-a
pr-b (separate): adds PP>1 support

Upstream PR

NVIDIA#4368

🤖 Generated with Claude Code

… (NMFW-17) COLOCATED mode replaces UNIFIED and covers both the legacy (no grid map) and heterogeneous TP/DP on shared ranks cases. Auto-detected from grid overlap. Core: - ColocatedBridgeCommunicator handles fan-in / fan-out / equal-DP with autograd-correct collectives. Fan-in uses all_gather_into_tensor into a pre-allocated buffer (no intermediate tensor list + torch.cat copy). - MimoModel._forward_all_modules applies the communicator after encoder forward when colocated comms are configured. - module_to_grid_map key validation runs regardless of colocated/non-colocated path so mismatched keys always raise. Tests: - test_mimo_colocated_communicator.py: rank mapping, slice info, fan-in/fan-out/equal-DP forward+backward correctness. - test_mimo_colocated_correctness.py: TransformerBlock-based multi-iteration correctness (9 checks x 3 iters). - test_mimo_colocated_e2e.py: end-to-end colocated VLM with heterogeneous TP/DP at PP=1. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

The fan-out backward was returning a locally zero-padded gradient with no collective, leaving each src rank with only its own dest rank's slice of the gradient and zeros elsewhere. Every other dest rank that consumed a different slice of the same src activation was dropped from the encoder param gradient. Fan-out forward is a local narrow (no comm), so per-rank autograd is correct in isolation — but logically the src rank's batch is shared across multiple dest ranks that each take a disjoint slice. The autograd-correct adjoint reconstructs the full src-batch gradient as the concatenation of slice gradients. We now build a process group per (src_dp_idx, dest_tp_idx) containing the `scale = dest_dp / src_dp` dest ranks that consume consecutive slots of the src rank's batch, and use `dist.all_gather_into_tensor` on the backward (symmetric with the fan-in forward). Rank order inside the group matches slot order, so the gathered tensor already has the correct layout. The previous test tolerance (atol=5e-3, rtol=0) was ~1000x looser than the discrepancy the bug produced on the toy scale used here (gradient ~1e-5, disagreement ~5e-6 on the samples handled by sibling llm_dp ranks), so the zero-pad bug passed the test silently. Tightened the input_grad check to atol=5e-6, rtol=1e-3, which reliably fails when the fix is reverted. Verified: with the fix, max|ci.grad - ri.grad| drops from ~1e-5 to ~2e-7 (floating-point noise) across all three fan_in / fan_out / equal cases. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

- Reject CP>1 in _validate_grids with a clear error (CP>1 support is a follow-up; the existing _gen_rank_enum(['tp']) enumeration would silently corrupt dp_idx otherwise). - Validate presence of 'tp' and 'dp' dims in both grids up-front. - Guard _get_fan_in_slice_info against non-src ranks (symmetric with _get_fan_out_slice_info). - Drop sorted() in group construction; ranks are already appended in slot order, which is what all_gather_into_tensor requires. Add a docstring spelling out the invariant. - Add destroy() to release all_gather_pg and fan_out_gather_pg. - Remove dead ctx.input_batch_size and _my_all_gather_group_idx. - Strengthen E2E test with cross-TP loss consistency and iter-over-iter change checks (catches silent corruption that leaves loss finite but training a no-op). - Add direct test for fan_out_gather_group_ranks structure + slot-order invariant. - Fix copyright year 2026 -> 2025 in new files. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

- Reject PP>1 on both src and dest grids (PP=1 scope only). - Replace float dp_scale_factor with integer fan_in_scale/fan_out_scale; drop fragile int(1/float) casts. - Switch to public HyperCommGrid.get_rank_enum; drop the PP filter and unused _get_rank_dim_coord helper that supported dest PP>1. - Pass dim_mapping={'b': 0, 'h': 1} from MimoModel so the communicator correctly slices the flattened (s*b, h) encoder output; document the flattened-tensor assumption on the class. - Drop redundant guards in _build_colocated_communicators; caller and _is_colocated already guarantee the invariants. - Remove unused is_equal_dp and the four PG getters; access all_gather_pg / fan_out_gather_pg directly from the autograd function. - Restore copyright year on config/role.py (2025, 2026). - Tests: destroy communicator PGs in teardowns to stay under NCCL's concurrent-communicator cap, add rank_offset>0 rank-mapping case, and add negative tests covering each _validate_grids raise path. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

Hardening in the communicator: - ``communicate()`` raises ``ValueError`` when the batch dim is not divisible by ``fan_out_scale``; ``get_slice_info()`` re-checks fan-in and fan-out on the backward-narrow path. Prevents silent sample truncation. - Document the TP-replication precondition on both the class and ``communicate()``. Add an opt-in debug check (``CHECK_TP_REPLICATION`` class flag or ``check_tp_replication`` ctor arg) that all-gathers across the src TP group and raises on any peer mismatch. New tests in ``test_mimo_colocated_communicator.py``: - ``TestBridgeGradients`` — bitwise-exact (``atol=0, rtol=0``) coverage of the four pure-data-movement paths plus equal-DP identity: * fan-in forward == ``torch.cat`` of sibling inputs (2 shape configs × 2 dim_mappings), * fan-in backward == ``grad_output.narrow`` at this rank's slot, * fan-out forward == ``input.narrow`` at this rank's slot, * fan-out backward == cat of every sibling's grad in slot order — the critical guard against zero-pad-without-gather, wrong slot order, double-counting, or missing siblings, * equal-DP fwd/bwd is a pure identity. - ``TestCommunicatePreconditions`` — fan-out communicate() and ``get_slice_info`` raise on non-divisible batch. - ``TestDestroy`` — ``destroy()`` nulls both PGs and is idempotent. - ``TestValidateGrids.test_pp_gt_1_rejected`` — dedicated coverage of the PP>1 guard on either grid. New tests in ``test_mimo_colocated_e2e.py``: - ``test_colocated_fan_out_8gpu`` — mirror of the fan-in e2e case with encoder TP4/DP2 and LLM TP2/DP4. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

yashaswikarnati · 2026-04-20T20:02:11Z

+            )
+
+    def _get_fan_out_slice_info(self, batch_size: int) -> SliceInfo:
+        if self.current_rank not in self.rank_to_dest_pos:


if self.current_rank not in self.rank_to_dest_pos: when is this condition true? all ranks must be in short ?

yashaswikarnati · 2026-04-20T20:10:24Z

+        """Return True if src DP > dest DP (encoder has more replicas)."""
+        return self.src_dp_size > self.dest_dp_size
+
+    def _assert_tp_replicated(self, tensor: torch.Tensor) -> None:


remove for now

yashaswikarnati · 2026-04-20T20:32:22Z

+            else:
+                self.role = RankRole.from_grid_map(mimo_config.module_to_grid_map, modality_names)


RankRole.from_grid_map and RankRole.colocated seem not consistent, can we have consistent interface, discuss some options

yashaswikarnati · 2026-04-20T20:33:23Z

+        self.colocated_comms = {}
        if mimo_config.module_to_grid_map:
-            self.role = RankRole.from_grid_map(mimo_config.module_to_grid_map, modality_names)
+            expected_keys = set(modality_names) | {MIMO_LANGUAGE_MODULE_KEY}


can we have this valdation as part of mimo model config post init? and not pollute boiler plate in main model constructor

yashaswikarnati · 2026-04-20T20:35:09Z

+    def _build_colocated_communicators(self):
+        """Build communicators for each encoder → language edge.
+
+        Encoder outputs reach the communicator after ``VisionModalitySubmodules``


why we have this doc string here

- Drop the redundant ``Scope: PP=1`` class-docstring line; tighten the docstring while still calling out the TP-replication precondition. - Remove the ``CHECK_TP_REPLICATION`` class flag, the ``check_tp_replication`` ctor arg, the ``_assert_tp_replicated`` helper, and its call in ``communicate()``. The precondition is documented; a collective equality check on every forward is not the right tool. - Unify ``_build_all_gather_groups`` and ``_build_fan_out_gather_groups`` behind a single ``_build_gather_groups`` helper. Both build the same shape of group (iterate one side's DP slot, gather ``scale`` sibling ranks per slot, per sibling TP shard) — only the endpoints differ. - Drop the ``current_rank not in rank_to_{src,dest}_pos`` guards in the slice-info helpers. Under the PP=1 colocated contract every participating rank is in both maps; the guard was dead. - Consolidate ``MimoModel.__init__`` role/bridge construction: the grid-map key validation now lives in ``MimoModelConfig.__post_init__``, and the ``RankRole.colocated`` / ``from_grid_map`` classmethods take a consistent ``modality_module_names`` argument (the language key is appended internally, matching ``from_grid_map``). - Drop the ``_build_colocated_communicators`` docstring — the body is self-explanatory. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

yashaswikarnati · 2026-04-20T20:53:24Z

-            self._make_config(encoder_in_grid=True, language_in_grid=True, pp_rank=1, pp_size=3)
+        # Stage info with PP on a non-colocated layout (encoder and language on
+        # different rank ranges, which routes through RankRole.from_grid_map).
+        world_size = dist.get_world_size()


the existing test should pass without passig in the grid? with default colocated and using process groups from parallel state? why did we add this?

yashaswikarnati · 2026-04-20T20:57:14Z

+_embedding_pg_cache: dict = {}
+
+
+def create_hypercomm_grid(offset=0, tp=1, cp=1, pp=1, dp=1):


tests/unit_tests/models/test_mimo_1f1b_schedule.py : can we reuse all the helpers and machinery from this other test we already have and not define these verbose tests all over again ?

yashaswikarnati · 2026-04-20T21:06:44Z

+H, NHEADS, SEQ, GBS = 1024, 8, 8, 8
+
+
+def _make_block(num_layers, dtype, pg):


can we reuse helper functions and actual module creations from tests/unit_tests/models/test_mimo_1f1b_schedule.py :

yashaswikarnati · 2026-04-20T21:07:57Z

+            dim_mapping={'s': 0, 'b': 1, 'h': 2},
+        )
+        _active_comms.append(comm)
+        ref_opt = torch.optim.SGD(list(ref_enc.parameters()) + list(ref_llm.parameters()), lr=lr)


can we use actual mimo optimizer which exercises the distributed optimizer path, which is what we care about

yashaswikarnati · 2026-04-20T21:11:35Z

+            ri = gi.clone().detach().requires_grad_(True)
+            reo = ref_enc(hidden_states=ri, attention_mask=None)
+            rlo = ref_llm(hidden_states=reo, attention_mask=None)
+            rlo.sum().backward()


can we use actual cross entropy loss? which is whats in megatron. and also resue infra fro other tests for data iterator fwd pass etc. we need to use actual modules as well. lets have consistent mimo model init across tests for colocated and non colocated and e2e and correctness tests

yashaswikarnati · 2026-04-20T21:14:56Z

+_DIM_MAPPING_IDS = ["sbh", "bsh"]
+
+
+class TestBridgeGradients:


how this relates to other correctness tests we added in other test file. is there any benifit of these particular tests

yashaswikarnati · 2026-04-20T21:46:46Z

+        # bridge whose DP indices map into one DP slot on the iterating side.
+        # Fan-in: iterate dest DP, gather src ranks; fan-out: iterate src DP,
+        # gather dest ranks. Same shape, different endpoints — share a builder.
+        if self.fan_in_scale > 1:


Question: do we need both fan in and fan out and self.fan_out_gather_group_ranks and also self.all_gather_pg,. only one can be true at a given time ? i wonder if we can simplify this abstraction.

…rectness test (NMFW-17) Bug: with colocated heterogeneous DP (enc_dp != llm_dp), the LLM's mean-CE loss implicitly divides per-sample encoder gradients by local_B_llm = B_full / llm_dp. The encoder's DDP then additionally divides by its own DP group size (enc_dp), giving a total divisor of ``local_B_llm * enc_dp`` instead of ``B_full`` — encoder grads are off by ``llm_dp / enc_dp`` in either direction (fan-in or fan-out). Fix: add ``DistributedDataParallelConfig.gradient_reduce_div_factor`` — when set, the divisor used in pre-reduction scaling becomes this value instead of ``dp_cp_group.size()``. The reduction group itself is unchanged; only the scaling math moves. Default is ``None`` (backwards- compatible). This mirrors the existing MoE ``expert_gradient_scaling_factor`` pattern: keep the collective local, retune the scalar. Wire-up: - In ``DistributedDataParallel.__init__`` compute ``effective_dp_size`` from the override (or fall back to ``dp_cp_group.size()``), and use it in both avg-in-collective and sum paths. Assertions parametrized so defaults still round-trip to ``1/dp_cp_group.size()``. - In ``test_mimo_colocated_e2e.get_mimo_model_colocated`` set ``gradient_reduce_div_factor=llm_dp`` on the encoder's DDP config via ``dataclasses.replace``; the LLM's DDP keeps the default. Correctness test rewrite: ``test_mimo_colocated_correctness.py`` now uses a real mean-CE loss and Megatron DDP on both encoder and LLM, and compares per-param ``main_grad`` (post-``finish_grad_sync``) against a single-GPU reference running the full batch with identical weights. The previous toy sum-loss test could not see this bug — sum() has no implicit ``1/N`` in the backward chain. Parametrized fan-in, fan-out, and equal-DP; encoder grad tolerance ``atol=5e-4`` for fp32. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

The unit-test file for ColocatedBridgeCommunicator had accumulated duplicates and model-level golden coverage that now belongs in ``test_mimo_colocated_correctness.py``. Keep this file lean and focused on bridge semantics. Dropped (all per SPARK's explicit keep/remove list): - ``TestRankMappings.test_rank_mappings`` ``equal`` and ``extreme`` parametrizations (equal-DP has stronger coverage in TestBridgeGradients; extreme adds no branch coverage beyond the 2x case). - ``TestAllGatherGroups`` extreme 8x parametrizations plus the two trivial ``no_*_gather`` null-inverse tests. - ``TestSliceInfo`` class in full — it over-fit a private helper. - ``TestValidateGrids.test_pp_gt_1_rejected`` — exact duplicate of the per-side PP guard tests. - ``TestCommunicatePreconditions.test_non_divisible_get_slice_info_fan_out`` — redundant with the public fan-out divisibility test. - ``TestBridgeGradients`` 8x parametrizations and the ``tp2_dp4`` equal-DP duplicate — same branch, bigger instance. - ``TestGolden.test_forward_backward_golden`` — heavy model-level golden test that belongs in ``test_mimo_colocated_correctness.py``. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

…ison (NMFW-17) Per SPARK's proposal, the correctness test now exercises the full distributed training objective and compares every shard-visible value against a single-GPU reference: - Adds a TP=1 vocab head (plain ``nn.Linear``) on both sides and wires the LLM block + head into a single ``LLMWithHead`` container so their params share one DDP reduction. - Replaces the prior per-LLM-rank ``F.cross_entropy(..., reduction='mean')`` with an exact global-mean CE: all-reduce ``(local_num, local_den)`` on the LLM DP group and divide. This is mathematically the full-batch mean; no implicit ``1/local_tokens`` factor to compensate for. - With global-mean CE the DP reduction must be a pure SUM on both sides (the per-token grad scalar is already ``1/global_den`` on every rank). Both encoder and LLM DDP configs now set ``gradient_reduce_div_factor=1``. - Reference: each rank independently runs the full batch on identical TP=1 weights (kept consistent via ``_avg_params``), collects ``param.grad``. Parallel: each rank runs its DP slice, ``finish_grad_sync``, then compares ``param.main_grad``. - Asserts match against the reference for: scalar loss, encoder input grads, every encoder param grad (TP-sharded), every LLM-block param grad (TP-sharded), vocab-head grads (TP=1), and post-SGD-step weights — shard-wise ``atol=5e-4`` in FP32 with deterministic env. - Parametrized fan-in and fan-out. The equal-DP case was dropped as SPARK noted — identity coverage lives in ``TestBridgeGradients.test_equal_dp_is_bitwise_identity_fwd_and_bwd``. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

The previous ``loss_func`` did ``output_tensor.float().sum()`` on the per- token CE from ``GPTModel.compute_language_model_loss`` — that is neither a mean nor a masked sum, and it sidesteps the heterogeneous-DP grad scaling bug entirely. This replaces it with the exact distributed equivalent of full-batch ``F.cross_entropy(..., reduction='mean')``: masked = output_tensor.float() * loss_mask.float() local_num = masked.sum() local_den = loss_mask.float().sum() dist.all_reduce(local_num, group=llm_dp_pg) dist.all_reduce(local_den, group=llm_dp_pg) loss = local_num / local_den.clamp_min(1.0) With that formulation the per-token grad scalar is ``1/global_den`` on every rank, so the DP reduction must be a pure SUM. Both encoder and LLM DDP configs now set ``gradient_reduce_div_factor=1``; the earlier ``gradient_reduce_div_factor=llm_dp`` override on the encoder was the right fix for a per-rank ``F.cross_entropy(mean)`` loss and is no longer correct under num+den. ``forward_step`` now threads the LLM DP group into ``loss_func`` via ``partial``. Module docstring updated to document the formulation and to point at ``test_mimo_colocated_correctness.py`` for the tighter shard-wise reference-vs-distributed gradient and weight comparisons. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

Per PR review comments 15, 16, 18: stop redefining helpers, fixtures, and model builders in the colocated test files and consume them from ``test_mimo_1f1b_schedule.py`` so all MimoModel integration tests go through the same init path. Shared surface (imported from ``test_mimo_1f1b_schedule``): * ``create_hypercomm_grid``, ``destroy_all_grids`` * ``create_all_embedding_groups``, ``get_pg_collection_with_embedding_groups`` * ``get_language_model_spec``, ``get_vision_submodules_spec`` * ``get_mimo_model`` * ``DataIterator`` ``get_mimo_model`` now takes an optional ``ddp_config`` parameter so colocated callers can pass a config with ``gradient_reduce_div_factor=1`` (required under the num+den mean-CE formulation — see per-file docstring). 1F1B tests keep their existing default config. ``test_mimo_colocated_e2e.py``: * Deleted the duplicate helpers (~500 lines of verbose copies). * ``loss_func`` and ``forward_step`` unchanged semantically (num+den masked-CE all-reduced on the LLM DP group only), now documented as using the actual Megatron CE from ``GPTModel``. * Added ``test_colocated_fan_in_grad_accumulation_8gpu`` with ``num_microbatches=4`` to exercise gradient accumulation in the scheduler — AXIOM's requested coverage for the DDP reduction path. ``test_mimo_colocated_correctness.py``: * Dropped the raw-TransformerBlock-plus-linear-head reference stack; it was structurally different from MimoModel and forced a second, ad-hoc init path. * Now exercises the **actual distributed optimizer** via ``get_mimo_optimizer`` with ``use_distributed_optimizer=True`` (comment 17). * Reuses ``forward_step`` from the e2e module so both files use the same CE loss construction. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

Comment 14 (``test_mimo_model.py``): drop the PP-stage block added to ``test_role_determination`` this PR. The role-stage logic is already covered at the unit level in ``test_mimo_role.py`` and at the integration level in ``test_mimo_1f1b_schedule.py``; having the same MockGrid dance here was duplicative. Point a short comment at those files for future readers. Comment 19 (``test_mimo_colocated_communicator.py``): expand the ``TestBridgeGradients`` docstring to explicitly call out why these tests exist alongside the model-level correctness tests. The bridge is pure data movement (narrow / all-gather, no FP compute), so its adjoint can be asserted bitwise exact (``rtol=0, atol=0``) — that is the local invariant the MimoModel-level tests cannot verify because they see GEMM reduction-order noise. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

Per SENTINEL review: 1. Runtime int-check on ``DistributedDataParallelConfig.gradient_reduce_div_factor`` — the field is typed ``Optional[int]`` but Python does not enforce that at runtime. A float or bool here would silently produce wrong-scaled grads. Raise ``ValueError`` at DDP construction if the override is set to anything other than a positive int. 2. ``config.calculate_per_token_loss=True`` forces ``scaling_factor=1.0`` unconditionally and does the final division externally via ``finalize_model_grads``. Combining that with a ``gradient_reduce_div_factor`` override is ambiguous and almost certainly a bug, so refuse the combination. (SENTINEL's "per-token loss bypass" concern from the earlier review.) AXIOM's three remaining items already land via earlier phases: * grad-accumulation test: ``test_colocated_fan_in_grad_accumulation_8gpu`` in test_mimo_colocated_e2e.py exercises ``num_microbatches=4``. * (num, den) all-reduce group scope: documented in ``loss_func`` as "LLM DP group only — never the full dp*tp group". * divide-by-zero guard: ``local_den.clamp_min(1.0)`` in the same ``loss_func``. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

Rewrite test_mimo_colocated_correctness.py to compare a heterogeneous-DP MimoModel against a TP=1, DP=world_size reference that receives the full global batch on every rank. Under the num+den mean CE and gradient_reduce_div_factor=1, both configurations yield the DP=1 gradient on each encoder shard, so one Adam step lands on identical TP-sliced weights (within bf16 precision). Key implementation details: * Both models use gradient_reduce_div_factor=1 so DDP pure-SUMs over DP. * Ref params are copied into the TP-sharded dist params (via tensor_split on partition_dim) BEFORE building the distributed optimizer, since DistributedOptimizer snapshots current .data into fp32 master weights at __init__. * Global batch is generated on rank 0 and broadcast; dist pre-slices per rank to match forward_step expectations, ref consumes the full batch. * Parametrized over fan-in (enc_tp=2, enc_dp=4, llm_tp=4, llm_dp=2) and fan-out (enc_tp=4, enc_dp=2, llm_tp=2, llm_dp=4). Closes AXIOM R1 — previously the test only asserted ordering invariants, not numerical equivalence to a DP=1 oracle, so a wrong scaling factor could pass silently. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>

Exercise the fail-fast guards added in 552d8c2 to DistributedDataParallel: * Non-positive ints (0, -1, -100) raise ValueError with "must be a positive int". * Non-ints (1.0, 2.5, "2") raise ValueError with the same message (the guard uses `type(x) is not int` so float/str/bool are all rejected). * Combination with calculate_per_token_loss=True raises ValueError — the per-token loss path pins scaling_factor=1.0 and divides externally, so layering an explicit div factor on top is ambiguous. * Sanity checks: div_factor=None (default) and div_factor=4 (valid) pass. Closes AXIOM R2. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>

Mirror test_colocated_fan_in_grad_accumulation_8gpu with the fan-out layout (enc_tp=4, enc_dp=2, llm_tp=2, llm_dp=4, num_microbatches=4). The LLM side sees the larger DP group and slices inputs per-rank; the encoder's smaller DP group accumulates grads across microbatches while receiving broadcast-derived hidden states from multiple LLM peers. Covering this direction explicitly catches accumulation bugs that only surface when the encoder-side DP group is the smaller of the two. Closes AXIOM R4. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>

The DP=1/TP=1 reference produced different forward-pass results due to different TP-parallel accumulation orders. Switch to an equal-DP reference that shares the same encoder TP layout as the heterogeneous-DP model under test. When enc_dp==llm_dp the bridge is BridgeDirection.equal (identity), so the reference exercises the same DDP scaling path without any bridge-induced reshaping. Compares post-step encoder weights shard-wise at rtol=atol=1e-3.

yashaswikarnati · 2026-04-21T15:41:36Z

+                f"src={self.src_grid.rank_offset}, dest={self.dest_grid.rank_offset}"
+            )
+
+        for name, grid in [("src", self.src_grid), ("dest", self.dest_grid)]:


instead of 3 different for loops over src and dest grid, we can move all these checks under single for loop above when we check for tp dp ?

yashaswikarnati · 2026-04-21T15:46:48Z

-        else:
-            self.role = RankRole.unified(modality_names + [MIMO_LANGUAGE_MODULE_KEY])
+        self.colocated_comms = {}
+        # Single dispatch point for both colocated and non-colocated layouts.


remove # Single dispatch point for both colocated and non-colocated layouts.

Under bf16 with bias/dropout, logits diverged ~1.0 between the heterogeneous-DP dist model and the equal-DP ref — masking whether the bridge itself was numerically correct. Thread fp32 + add_bias_linear=False + attention/hidden dropout=0 through the test helpers, and add a post-bridge hidden-state oracle so the bridge output is checked directly (bit-exact) in addition to logits/grads/weights. Diff prints now include ref_max/ref_p95/ref_mean/rel_max so small absolute diffs on small-scale refs don't get flagged as clean. Drop the superseded e2e test. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>

Revert the DDP change that added ``gradient_reduce_div_factor`` and instead drive colocated hetero-DP grad scaling entirely from MIMO test code. Sub-model ``TransformerConfig``s now set ``calculate_per_token_loss=True`` (pins DDP ``scaling_factor=1.0``), the loss_func emits ``(local_sum, local_num_tokens, log_dict)``, and the MIMO ``finalize_grads_func`` lifts the ``num_tokens`` all-reduce over the LLM dp_cp group then uniformly calls ``scale_gradients(1/N_global)`` on both encoder and language sides. Semantically this is also an upgrade: the default mean-of-per-rank-means path is biased for variable-token VLM loss masks, whereas the per-token path gives the true global per-token mean. Remove: ``DistributedDataParallelConfig.gradient_reduce_div_factor``, its validator block, the two ``effective_dp_size`` call sites in DDP, and ``TestGradientReduceDivFactorValidation``. Correctness oracle (8 parametrizations, mbs1/mbs4 x uniform/asymmetric x fan_in/fan_out) passes; 1f1b MIMO test (5 passed, 2 skipped on GPU count); colocated communicator test (24 passed) — all green on 8 GPUs. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>

Address review agent concerns from PR NVIDIA#4368: 1. MimoModel.destroy() releases ColocatedBridgeCommunicator subgroups. Without it, NCCL leaks concurrent communicators across long-lived or repeatedly-rebuilt models. Tests now call destroy() before destroy_all_grids() via try/finally and setup/teardown hooks. 2. finalize_grads_func in the correctness test forwards force_all_reduce to the per-side finalize_model_grads() so PP grad-sync semantics aren't silently dropped if the schedule ever exercises them here. Verified on cog (cw-dfw): - test_mimo_colocated_correctness (8 gpu): 8 passed in 4:43 - test_mimo_colocated_communicator (8 gpu): 24 passed in 14s - test_mimo_1f1b_schedule::test_full_pp_8gpu: passed in 20s Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>

yashaswikarnati · 2026-04-22T01:06:17Z

-                        assert (gradient_scaling_factor == 1) or (
-                            gradient_scaling_factor
-                            == (self.expt_dp_group.size() / self.dp_cp_group.size())
+                        # For non-expert parameters, gradient_scaling_factor is 1.0.


unwanted diff

yashaswikarnati · 2026-04-22T01:09:57Z

@@ -1,11 +1,11 @@
-# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved.
+# Copyright (c) 2025, 2026, NVIDIA CORPORATION. All rights reserved.


dont change the copy right for this

yashaswikarnati · 2026-04-22T01:11:47Z

+    ) -> 'RankRole':
+        """Construct a RankRole from modality names and an optional grid map.
+
+        This is the unified entry point — callers pass the same first


we dont need big doc string

yashaswikarnati · 2026-04-22T01:15:37Z

+              module this rank participates in.
+        """
+        if module_to_grid_map is None or cls._all_grids_colocated(module_to_grid_map):
+            return cls.colocated(modality_module_names)


can we cleanly expose a api that covers both colocated and non colocated? and we dont have to deal with from gird_map ? so is module to grid map available any way for both? do we also need modality module names expicitly or this is already avail in module to grid map.

yashaswikarnati · 2026-04-22T01:19:20Z

        for p1, p2 in zip(ddp_model1.parameters(), ddp_model2.parameters()):
            if hasattr(p1, 'main_grad') and hasattr(p2, 'main_grad'):
                testing.assert_close(p1.main_grad, p2.main_grad, rtol=0, atol=0)
+


these changes arre not required ?

yashaswikarnati · 2026-04-22T01:20:59Z

-        )
-        assert model_pp.role.is_first_stage("images") is False
-        assert model_pp.role.is_last_stage("images") is False
+        # Non-colocated / PP role-stage coverage lives in tests/unit_tests/


do we need this verbose comment? why diid we modify this filee?

yashaswikarnati · 2026-04-22T01:21:24Z

@@ -0,0 +1,560 @@
+# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.


new files copy right should be 2026

yashaswikarnati · 2026-04-22T01:29:57Z

+# ── Test 3d: destroy() releases PGs ──────────────────────────────────────────
+
+
+class TestDestroy:


we can trim no so important tests from here

yashaswikarnati · 2026-04-22T01:34:03Z

+    return captures, handle
+
+
+def _register_llm_input_capture(mimo_model):


should we also check the inputs to the llm, combined embeddings as encoder tp is same the outputs after bridge communication must statisfy the invariant of no drift for different llm tp/dp

Applies mechanical fixes and design-level refactors from the PR review: - Revert DDP file + DDP test to origin/main (all unwanted diff removed). - Revert role.py copyright year; shrink verbose RankRole.build docstring. - Collapse RankRole API: _colocated/_from_grid_map are private; build() is the single public entry point. Drop the duplicate key validation in _from_grid_map since MimoModelConfig.__post_init__ already validates. - Remove the "Single dispatch point" comment in MimoModel.__init__. - Restore PP-stage assertion in test_mimo_model.py; drop pointer comment. - Bump new test file copyright to 2026. - Merge _validate_grids for-loops into a single pass; collapse three pp/cp parametrizations into one test. - Extract build_no_sync_func helper in test_mimo_1f1b_schedule.py and reuse it from the colocated-correctness test to cut duplication. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>

PR #10 replaced the ad-hoc gradient_reduce_div_factor DDP knob with calculate_per_token_loss=True on both sub-model configs plus a custom finalize hook that divides grads by the global valid-token count. The three-phase PP schedule now has to forward the schedule's total_num_tokens to the deferred finalize call, otherwise the hook's assertion fails and per-token normalization never happens on the encoder/LLM grads. * _loss_func now returns the 3-tuple (local_sum, local_num_tokens, log_dict) contract the schedule expects when per-token loss is on. * _deferred_finalize swaps the finalize hook with a capturing stub that records the num_tokens the inner schedule would have passed; after Phase 3, we invoke the original finalize with the captured value. test_mimo_colocated_pp: adopt per-token-loss wiring, add PP broadcast _wire_training_hooks from the PR #10 correctness test only all-reduces num_tokens over the LLM DP group. With LLM PP>1, non-last PP stages see num_tokens=0 from the inner schedule (loss runs only on the last stage), so the DP sum would land at N_last_stage instead of N_global and encoder/LLM grads would end up scaled differently per PP stage. _wire_pp_training_hooks broadcasts num_tokens from the last LLM PP rank first, then all-reduces across DP — every rank arrives at the same N_global. The PP test also drops the removed gradient_reduce_div_factor kwarg, switches both models to fp32 / no-bias / no-dropout for exact comparison, and uses the 3-tuple loss shape on the ref forward path.

…comment Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>

Apply black formatting to PR-changed files and add the missing copyright header to megatron/core/models/mimo/comm/__init__.py. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>

…W-17) Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>

The non-colocated test grids only carry 'pp' (or no) dim_names, so the auto-detected colocated path tripped _validate_grids requiring 'tp'. Early-return when any grid is missing TP/DP topology so non-colocated configs work, and assert colocated_comms stays empty in the test. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

`TestMimoModelNonColocated` exists to exercise the non-colocated path, but `_make_config(True, True)` made both grids share rank_offset=0 and size=world_size, which is the only signal `_all_grids_colocated` looks at — so role.build dispatched to `_colocated` and forced is_first_stage= True for every module. The 4th sub-case of test_role_determination only asserts about `images`, so flip language_in_grid to False; that gives the language grid a distinct rank_offset, role.build takes `_from_grid_map`, and the PP-stage info from the encoder grid survives. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

yashaswikarnati and others added 5 commits April 17, 2026 12:08