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

docs/user-guide/features/fine_grained_activation_offloading.md

@@ -13,7 +13,7 @@ Contributed in collaboration with RedNote.
 
 Memory is often the limiting factor for very large sparse MoE models such as DeepSeek-V3 and Qwen3-235B. Fine-grained recomputation lowers activation memory at the cost of extra compute. Offloading can use host-device bandwidth so that reload overlaps compute and keeps overhead small in many setups. Fine-grained activation offloading moves activations at module granularity so you can tune how much activation memory leaves the device and adjust training throughput.
 
-Supported offloading modules are `"attn_norm"`, `"core_attn"`, `"attn_proj"`, `"mlp_norm"`, `"expert_fc1"`, and `"moe_act"`. They can be combined with fine-grained recomputation to free almost all activations for a transformer layer on the device.
+Supported offloading modules are `"attn_norm"`, `"qkv_linear"`, `"core_attn"`, `"attn_proj"`, `"mlp_norm"`, `"expert_fc1"`, `"moe_act"`, and `"fused_group_mlp"`. They can be combined with fine-grained recomputation to free almost all activations for a transformer layer on the device. `fused_group_mlp` requires `--use-transformer-engine-op-fuser` and offloads the whole fused grouped MLP, so it cannot be combined with `expert_fc1` or `moe_act`.
 
 ## Features
 
@@ -33,7 +33,7 @@ Supported offloading modules are `"attn_norm"`, `"core_attn"`, `"attn_proj"`, `"
 --fine-grained-activation-offloading
 
 # Modules whose inputs are offloaded (refer to your training script for list or delimiter syntax).
-# Choices: "attn_norm", "core_attn", "attn_proj", "mlp_norm", "expert_fc1", "moe_act".
+# Choices: "attn_norm", "qkv_linear", "core_attn", "attn_proj", "mlp_norm", "expert_fc1", "moe_act", "fused_group_mlp".
 --offload-modules expert_fc1
 ```
 

docs/user-guide/features/paged_stash.md

@@ -21,7 +21,7 @@ Whenever `moe_expert_rank_capacity_factor` is set, a **runner** wraps forward-ba
 
 ## Prerequisites
 
-HybridEP + TE fused grouped experts are required whenever `moe_expert_rank_capacity_factor` is set. With `moe_paged_stash` enabled: capacity factor must be set; no `cpu_offloading`; `offload_modules` must not include `expert_fc1` or `moe_act`. The runner is active whenever capacity factor is set (even without `--moe-paged-stash`) for over-budget reruns; stash overflow is checked only when paged stashing is on.
+HybridEP + TE fused grouped experts are required whenever `moe_expert_rank_capacity_factor` is set. With `moe_paged_stash` enabled: capacity factor must be set; no `cpu_offloading`; `offload_modules` must not include `expert_fc1`, `moe_act`, or `fused_group_mlp`. The runner is active whenever capacity factor is set (even without `--moe-paged-stash`) for over-budget reruns; stash overflow is checked only when paged stashing is on.
 
 ## Configuration
 

megatron/core/pipeline_parallel/fine_grained_activation_offload.py

@@ -1,4 +1,4 @@
-# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 
 from collections import defaultdict, deque
 from contextlib import nullcontext
@@ -24,6 +24,22 @@ def debug_rank(message):
         print(message)
 
 
+def _te_do_not_offload(tensor):
+    """Return whether TE marked a tensor-like object as non-offloadable."""
+    if getattr(tensor, "_TE_do_not_offload", False):
+        return True
+    if not hasattr(tensor, "get_data_tensors"):
+        return False
+    try:
+        data_tensors = tensor.get_data_tensors()
+    except Exception:  # pragma: no cover - best effort for third-party tensor wrappers
+        return False
+    return any(
+        data_tensor is not None and getattr(data_tensor, "_TE_do_not_offload", False)
+        for data_tensor in data_tensors
+    )
+
+
 def print_offload_summary_table(total_offload_bytes: Dict[str, int]):
     """
     Print an ASCII table summarizing offload bytes across all ranks.
@@ -343,9 +359,9 @@ def __init__(self, name):
         self.total_offload_bytes = 0
         self.total_tensor_count = 0
         # Using memory pool is for the compatibility with cuda graph.
-        # Shapes of tensors for expert_fc1 and moe_act are not known in advance,
+        # Shapes of tensors for MoE activation offload groups are not known in advance,
         # so we do not use CPU pool for them.
-        if name == "expert_fc1" or name == "moe_act":
+        if name in ("expert_fc1", "moe_act", "fused_group_mlp"):
             self.use_cpu_pool = False
         else:
             self.use_cpu_pool = True
@@ -732,7 +748,7 @@ def on_get_saved_tensor(self, saved_state: Any) -> torch.Tensor:
         Hook called when autograd retrieves a saved tensor during backward pass.
         Returns the actual tensor (potentially reloading from CPU).
         """
