Add ColocatedBridgeCommunicator for heterogeneous TP/DP MIMO training (NMFW-17) - #4368
Merged
yashaswikarnati merged 29 commits intoApr 28, 2026
Conversation
… (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>
2 tasks
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>
- 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>
…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>
Per PR review comments 4, 5, 7, 10: the communicator carried two of
everything — ``all_gather_pg`` and ``fan_out_gather_pg``; the matching
``*_group_ranks`` lists; ``fan_in_scale`` and ``fan_out_scale`` — even
though fan-in and fan-out are mutually exclusive. Collapse into a
single ``direction: BridgeDirection {FAN_IN, FAN_OUT, EQUAL}`` with one
``gather_pg``, one ``gather_group_ranks``, and one integer ``scale``.
The autograd Function and ``get_slice_info`` pick behavior off
``direction``; ``is_fan_in()`` / ``is_fan_out()`` stay as thin
predicates. ``_all_gather_along_batch_dim`` now lives as a module-level
helper instead of being inlined twice (once in ``forward``, once in
``backward``).
Also addresses comment 6: drop the "Silent truncation on non-divisible
batches is a correctness bug that produced one-off mis-slicing in early
versions" line — the raise itself is self-explanatory.
Test file updated to the renamed attributes:
* ``gather_pg`` / ``gather_group_ranks`` / ``scale`` replace the four
old names everywhere they appeared.
* Divisibility error messages now match ``not divisible by fan_(in|out)``.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Per PR review comment 11: ``RankRole.colocated(modality_module_names)`` and ``RankRole.from_grid_map(module_to_grid_map, modality_module_names)`` had asymmetric signatures despite being the two branches of the same dispatch. Add a single ``RankRole.build(modality_module_names, module_to_grid_map=None)`` that picks between them. The legacy factories stay as the actual implementations (callers that know exactly which branch they want can still hit them directly), but the canonical entry point is ``build``. Colocated-detection logic (``rank_offset`` + ``size`` agreement across grids) moved onto ``RankRole`` as ``_all_grids_colocated`` so both the dispatcher and future callers can reuse it. ``MimoModel._is_colocated`` was the only user and has been removed; ``MimoModel.__init__`` now dispatches once via ``RankRole.build`` and branches on ``self.role.mode`` for bridge construction. Phase C review items 12 (``MimoModelConfig.__post_init__`` key-validation) and 13 (drop stray docstring on ``_build_colocated_communicators``) were completed in earlier commits; verified still in place. 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.
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>
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>
yashaswikarnati
commented
Apr 23, 2026
…comment Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
kvareddy
approved these changes
Apr 27, 2026
Contributor
Author
|
/ok to test 5073ad0 |
svcnvidia-nemo-ci
added
Approved
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>
Contributor
Author
|
/ok to test 189a9eb |
…W-17) Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
Contributor
Author
|
/ok to test 58a7687 |
yashaswikarnati
enabled auto-merge
April 27, 2026 17:54
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>
yashaswikarnati
force-pushed
the
ykarnati/nmfw-17-colocated-bridge-pp1
branch
from
April 27, 2026 22:43
44863d3 to
1d4a117
Compare
Contributor
Author
|
/ok to test 1d4a117 |
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Contributor
Author
|
/ok to test 641a861 |
Contributor
|
🔄 Merge queue validation started! You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/25028887222 |
yangbofun
pushed a commit
to xlm-research/Megatron-LM
that referenced
this pull request
May 22, 2026
… (NMFW-17) (NVIDIA#4368) Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
yhgalaxy
pushed a commit
to yhgalaxy/Megatron-LM
that referenced
this pull request
Jun 17, 2026
… (NMFW-17) (NVIDIA#4368) Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com> Signed-off-by: yhgalaxy <yhgalaxy@outlook.com>
jon-barker
pushed a commit
to jon-barker/Megatron-LM
that referenced
this pull request
Jul 10, 2026
… (NMFW-17) (NVIDIA#4368) Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com> Signed-off-by: Jon Barker <jbarker@aws-cmh-slurm-1-vscode-02.cm.cluster>
terminator123
pushed a commit
to 021ai/Megatron-LM
that referenced
this pull request
Aug 3, 2026
… (NMFW-17) (NVIDIA#4368) Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
4 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Adds
ColocatedBridgeCommunicator— the core primitive for heterogeneous TP/DP MIMO training where the encoder and language model share the same rank pool but use different tensor-parallel/data-parallel layouts. PP=1 scope only; PP>1 support is a stacked follow-up PR.ModuleLayout.UNIFIED→ModuleLayout.COLOCATEDand auto-detects colocated layouts from grid overlap (samerank_offset+size)._ColocatedCommunicateautograd function.dist.all_gather_into_tensorinto a pre-allocated buffer (handlesbatch_dim != 0viamovedim), avoiding the Python list +dist.all_gather+torch.catcopy.MimoModel._forward_all_modulesapplies the communicator on encoder embeddings before combining them with text embeddings.module_to_grid_mapkey validation is hoisted so mismatched keys always raise, independent of colocated/non-colocated path.Test plan
tests/unit_tests/models/test_mimo_colocated_communicator.py— rank mapping, slice info, fan-in/fan-out/equal-DP forward+backward correctness (11 tests).tests/unit_tests/models/test_mimo_colocated_correctness.py— TransformerBlock-based multi-iteration correctness, 9 checks × 3 iterations, across fan-in/fan-out/equal variants.tests/unit_tests/models/test_mimo_colocated_e2e.py— end-to-end colocated VLM with heterogeneous TP/DP at PP=1, full optimizer + DDP + finalize_model_grads path.tests/unit_tests/models/test_mimo_model.py— enum rename + updates to keep the PP role-determination test truly non-colocated (distinctrank_offsets).uv run python -m torch.distributed.run --nproc_per_node=8 -m pytest ....Run commands
```bash
uv run python -m torch.distributed.run --nproc_per_node=8 -m pytest tests/unit_tests/models/test_mimo_colocated_communicator.py -v
uv run python -m torch.distributed.run --nproc_per_node=8 -m pytest tests/unit_tests/models/test_mimo_colocated_correctness.py -v
uv run python -m torch.distributed.run --nproc_per_node=8 -m pytest tests/unit_tests/models/test_mimo_colocated_e2e.py -v
uv run python -m torch.distributed.run --nproc_per_node=8 -m pytest tests/unit_tests/models/test_mimo_model.py -v
```
Follow-ups
🤖 Generated with Claude Code