fix(moe): preserve MLP dispatch through checkpoint wrappers

[akoumpa]

What does this PR do ?

Fix GPT-OSS and Qwen3-MoE MLP dispatch after activation checkpointing wraps the decoder block's mlp module.

Root cause

The decoder blocks selected the dense or MoE call signature with a forward-time isinstance(self.mlp, ...) check. Activation checkpointing replaces self.mlp with a CheckpointWrapper before FSDP2 sharding, so the wrapped module is no longer recognized as MLP or MoE and the MoE assertion fails on the first forward pass.

The block already knows whether it is dense or MoE during construction. This change records that role and uses it for dispatch, independent of later wrappers.

Changelog

• Preserve the GPT-OSS and Qwen3-MoE layer role through checkpoint wrapping.
• Dispatch the dense and MoE call signatures from the construction-time role.
• Add CPU regression coverage for checkpoint-wrapped GPT-OSS and Qwen3-MoE dispatch.
• Update existing tests that replace a block's MLP to set the corresponding layer role.

Validation

• uv run --group test pytest -q tests/unit_tests/models/gpt_oss/test_gptoss_model.py tests/unit_tests/models/qwen3_moe/test_qwen3_moe_model.py tests/unit_tests/models/test_moe_checkpoint_wrapper_dispatch.py
  • 13 passed, 39 skipped (GPU-gated)
• Ruff format and lint checks passed.
• GPT-OSS 20B: 50 training steps and checkpoint robustness passed.
• Qwen3-MoE two-layer activation-checkpoint proxy: 50 training steps passed.
• Qwen3-MoE 30B completed all 50 activation-checkpointed training steps, then failed a separate HF checkpoint reload KL threshold check (0.332 > 0.01). This change adds no parameters or state-dict keys.

nemo-ci pipeline references

• Scoped GPT-OSS + Qwen3-MoE parent pipeline
• Scoped GPT-OSS + Qwen3-MoE leaf pipeline
• Qwen3-MoE proxy parent pipeline
• Qwen3-MoE proxy leaf pipeline

Before your PR is "Ready for review"

Pre checks:

• Read and followed the contributor guidelines.
• Added regression tests.
• Documentation is not required for this internal dispatch fix.

Additional Information

• Linear AM-370

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[akoumpa]

/ok to test dabb4e1

[akoumpa]

/ok to test 26c8ed5

[akoumpa]

/ok to test 439295d