diff --git a/docs/user-guide/features/fine_grained_activation_offloading.md b/docs/user-guide/features/fine_grained_activation_offloading.md index 66a4abc8643..915926a6b9b 100644 --- a/docs/user-guide/features/fine_grained_activation_offloading.md +++ b/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 ``` diff --git a/docs/user-guide/features/paged_stash.md b/docs/user-guide/features/paged_stash.md index 4b7d807ace2..b5b97144905 100644 --- a/docs/user-guide/features/paged_stash.md +++ b/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 diff --git a/megatron/core/pipeline_parallel/fine_grained_activation_offload.py b/megatron/core/pipeline_parallel/fine_grained_activation_offload.py index 0c2e92ae7bf..e5c82876516 100644 --- a/megatron/core/pipeline_parallel/fine_grained_activation_offload.py +++ b/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,8 +868,9 @@ 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, ( @@ -861,7 +878,16 @@ def tensor_push(self, tensor): 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 diff --git a/megatron/core/transformer/moe/README.md b/megatron/core/transformer/moe/README.md index 2f731fcb6f0..57a657d0a61 100644 --- a/megatron/core/transformer/moe/README.md +++ b/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` diff --git a/megatron/core/transformer/moe/experts.py b/megatron/core/transformer/moe/experts.py index fad6966d0a8..bde1737fce0 100644 --- a/megatron/core/transformer/moe/experts.py +++ b/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: diff --git a/megatron/core/transformer/transformer_config.py b/megatron/core/transformer/transformer_config.py index ab84abd3a17..c04b6aba645 100644 --- a/megatron/core/transformer/transformer_config.py +++ b/megatron/core/transformer/transformer_config.py @@ -1132,7 +1132,7 @@ 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. @@ -1140,6 +1140,7 @@ class TransformerConfig(ModelParallelConfig): "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 diff --git a/tests/unit_tests/pipeline_parallel/test_fine_grained_activation_offloading.py b/tests/unit_tests/pipeline_parallel/test_fine_grained_activation_offloading.py index 4b68d4a48b5..515f6a01ddf 100644 --- a/tests/unit_tests/pipeline_parallel/test_fine_grained_activation_offloading.py +++ b/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,