-        debug_rank(f"----on_get_saved_tensor {saved_state}")
+        debug_rank("----on_get_saved_tensor")
         return self.cur_backward_chunk().tensor_pop(saved_state)
 
 
@@ -852,16 +868,26 @@ def find_next_group(self, name=None):
         assert name is not None, "Name is required"
         return self.find_group_with_name(name, self._offloaded_group_index)
 
-    def tensor_push(self, tensor):
-        """Push tensor to the offload handler."""
+    @staticmethod
+    def _can_manage_tensor_for_offload(tensor):
+        """Return whether the tensor can be managed by activation offload hooks."""
         torch_stray_tensor = isinstance(
             tensor,
             (
                 torch._subclasses.fake_tensor.FakeTensor,
                 torch._subclasses.functional_tensor.FunctionalTensor,
             ),
         )
-        assert not torch_stray_tensor, "Stray tensor should not be offloaded"
+        return (
+            not isinstance(tensor, torch.nn.Parameter)
+            and not torch_stray_tensor
+            and tensor.device.type == "cuda"
+        )
+
+    def tensor_push(self, tensor):
+        """Push tensor to the offload handler."""
+        if not self._can_manage_tensor_for_offload(tensor):
+            return tensor
 
         # Assign unique tag based on group index and position within group
         tensor_tag = (self._offloaded_group_index, self._tensor_count_current_group)
@@ -872,6 +898,9 @@ def tensor_push(self, tensor):
 
     def tensor_pop(self, tensor_tag):
         """Pop tensor from the offload handler."""
+        if isinstance(tensor_tag, torch.Tensor):
+            debug_rank(f"--------tensor_pop passthrough tensor {tensor_tag.shape}")
+            return tensor_tag
         debug_rank(f"--------tensor_pop {tensor_tag}")
         group_id, idx = tensor_tag
         tensor = self.offload_groups[group_id - 1].pop_tensor(tensor_tag)
@@ -886,6 +915,10 @@ def tensor_need_offloading_checker(self, tensor):
         debug_rank(
             f"tensor_need_offloading_checker {getattr(tensor, 'offloading_activation', None)}"
         )
+        if not self._can_manage_tensor_for_offload(tensor):
+            return False
+        if _te_do_not_offload(tensor):
+            return False
         if tensor.numel() < self.min_offloaded_tensor_size:
             return False
         # Respect tensor's offload preference if specified

megatron/core/transformer/moe/README.md

@@ -386,7 +386,7 @@ Unlike recomputation (which trades compute for memory), offloading trades **GPU-
 **Usage**
 ```bash
 --fine-grained-activation-offloading
---offload-modules expert_fc1 moe_act # Choices: attn_norm, core_attn, attn_proj, mlp_norm, expert_fc1, moe_act
+--offload-modules expert_fc1 moe_act # Choices: attn_norm, qkv_linear, core_attn, attn_proj, mlp_norm, expert_fc1, moe_act, fused_group_mlp
 ```
 
 For more details, see `docs/user-guide/features/fine_grained_activation_offloading.md`

megatron/core/transformer/moe/experts.py

@@ -250,6 +250,11 @@ def __init__(
             and "moe_act" in self.config.offload_modules
         )
 
