diff --git a/nemo_automodel/components/moe/parallelizer.py b/nemo_automodel/components/moe/parallelizer.py index de3954337c..2a1c97e14c 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 if bool(getattr(block, "_nemo_disable_activation_checkpointing", False)): logger.info("Skipping activation checkpointing for model-owned eager block %s", layer_id) continue @@ -595,26 +620,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 bool(getattr(block, "_nemo_disable_activation_checkpointing", False)): logger.info("Skipping activation checkpointing for model-owned eager block %s", layer_id) diff --git a/tests/unit_tests/moe/test_parallelizer.py b/tests/unit_tests/moe/test_parallelizer.py index 039bf2e36d..12bc1126d0 100644 --- a/tests/unit_tests/moe/test_parallelizer.py +++ b/tests/unit_tests/moe/test_parallelizer.py @@ -707,16 +707,24 @@ def fake_wrapper(block, preserve_rng_state, determinism_check=None, context_fn=N assert len(model.layers.registered) == 2 -def test_apply_ac_skips_model_owned_eager_block(monkeypatch): +@pytest.mark.parametrize("selective", [False, True]) +def test_apply_ac_skips_model_owned_eager_block(monkeypatch, selective): P = _import_parallelizer_with_stubs(monkeypatch) + if selective: + 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 = "_nemo_selective_ac" + 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) + eager_block = DummyBlock() eager_block._nemo_disable_activation_checkpointing = True - wrapped_block = object() + wrapped_block = types.SimpleNamespace() wrapper_mock = MagicMock(return_value=wrapped_block) monkeypatch.setattr(P, "ptd_checkpoint_wrapper", wrapper_mock) model = DummyModel([DummyBlock(), eager_block]) - P.apply_ac(model, ignore_router=True, hidden_size=7168, num_experts=256) + P.apply_ac(model, ignore_router=True, hidden_size=7168, num_experts=256, selective=selective) wrapper_mock.assert_called_once() assert model.layers.registered == {"0": wrapped_block} @@ -2756,6 +2764,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) # ============================================================================