diff --git a/examples/mimo/training/grad_sync.py b/examples/mimo/training/grad_sync.py new file mode 100644 index 00000000000..9ac6a495aa5 --- /dev/null +++ b/examples/mimo/training/grad_sync.py @@ -0,0 +1,193 @@ +# 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 on the LLM (last PP stage, TP rank 0) coordinate that sums + 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 + 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 _token_source_global_rank(language_grid) -> int: + """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) + for line in pp_lines: + if min_rank in line: + return int(line[-1]) + raise RuntimeError( + f"Could not derive token-source global rank from language grid 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 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.broadcast(global_num_tokens, src=src_global_rank) + 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) + # 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) + ) + + 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, src_global_rank) + 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..0a08e6d93f2 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,21 @@ 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( + grids=module_to_grid_map, + 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 +681,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 +701,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..747b66a815a 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 @@ -164,88 +166,24 @@ def _set_deterministic_env(): os.environ.pop('NVTE_UNFUSED_ATTN', None) -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." - ) +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. - # 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( + grids=module_to_grid_map, + 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( @@ -990,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, @@ -1009,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, @@ -1044,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 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