+        self.offload_fused_group_mlp = (
+            self.config.fine_grained_activation_offloading
+            and "fused_group_mlp" in self.config.offload_modules
+        )
+
         self.activation_recompute = (
             self.config.recompute_granularity == 'selective'
             and "moe_act" in self.config.recompute_modules
@@ -329,8 +334,8 @@ def _is_fused_impl_supported(self) -> bool:
         # Check for unsupported features
         if self.tp_group.size() > 1:
             return False  # Tensor parallelism is not supported
-        if self.offload_expert_fc1 or self.offload_moe_act:
-            return False  # Fine-grained activation offloading is not supported
+        if getattr(self, "offload_expert_fc1", False) or getattr(self, "offload_moe_act", False):
+            return False  # Selective expert_fc1/moe_act offload is only supported unfused.
         if self.config.moe_apply_probs_on_input:
             return False  # Pre-multiplying probs is not supported
 
@@ -603,13 +608,25 @@ def _fused_forward(
             )
         else:
             stash_context = nullcontext()
-        with stash_context:
-            # Call fused impl
-            output = ops(
-                permuted_local_hidden_states,
-                tokens_per_expert,  # FC1
-                permuted_probs,  # Scaled SwiGLU
-                tokens_per_expert,  # FC2
+        fine_grained_activation_offloading = getattr(self, "offload_fused_group_mlp", False)
+        offload_name = "fused_group_mlp"
+        with off_interface(
+            fine_grained_activation_offloading, permuted_local_hidden_states, offload_name
+        ) as permuted_local_hidden_states:
+            forced_released_tensors = (
+                [permuted_local_hidden_states] if fine_grained_activation_offloading else []
+            )
+            with stash_context:
+                # Call fused impl
+                output = ops(
+                    permuted_local_hidden_states,
+                    tokens_per_expert,  # FC1
+                    permuted_probs,  # Scaled activation
+                    tokens_per_expert,  # FC2
+                )
+        if fine_grained_activation_offloading:
+            output = off_interface.group_commit(
+                output, name=offload_name, forced_released_tensors=forced_released_tensors
             )
         # Remove padding if needed
         if unpadded_tokens_per_expert is not None:

megatron/core/transformer/transformer_config.py

@@ -1132,14 +1132,15 @@ class TransformerConfig(ModelParallelConfig):
     offload_modules: Optional[list[str]] = field(default_factory=list)
     """The submodules to offload its input.
     choices: "attn_norm", "qkv_linear", "core_attn", "attn_proj",
-             "mlp_norm", "expert_fc1", "moe_act".
+             "mlp_norm", "expert_fc1", "moe_act", "fused_group_mlp".
     "attn_norm": offload the input of the normalization in the attention part.
     "qkv_linear": offload the input of the qkv linear part.
     "core_attn": offload the input of the core attention part.
     "attn_proj": offload the input of the attn linear projection part.
     "mlp_norm": offload the input of the normalization in the mlp part.
     "expert_fc1": offload the input of the expert fc1 part.
     "moe_act": offload the input of the moe act part.
+    "fused_group_mlp": offload the input of the whole fused grouped MLP.
     """
     min_offloaded_tensor_size: int = 1024 * 1024
     """The minimum size of the tensor to be offloaded."""
@@ -1681,6 +1682,7 @@ def __post_init__(self):
                 "core_attn",
                 "attn_proj",
                 "expert_fc1",
+                "fused_group_mlp",
                 "moe_act",
                 "attn_norm",
                 "mlp_norm",
@@ -1697,6 +1699,15 @@ def __post_init__(self):
                     "because the input of attn_proj is the output of core_attn, "
                     "which is needed in core_attn.backward()."
                 )
+            if "fused_group_mlp" in self.offload_modules:
+                if not self.use_transformer_engine_op_fuser:
+                    raise ValueError("fused_group_mlp requires use_transformer_engine_op_fuser.")
+                moe_partial_offload = {"expert_fc1", "moe_act"} & set(self.offload_modules)
+                if moe_partial_offload:
+                    raise ValueError(
+                        "fused_group_mlp offloads the whole fused grouped MLP and cannot be "
+                        f"combined with expert_fc1 or moe_act. Remove: {moe_partial_offload}"
+                    )
         if self.moe_paged_stash:
             if self.cpu_offloading:
                 raise ValueError("moe_paged_stash cannot be enabled with cpu_offloading.")
@@ -1705,11 +1716,14 @@ def __post_init__(self):
                     "moe_paged_stash requires moe_expert_rank_capacity_factor to be set; "
                     "there is no need to use paged stashing without it."
                 )
