From 7e822a30fd49512aa6beb31ce112e4e8163b272c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Piotr=20=C5=BBelasko?= Date: Tue, 25 Aug 2026 09:32:48 -0700 Subject: [PATCH] fix(mtp): checkpoint repeated dense blocks MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Piotr Żelasko --- nemo_automodel/components/moe/parallelizer.py | 49 +++++++++------- tests/unit_tests/moe/test_parallelizer.py | 56 +++++++++++++++++++ 2 files changed, 86 insertions(+), 19 deletions(-) diff --git a/nemo_automodel/components/moe/parallelizer.py b/nemo_automodel/components/moe/parallelizer.py index e95ed49221..c9af42c781 100644 --- a/nemo_automodel/components/moe/parallelizer.py +++ b/nemo_automodel/components/moe/parallelizer.py @@ -144,6 +144,23 @@ def _get_moe_module(block: nn.Module) -> MoE | None: return module +def _repeated_mtp_moe_block_ids(model: nn.Module) -> set[int]: + """Return weight-tied MTP blocks whose experts cannot be recomputed safely. + + A repeated MTP depth may contain multiple physical sublayers. Only MoE + sublayers own the EP-sharded expert parameter group whose second FSDP2 + checkpoint recompute is unsafe; attention, MLP, and Mamba sublayers remain + eligible for activation checkpointing. + """ + mtp_module = getattr(model, "mtp", None) + if mtp_module is None or not hasattr(mtp_module, "layers"): + return set() + mtp_repeated = bool(getattr(getattr(mtp_module, "mtp_config", None), "use_repeated_layer", False)) + if not mtp_repeated: + return set() + return {id(block) for block in mtp_module.layers.children() if _get_moe_module(block) is not None} + + def _preserve_gate_load_during_recompute( block: nn.Module, context_fn: Callable[[], tuple[AbstractContextManager, AbstractContextManager]] | None = None, @@ -484,6 +501,12 @@ def apply_ac( scopes = normalize_activation_checkpointing_scope(activation_checkpointing_scope) checkpoint_decoder = "all" in scopes or "language" in scopes + repeated_mtp_moe_block_ids = _repeated_mtp_moe_block_ids(model) if checkpoint_decoder else set() + if repeated_mtp_moe_block_ids: + logger.info( + "Skipping activation checkpointing on %d weight-tied MTP MoE block(s)", + len(repeated_mtp_moe_block_ids), + ) if checkpoint_decoder and not selective and not ignore_router: logger.warning( "Activation checkpointing is enabled with ignore_router_for_ac=False. The MoE " @@ -510,6 +533,8 @@ def apply_ac( selective_context_fn, ) for parent_layers, layer_id, block in iter_transformer_and_mtp_blocks(model): + if id(block) in repeated_mtp_moe_block_ids: + continue block = ptd_checkpoint_wrapper( block, preserve_rng_state=True, @@ -592,26 +617,12 @@ def selective_checkpointing_context_fn(): def _with_attention_backend_snapshot(context_fn=None): return functools.partial(transformer_engine_attention_backend_snapshot_context_fn, context_fn) - # Weight-tied (use_repeated_layer) MTP head blocks must NOT be activation - # checkpointed: the single physical block is recomputed once per MTP depth in - # backward, and FSDP2 cannot re-unshard the *shared* EP-sharded experts param - # group on the 2nd+ recompute (the 1st recompute's post_backward reshards it, and - # the 2nd recompute's pre_forward unshard does not re-gather it) -> the experts - # weight is read in the resharded Shard(1) state and grouped_gemm raises - # "Expected hidden_in == a.size(1)". The MTP head is tiny (1 physical block), so - # skipping its recompute costs negligible activation memory. Non-tied MTP heads - # (each physical block recomputed exactly once) are unaffected and keep AC. - mtp_module = getattr(model, "mtp", None) - mtp_block_ids: set[int] = set() - mtp_repeated = False - if mtp_module is not None and hasattr(mtp_module, "layers"): - mtp_block_ids = {id(b) for b in mtp_module.layers.children()} - mtp_repeated = bool(getattr(getattr(mtp_module, "mtp_config", None), "use_repeated_layer", False)) - if mtp_repeated and mtp_block_ids: - logger.info("Skipping activation checkpointing on %d weight-tied MTP head block(s)", len(mtp_block_ids)) - for parent_layers, layer_id, block in iter_transformer_and_mtp_blocks(model): - if mtp_repeated and id(block) in mtp_block_ids: + # A weight-tied MoE block is recomputed once per logical MTP depth. + # FSDP2 cannot re-unshard its shared EP-sharded experts group after the + # first recompute, so leave only those blocks uncheckpointed. Repeated + # dense/attention/Mamba blocks have no such expert group and keep AC. + if id(block) in repeated_mtp_moe_block_ids: continue if ignore_router: block = ptd_checkpoint_wrapper( diff --git a/tests/unit_tests/moe/test_parallelizer.py b/tests/unit_tests/moe/test_parallelizer.py index d4770b5793..755db183cf 100644 --- a/tests/unit_tests/moe/test_parallelizer.py +++ b/tests/unit_tests/moe/test_parallelizer.py @@ -341,6 +341,14 @@ class GroupedExpertsTE: experts_stub.GroupedExpertsTE = GroupedExpertsTE monkeypatch.setitem(sys.modules, "nemo_automodel.components.moe.experts", experts_stub) + mok_experts_stub = types.ModuleType("nemo_automodel.components.moe.mok_experts") + + class GroupedExpertsMoK: + pass + + mok_experts_stub.GroupedExpertsMoK = GroupedExpertsMoK + monkeypatch.setitem(sys.modules, "nemo_automodel.components.moe.mok_experts", mok_experts_stub) + def _import_parallelizer_with_stubs(monkeypatch): import importlib @@ -350,6 +358,7 @@ def _import_parallelizer_with_stubs(monkeypatch): "nemo_automodel.components.moe.parallelizer", "nemo_automodel.components.moe.layers", "nemo_automodel.components.moe.experts", + "nemo_automodel.components.moe.mok_experts", "nemo_automodel.components.distributed.pipelining", "nemo_automodel.components.distributed.pipelining.config", "nemo_automodel.components.distributed.pipelining.hf_utils", @@ -2644,6 +2653,53 @@ def fake_wrapper(block, preserve_rng_state, determinism_check=None, context_fn=N assert getattr(w, sentinel_flag, False) is True +@pytest.mark.parametrize("selective", [False, True]) +def test_apply_ac_repeated_mtp_checkpoints_dense_blocks_and_skips_moe(monkeypatch, selective): + """Weight-tied MTP AC skips only blocks with the unsafe shared experts group.""" + P = _import_parallelizer_with_stubs(monkeypatch) + monkeypatch.setattr(P, "MoE", DummyMoE) + + if selective: + sentinel_flag = "_nemo_selective_ac" + dense_stub = types.ModuleType("nemo_automodel.components.distributed.activation_checkpointing") + dense_stub.make_selective_checkpoint_context_fn = MagicMock(return_value=object()) + dense_stub.SELECTIVE_AC_WRAPPER_FLAG = sentinel_flag + dense_stub.transformer_engine_attention_backend_snapshot_context_fn = lambda context_fn=None: context_fn + monkeypatch.setitem(sys.modules, "nemo_automodel.components.distributed.activation_checkpointing", dense_stub) + + wrapped = [] + + class _Wrapper: + def __init__(self, block): + self.block = block + + def fake_wrapper(block, **_kwargs): + wrapper = _Wrapper(block) + wrapped.append(wrapper) + return wrapper + + monkeypatch.setattr(P, "ptd_checkpoint_wrapper", MagicMock(side_effect=fake_wrapper)) + + backbone = DummyBlock(mlp=object()) + mtp_attention = DummyBlock(mlp=object()) + mtp_moe = DummyBlock(mlp=DummyMoE()) + model = types.SimpleNamespace( + model=DummyModel([backbone]), + mtp=types.SimpleNamespace( + layers=LayerContainer([mtp_attention, mtp_moe]), + mtp_config=types.SimpleNamespace(use_repeated_layer=True), + ), + ) + + P.apply_ac(model, selective=selective, hidden_size=8, num_experts=4) + + assert [wrapper.block for wrapper in wrapped] == [backbone, mtp_attention] + assert set(model.model.layers.registered) == {"0"} + assert set(model.mtp.layers.registered) == {"0"} + if selective: + assert all(getattr(wrapper, sentinel_flag, False) for wrapper in wrapped) + + # ============================================================================ # Tests for block.moe attribute handling (Step3p5 style models) # ============================================================================