[feat] Support fine-grained activation offloading in fused group mlp

[lhb8125]

Summary

This PR updates fine-grained activation offloading for the Transformer Engine fused grouped MLP path after the TE-side selective offload API was reverted. Instead of relying on per-op TE activation-offload opt-out markers, Megatron now exposes a separate offload module, fused_group_mlp, that offloads the whole fused grouped MLP as one group.

Changes:

• Adds fused_group_mlp to TransformerConfig.offload_modules.
• Requires use_transformer_engine_op_fuser=True when fused_group_mlp is selected.
• Rejects mixing fused_group_mlp with partial MoE offload modules (expert_fc1 / moe_act).
• Wraps the fused grouped MLP forward with Megatron's fine-grained offload group interface and commits the output as fused_group_mlp.
• Keeps expert_fc1 / moe_act selective offload on the unfused grouped MLP path only.
• Treats fused_group_mlp like other dynamic MoE activation offload groups by disabling the CPU tensor pool for it.
• Skips non-manageable tensors in the activation offload manager (Parameter, fake/functional tensors, and non-CUDA tensors), while still respecting _TE_do_not_offload and offloading_activation=False.
• Updates docs and unit coverage for the new module.

Testing

• git diff --check
• python3.12 -m py_compile megatron/core/pipeline_parallel/fine_grained_activation_offload.py megatron/core/transformer/moe/experts.py megatron/core/transformer/transformer_config.py tests/unit_tests/pipeline_parallel/test_fine_grained_activation_offloading.py
• PATH=.venv/bin:$PATH CHECK_ONLY=true BASE_REF=main bash tools/autoformat.sh
• Local DCO verification: all 19 commits have Signed-off-by: hongbinl <hongbinl@nvidia.com>

Notes:

• Local pytest was not run in this environment because the required PyTorch/test stack is not installed here.

[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.

[github-actions]

This PR has been automatically converted to draft because all PRs must start as drafts.

When you are ready for review, click Ready for Review to begin the review process. This will:

1. Add the oncall reviewer (optional reviewer)
2. Add required review teams based on your changes

See the contribution guide for more details.

[lhb8125]

/ok to test 950455b

[lhb8125]

/ok to test 577bb49

[lhb8125]

/ok to test ebf631c

[lhb8125]

/ok to test 2d6f25f

[lhb8125]

/ok to test 8014953

[lhb8125]

/ok to test b389f3e

[guihong-nv]

/claude review

[claude]

LGTM

[fanshiqing]

@lhb8125 It seems a bug:
https://github.com/NVIDIA/TransformerEngine/blob/main/transformer_engine/pytorch/ops/fused/grouped_mlp.py#L1455

if cpu_offloading:
    activation_tensors = [
        t for t in (grouped_fc1_x, activation_in, saved_grouped_fc2_x) if t is not None
    ]
    start_offload(*activation_tensors)
    mark_activation_offload(*activation_tensors)

Mcore's side fine-grained offloading modules (expert-fc1, moe_act) can not control the above behavior in TE, any of these two is effective will make TE to offload expert-fc1 AND moe_act at the same time.

[fanshiqing]

@lhb8125 It seems a bug: https://github.com/NVIDIA/TransformerEngine/blob/main/transformer_engine/pytorch/ops/fused/grouped_mlp.py#L1455

if cpu_offloading:
    activation_tensors = [
        t for t in (grouped_fc1_x, activation_in, saved_grouped_fc2_x) if t is not None
    ]
    start_offload(*activation_tensors)
    mark_activation_offload(*activation_tensors)

Mcore's side fine-grained offloading modules (expert-fc1, moe_act) can not control the above behavior in TE, any of these two is effective will make TE to offload expert-fc1 AND moe_act at the same time.

Confirmed, while the TE side corresponding MR is already merged, draft a quick fix, please help review: NVIDIA/TransformerEngine#3128

[lhb8125]

/ok to test c6dd285

[svcnvidia-nemo-ci]

🔄 Merge queue validation started!

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