-            moe_offload_conflict = {"expert_fc1", "moe_act"} & set(self.offload_modules)
+            moe_offload_conflict = {"expert_fc1", "moe_act", "fused_group_mlp"} & set(
+                self.offload_modules
+            )
             if moe_offload_conflict:
                 raise ValueError(
                     "When moe_paged_stash is enabled, offload_modules must not include "
-                    f"expert_fc1 or moe_act (paged stash covers those activations). "
+                    f"expert_fc1, moe_act, or fused_group_mlp "
+                    f"(paged stash covers those activations). "
                     f"Remove: {moe_offload_conflict}"
                 )
 
@@ -2366,10 +2380,22 @@ def _scope_to_str(s):
                         )
 
             if self.fine_grained_activation_offloading:
-                assert self.cuda_graph_impl in ("transformer_engine", "full_iteration"), (
+                offload_modules = set(self.offload_modules or [])
+                local_partial_moe_offload = (
+                    self.cuda_graph_impl == "local"
+                    and bool(offload_modules)
+                    and offload_modules <= {"expert_fc1", "moe_act", "fused_group_mlp"}
+                    and CudaGraphModule.moe not in self.cuda_graph_modules
+                )
+                assert (
+                    self.cuda_graph_impl in ("transformer_engine", "full_iteration")
+                    or local_partial_moe_offload
+                ), (
                     "fine-grained activation offloading is only supported with "
                     "transformer_engine CUDA graph implementation or local CUDA graph "
-                    "implementation with full_iteration scope."
+                    "implementation with full_iteration scope. Local partial CUDA graphs "
+                    "are supported only for expert_fc1, moe_act, or fused_group_mlp "
+                    "offload when the full MoE module is not captured."
                 )
                 assert (
                     CudaGraphModule.moe not in self.cuda_graph_modules

tests/unit_tests/pipeline_parallel/test_fine_grained_activation_offloading.py

@@ -10,6 +10,7 @@
 
 from megatron.core.models.gpt.gpt_layer_specs import get_gpt_layer_with_transformer_engine_spec
 from megatron.core.models.gpt.gpt_model import GPTModel
+from megatron.core.pipeline_parallel.fine_grained_activation_offload import ChunkOffloadHandler
 from megatron.core.pipeline_parallel.fine_grained_activation_offload import (
     FineGrainedActivationOffloadingInterface as off_interface,
 )
@@ -32,6 +33,52 @@ def _reset_cuda_memory() -> None:
         torch.cuda.synchronize()
 
 
+def _make_chunk_handler_for_offload_checker(min_offloaded_tensor_size: int = 1):
+    handler = ChunkOffloadHandler.__new__(ChunkOffloadHandler)
+    handler.min_offloaded_tensor_size = min_offloaded_tensor_size
+    return handler
+
+
+def test_chunk_offload_handler_skips_non_offloadable_tensor_types():
+    handler = _make_chunk_handler_for_offload_checker()
+
+    cpu_tensor = torch.empty(1024)
+    assert not handler.tensor_need_offloading_checker(cpu_tensor)
+    assert handler.tensor_push(cpu_tensor) is cpu_tensor
+    assert handler.tensor_pop(cpu_tensor) is cpu_tensor
+
+    parameter = torch.nn.Parameter(torch.empty(1024))
+    assert not handler.tensor_need_offloading_checker(parameter)
+    assert handler.tensor_push(parameter) is parameter
+    assert handler.tensor_pop(parameter) is parameter
+
+    try:
+        from torch._subclasses.fake_tensor import FakeTensorMode
+    except ImportError:
+        pytest.skip("FakeTensorMode is not available in this PyTorch version.")
+
+    with FakeTensorMode():
+        fake_tensor = torch.empty(1024, device="cuda")
+    assert not handler.tensor_need_offloading_checker(fake_tensor)
+    assert handler.tensor_push(fake_tensor) is fake_tensor
+    assert handler.tensor_pop(fake_tensor) is fake_tensor
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required for offload check.")
+def test_chunk_offload_handler_respects_tensor_offloading_activation_opt_out():
+    handler = _make_chunk_handler_for_offload_checker()
+
+    tensor = torch.empty(1024, device="cuda")
+    assert handler.tensor_need_offloading_checker(tensor)
+
+    tensor._TE_do_not_offload = True
+    assert not handler.tensor_need_offloading_checker(tensor)
+
+    tensor = torch.empty(1024, device="cuda")
+    tensor.offloading_activation = False
+    assert not handler.tensor_need_offloading_checker(tensor)
+
+
 def _build_gpt_model(
     *,
     seed: int,