From 19364dbe76d19703e872d84bd8cd89708b5fd008 Mon Sep 17 00:00:00 2001 From: ykarnati Date: Wed, 10 Jun 2026 22:05:35 -0700 Subject: [PATCH 1/6] mimo: add dual gradient finalization with vision participation correction Add examples/mimo/training/grad_sync.py with configure_grad_sync, which installs a finalize_model_grads_func that finalizes the language model and each modality submodule over its own per-module ProcessGroupCollection via megatron.core.distributed.finalize_model_grads. The common path is the per-token mean: with calculate_per_token_loss=True the DDP gradient scaling factor is 1.0 (pure SUM), and grads are externally scaled by 1/N_global over the LLM DP group. Add mark_modality_participation / reset_modality_participation / _vision_participation_count for the vision partial-participation correction, guarded (along with the cross-grid end-of-iter barrier) on the non-colocated role layout. The no_sync / zero-grad-buffer / active-DDP-iteration helpers are deferred to the MM4 step PR since they depend on runtime DDP-wrap helpers not present here. Add a real 8-GPU test that builds a colocated MimoModel via the existing get_mimo_model helper, wires configure_grad_sync, runs one forward/backward plus finalize, and asserts the encoder grads are finalized over the vision DP group as the global per-token mean (DP-replica-invariant), plus a focused non-colocated check of the vision participation correction over a grid-derived process group. Co-Authored-By: Claude Opus 4.8 (1M context) Signed-off-by: ykarnati --- examples/mimo/training/grad_sync.py | 170 ++++++++++++++++++ .../models/mimo/test_mimo_1f1b_schedule.py | 54 +++--- .../mimo/test_mimo_colocated_correctness.py | 97 ++-------- .../models/mimo/test_mimo_grad_sync.py | 74 ++++++++ 4 files changed, 294 insertions(+), 101 deletions(-) create mode 100644 examples/mimo/training/grad_sync.py create mode 100644 tests/unit_tests/models/mimo/test_mimo_grad_sync.py diff --git a/examples/mimo/training/grad_sync.py b/examples/mimo/training/grad_sync.py new file mode 100644 index 00000000000..406bd602b43 --- /dev/null +++ b/examples/mimo/training/grad_sync.py @@ -0,0 +1,170 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +"""Dual gradient finalization for MIMO training on the stock Megatron loop.""" + +from __future__ import annotations + +import torch +import torch.distributed as dist + +from examples.mimo.training.topology import HeteroTopology +from megatron.core.distributed.finalize_model_grads import finalize_model_grads +from megatron.core.models.mimo.config.role import MIMO_LANGUAGE_MODULE_KEY +from megatron.core.models.mimo.model.base import MimoModel +from megatron.core.pipeline_parallel.utils import is_pp_last_stage + +# Sentinel set per modality submodule when this rank had that modality's input this step. +_PARTICIPATED_ATTR = "_mimo_rank_processed_input" + + +def _has_modality_input(value) -> bool: + """Whether this rank received this modality's input this step. + + The batch omits a modality's key when absent, so ``value`` is None (not present) or a + non-empty nested dict (present); an empty tensor also counts as absent. + """ + if isinstance(value, torch.Tensor): + return value.numel() > 0 + return bool(value) + + +def mark_modality_participation(mimo_model: MimoModel, batch) -> None: + """Tag each modality submodule with whether this rank had that modality's input this step. + + Reads ``batch["modality_inputs"]`` (keyed by modality name) so the flag is per modality + rather than vision-specific. + """ + modality_inputs = batch.get("modality_inputs", {}) if isinstance(batch, dict) else {} + for name, submodule in mimo_model.modality_submodules.items(): + if submodule is not None: + setattr(submodule, _PARTICIPATED_ATTR, _has_modality_input(modality_inputs.get(name))) + + +def reset_modality_participation(mimo_model: MimoModel) -> None: + """Clear per-step participation flags at the top of each train step.""" + for submodule in mimo_model.modality_submodules.values(): + if submodule is not None: + setattr(submodule, _PARTICIPATED_ATTR, False) + + +def _vision_participation_count(submodule, vision_dp_group) -> float: + """Number of vision-DP ranks that processed image input this step.""" + val = 1.0 if getattr(submodule, _PARTICIPATED_ATTR, False) else 0.0 + indicator = torch.tensor([val], dtype=torch.float32, device="cuda") + dist.all_reduce(indicator, op=dist.ReduceOp.SUM, group=vision_dp_group) + return float(indicator.item()) + + +def _is_pg_member(pg) -> bool: + """Whether the current rank belongs to ``pg`` (defensive; -1 for non-members).""" + return pg is not None and dist.get_rank(group=pg) >= 0 + + +def _is_token_source_rank(language_pg) -> bool: + """Whether this rank is the single LLM coordinate that owns the global token count. + + Sourcing the count from one coordinate (last PP stage, TP rank 0) avoids + double-counting it across TP/PP replicas. In non-colocated, language_pg is the LLM + collection seen on every rank, so encoder-grid ranks reach here with non-member + (None) pp/tp groups — the _is_pg_member guards short-circuit so they are never the + source. + """ + if language_pg is None: + return False + pp = getattr(language_pg, "pp", None) + tp = getattr(language_pg, "tp", None) + return ( + _is_pg_member(pp) + and _is_pg_member(tp) + and is_pp_last_stage(pp) + and dist.get_rank(group=tp) == 0 + ) + + +def _global_token_count(num_tokens, language_pg) -> float: + """Total non-padded tokens in the global batch, visible on every rank. + + Only the LLM token-source ranks contribute: they sum over the LLM DP/CP group; + a world MAX then publishes that N_global to every rank, including the + non-colocated encoder grid (where ``language_pg`` is None and the count is 0). + """ + global_num_tokens = torch.zeros(1, dtype=torch.float32, device="cuda") + if _is_token_source_rank(language_pg): + token_count = num_tokens.to(dtype=torch.float32).sum().view(1) + dist.all_reduce(token_count, group=language_pg.dp_cp, op=dist.ReduceOp.SUM) + if dist.get_rank(group=language_pg.dp_cp) == 0: + global_num_tokens.copy_(token_count) + dist.all_reduce(global_num_tokens, op=dist.ReduceOp.MAX) + return float(global_num_tokens.item()) + + +def configure_grad_sync(args, mimo_model: MimoModel, topology: HeteroTopology) -> None: + """Configure per-module gradient finalization: each module finalizes over its own groups. + + The encoder and LLM have decoupled parallelism (separate grids), so each reduces its + gradients over its own process-group collection; both then divide by one shared + per-token mean (N_global). + + MimoModel structure (each a separately DDP-wrapped module on its own grid):: + + MimoModel + ├─ language_model (LLM) -> own process groups + └─ modality_submodules[*] (encoders) -> own process groups + """ + module_pgs = topology.module_pgs + language_pg = module_pgs.get(MIMO_LANGUAGE_MODULE_KEY) + correct_vision_grad = bool( + getattr(args, "correct_encoder_grad_for_partial_participation", False) + ) + + def finalize_grads_func(_model_list, num_tokens, force_all_reduce=False, **_kwargs): + # calculate_per_token_loss=True => DDP gradient_scaling_factor 1.0 (pure SUM), + # so the per-token mean is applied here by dividing every shard by N_global. + assert num_tokens is not None, ( + "MIMO grad sync expects calculate_per_token_loss=True so the schedule " + "forwards total_num_tokens; got None." + ) + + # N_global is the global token count, published to every rank (including the + # non-colocated encoder grid) so both modules divide by the same per-token mean. + n_global = _global_token_count(num_tokens, language_pg) + inv = 1.0 / n_global if n_global > 0 else 0.0 + + if mimo_model.language_model is not None: + finalize_model_grads( + [mimo_model.language_model], + num_tokens=None, + pg_collection=language_pg, + force_all_reduce=force_all_reduce, + ) + if inv != 0.0: + mimo_model.language_model.scale_gradients(inv) + + for name, submodule in mimo_model.modality_submodules.items(): + if submodule is None: + continue + vision_pg = module_pgs.get(name) + finalize_model_grads( + [submodule], + num_tokens=None, + pg_collection=vision_pg, + force_all_reduce=force_all_reduce, + ) + + vision_scale = inv + if correct_vision_grad and vision_pg is not None and vision_pg.dp is not None: + vision_dp_group = vision_pg.dp + if _is_pg_member(vision_dp_group): + vision_dp_size = dist.get_world_size(vision_dp_group) + if vision_dp_size > 1: + participation = _vision_participation_count(submodule, vision_dp_group) + if 0.0 < participation < vision_dp_size: + vision_scale *= vision_dp_size / participation + + if vision_scale != 0.0: + submodule.scale_gradients(vision_scale) + + mimo_model.config.finalize_model_grads_func = finalize_grads_func + # The schedule always calls grad_scale_func with a Tensor loss; the per-token + # mean is applied in finalize_grads_func, so no extra scaling is needed here. + mimo_model.config.grad_scale_func = lambda loss: loss diff --git a/tests/unit_tests/models/mimo/test_mimo_1f1b_schedule.py b/tests/unit_tests/models/mimo/test_mimo_1f1b_schedule.py index 64824898927..a91d1741868 100644 --- a/tests/unit_tests/models/mimo/test_mimo_1f1b_schedule.py +++ b/tests/unit_tests/models/mimo/test_mimo_1f1b_schedule.py @@ -9,6 +9,7 @@ import logging from contextlib import ExitStack, contextmanager from functools import partial +from types import SimpleNamespace import pytest import torch @@ -16,8 +17,8 @@ from packaging import version import megatron.core.pipeline_parallel.schedules as schedule +from examples.mimo.training.grad_sync import configure_grad_sync from megatron.core.distributed import DistributedDataParallel, DistributedDataParallelConfig -from megatron.core.distributed.finalize_model_grads import finalize_model_grads from megatron.core.hyper_comm_grid import HyperCommGrid from megatron.core.models.gpt.gpt_layer_specs import get_gpt_layer_with_transformer_engine_spec from megatron.core.models.gpt.gpt_model import GPTModel @@ -568,7 +569,12 @@ def run_mimo_1f1b_test( micro_batch_size=2, num_microbatches=4, ): - """Run MIMO model through 1F1B schedule and verify.""" + """Run MIMO model through 1F1B schedule and verify. + + Uses the production examples/mimo configure_grad_sync (calculate_per_token_loss=True) + as the grad-finalization hook, exercising its cross-grid token sourcing + N_global + broadcast on this non-colocated topology. + """ # Clear NVTE env vars that the conftest set_env fixture sets to '0'. # GPTModel (LanguageModule) asserts these are unset or match the attention backend. import os @@ -599,26 +605,20 @@ def run_mimo_1f1b_test( num_layers=num_layers, vocab_size=vocab_size, seq_len=seq_length, + per_token_loss=True, ) - no_sync_func = build_no_sync_func(mimo_model) + mimo_model.config.no_sync_func = build_no_sync_func(mimo_model) - def finalize_grads_func(*args, **kwargs): - if mimo_model.language_model is not None: - finalize_model_grads( - [mimo_model.language_model], num_tokens=None, pg_collection=language_pg - ) - for submodule in mimo_model.modality_submodules.values(): - if submodule is not None: - finalize_model_grads([submodule], num_tokens=None, pg_collection=vision_pg) - - mimo_model.config.no_sync_func = no_sync_func - mimo_model.config.finalize_model_grads_func = finalize_grads_func - mimo_model.config.grad_scale_func = lambda loss: ( - torch.tensor(loss, dtype=torch.float32, device='cuda', requires_grad=True) - if isinstance(loss, (int, float)) - else loss + # Use the production grad-sync hook (finalize per module over its own groups + + # cross-grid N_global per-token mean) for every config. + grad_sync_topology = SimpleNamespace( + module_pgs={ + MIMO_LANGUAGE_MODULE_KEY: language_pg, + **{name: vision_pg for name in mimo_model.modality_submodules}, + } ) + configure_grad_sync(SimpleNamespace(), mimo_model, grad_sync_topology) # Create optimizer opt_config = OptimizerConfig( @@ -680,8 +680,17 @@ def finalize_grads_func(*args, **kwargs): def step_func(data_iterator, model): def loss_func(loss_mask, output_tensor): + # calculate_per_token_loss=True: the schedule expects a + # (loss_sum, num_tokens, loss_dict) triple, with num_tokens an int tensor. + def _ret(loss, num_tokens, reduced): + return loss, num_tokens, {'loss_reduced': reduced} + + zero = torch.tensor(0.0, device='cuda', requires_grad=True) + # num_tokens must be an int tensor: the schedule accumulates it into an + # int total_num_tokens when calculate_per_token_loss=True. + one = torch.tensor(1, device='cuda', dtype=torch.int) if output_tensor is None: - return torch.tensor(0.0, device='cuda', requires_grad=True), {'loss_reduced': 0.0} + return _ret(zero, one, 0.0) if isinstance(output_tensor, dict): output = output_tensor.get( @@ -691,10 +700,13 @@ def loss_func(loss_mask, output_tensor): output = output_tensor if output is None: - return torch.tensor(0.0, device='cuda', requires_grad=True), {'loss_reduced': 0.0} + return _ret(zero, one, 0.0) loss = output.float().sum() - return loss, {'loss_reduced': loss} + num_tokens = ( + loss_mask.sum().to(torch.int).clamp(min=1) if loss_mask is not None else one + ) + return _ret(loss, num_tokens, loss) batch = next(data_iterator) if data_iterator is not None else {'input_ids': None} output_tensor, loss_mask = model(**batch) diff --git a/tests/unit_tests/models/mimo/test_mimo_colocated_correctness.py b/tests/unit_tests/models/mimo/test_mimo_colocated_correctness.py index 71ff13ec557..879e93e7a5b 100644 --- a/tests/unit_tests/models/mimo/test_mimo_colocated_correctness.py +++ b/tests/unit_tests/models/mimo/test_mimo_colocated_correctness.py @@ -52,6 +52,7 @@ import os from functools import partial +from types import SimpleNamespace import pytest import torch @@ -59,8 +60,9 @@ from packaging import version import megatron.core.pipeline_parallel.schedules as schedule +from examples.mimo.training.grad_sync import configure_grad_sync from megatron.core.distributed import DistributedDataParallelConfig -from megatron.core.distributed.finalize_model_grads import finalize_model_grads +from megatron.core.models.mimo.config.role import MIMO_LANGUAGE_MODULE_KEY from megatron.core.models.mimo.optimizer import get_mimo_optimizer from megatron.core.optimizer.optimizer_config import OptimizerConfig from megatron.core.transformer.enums import ModelType @@ -165,87 +167,22 @@ def _set_deterministic_env(): def _wire_training_hooks(mimo_model, language_pg, vision_pg): - """Attach no_sync / finalize_grads / grad_scale hooks to a MimoModel. - - The finalize hook implements the heterogeneous-DP grad-scaling story - without touching ``DistributedDataParallel``. Both sub-model configs - set ``calculate_per_token_loss=True``, so both DDPs pure-SUM across - their own DP group (``gradient_scaling_factor=1.0``). After backward - and DDP reduce, every rank's ``main_grad`` holds the un-normalized - full-batch sum of per-token gradients. - - This hook then: - 1. all-reduces the schedule's ``total_num_tokens`` across the LLM - DP group to obtain ``N_global`` (total valid tokens in the global - batch). Since ranks are colocated, every rank now knows - ``N_global``. - 2. Calls ``finalize_model_grads(num_tokens=None)`` per side — runs - the usual DDP grad finish + layernorm/embedding AR work without - letting the built-in divisor path fire. - 3. Calls ``scale_gradients(1/N_global)`` on each side — lands the - true global per-token mean uniformly on encoder and LLM grads. - - Note: encoder has no loss_func (so nothing emits a per-encoder-DP - ``num_tokens`` to feed ``finalize_model_grads``' internal all-reduce). - Doing the all-reduce once ourselves and calling ``scale_gradients`` - directly avoids engineering a fictitious per-encoder-rank count whose - sum happens to equal ``N_global``. - """ - - no_sync_func = build_no_sync_func(mimo_model) - - def finalize_grads_func(model_list, num_tokens, force_all_reduce=False, **kwargs): - # Schedule passes the per-rank sum-across-microbatches of what the - # loss_func returned. Because loss_func runs only on the LLM side, - # this is the LLM-local token count. - assert num_tokens is not None, ( - "finalize_grads_func expects calculate_per_token_loss=True on the " - "TransformerConfig so the schedule forwards total_num_tokens; got None." - ) + """Attach no_sync plus the production grad-sync hooks to a MimoModel. - # Phase 1: lift the all-reduce. After this, every rank (including - # encoder-only replicas) has N_global = total non-padded tokens in - # the global batch. - llm_dp_pg = language_pg.dp_cp if language_pg.dp_cp is not None else language_pg.dp - dist.all_reduce(num_tokens, group=llm_dp_pg, op=dist.ReduceOp.SUM) - n_global = num_tokens.item() - - # Phase 2: per-side DDP finish without built-in num_tokens scaling. - # Forward ``force_all_reduce`` so PP grad-sync semantics (if ever - # exercised here) aren't silently dropped. - if mimo_model.language_model is not None: - finalize_model_grads( - [mimo_model.language_model], - num_tokens=None, - pg_collection=language_pg, - force_all_reduce=force_all_reduce, - ) - for submodule in mimo_model.modality_submodules.values(): - if submodule is not None: - finalize_model_grads( - [submodule], - num_tokens=None, - pg_collection=vision_pg, - force_all_reduce=force_all_reduce, - ) - - # Phase 3: uniform divide by N_global. Guard div-by-zero for the - # degenerate fully-masked batch. - if n_global > 0: - inv = 1.0 / n_global - if mimo_model.language_model is not None: - mimo_model.language_model.scale_gradients(inv) - for submodule in mimo_model.modality_submodules.values(): - if submodule is not None: - submodule.scale_gradients(inv) - - mimo_model.config.no_sync_func = no_sync_func - mimo_model.config.finalize_model_grads_func = finalize_grads_func - mimo_model.config.grad_scale_func = lambda loss: ( - torch.tensor(loss, dtype=torch.float32, device='cuda', requires_grad=True) - if isinstance(loss, (int, float)) - else loss + Delegates the finalize/grad-scale wiring to ``configure_grad_sync`` (the real + examples/mimo path), so this test's dp1-reference assertions validate that + production hook directly. ``configure_grad_sync`` implements the same per-token + mean: all-reduce ``total_num_tokens`` over the LLM DP group to get ``N_global``, + finalize each submodule over its own group, then ``scale_gradients(1/N_global)``. + """ + mimo_model.config.no_sync_func = build_no_sync_func(mimo_model) + topology = SimpleNamespace( + module_pgs={ + MIMO_LANGUAGE_MODULE_KEY: language_pg, + **{name: vision_pg for name in mimo_model.modality_submodules}, + } ) + configure_grad_sync(SimpleNamespace(), mimo_model, topology) def _generate_and_broadcast_global_batches( diff --git a/tests/unit_tests/models/mimo/test_mimo_grad_sync.py b/tests/unit_tests/models/mimo/test_mimo_grad_sync.py new file mode 100644 index 00000000000..33eaa88e907 --- /dev/null +++ b/tests/unit_tests/models/mimo/test_mimo_grad_sync.py @@ -0,0 +1,74 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +"""Real-distributed test for the grad_sync vision partial-participation correction. + +The dual-finalize per-token-mean path is validated end-to-end by +test_mimo_colocated_correctness (which wires configure_grad_sync into its +dp1-reference oracle). This file covers the participation-count helper directly +on grid-derived process groups (no parallel_state). +""" + +from types import SimpleNamespace + +import pytest +import torch +import torch.distributed as dist + +from examples.mimo.training.grad_sync import ( + _vision_participation_count, + mark_modality_participation, + reset_modality_participation, +) +from tests.unit_tests.models.mimo.test_mimo_1f1b_schedule import ( + create_hypercomm_grid, + destroy_all_grids, +) +from tests.unit_tests.test_utilities import Utils + + +class TestVisionParticipation: + @classmethod + def setup_class(cls): + Utils.initialize_distributed() + cls.world_size = dist.get_world_size() + + @classmethod + def teardown_class(cls): + Utils.destroy_model_parallel() + + def teardown_method(self): + destroy_all_grids() + + def test_vision_participation_correction(self): + """Partial participation: text-only ranks upscale present ranks. + + With only some DP ranks holding image input, the participation count is + < dp_size and the correction factor dp_size/participation is applied. + """ + if self.world_size != 8: + pytest.skip(f"Requires 8 GPUs, got {self.world_size}") + + grid = create_hypercomm_grid(offset=0, tp=1, cp=1, pp=1, dp=self.world_size) + vision_dp = grid.get_pg("dp") + dp_size = dist.get_world_size(vision_dp) + + submodule = SimpleNamespace() + fake_model = SimpleNamespace(modality_submodules={"images": submodule}) + + rank = dist.get_rank(vision_dp) + has_image = rank < dp_size // 2 + batch = ( + {"modality_inputs": {"images": {"hidden_states": torch.ones(1, device="cuda")}}} + if has_image + else {"modality_inputs": {}} + ) + reset_modality_participation(fake_model) + mark_modality_participation(fake_model, batch) + + count = _vision_participation_count(submodule, vision_dp) + assert count == float(dp_size // 2) + factor = dp_size / count + assert factor == pytest.approx(2.0) + + reset_modality_participation(fake_model) + assert getattr(submodule, "_mimo_rank_processed_input") is False From f0cd889a4babff99b6717bc5702219e948a8e9ba Mon Sep 17 00:00:00 2001 From: ykarnati Date: Tue, 16 Jun 2026 14:44:24 -0700 Subject: [PATCH 2/6] mimo: publish N_global via explicit broadcast from derived source rank Replace the world MAX-allreduce that published the global token count with an explicit dist.broadcast from the token-source rank. The source's global rank is derived statically from the language grid metadata (shape/dim_names/rank_offset), which is identical on every rank, so the non-colocated encoder grid -- which is not a member of any LLM process group -- can still name the broadcast root. The (pp_last, tp0) ranks still all-reduce the count over DP/CP collectively; DP/CP rank 0 holds the result and is exactly the derived broadcast root, so the broadcast is well defined and the collective cannot hang. Signed-off-by: ykarnati --- examples/mimo/training/grad_sync.py | 59 ++++++++++++++++++++++------- 1 file changed, 46 insertions(+), 13 deletions(-) diff --git a/examples/mimo/training/grad_sync.py b/examples/mimo/training/grad_sync.py index 406bd602b43..267ae15d35b 100644 --- a/examples/mimo/training/grad_sync.py +++ b/examples/mimo/training/grad_sync.py @@ -61,13 +61,15 @@ def _is_pg_member(pg) -> bool: def _is_token_source_rank(language_pg) -> bool: - """Whether this rank is the single LLM coordinate that owns the global token count. - - Sourcing the count from one coordinate (last PP stage, TP rank 0) avoids - double-counting it across TP/PP replicas. In non-colocated, language_pg is the LLM - collection seen on every rank, so encoder-grid ranks reach here with non-member - (None) pp/tp groups — the _is_pg_member guards short-circuit so they are never the - source. + """Whether this rank is on the LLM (last PP stage, TP rank 0) coordinate that sums + the global token count over the DP/CP group. + + These ranks (one per DP/CP position) collectively all-reduce the count over DP/CP; + only DP/CP rank 0 then holds the value and serves as the broadcast root. Sourcing + from this single coordinate avoids double-counting across TP/PP replicas. In + non-colocated, language_pg is the LLM collection seen on every rank, so encoder-grid + ranks reach here with non-member (None) pp/tp groups — the _is_pg_member guards + short-circuit so they never participate. """ if language_pg is None: return False @@ -81,20 +83,46 @@ def _is_token_source_rank(language_pg) -> bool: ) -def _global_token_count(num_tokens, language_pg) -> float: +def _token_source_global_rank(language_grid) -> int: + """Global (default-group) rank of the single LLM token-source coordinate. + + The source is the language coordinate (tp=0, cp=0, dp=0, pp=last). It is derived + statically from grid metadata (shape/dim_names/rank_offset) that is identical on + every rank, so even the non-colocated encoder grid — which is not a member of any + LLM group — can name this global rank as the broadcast root. + + ``get_rank_enum("pp")`` yields the exact PP rank lists ``create_pg`` uses. The grid's + ``rank_offset`` is always the all-zero coordinate (tp=0, cp=0, dp=0, pp=0), so the PP + line that starts at ``rank_offset`` is the (tp=0, cp=0, dp=0) line and its last entry + is the (pp=last) source rank. + """ + pp_lines = language_grid.get_rank_enum("pp") + for line in pp_lines: + if line[0] == language_grid.rank_offset: + return int(line[-1]) + raise RuntimeError( + f"Could not derive token-source global rank from language grid " + f"(rank_offset={language_grid.rank_offset}, pp_lines={pp_lines})" + ) + + +def _global_token_count(num_tokens, language_pg, src_global_rank) -> float: """Total non-padded tokens in the global batch, visible on every rank. - Only the LLM token-source ranks contribute: they sum over the LLM DP/CP group; - a world MAX then publishes that N_global to every rank, including the - non-colocated encoder grid (where ``language_pg`` is None and the count is 0). + Only the LLM token-source rank computes the count by summing over the LLM DP/CP + group; it then broadcasts that N_global from its global rank to every rank in the + world (including the non-colocated encoder grid, where ``language_pg`` is None) so + both modules divide by the same per-token mean. """ global_num_tokens = torch.zeros(1, dtype=torch.float32, device="cuda") if _is_token_source_rank(language_pg): + # Collective over DP/CP: every (pp_last, tp0) rank participates so the all-reduce + # does not hang; only DP/CP rank 0 keeps the result and is the broadcast root. token_count = num_tokens.to(dtype=torch.float32).sum().view(1) dist.all_reduce(token_count, group=language_pg.dp_cp, op=dist.ReduceOp.SUM) if dist.get_rank(group=language_pg.dp_cp) == 0: global_num_tokens.copy_(token_count) - dist.all_reduce(global_num_tokens, op=dist.ReduceOp.MAX) + dist.broadcast(global_num_tokens, src=src_global_rank) return float(global_num_tokens.item()) @@ -113,6 +141,11 @@ def configure_grad_sync(args, mimo_model: MimoModel, topology: HeteroTopology) - """ module_pgs = topology.module_pgs language_pg = module_pgs.get(MIMO_LANGUAGE_MODULE_KEY) + # Static (collective-free) derivation of the token source's global rank, used as the + # broadcast root that publishes N_global to every rank — including the non-colocated + # encoder grid, which belongs to no LLM group and therefore cannot name the root via + # any LLM process group. + src_global_rank = _token_source_global_rank(topology.grids[MIMO_LANGUAGE_MODULE_KEY]) correct_vision_grad = bool( getattr(args, "correct_encoder_grad_for_partial_participation", False) ) @@ -127,7 +160,7 @@ def finalize_grads_func(_model_list, num_tokens, force_all_reduce=False, **_kwar # N_global is the global token count, published to every rank (including the # non-colocated encoder grid) so both modules divide by the same per-token mean. - n_global = _global_token_count(num_tokens, language_pg) + n_global = _global_token_count(num_tokens, language_pg, src_global_rank) inv = 1.0 / n_global if n_global > 0 else 0.0 if mimo_model.language_model is not None: From a2e26fc976040b067af8de33a0b79f47545e679f Mon Sep 17 00:00:00 2001 From: ykarnati Date: Tue, 16 Jun 2026 15:07:38 -0700 Subject: [PATCH 3/6] mimo: derive token-source rank from enumeration, not rank_offset Read the broadcast root directly from the grid's PP rank enumeration instead of assuming rank_offset is the all-zero coordinate. get_rank_enum("pp") is the grid's authoritative coordinate->global-rank mapping and already accounts for rank_offset; the source line is the PP line containing the grid's global minimum rank (the (tp=0, cp=0, dp=0) coordinate), and its last entry is the (pp=last) source rank. Signed-off-by: ykarnati --- examples/mimo/training/grad_sync.py | 24 +++++++++++++----------- 1 file changed, 13 insertions(+), 11 deletions(-) diff --git a/examples/mimo/training/grad_sync.py b/examples/mimo/training/grad_sync.py index 267ae15d35b..5eceac254ac 100644 --- a/examples/mimo/training/grad_sync.py +++ b/examples/mimo/training/grad_sync.py @@ -87,22 +87,24 @@ def _token_source_global_rank(language_grid) -> int: """Global (default-group) rank of the single LLM token-source coordinate. The source is the language coordinate (tp=0, cp=0, dp=0, pp=last). It is derived - statically from grid metadata (shape/dim_names/rank_offset) that is identical on - every rank, so even the non-colocated encoder grid — which is not a member of any - LLM group — can name this global rank as the broadcast root. - - ``get_rank_enum("pp")`` yields the exact PP rank lists ``create_pg`` uses. The grid's - ``rank_offset`` is always the all-zero coordinate (tp=0, cp=0, dp=0, pp=0), so the PP - line that starts at ``rank_offset`` is the (tp=0, cp=0, dp=0) line and its last entry - is the (pp=last) source rank. + statically from the grid's rank enumeration (identical on every rank), so even the + non-colocated encoder grid — which is not a member of any LLM group — can name this + global rank as the broadcast root. + + ``get_rank_enum("pp")`` returns the authoritative PP rank lists that ``create_pg`` + uses for this grid; each list is one (tp, cp, dp) coordinate ordered along PP and + already accounts for ``rank_offset``. Rather than assume which list is the all-zero + coordinate, the source line is read directly from the enumeration: the global minimum + rank belongs to (tp=0, cp=0, dp=0), so its PP line is the source line and that line's + last entry is the (pp=last) source rank. """ pp_lines = language_grid.get_rank_enum("pp") + min_rank = min(rank for line in pp_lines for rank in line) for line in pp_lines: - if line[0] == language_grid.rank_offset: + if min_rank in line: return int(line[-1]) raise RuntimeError( - f"Could not derive token-source global rank from language grid " - f"(rank_offset={language_grid.rank_offset}, pp_lines={pp_lines})" + f"Could not derive token-source global rank from language grid pp_lines={pp_lines}" ) From 75872f2e8718c8967501340a158d9150b400bc65 Mon Sep 17 00:00:00 2001 From: ykarnati Date: Tue, 16 Jun 2026 15:10:48 -0700 Subject: [PATCH 4/6] mimo: condense grad-sync comments Tighten the broadcast-root comment and trim the over-explained _token_source_global_rank and _is_token_source_rank docstrings while keeping the load-bearing rationale. Signed-off-by: ykarnati --- examples/mimo/training/grad_sync.py | 38 ++++++++++------------------- 1 file changed, 13 insertions(+), 25 deletions(-) diff --git a/examples/mimo/training/grad_sync.py b/examples/mimo/training/grad_sync.py index 5eceac254ac..9ac6a495aa5 100644 --- a/examples/mimo/training/grad_sync.py +++ b/examples/mimo/training/grad_sync.py @@ -62,14 +62,11 @@ def _is_pg_member(pg) -> bool: def _is_token_source_rank(language_pg) -> bool: """Whether this rank is on the LLM (last PP stage, TP rank 0) coordinate that sums - the global token count over the DP/CP group. - - These ranks (one per DP/CP position) collectively all-reduce the count over DP/CP; - only DP/CP rank 0 then holds the value and serves as the broadcast root. Sourcing - from this single coordinate avoids double-counting across TP/PP replicas. In - non-colocated, language_pg is the LLM collection seen on every rank, so encoder-grid - ranks reach here with non-member (None) pp/tp groups — the _is_pg_member guards - short-circuit so they never participate. + the global token count over DP/CP. + + Sourcing from this single coordinate avoids double-counting across TP/PP replicas. + The _is_pg_member guards short-circuit encoder-grid ranks (non-member pp/tp groups) + so they never participate. """ if language_pg is None: return False @@ -84,19 +81,12 @@ def _is_token_source_rank(language_pg) -> bool: def _token_source_global_rank(language_grid) -> int: - """Global (default-group) rank of the single LLM token-source coordinate. - - The source is the language coordinate (tp=0, cp=0, dp=0, pp=last). It is derived - statically from the grid's rank enumeration (identical on every rank), so even the - non-colocated encoder grid — which is not a member of any LLM group — can name this - global rank as the broadcast root. - - ``get_rank_enum("pp")`` returns the authoritative PP rank lists that ``create_pg`` - uses for this grid; each list is one (tp, cp, dp) coordinate ordered along PP and - already accounts for ``rank_offset``. Rather than assume which list is the all-zero - coordinate, the source line is read directly from the enumeration: the global minimum - rank belongs to (tp=0, cp=0, dp=0), so its PP line is the source line and that line's - last entry is the (pp=last) source rank. + """Global rank of the single LLM token-source coordinate (tp=0, cp=0, dp=0, pp=last). + + Derived statically from ``get_rank_enum("pp")`` (the grid's authoritative rank + enumeration, identical on every rank), so encoder-grid ranks in no LLM group can name + it. The global minimum rank is (tp=0, cp=0, dp=0), so its PP line is the source line + and that line's last entry is the (pp=last) source rank. """ pp_lines = language_grid.get_rank_enum("pp") min_rank = min(rank for line in pp_lines for rank in line) @@ -143,10 +133,8 @@ def configure_grad_sync(args, mimo_model: MimoModel, topology: HeteroTopology) - """ module_pgs = topology.module_pgs language_pg = module_pgs.get(MIMO_LANGUAGE_MODULE_KEY) - # Static (collective-free) derivation of the token source's global rank, used as the - # broadcast root that publishes N_global to every rank — including the non-colocated - # encoder grid, which belongs to no LLM group and therefore cannot name the root via - # any LLM process group. + # Broadcast root for N_global; derived statically so encoder-grid ranks (in no LLM + # group) can still name it. src_global_rank = _token_source_global_rank(topology.grids[MIMO_LANGUAGE_MODULE_KEY]) correct_vision_grad = bool( getattr(args, "correct_encoder_grad_for_partial_participation", False) From 7457d0824427571f6aa112a365c4f8e048ffbc67 Mon Sep 17 00:00:00 2001 From: ykarnati Date: Mon, 22 Jun 2026 12:50:27 -0700 Subject: [PATCH 5/6] Update MIMO grad-sync test topology Signed-off-by: ykarnati --- .../unit_tests/models/mimo/test_mimo_1f1b_schedule.py | 1 + .../models/mimo/test_mimo_colocated_correctness.py | 11 ++++++----- 2 files changed, 7 insertions(+), 5 deletions(-) diff --git a/tests/unit_tests/models/mimo/test_mimo_1f1b_schedule.py b/tests/unit_tests/models/mimo/test_mimo_1f1b_schedule.py index a91d1741868..c573e61479c 100644 --- a/tests/unit_tests/models/mimo/test_mimo_1f1b_schedule.py +++ b/tests/unit_tests/models/mimo/test_mimo_1f1b_schedule.py @@ -613,6 +613,7 @@ def run_mimo_1f1b_test( # Use the production grad-sync hook (finalize per module over its own groups + # cross-grid N_global per-token mean) for every config. grad_sync_topology = SimpleNamespace( + grids=module_to_grid_map, module_pgs={ MIMO_LANGUAGE_MODULE_KEY: language_pg, **{name: vision_pg for name in mimo_model.modality_submodules}, diff --git a/tests/unit_tests/models/mimo/test_mimo_colocated_correctness.py b/tests/unit_tests/models/mimo/test_mimo_colocated_correctness.py index 879e93e7a5b..ccaf791bea6 100644 --- a/tests/unit_tests/models/mimo/test_mimo_colocated_correctness.py +++ b/tests/unit_tests/models/mimo/test_mimo_colocated_correctness.py @@ -166,7 +166,7 @@ def _set_deterministic_env(): os.environ.pop('NVTE_UNFUSED_ATTN', None) -def _wire_training_hooks(mimo_model, language_pg, vision_pg): +def _wire_training_hooks(mimo_model, module_to_grid_map, language_pg, vision_pg): """Attach no_sync plus the production grad-sync hooks to a MimoModel. Delegates the finalize/grad-scale wiring to ``configure_grad_sync`` (the real @@ -177,6 +177,7 @@ def _wire_training_hooks(mimo_model, language_pg, vision_pg): """ mimo_model.config.no_sync_func = build_no_sync_func(mimo_model) topology = SimpleNamespace( + grids=module_to_grid_map, module_pgs={ MIMO_LANGUAGE_MODULE_KEY: language_pg, **{name: vision_pg for name in mimo_model.modality_submodules}, @@ -927,7 +928,7 @@ def test_dist_matches_dp1_reference_post_step_weights( # Build dist first (heterogeneous TP/DP). torch.manual_seed(12345) - dist_mimo, _, _, dist_language_pg, dist_vision_pg = get_mimo_model( + dist_mimo, dist_module_to_grid_map, _, dist_language_pg, dist_vision_pg = get_mimo_model( encoder_name=encoder_name, encoder_grid=dist_enc_grid, llm_grid=dist_llm_grid, @@ -946,7 +947,7 @@ def test_dist_matches_dp1_reference_post_step_weights( # Reference with equal-DP uniform (enc_tp == llm_tp, enc_dp == llm_dp). torch.manual_seed(12345) - ref_mimo, _, _, ref_language_pg, ref_vision_pg = get_mimo_model( + ref_mimo, ref_module_to_grid_map, _, ref_language_pg, ref_vision_pg = get_mimo_model( encoder_name=encoder_name, encoder_grid=ref_enc_grid, llm_grid=ref_llm_grid, @@ -981,8 +982,8 @@ def test_dist_matches_dp1_reference_post_step_weights( dist_llm_grid.get_pg("tp"), ) - _wire_training_hooks(dist_mimo, dist_language_pg, dist_vision_pg) - _wire_training_hooks(ref_mimo, ref_language_pg, ref_vision_pg) + _wire_training_hooks(dist_mimo, dist_module_to_grid_map, dist_language_pg, dist_vision_pg) + _wire_training_hooks(ref_mimo, ref_module_to_grid_map, ref_language_pg, ref_vision_pg) # Distributed optimizers snapshot current param.data into fp32 master # weights at __init__, so both must be built AFTER the ref-to-dist From cadd331400dac574f2e079959276646efe34c0e0 Mon Sep 17 00:00:00 2001 From: ykarnati Date: Mon, 22 Jun 2026 13:10:37 -0700 Subject: [PATCH 6/6] Format MIMO grad-sync topology tests Signed-off-by: ykarnati --- tests/unit_tests/models/mimo/test_mimo_1f1b_schedule.py | 2 +- tests/unit_tests/models/mimo/test_mimo_colocated_correctness.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/unit_tests/models/mimo/test_mimo_1f1b_schedule.py b/tests/unit_tests/models/mimo/test_mimo_1f1b_schedule.py index c573e61479c..0a08e6d93f2 100644 --- a/tests/unit_tests/models/mimo/test_mimo_1f1b_schedule.py +++ b/tests/unit_tests/models/mimo/test_mimo_1f1b_schedule.py @@ -617,7 +617,7 @@ def run_mimo_1f1b_test( module_pgs={ MIMO_LANGUAGE_MODULE_KEY: language_pg, **{name: vision_pg for name in mimo_model.modality_submodules}, - } + }, ) configure_grad_sync(SimpleNamespace(), mimo_model, grad_sync_topology) diff --git a/tests/unit_tests/models/mimo/test_mimo_colocated_correctness.py b/tests/unit_tests/models/mimo/test_mimo_colocated_correctness.py index ccaf791bea6..747b66a815a 100644 --- a/tests/unit_tests/models/mimo/test_mimo_colocated_correctness.py +++ b/tests/unit_tests/models/mimo/test_mimo_colocated_correctness.py @@ -181,7 +181,7 @@ def _wire_training_hooks(mimo_model, module_to_grid_map, language_pg, vision_pg) module_pgs={ MIMO_LANGUAGE_MODULE_KEY: language_pg, **{name: vision_pg for name in mimo_model.modality_submodules}, - } + }, ) configure_grad_sync(SimpleNamespace(), mimo_model, topology)