Feat/selective offload on srelu fuser

[lhb8125]

Description

Please include a brief summary of the changes, relevant motivation and context.

Fixes # (issue)

Type of change

• Documentation change (change only to the documentation, either a fix or a new content)
• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Breaking change (fix or feature that would cause existing functionality to not work as expected)
• Infra/Build change
• Code refactoring

Changes

Please list the changes introduced in this PR:

• Change A
• Change B

Checklist:

• I have read and followed the contributing guidelines
• The functionality is complete
• I have commented my code, particularly in hard-to-understand areas
• I have made corresponding changes to the documentation
• My changes generate no new warnings
• I have added tests that prove my fix is effective or that my feature works
• New and existing unit tests pass locally with my changes

[greptile-apps]

Greptile Summary

This PR adds selective CPU-offload control to the SReLU fuser path of _ForwardGroupedMLP_CuTeGEMMBase. When is_cpu_offload_enabled() returns true, weights and saved FC2 intermediate tensors are unconditionally kept resident on GPU, while the FC1 input activation and MoE activation tensors can optionally be kept resident via a per-op flag.

• GroupedTensorStorage.get_data_tensors() is introduced to expose component tensors to the V1 offload path's mark_activation_offload helper.
• grouped_fc2_x in the SFD kernel path is promoted from GroupedTensor to GroupedTensorStorage, and fc1_input_quantizer.internal = True is set unconditionally so both input branches produce GroupedTensorStorage.
• mark_not_offload calls are wired around save_for_backward for weights, FC2 input, FC1 input, and MoE activation tensors, gated on new per-op no_offload_activation attributes.

Confidence Score: 3/5

Not ready to merge without clarification on two forward_grouped_mlp.py changes.

The no_offload_activation attribute controls whether critical tensors are kept on GPU, but it is never set anywhere in the codebase and has no public API, making the selective-offload opt-out path unreachable without undocumented monkey-patching. Separately, fc1_input_quantizer.internal = True is set unconditionally on every forward pass, changing grouped_fc1_x from a GroupedTensor to a GroupedTensorStorage regardless of whether offloading is active, which could silently break downstream consumers expecting a torch.Tensor subclass.

transformer_engine/pytorch/ops/fused/forward_grouped_mlp.py — specifically the no_offload_activation gating logic (lines 680-681) and the unconditional fc1_input_quantizer.internal = True assignment (line 302).

Important Files Changed

Filename	Overview
transformer_engine/pytorch/ops/fused/forward_grouped_mlp.py	Adds selective offload logic: fc1_input_quantizer.internal is set unconditionally (not just during offloading), grouped_fc2_x is changed from GroupedTensor to GroupedTensorStorage in the SFD path, mark_not_offload is called for weights/activations, and a new no_offload_activation attribute controls per-tensor opt-out — but this attribute is never set or documented anywhere in the repo.
transformer_engine/pytorch/tensor/storage/grouped_tensor_storage.py	Adds get_data_tensors() that mirrors the field list in prepare_for_saving/restore_from_saved, enabling the V1 offload path to enumerate component tensors for marking.

Sequence Diagram

sequenceDiagram
    participant Caller as Megatron / Caller
    participant FwdMLP as _ForwardGroupedMLP_CuTeGEMMBase.forward
    participant Offload as cpu_offload (V1 or V2)
    participant Ctx as fc1_ctx / activation_ctx / fc2_ctx

    Caller->>FwdMLP: forward(input_, ...)
    FwdMLP->>FwdMLP: Build grouped_fc1_x (GroupedTensorStorage)
    FwdMLP->>FwdMLP: GEMM to activation_in, grouped_fc2_x (GroupedTensorStorage)
    FwdMLP->>Offload: is_cpu_offload_enabled()?
    alt "cpu_offloading = True"
        FwdMLP->>Offload: mark_not_offload fc1_weight_tensors
        opt "no_offload_fc1_activation=True"
            FwdMLP->>Offload: mark_not_offload grouped_fc1_x
        end
        FwdMLP->>Ctx: fc1_ctx.save_for_backward
        opt "no_offload_moe_activation=True"
            FwdMLP->>Offload: mark_not_offload activation_in scales
        end
        FwdMLP->>Ctx: activation_ctx.save_for_backward
        opt saved_grouped_fc2_x not None
            FwdMLP->>Offload: mark_not_offload saved_grouped_fc2_x
        end
        FwdMLP->>Offload: mark_not_offload fc2_weight_tensors
        FwdMLP->>Ctx: fc2_ctx.save_for_backward
    end
    FwdMLP-->>Caller: fc2_out

Loading

_{Reviews (9): Last reviewed commit: "Support selective offload for fused grou..." | Re-trigger Greptile}

[timmoon10]

Overall looks good, but with one design suggestion.

Followup tasks after merging this PR:

• Enable activation checkpointing in the unfused grouped linear op.
• Update activation checkpointing to support v2 infrastructure from #1762, which is opt-out rather than opt-in.

[lhb8125]

/te-ci pytorch L1

[timmoon10]

Fine, but we should avoid relying on these in Mcore. They should be considered as debugging tools and we don't make guarantees to maintain them.

If we do want to control this in Mcore, we should do it through proper, user-facing APIs. We can add it as an arg to the unfused ops, but we should also make sure the unfused ops actually respect it too.

[timmoon10]

Now that we set the input quantizer with .internal = True, isn't it redundant to repack grouped_fc1_x into a GroupedTensorStorage?