Paged stashing no longer assumes single layer config

[Phlip79]

• I, the PR author, have personally reviewed every line of this PR.

What does this PR do?

Design doc (internal-only).

#6313 changes the layer logic so that each layer type has its own config and megatron/core/models/hybrid/hybrid_model.py uses a list of configs instead of a list of symbols with a single config.

This PR removes PagedStashRunner's assumption that every MoE module shares its root model config, without changing existing model behavior.

PagedStashRunner currently toggles moe_paged_stash only on the training and root model configs. This works while every MoE module aliases one of those configs, but it does not support a module owning a distinct config.

This change:

• keeps the existing structural MoE discovery, ordering, overflow detection, capacity reset, and retry behavior unchanged;
• separately collects the configs used by MoE modules, deduplicated by identity;
• disables paged stashing on every collected config for the retry; and
• restores each config's own original moe_paged_stash value afterward.

For existing models whose MoE modules share their root model config, identity deduplication makes this behavior equivalent to the current implementation.

This PR is intentionally independent of, and intended to merge before, #6313. It does not modify #6313 or depend on its per-layer config changes.

[copy-pr-bot]

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

[Phlip79]

/ok to test cb6668d

[Phlip79]

/ok to test 99bd9c8

[fanshiqing]

/claude review

[claude]

Light review. The core change is correct: per-config save/restore replaces the single self.config.moe_paged_stash snapshot, identity dedup keeps shared-root-config models behaviorally identical, and the new tests cover both the shared-config equivalence case and the distinct-per-config disable/restore case (including a config starting at False, which the old code silently clobbered). Two notes, neither blocking.

1. moe_layers and the tracked-config set now disagree for nested-MTP models.

The new modules() walk finds MoE modules anywhere under the decoder root, including a HybridStack MTP inner block. The existing structural discovery does not — getattr(transformer_layer, "mlp", None) is None when mtp_model_layer is a HybridStack. The new test encodes this asymmetry:

assert run.runner.moe_layers == [run.decoder_moe]          # nested MTP MoE excluded
assert [id(c) for c in run.runner._configs_to_sync_moe_paged_stash] == [
    id(training_config), id(model_config), id(decoder_moe_config), id(mtp_moe_config),
]                                                           # but its config IS tracked

Consequence for a hybrid nested-MTP model: its moe_paged_stash is now correctly disabled for the retry, but since it is absent from moe_layers, check_moe_overflow never polls its check_over_budget(), and prepare_for_rerun never clears its moe_expert_rank_capacity_factor, calls reset_over_budget(), or invalidate_ep_bootstrap(). That layer can go over budget without triggering a retry, and keeps its static budget across one.

This is pre-existing (structural discovery was already blind to HybridStack MTP layers) and the PR explicitly leaves discovery unchanged, so deferring is fine. But the walk you just added identifies exactly the modules the structural loops miss, so a follow-up could collapse both onto it:

        for module in model_with_decoder.modules():
            token_dispatcher = getattr(module, "token_dispatcher", None)
            if token_dispatcher is None or not hasattr(token_dispatcher, "check_over_budget"):
                continue
            _track_cfg(getattr(module, "config", None))
            self.moe_layers.append(module)   # replaces the two structural loops below

Caveat if you do: modules() yields in registration order and would need id()-dedup like _track_cfg, since a chunk registering one MoE module through two attribute paths would otherwise append it twice.

2. The hasattr(c, "moe_paged_stash") guard is untested and can mask a miss.

For any real TransformerConfig the attribute is always present (megatron/core/transformer/transformer_config.py:1278), so the guard only fires for a non-config object — where a silent skip means that module keeps paged stash enabled through the retry and re-overflows. Consider a test passing a config-like object without the attribute and asserting it is skipped, or drop the guard so a structural surprise surfaces as an AttributeError rather than a silent skip.

Test placement matches the suite and is already covered by the tests/unit_tests/transformer/moe/**/*.py bucket in tests/test_utils/recipes/h100/unit-tests.yaml, so no recipe change is needed. Nice that these are CPU-only and need no process group.

[vasunvidia]

This looks good to me.

[svcnvidia-nemo-ci]

🔄 Merge queue validation started!

You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/31670873122