From 50f5dd621040703d51e1abe0cdfde50b0749477f Mon Sep 17 00:00:00 2001 From: hongbinl Date: Wed, 27 May 2026 00:01:25 -0700 Subject: [PATCH 01/26] Support selective offload for TE fused grouped MLP Signed-off-by: hongbinl --- .../fine_grained_activation_offload.py | 4 +- megatron/core/transformer/moe/experts.py | 40 ++++++++++++++----- .../core/transformer/transformer_config.py | 16 +++++++- 3 files changed, 48 insertions(+), 12 deletions(-) diff --git a/megatron/core/pipeline_parallel/fine_grained_activation_offload.py b/megatron/core/pipeline_parallel/fine_grained_activation_offload.py index 0c2e92ae7bf..a3cf97ea1bf 100644 --- a/megatron/core/pipeline_parallel/fine_grained_activation_offload.py +++ b/megatron/core/pipeline_parallel/fine_grained_activation_offload.py @@ -345,7 +345,7 @@ def __init__(self, name): # Using memory pool is for the compatibility with cuda graph. # Shapes of tensors for expert_fc1 and moe_act 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", "expert_fc1_moe_act"): self.use_cpu_pool = False else: self.use_cpu_pool = True @@ -886,6 +886,8 @@ def tensor_need_offloading_checker(self, tensor): debug_rank( f"tensor_need_offloading_checker {getattr(tensor, 'offloading_activation', None)}" ) + if getattr(tensor, "_TE_do_not_offload", False): + 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/experts.py b/megatron/core/transformer/moe/experts.py index fad6966d0a8..966f52ff99c 100644 --- a/megatron/core/transformer/moe/experts.py +++ b/megatron/core/transformer/moe/experts.py @@ -329,8 +329,6 @@ 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 self.config.moe_apply_probs_on_input: return False # Pre-multiplying probs is not supported @@ -415,6 +413,9 @@ def _make_fused_ops(self) -> torch.nn.Module: single_grouped_bias=fc1_single_grouped_bias, delay_wgrad_compute=fc1_delay_wgrad_compute, ) + fused_activation_offloading = self.offload_expert_fc1 or self.offload_moe_act + if fused_activation_offloading: + op.activation_offloading = self.offload_expert_fc1 # Copy the weights from GroupedLinear module to GroupedLinear op. if fc1_single_grouped_weight: @@ -489,6 +490,8 @@ def _make_fused_ops(self) -> torch.nn.Module: "_make_fused_ops expected SwiGLU, quick_gelu, or weighted squared_relu; " "call _is_fused_impl_supported() before constructing fused ops." ) + if fused_activation_offloading: + op.activation_offloading = self.offload_moe_act ops.append(op) # FC2 @@ -603,13 +606,32 @@ 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 + fused_activation_offloading = self.offload_expert_fc1 or self.offload_moe_act + offload_name = "_".join( + name + for name, enabled in ( + ("expert_fc1", self.offload_expert_fc1), + ("moe_act", self.offload_moe_act), + ) + if enabled + ) + with off_interface( + fused_activation_offloading, permuted_local_hidden_states, offload_name + ) as permuted_local_hidden_states: + forced_released_tensors = ( + [permuted_local_hidden_states] if self.offload_expert_fc1 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 fused_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..3ba0dcff0fa 100644 --- a/megatron/core/transformer/transformer_config.py +++ b/megatron/core/transformer/transformer_config.py @@ -2366,10 +2366,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"} + 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 offload when the full MoE " + "module is not captured." ) assert ( CudaGraphModule.moe not in self.cuda_graph_modules From 7111bc74243eb17ff89da814332081047c37bdc5 Mon Sep 17 00:00:00 2001 From: hongbinl Date: Wed, 27 May 2026 01:54:49 -0700 Subject: [PATCH 02/26] Rename TE fine-grained offload marker Signed-off-by: hongbinl --- megatron/core/transformer/moe/experts.py | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/megatron/core/transformer/moe/experts.py b/megatron/core/transformer/moe/experts.py index 966f52ff99c..40ebfcc4b6c 100644 --- a/megatron/core/transformer/moe/experts.py +++ b/megatron/core/transformer/moe/experts.py @@ -413,9 +413,9 @@ def _make_fused_ops(self) -> torch.nn.Module: single_grouped_bias=fc1_single_grouped_bias, delay_wgrad_compute=fc1_delay_wgrad_compute, ) - fused_activation_offloading = self.offload_expert_fc1 or self.offload_moe_act - if fused_activation_offloading: - op.activation_offloading = self.offload_expert_fc1 + fine_grained_activation_offloading = self.offload_expert_fc1 or self.offload_moe_act + if fine_grained_activation_offloading: + op.fine_grained_activation_offloading = self.offload_expert_fc1 # Copy the weights from GroupedLinear module to GroupedLinear op. if fc1_single_grouped_weight: @@ -490,8 +490,8 @@ def _make_fused_ops(self) -> torch.nn.Module: "_make_fused_ops expected SwiGLU, quick_gelu, or weighted squared_relu; " "call _is_fused_impl_supported() before constructing fused ops." ) - if fused_activation_offloading: - op.activation_offloading = self.offload_moe_act + if fine_grained_activation_offloading: + op.fine_grained_activation_offloading = self.offload_moe_act ops.append(op) # FC2 @@ -606,7 +606,7 @@ def _fused_forward( ) else: stash_context = nullcontext() - fused_activation_offloading = self.offload_expert_fc1 or self.offload_moe_act + fine_grained_activation_offloading = self.offload_expert_fc1 or self.offload_moe_act offload_name = "_".join( name for name, enabled in ( @@ -616,7 +616,7 @@ def _fused_forward( if enabled ) with off_interface( - fused_activation_offloading, permuted_local_hidden_states, offload_name + fine_grained_activation_offloading, permuted_local_hidden_states, offload_name ) as permuted_local_hidden_states: forced_released_tensors = ( [permuted_local_hidden_states] if self.offload_expert_fc1 else [] @@ -629,7 +629,7 @@ def _fused_forward( permuted_probs, # Scaled activation tokens_per_expert, # FC2 ) - if fused_activation_offloading: + if fine_grained_activation_offloading: output = off_interface.group_commit( output, name=offload_name, forced_released_tensors=forced_released_tensors ) From 52cb66cfebd382d2baa6cabb533e7925512286e8 Mon Sep 17 00:00:00 2001 From: hongbinl Date: Wed, 27 May 2026 07:10:52 -0700 Subject: [PATCH 03/26] Simplify fused grouped MLP offload attrs Signed-off-by: hongbinl --- megatron/core/transformer/moe/experts.py | 7 ++----- 1 file changed, 2 insertions(+), 5 deletions(-) diff --git a/megatron/core/transformer/moe/experts.py b/megatron/core/transformer/moe/experts.py index 40ebfcc4b6c..95f461780dd 100644 --- a/megatron/core/transformer/moe/experts.py +++ b/megatron/core/transformer/moe/experts.py @@ -413,9 +413,7 @@ def _make_fused_ops(self) -> torch.nn.Module: single_grouped_bias=fc1_single_grouped_bias, delay_wgrad_compute=fc1_delay_wgrad_compute, ) - fine_grained_activation_offloading = self.offload_expert_fc1 or self.offload_moe_act - if fine_grained_activation_offloading: - op.fine_grained_activation_offloading = self.offload_expert_fc1 + op.fine_grained_activation_offloading = self.offload_expert_fc1 # Copy the weights from GroupedLinear module to GroupedLinear op. if fc1_single_grouped_weight: @@ -490,8 +488,7 @@ def _make_fused_ops(self) -> torch.nn.Module: "_make_fused_ops expected SwiGLU, quick_gelu, or weighted squared_relu; " "call _is_fused_impl_supported() before constructing fused ops." ) - if fine_grained_activation_offloading: - op.fine_grained_activation_offloading = self.offload_moe_act + op.fine_grained_activation_offloading = self.offload_moe_act ops.append(op) # FC2 From 531a4368882ab9d53a4a64f578bbd8facd8521a1 Mon Sep 17 00:00:00 2001 From: hongbinl Date: Sun, 31 May 2026 22:59:58 -0700 Subject: [PATCH 04/26] Gate fused grouped MLP offload on TE 2.17 Signed-off-by: hongbinl --- .../core/extensions/transformer_engine.py | 5 ++ megatron/core/transformer/moe/experts.py | 26 ++++--- .../transformer/moe/test_grouped_mlp.py | 71 +++++++++++++++++++ 3 files changed, 91 insertions(+), 11 deletions(-) diff --git a/megatron/core/extensions/transformer_engine.py b/megatron/core/extensions/transformer_engine.py index b84565dd1f3..b6ea8e44b31 100644 --- a/megatron/core/extensions/transformer_engine.py +++ b/megatron/core/extensions/transformer_engine.py @@ -86,6 +86,11 @@ _TE_CONFIG_TYPE_KEY = "transformer_engine_config_type" +def fused_grouped_mlp_activation_offload_supported() -> bool: + """Return whether TE fused grouped MLP supports selective activation offload markers.""" + return HAVE_TE and is_te_min_version("2.17") + + class TransformerEngineConfigType(enum.Enum): """Configuration object types in config dictionary""" diff --git a/megatron/core/transformer/moe/experts.py b/megatron/core/transformer/moe/experts.py index 95f461780dd..fe4b8e0bfe4 100644 --- a/megatron/core/transformer/moe/experts.py +++ b/megatron/core/transformer/moe/experts.py @@ -18,7 +18,10 @@ from megatron.core.activations import squared_relu from megatron.core.dist_checkpointing.mapping import ShardedStateDict from megatron.core.dist_checkpointing.utils import replace_prefix_for_sharding -from megatron.core.extensions.transformer_engine import HAVE_TE +from megatron.core.extensions.transformer_engine import ( + HAVE_TE, + fused_grouped_mlp_activation_offload_supported, +) from megatron.core.fusions.fused_bias_geglu import quick_gelu, weighted_bias_quick_geglu_impl from megatron.core.fusions.fused_bias_swiglu import weighted_bias_swiglu_impl from megatron.core.fusions.fused_weighted_squared_relu import weighted_squared_relu_impl @@ -329,6 +332,10 @@ 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 ( + getattr(self, "offload_expert_fc1", False) or getattr(self, "offload_moe_act", False) + ) and not fused_grouped_mlp_activation_offload_supported(): + return False # TE fused grouped MLP offload markers require TE >= 2.17. if self.config.moe_apply_probs_on_input: return False # Pre-multiplying probs is not supported @@ -413,7 +420,7 @@ def _make_fused_ops(self) -> torch.nn.Module: single_grouped_bias=fc1_single_grouped_bias, delay_wgrad_compute=fc1_delay_wgrad_compute, ) - op.fine_grained_activation_offloading = self.offload_expert_fc1 + op.fine_grained_activation_offloading = getattr(self, "offload_expert_fc1", False) # Copy the weights from GroupedLinear module to GroupedLinear op. if fc1_single_grouped_weight: @@ -488,7 +495,7 @@ def _make_fused_ops(self) -> torch.nn.Module: "_make_fused_ops expected SwiGLU, quick_gelu, or weighted squared_relu; " "call _is_fused_impl_supported() before constructing fused ops." ) - op.fine_grained_activation_offloading = self.offload_moe_act + op.fine_grained_activation_offloading = getattr(self, "offload_moe_act", False) ops.append(op) # FC2 @@ -603,21 +610,18 @@ def _fused_forward( ) else: stash_context = nullcontext() - fine_grained_activation_offloading = self.offload_expert_fc1 or self.offload_moe_act + offload_expert_fc1 = getattr(self, "offload_expert_fc1", False) + offload_moe_act = getattr(self, "offload_moe_act", False) + fine_grained_activation_offloading = offload_expert_fc1 or offload_moe_act offload_name = "_".join( name - for name, enabled in ( - ("expert_fc1", self.offload_expert_fc1), - ("moe_act", self.offload_moe_act), - ) + for name, enabled in (("expert_fc1", offload_expert_fc1), ("moe_act", offload_moe_act)) if enabled ) 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 self.offload_expert_fc1 else [] - ) + forced_released_tensors = [permuted_local_hidden_states] if offload_expert_fc1 else [] with stash_context: # Call fused impl output = ops( diff --git a/tests/unit_tests/transformer/moe/test_grouped_mlp.py b/tests/unit_tests/transformer/moe/test_grouped_mlp.py index f6cef08b1b6..b1ffadfdc7c 100644 --- a/tests/unit_tests/transformer/moe/test_grouped_mlp.py +++ b/tests/unit_tests/transformer/moe/test_grouped_mlp.py @@ -631,6 +631,40 @@ def test_is_fused_impl_supported_requires_scaled_srelu_op(monkeypatch): assert module._is_fused_impl_supported() is False +def test_fused_grouped_mlp_activation_offload_requires_te_217(monkeypatch): + import megatron.core.extensions.transformer_engine as te_ext + + checked_versions = [] + + def fake_is_te_min_version(version): + checked_versions.append(version) + return False + + monkeypatch.setattr(te_ext, "HAVE_TE", True) + monkeypatch.setattr(te_ext, "is_te_min_version", fake_is_te_min_version) + + assert te_ext.fused_grouped_mlp_activation_offload_supported() is False + assert checked_versions == ["2.17"] + + +def test_is_fused_impl_supported_rejects_offload_without_te_217(monkeypatch): + fake_te, FakeGroupedLinear = _make_fake_te_namespace() + monkeypatch.setattr(experts_module, "te", fake_te) + monkeypatch.setattr(experts_module, "HAVE_TE", True) + monkeypatch.setattr(experts_module, "is_te_min_version", lambda _: True) + monkeypatch.setattr( + experts_module, "fused_grouped_mlp_activation_offload_supported", lambda: False + ) + _install_fake_te_ops_modules(monkeypatch, fake_te) + + module = _make_fused_impl_support_module( + FakeGroupedLinear, activation_func=F.silu, gated_linear_unit=True + ) + module.offload_expert_fc1 = True + + assert module._is_fused_impl_supported() is False + + def test_make_fused_ops_attaches_single_grouped_bias_for_fc1(monkeypatch): """single_grouped_bias=True → bias attached as `bias` (not `bias{idx}`).""" fake_te, FakeGroupedLinear = _make_fake_te_namespace() @@ -670,6 +704,43 @@ def test_make_fused_ops_attaches_single_grouped_bias_for_fc1(monkeypatch): ), "bias should not be split into bias{idx} when single_grouped_bias=True" +def test_make_fused_ops_marks_fc1_and_activation_for_offload(monkeypatch): + fake_te, FakeGroupedLinear = _make_fake_te_namespace() + monkeypatch.setattr(experts_module, "te", fake_te) + + module = TEGroupedMLP.__new__(TEGroupedMLP) + torch.nn.Module.__init__(module) + module.config = SimpleNamespace( + moe_mlp_glu_interleave_size=2, + delay_wgrad_compute=False, + activation_func_clamp_value=None, + activation_func=F.silu, + gated_linear_unit=True, + ) + module.activation_func = F.silu + module.activation_recompute = False + module.offload_expert_fc1 = True + module.offload_moe_act = True + common = dict( + device="cuda", + dtype=torch.bfloat16, + accumulate_into_main_grad=False, + single_grouped_weight=False, + ) + module.linear_fc1 = FakeGroupedLinear(2, 4, 8, bias=False, **common) + module.linear_fc2 = FakeGroupedLinear(2, 8, 4, bias=False, **common) + module.linear_fc1.weight0 = torch.nn.Parameter(torch.ones(8, 4)) + module.linear_fc1.weight1 = torch.nn.Parameter(torch.ones(8, 4)) + module.linear_fc2.weight0 = torch.nn.Parameter(torch.ones(4, 8)) + module.linear_fc2.weight1 = torch.nn.Parameter(torch.ones(4, 8)) + + ops = module._make_fused_ops() + + assert ops[0].fine_grained_activation_offloading is True + assert ops[1].fine_grained_activation_offloading is True + assert not hasattr(ops[2], "fine_grained_activation_offloading") + + def test_backward_dw_dispatches_fused_children_in_fc2_then_fc1_order(): """delay_wgrad_compute=True (via wrapper) → backward_dw calls fused [2] then [0], then triggers the wrapper-side wgrad/reduce hooks (PR 4311) so DDP's reduce-scatter From faae08530c37604868b9265b7e124cbc423be25e Mon Sep 17 00:00:00 2001 From: hongbinl Date: Mon, 1 Jun 2026 03:53:36 -0700 Subject: [PATCH 05/26] Refine fused grouped MLP offload grouping Signed-off-by: hongbinl --- .../fine_grained_activation_offload.py | 37 +++++++++------- megatron/core/transformer/moe/experts.py | 18 +++++--- .../transformer/moe/test_grouped_mlp.py | 42 ++++++++++++++++++- 3 files changed, 74 insertions(+), 23 deletions(-) diff --git a/megatron/core/pipeline_parallel/fine_grained_activation_offload.py b/megatron/core/pipeline_parallel/fine_grained_activation_offload.py index a3cf97ea1bf..d91d1378212 100644 --- a/megatron/core/pipeline_parallel/fine_grained_activation_offload.py +++ b/megatron/core/pipeline_parallel/fine_grained_activation_offload.py @@ -334,7 +334,7 @@ class OffloadTensorGroup: A group of tensors to be offloaded together. """ - def __init__(self, name): + def __init__(self, name, use_cpu_pool: Optional[bool] = None): self._name = name self._tensors = {} self._offload_event = torch.cuda.Event() @@ -342,13 +342,11 @@ def __init__(self, name): self.offload = True 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, - # so we do not use CPU pool for them. - if name in ("expert_fc1", "moe_act", "expert_fc1_moe_act"): - self.use_cpu_pool = False - else: - self.use_cpu_pool = True + # Using memory pool is for the compatibility with cuda graph. Shapes of tensors for + # expert_fc1 and moe_act are not known in advance, so do not use CPU pool for them. + if use_cpu_pool is None: + use_cpu_pool = name not in ("expert_fc1", "moe_act") + self.use_cpu_pool = use_cpu_pool def push_tensor(self, tag, tensor): """Push a tensor to the group.""" @@ -1060,7 +1058,7 @@ def on_group_commit_backward(self, name): self._reloading_group.remove(reloading_group) break - def on_group_start_forward(self, name): + def on_group_start_forward(self, name, use_cpu_pool: Optional[bool] = None): """ Called at the start of a layer group's forward pass. Increments group index and prepares for offloading. @@ -1070,7 +1068,7 @@ def on_group_start_forward(self, name): debug_rank(f"--on_group_start_forward {name}") self._offloaded_group_index = self._offloaded_group_index + 1 if self.is_warmup: - self.offload_groups.append(OffloadTensorGroup(name)) + self.offload_groups.append(OffloadTensorGroup(name, use_cpu_pool=use_cpu_pool)) self._max_group_size = max(self._max_group_size, self._offloaded_group_index) debug_rank(f"max group size {self._max_group_size}") else: @@ -1194,12 +1192,12 @@ class FineGrainedOffloadingGroupStartFunction(torch.autograd.Function): """ @staticmethod - def forward(ctx, tensor, cpu_offload_handler, name): + def forward(ctx, tensor, cpu_offload_handler, name, use_cpu_pool): # pylint: disable=missing-function-docstring ctx.cpu_offload_handler = cpu_offload_handler debug_rank("FineGrainedOffloadingGroupStartFunction forward") - cpu_offload_handler.on_group_start_forward(name) + cpu_offload_handler.on_group_start_forward(name, use_cpu_pool=use_cpu_pool) # return the identical tensor return tensor @@ -1212,12 +1210,14 @@ def backward(ctx, grad_output): return grad_output, None, None, None -def fine_grained_offloading_group_start(tensor, name=None): +def fine_grained_offloading_group_start(tensor, name=None, use_cpu_pool: Optional[bool] = None): """Mark the start of a layer group and prepare for offload/reload.""" cur_forward_chunk = PipelineOffloadManager.get_instance().pop_forward_chunk(name=name) if cur_forward_chunk is None: return tensor - return FineGrainedOffloadingGroupStartFunction.apply(tensor, cur_forward_chunk, name) + return FineGrainedOffloadingGroupStartFunction.apply( + tensor, cur_forward_chunk, name, use_cpu_pool + ) def fine_grained_offloading_forward_record(event: torch.cuda.Event) -> None: @@ -1256,15 +1256,20 @@ def fine_grained_offloading_backward_record(tensor, event: torch.cuda.Event) -> class FineGrainedActivationOffloadingInterface: """Interface for fine-grained activation offloading.""" - def __init__(self, offload: bool, tensor: torch.Tensor, name: str): + def __init__( + self, offload: bool, tensor: torch.Tensor, name: str, use_cpu_pool: Optional[bool] = None + ): self.offload = offload self.tensor = tensor self.name = name + self.use_cpu_pool = use_cpu_pool def __enter__(self): """Enter context manager to enable activation offloading hooks.""" if self.offload: - self.tensor = fine_grained_offloading_group_start(self.tensor, self.name) + self.tensor = fine_grained_offloading_group_start( + self.tensor, self.name, self.use_cpu_pool + ) PipelineOffloadManager.get_instance().__enter__() return self.tensor diff --git a/megatron/core/transformer/moe/experts.py b/megatron/core/transformer/moe/experts.py index fe4b8e0bfe4..0f6af4ca293 100644 --- a/megatron/core/transformer/moe/experts.py +++ b/megatron/core/transformer/moe/experts.py @@ -613,13 +613,19 @@ def _fused_forward( offload_expert_fc1 = getattr(self, "offload_expert_fc1", False) offload_moe_act = getattr(self, "offload_moe_act", False) fine_grained_activation_offloading = offload_expert_fc1 or offload_moe_act - offload_name = "_".join( - name - for name, enabled in (("expert_fc1", offload_expert_fc1), ("moe_act", offload_moe_act)) - if enabled - ) + if offload_expert_fc1 and offload_moe_act: + offload_name = "grouped_mlp" + elif offload_expert_fc1: + offload_name = "expert_fc1" + elif offload_moe_act: + offload_name = "moe_act" + else: + offload_name = "" with off_interface( - fine_grained_activation_offloading, permuted_local_hidden_states, offload_name + fine_grained_activation_offloading, + permuted_local_hidden_states, + offload_name, + use_cpu_pool=False, ) as permuted_local_hidden_states: forced_released_tensors = [permuted_local_hidden_states] if offload_expert_fc1 else [] with stash_context: diff --git a/tests/unit_tests/transformer/moe/test_grouped_mlp.py b/tests/unit_tests/transformer/moe/test_grouped_mlp.py index b1ffadfdc7c..3c03d113956 100644 --- a/tests/unit_tests/transformer/moe/test_grouped_mlp.py +++ b/tests/unit_tests/transformer/moe/test_grouped_mlp.py @@ -150,12 +150,33 @@ def register_forward_pre_hook(self, hook): assert hasattr(ops, "forward_pre_hook") -def test_fused_forward_caches_ops_and_forwards_expected_arguments(): +def test_fused_forward_caches_ops_and_forwards_expected_arguments(monkeypatch): class FakeFusedOps: def __call__(self, hidden_states, fc1_tokens, probs, fc2_tokens): self.args = (hidden_states, fc1_tokens, probs, fc2_tokens) return hidden_states + 1 + class FakeOffloadInterface: + enter_calls = [] + commit_calls = [] + + def __init__(self, offload, tensor, name, use_cpu_pool=None): + self.tensor = tensor + self.enter_calls.append((offload, tensor, name, use_cpu_pool)) + + def __enter__(self): + return self.tensor + + def __exit__(self, *args): + return None + + @staticmethod + def group_commit(tensor, name, forced_released_tensors=None, delay_offload=False): + FakeOffloadInterface.commit_calls.append( + (tensor, name, forced_released_tensors, delay_offload) + ) + return tensor + module = TEGroupedMLP.__new__(TEGroupedMLP) # `_fused_forward` calls `skip_routed_expert_padding(config)` (added by PR 4071), which # reads `moe_token_dispatcher_type` and `moe_flex_dispatcher_backend` after the @@ -171,9 +192,13 @@ def __call__(self, hidden_states, fc1_tokens, probs, fc2_tokens): module._fused_ops = None fused_ops = FakeFusedOps() module._make_fused_ops = lambda: fused_ops + module.offload_expert_fc1 = True + module.offload_moe_act = True + module.activation_recompute = False hidden_states = torch.zeros(2, 4) tokens_per_expert = torch.tensor([1, 1]) probs = torch.ones(2) + monkeypatch.setattr(experts_module, "off_interface", FakeOffloadInterface) output = module._fused_forward(hidden_states, tokens_per_expert, probs) @@ -183,6 +208,21 @@ def __call__(self, hidden_states, fc1_tokens, probs, fc2_tokens): assert fused_ops.args[1] is tokens_per_expert assert fused_ops.args[2] is probs assert fused_ops.args[3] is tokens_per_expert + assert len(FakeOffloadInterface.enter_calls) == 1 + offload, tensor, name, use_cpu_pool = FakeOffloadInterface.enter_calls[0] + assert offload is True + assert tensor is hidden_states + assert name == "grouped_mlp" + assert use_cpu_pool is False + assert len(FakeOffloadInterface.commit_calls) == 1 + committed_tensor, name, forced_released_tensors, delay_offload = ( + FakeOffloadInterface.commit_calls[0] + ) + assert committed_tensor is output + assert name == "grouped_mlp" + assert len(forced_released_tensors) == 1 + assert forced_released_tensors[0] is hidden_states + assert delay_offload is False def test_apply_bias_returns_input_unchanged_when_bias_is_none(): From 7739185e93f422528e930cf14e2bd21e4209364b Mon Sep 17 00:00:00 2001 From: hongbinl Date: Mon, 1 Jun 2026 06:25:05 -0700 Subject: [PATCH 06/26] Revert "Refine fused grouped MLP offload grouping" This reverts commit a76a38eb67afb539e551dab79911832f3af2c21c. Signed-off-by: hongbinl --- .../fine_grained_activation_offload.py | 37 +++++++--------- megatron/core/transformer/moe/experts.py | 18 +++----- .../transformer/moe/test_grouped_mlp.py | 42 +------------------ 3 files changed, 23 insertions(+), 74 deletions(-) diff --git a/megatron/core/pipeline_parallel/fine_grained_activation_offload.py b/megatron/core/pipeline_parallel/fine_grained_activation_offload.py index d91d1378212..a3cf97ea1bf 100644 --- a/megatron/core/pipeline_parallel/fine_grained_activation_offload.py +++ b/megatron/core/pipeline_parallel/fine_grained_activation_offload.py @@ -334,7 +334,7 @@ class OffloadTensorGroup: A group of tensors to be offloaded together. """ - def __init__(self, name, use_cpu_pool: Optional[bool] = None): + def __init__(self, name): self._name = name self._tensors = {} self._offload_event = torch.cuda.Event() @@ -342,11 +342,13 @@ def __init__(self, name, use_cpu_pool: Optional[bool] = None): self.offload = True 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, so do not use CPU pool for them. - if use_cpu_pool is None: - use_cpu_pool = name not in ("expert_fc1", "moe_act") - self.use_cpu_pool = use_cpu_pool + # Using memory pool is for the compatibility with cuda graph. + # Shapes of tensors for expert_fc1 and moe_act are not known in advance, + # so we do not use CPU pool for them. + if name in ("expert_fc1", "moe_act", "expert_fc1_moe_act"): + self.use_cpu_pool = False + else: + self.use_cpu_pool = True def push_tensor(self, tag, tensor): """Push a tensor to the group.""" @@ -1058,7 +1060,7 @@ def on_group_commit_backward(self, name): self._reloading_group.remove(reloading_group) break - def on_group_start_forward(self, name, use_cpu_pool: Optional[bool] = None): + def on_group_start_forward(self, name): """ Called at the start of a layer group's forward pass. Increments group index and prepares for offloading. @@ -1068,7 +1070,7 @@ def on_group_start_forward(self, name, use_cpu_pool: Optional[bool] = None): debug_rank(f"--on_group_start_forward {name}") self._offloaded_group_index = self._offloaded_group_index + 1 if self.is_warmup: - self.offload_groups.append(OffloadTensorGroup(name, use_cpu_pool=use_cpu_pool)) + self.offload_groups.append(OffloadTensorGroup(name)) self._max_group_size = max(self._max_group_size, self._offloaded_group_index) debug_rank(f"max group size {self._max_group_size}") else: @@ -1192,12 +1194,12 @@ class FineGrainedOffloadingGroupStartFunction(torch.autograd.Function): """ @staticmethod - def forward(ctx, tensor, cpu_offload_handler, name, use_cpu_pool): + def forward(ctx, tensor, cpu_offload_handler, name): # pylint: disable=missing-function-docstring ctx.cpu_offload_handler = cpu_offload_handler debug_rank("FineGrainedOffloadingGroupStartFunction forward") - cpu_offload_handler.on_group_start_forward(name, use_cpu_pool=use_cpu_pool) + cpu_offload_handler.on_group_start_forward(name) # return the identical tensor return tensor @@ -1210,14 +1212,12 @@ def backward(ctx, grad_output): return grad_output, None, None, None -def fine_grained_offloading_group_start(tensor, name=None, use_cpu_pool: Optional[bool] = None): +def fine_grained_offloading_group_start(tensor, name=None): """Mark the start of a layer group and prepare for offload/reload.""" cur_forward_chunk = PipelineOffloadManager.get_instance().pop_forward_chunk(name=name) if cur_forward_chunk is None: return tensor - return FineGrainedOffloadingGroupStartFunction.apply( - tensor, cur_forward_chunk, name, use_cpu_pool - ) + return FineGrainedOffloadingGroupStartFunction.apply(tensor, cur_forward_chunk, name) def fine_grained_offloading_forward_record(event: torch.cuda.Event) -> None: @@ -1256,20 +1256,15 @@ def fine_grained_offloading_backward_record(tensor, event: torch.cuda.Event) -> class FineGrainedActivationOffloadingInterface: """Interface for fine-grained activation offloading.""" - def __init__( - self, offload: bool, tensor: torch.Tensor, name: str, use_cpu_pool: Optional[bool] = None - ): + def __init__(self, offload: bool, tensor: torch.Tensor, name: str): self.offload = offload self.tensor = tensor self.name = name - self.use_cpu_pool = use_cpu_pool def __enter__(self): """Enter context manager to enable activation offloading hooks.""" if self.offload: - self.tensor = fine_grained_offloading_group_start( - self.tensor, self.name, self.use_cpu_pool - ) + self.tensor = fine_grained_offloading_group_start(self.tensor, self.name) PipelineOffloadManager.get_instance().__enter__() return self.tensor diff --git a/megatron/core/transformer/moe/experts.py b/megatron/core/transformer/moe/experts.py index 0f6af4ca293..fe4b8e0bfe4 100644 --- a/megatron/core/transformer/moe/experts.py +++ b/megatron/core/transformer/moe/experts.py @@ -613,19 +613,13 @@ def _fused_forward( offload_expert_fc1 = getattr(self, "offload_expert_fc1", False) offload_moe_act = getattr(self, "offload_moe_act", False) fine_grained_activation_offloading = offload_expert_fc1 or offload_moe_act - if offload_expert_fc1 and offload_moe_act: - offload_name = "grouped_mlp" - elif offload_expert_fc1: - offload_name = "expert_fc1" - elif offload_moe_act: - offload_name = "moe_act" - else: - offload_name = "" + offload_name = "_".join( + name + for name, enabled in (("expert_fc1", offload_expert_fc1), ("moe_act", offload_moe_act)) + if enabled + ) with off_interface( - fine_grained_activation_offloading, - permuted_local_hidden_states, - offload_name, - use_cpu_pool=False, + fine_grained_activation_offloading, permuted_local_hidden_states, offload_name ) as permuted_local_hidden_states: forced_released_tensors = [permuted_local_hidden_states] if offload_expert_fc1 else [] with stash_context: diff --git a/tests/unit_tests/transformer/moe/test_grouped_mlp.py b/tests/unit_tests/transformer/moe/test_grouped_mlp.py index 3c03d113956..b1ffadfdc7c 100644 --- a/tests/unit_tests/transformer/moe/test_grouped_mlp.py +++ b/tests/unit_tests/transformer/moe/test_grouped_mlp.py @@ -150,33 +150,12 @@ def register_forward_pre_hook(self, hook): assert hasattr(ops, "forward_pre_hook") -def test_fused_forward_caches_ops_and_forwards_expected_arguments(monkeypatch): +def test_fused_forward_caches_ops_and_forwards_expected_arguments(): class FakeFusedOps: def __call__(self, hidden_states, fc1_tokens, probs, fc2_tokens): self.args = (hidden_states, fc1_tokens, probs, fc2_tokens) return hidden_states + 1 - class FakeOffloadInterface: - enter_calls = [] - commit_calls = [] - - def __init__(self, offload, tensor, name, use_cpu_pool=None): - self.tensor = tensor - self.enter_calls.append((offload, tensor, name, use_cpu_pool)) - - def __enter__(self): - return self.tensor - - def __exit__(self, *args): - return None - - @staticmethod - def group_commit(tensor, name, forced_released_tensors=None, delay_offload=False): - FakeOffloadInterface.commit_calls.append( - (tensor, name, forced_released_tensors, delay_offload) - ) - return tensor - module = TEGroupedMLP.__new__(TEGroupedMLP) # `_fused_forward` calls `skip_routed_expert_padding(config)` (added by PR 4071), which # reads `moe_token_dispatcher_type` and `moe_flex_dispatcher_backend` after the @@ -192,13 +171,9 @@ def group_commit(tensor, name, forced_released_tensors=None, delay_offload=False module._fused_ops = None fused_ops = FakeFusedOps() module._make_fused_ops = lambda: fused_ops - module.offload_expert_fc1 = True - module.offload_moe_act = True - module.activation_recompute = False hidden_states = torch.zeros(2, 4) tokens_per_expert = torch.tensor([1, 1]) probs = torch.ones(2) - monkeypatch.setattr(experts_module, "off_interface", FakeOffloadInterface) output = module._fused_forward(hidden_states, tokens_per_expert, probs) @@ -208,21 +183,6 @@ def group_commit(tensor, name, forced_released_tensors=None, delay_offload=False assert fused_ops.args[1] is tokens_per_expert assert fused_ops.args[2] is probs assert fused_ops.args[3] is tokens_per_expert - assert len(FakeOffloadInterface.enter_calls) == 1 - offload, tensor, name, use_cpu_pool = FakeOffloadInterface.enter_calls[0] - assert offload is True - assert tensor is hidden_states - assert name == "grouped_mlp" - assert use_cpu_pool is False - assert len(FakeOffloadInterface.commit_calls) == 1 - committed_tensor, name, forced_released_tensors, delay_offload = ( - FakeOffloadInterface.commit_calls[0] - ) - assert committed_tensor is output - assert name == "grouped_mlp" - assert len(forced_released_tensors) == 1 - assert forced_released_tensors[0] is hidden_states - assert delay_offload is False def test_apply_bias_returns_input_unchanged_when_bias_is_none(): From ccb6da3c09a58e2f046ce4c3c418a879ea65a8bb Mon Sep 17 00:00:00 2001 From: hongbinl Date: Tue, 2 Jun 2026 02:23:36 -0700 Subject: [PATCH 07/26] Use opt-out TE grouped MLP offload markers Signed-off-by: hongbinl --- megatron/core/transformer/moe/experts.py | 4 ++-- .../transformer/moe/test_grouped_mlp.py | 23 ++++++++++++++----- 2 files changed, 19 insertions(+), 8 deletions(-) diff --git a/megatron/core/transformer/moe/experts.py b/megatron/core/transformer/moe/experts.py index fe4b8e0bfe4..35e2bf761a9 100644 --- a/megatron/core/transformer/moe/experts.py +++ b/megatron/core/transformer/moe/experts.py @@ -420,7 +420,7 @@ def _make_fused_ops(self) -> torch.nn.Module: single_grouped_bias=fc1_single_grouped_bias, delay_wgrad_compute=fc1_delay_wgrad_compute, ) - op.fine_grained_activation_offloading = getattr(self, "offload_expert_fc1", False) + op.no_offload_expert_fc1 = not getattr(self, "offload_expert_fc1", False) # Copy the weights from GroupedLinear module to GroupedLinear op. if fc1_single_grouped_weight: @@ -495,7 +495,7 @@ def _make_fused_ops(self) -> torch.nn.Module: "_make_fused_ops expected SwiGLU, quick_gelu, or weighted squared_relu; " "call _is_fused_impl_supported() before constructing fused ops." ) - op.fine_grained_activation_offloading = getattr(self, "offload_moe_act", False) + op.no_offload_moe_act = not getattr(self, "offload_moe_act", False) ops.append(op) # FC2 diff --git a/tests/unit_tests/transformer/moe/test_grouped_mlp.py b/tests/unit_tests/transformer/moe/test_grouped_mlp.py index b1ffadfdc7c..73fccb3e716 100644 --- a/tests/unit_tests/transformer/moe/test_grouped_mlp.py +++ b/tests/unit_tests/transformer/moe/test_grouped_mlp.py @@ -704,7 +704,17 @@ def test_make_fused_ops_attaches_single_grouped_bias_for_fc1(monkeypatch): ), "bias should not be split into bias{idx} when single_grouped_bias=True" -def test_make_fused_ops_marks_fc1_and_activation_for_offload(monkeypatch): +@pytest.mark.parametrize( + "offload_expert_fc1, offload_moe_act, expected_no_offload_expert_fc1, expected_no_offload_moe_act", + [(True, True, False, False), (True, False, False, True), (False, True, True, False)], +) +def test_make_fused_ops_sets_fine_grained_offload_opt_out_attrs( + monkeypatch, + offload_expert_fc1, + offload_moe_act, + expected_no_offload_expert_fc1, + expected_no_offload_moe_act, +): fake_te, FakeGroupedLinear = _make_fake_te_namespace() monkeypatch.setattr(experts_module, "te", fake_te) @@ -719,8 +729,8 @@ def test_make_fused_ops_marks_fc1_and_activation_for_offload(monkeypatch): ) module.activation_func = F.silu module.activation_recompute = False - module.offload_expert_fc1 = True - module.offload_moe_act = True + module.offload_expert_fc1 = offload_expert_fc1 + module.offload_moe_act = offload_moe_act common = dict( device="cuda", dtype=torch.bfloat16, @@ -736,9 +746,10 @@ def test_make_fused_ops_marks_fc1_and_activation_for_offload(monkeypatch): ops = module._make_fused_ops() - assert ops[0].fine_grained_activation_offloading is True - assert ops[1].fine_grained_activation_offloading is True - assert not hasattr(ops[2], "fine_grained_activation_offloading") + assert ops[0].no_offload_expert_fc1 is expected_no_offload_expert_fc1 + assert ops[1].no_offload_moe_act is expected_no_offload_moe_act + assert not hasattr(ops[2], "no_offload_expert_fc1") + assert not hasattr(ops[2], "no_offload_moe_act") def test_backward_dw_dispatches_fused_children_in_fc2_then_fc1_order(): From c4c10c61cfd03133df599463b350696a7062078d Mon Sep 17 00:00:00 2001 From: hongbinl Date: Tue, 2 Jun 2026 02:41:38 -0700 Subject: [PATCH 08/26] Set generic TE activation offload opt-out marker Signed-off-by: hongbinl --- megatron/core/transformer/moe/experts.py | 4 ++-- .../unit_tests/transformer/moe/test_grouped_mlp.py | 13 ++++++------- 2 files changed, 8 insertions(+), 9 deletions(-) diff --git a/megatron/core/transformer/moe/experts.py b/megatron/core/transformer/moe/experts.py index 35e2bf761a9..99efbd91c4b 100644 --- a/megatron/core/transformer/moe/experts.py +++ b/megatron/core/transformer/moe/experts.py @@ -420,7 +420,7 @@ def _make_fused_ops(self) -> torch.nn.Module: single_grouped_bias=fc1_single_grouped_bias, delay_wgrad_compute=fc1_delay_wgrad_compute, ) - op.no_offload_expert_fc1 = not getattr(self, "offload_expert_fc1", False) + op.no_offload_activation = not getattr(self, "offload_expert_fc1", False) # Copy the weights from GroupedLinear module to GroupedLinear op. if fc1_single_grouped_weight: @@ -495,7 +495,7 @@ def _make_fused_ops(self) -> torch.nn.Module: "_make_fused_ops expected SwiGLU, quick_gelu, or weighted squared_relu; " "call _is_fused_impl_supported() before constructing fused ops." ) - op.no_offload_moe_act = not getattr(self, "offload_moe_act", False) + op.no_offload_activation = not getattr(self, "offload_moe_act", False) ops.append(op) # FC2 diff --git a/tests/unit_tests/transformer/moe/test_grouped_mlp.py b/tests/unit_tests/transformer/moe/test_grouped_mlp.py index 73fccb3e716..b5fce7d0c53 100644 --- a/tests/unit_tests/transformer/moe/test_grouped_mlp.py +++ b/tests/unit_tests/transformer/moe/test_grouped_mlp.py @@ -705,15 +705,15 @@ def test_make_fused_ops_attaches_single_grouped_bias_for_fc1(monkeypatch): @pytest.mark.parametrize( - "offload_expert_fc1, offload_moe_act, expected_no_offload_expert_fc1, expected_no_offload_moe_act", + "offload_expert_fc1, offload_moe_act, expected_fc1_no_offload, expected_act_no_offload", [(True, True, False, False), (True, False, False, True), (False, True, True, False)], ) def test_make_fused_ops_sets_fine_grained_offload_opt_out_attrs( monkeypatch, offload_expert_fc1, offload_moe_act, - expected_no_offload_expert_fc1, - expected_no_offload_moe_act, + expected_fc1_no_offload, + expected_act_no_offload, ): fake_te, FakeGroupedLinear = _make_fake_te_namespace() monkeypatch.setattr(experts_module, "te", fake_te) @@ -746,10 +746,9 @@ def test_make_fused_ops_sets_fine_grained_offload_opt_out_attrs( ops = module._make_fused_ops() - assert ops[0].no_offload_expert_fc1 is expected_no_offload_expert_fc1 - assert ops[1].no_offload_moe_act is expected_no_offload_moe_act - assert not hasattr(ops[2], "no_offload_expert_fc1") - assert not hasattr(ops[2], "no_offload_moe_act") + assert ops[0].no_offload_activation is expected_fc1_no_offload + assert ops[1].no_offload_activation is expected_act_no_offload + assert not hasattr(ops[2], "no_offload_activation") def test_backward_dw_dispatches_fused_children_in_fc2_then_fc1_order(): From 8d7a05b4b121a2351900f487b60e9abc63b8867b Mon Sep 17 00:00:00 2001 From: hongbinl Date: Tue, 2 Jun 2026 06:20:07 -0700 Subject: [PATCH 09/26] Fix integrity manifest direct test race Signed-off-by: hongbinl --- .../dist_checkpointing/test_integrity.py | 24 ++++++++++++------- 1 file changed, 15 insertions(+), 9 deletions(-) diff --git a/tests/unit_tests/dist_checkpointing/test_integrity.py b/tests/unit_tests/dist_checkpointing/test_integrity.py index e87af62af93..10aaa4110ae 100644 --- a/tests/unit_tests/dist_checkpointing/test_integrity.py +++ b/tests/unit_tests/dist_checkpointing/test_integrity.py @@ -64,9 +64,11 @@ def test_save_verify_integrity_manifest_directly(self, init_model_parallel, tmp_ tmp_path_dist_ckpt / 'test_save_integrity_manifest_directly', sync=True ) as ckpt_dir: metadata_file = Path(ckpt_dir / "metadata.json") - with open(metadata_file, "w") as f: - data = {"test_metadata": 1} - json.dump(data, f) + if torch.distributed.get_rank() == 0: + with open(metadata_file, "w") as f: + data = {"test_metadata": 1} + json.dump(data, f) + torch.distributed.barrier() if torch.distributed.get_rank() == 0: save_integrity_manifest(ckpt_dir) @@ -89,17 +91,21 @@ def test_save_verify_integrity_manifest_error(self, init_model_parallel, tmp_pat ) as ckpt_dir: metadata_file = Path(ckpt_dir / "metadata.json") - with open(metadata_file, "w") as f: - data = {"test_metadata": 1} - json.dump(data, f) + if torch.distributed.get_rank() == 0: + with open(metadata_file, "w") as f: + data = {"test_metadata": 1} + json.dump(data, f) + torch.distributed.barrier() if torch.distributed.get_rank() == 0: save_integrity_manifest(ckpt_dir) torch.distributed.barrier() - with open(metadata_file, "w") as f: - data = {"test_metadata": 11} - json.dump(data, f) + if torch.distributed.get_rank() == 0: + with open(metadata_file, "w") as f: + data = {"test_metadata": 11} + json.dump(data, f) + torch.distributed.barrier() # CheckpointingException, hash mismatch with pytest.raises(CheckpointingException): From 061a0dd8ed5afb48d3ed669240154f60eb25d0ca Mon Sep 17 00:00:00 2001 From: hongbinl Date: Mon, 8 Jun 2026 01:00:13 -0700 Subject: [PATCH 10/26] Sync fused grouped MLP offload with TE API Signed-off-by: hongbinl --- .../core/extensions/transformer_engine.py | 2 +- megatron/core/transformer/moe/experts.py | 4 +--- .../transformer/moe/test_grouped_mlp.py | 20 +++++-------------- 3 files changed, 7 insertions(+), 19 deletions(-) diff --git a/megatron/core/extensions/transformer_engine.py b/megatron/core/extensions/transformer_engine.py index b6ea8e44b31..122fdc34b23 100644 --- a/megatron/core/extensions/transformer_engine.py +++ b/megatron/core/extensions/transformer_engine.py @@ -87,7 +87,7 @@ def fused_grouped_mlp_activation_offload_supported() -> bool: - """Return whether TE fused grouped MLP supports selective activation offload markers.""" + """Return whether TE fused grouped MLP handles fine-grained activation offload.""" return HAVE_TE and is_te_min_version("2.17") diff --git a/megatron/core/transformer/moe/experts.py b/megatron/core/transformer/moe/experts.py index 99efbd91c4b..f63f1217479 100644 --- a/megatron/core/transformer/moe/experts.py +++ b/megatron/core/transformer/moe/experts.py @@ -335,7 +335,7 @@ def _is_fused_impl_supported(self) -> bool: if ( getattr(self, "offload_expert_fc1", False) or getattr(self, "offload_moe_act", False) ) and not fused_grouped_mlp_activation_offload_supported(): - return False # TE fused grouped MLP offload markers require TE >= 2.17. + return False # TE fused grouped MLP offload support requires TE >= 2.17. if self.config.moe_apply_probs_on_input: return False # Pre-multiplying probs is not supported @@ -420,7 +420,6 @@ def _make_fused_ops(self) -> torch.nn.Module: single_grouped_bias=fc1_single_grouped_bias, delay_wgrad_compute=fc1_delay_wgrad_compute, ) - op.no_offload_activation = not getattr(self, "offload_expert_fc1", False) # Copy the weights from GroupedLinear module to GroupedLinear op. if fc1_single_grouped_weight: @@ -495,7 +494,6 @@ def _make_fused_ops(self) -> torch.nn.Module: "_make_fused_ops expected SwiGLU, quick_gelu, or weighted squared_relu; " "call _is_fused_impl_supported() before constructing fused ops." ) - op.no_offload_activation = not getattr(self, "offload_moe_act", False) ops.append(op) # FC2 diff --git a/tests/unit_tests/transformer/moe/test_grouped_mlp.py b/tests/unit_tests/transformer/moe/test_grouped_mlp.py index b5fce7d0c53..21af088e258 100644 --- a/tests/unit_tests/transformer/moe/test_grouped_mlp.py +++ b/tests/unit_tests/transformer/moe/test_grouped_mlp.py @@ -704,17 +704,7 @@ def test_make_fused_ops_attaches_single_grouped_bias_for_fc1(monkeypatch): ), "bias should not be split into bias{idx} when single_grouped_bias=True" -@pytest.mark.parametrize( - "offload_expert_fc1, offload_moe_act, expected_fc1_no_offload, expected_act_no_offload", - [(True, True, False, False), (True, False, False, True), (False, True, True, False)], -) -def test_make_fused_ops_sets_fine_grained_offload_opt_out_attrs( - monkeypatch, - offload_expert_fc1, - offload_moe_act, - expected_fc1_no_offload, - expected_act_no_offload, -): +def test_make_fused_ops_does_not_set_deprecated_offload_opt_out_attrs(monkeypatch): fake_te, FakeGroupedLinear = _make_fake_te_namespace() monkeypatch.setattr(experts_module, "te", fake_te) @@ -729,8 +719,8 @@ def test_make_fused_ops_sets_fine_grained_offload_opt_out_attrs( ) module.activation_func = F.silu module.activation_recompute = False - module.offload_expert_fc1 = offload_expert_fc1 - module.offload_moe_act = offload_moe_act + module.offload_expert_fc1 = True + module.offload_moe_act = True common = dict( device="cuda", dtype=torch.bfloat16, @@ -746,8 +736,8 @@ def test_make_fused_ops_sets_fine_grained_offload_opt_out_attrs( ops = module._make_fused_ops() - assert ops[0].no_offload_activation is expected_fc1_no_offload - assert ops[1].no_offload_activation is expected_act_no_offload + assert not hasattr(ops[0], "no_offload_activation") + assert not hasattr(ops[1], "no_offload_activation") assert not hasattr(ops[2], "no_offload_activation") From db02c4781881d690a27f1a10dbb83482a5589b87 Mon Sep 17 00:00:00 2001 From: hongbinl Date: Tue, 9 Jun 2026 03:26:38 -0700 Subject: [PATCH 11/26] Use TE op offload opt-out API Signed-off-by: hongbinl --- .../core/extensions/transformer_engine.py | 7 ++- megatron/core/transformer/moe/experts.py | 15 +++++- .../transformer/moe/test_grouped_mlp.py | 54 +++++++++++++++---- 3 files changed, 64 insertions(+), 12 deletions(-) diff --git a/megatron/core/extensions/transformer_engine.py b/megatron/core/extensions/transformer_engine.py index 122fdc34b23..11f8d0a2a48 100644 --- a/megatron/core/extensions/transformer_engine.py +++ b/megatron/core/extensions/transformer_engine.py @@ -88,7 +88,12 @@ def fused_grouped_mlp_activation_offload_supported() -> bool: """Return whether TE fused grouped MLP handles fine-grained activation offload.""" - return HAVE_TE and is_te_min_version("2.17") + if not HAVE_TE or not is_te_min_version("2.17"): + return False + grouped_linear_cls = getattr(getattr(te.pytorch, "ops", None), "GroupedLinear", None) + return grouped_linear_cls is not None and hasattr( + grouped_linear_cls, "disable_cpu_offloading" + ) class TransformerEngineConfigType(enum.Enum): diff --git a/megatron/core/transformer/moe/experts.py b/megatron/core/transformer/moe/experts.py index f63f1217479..c07b3bfe0d3 100644 --- a/megatron/core/transformer/moe/experts.py +++ b/megatron/core/transformer/moe/experts.py @@ -335,7 +335,7 @@ def _is_fused_impl_supported(self) -> bool: if ( getattr(self, "offload_expert_fc1", False) or getattr(self, "offload_moe_act", False) ) and not fused_grouped_mlp_activation_offload_supported(): - return False # TE fused grouped MLP offload support requires TE >= 2.17. + return False # TE fused grouped MLP selective offload support requires a TE opt-out API. if self.config.moe_apply_probs_on_input: return False # Pre-multiplying probs is not supported @@ -432,6 +432,7 @@ def _make_fused_ops(self) -> torch.nn.Module: setattr(op, f"bias{idx}", getattr(self.linear_fc1, f"bias{idx}")) if self.linear_fc1.use_bias and fc1_single_grouped_bias: setattr(op, "bias", getattr(self.linear_fc1, "bias")) + fc1_op = op ops.append(op) # Activation and post-multiply probs (SwiGLU, clamped quick-GeGLU, or SReLU) @@ -494,6 +495,7 @@ def _make_fused_ops(self) -> torch.nn.Module: "_make_fused_ops expected SwiGLU, quick_gelu, or weighted squared_relu; " "call _is_fused_impl_supported() before constructing fused ops." ) + activation_op = op ops.append(op) # FC2 @@ -521,8 +523,19 @@ def _make_fused_ops(self) -> torch.nn.Module: setattr(op, f"bias{idx}", getattr(self.linear_fc2, f"bias{idx}")) if self.linear_fc2.use_bias and fc2_single_grouped_bias: setattr(op, "bias", getattr(self.linear_fc2, "bias")) + fc2_op = op ops.append(op) + fine_grained_activation_offloading = getattr( + self, "offload_expert_fc1", False + ) or getattr(self, "offload_moe_act", False) + if fine_grained_activation_offloading: + if not getattr(self, "offload_expert_fc1", False): + fc1_op.disable_cpu_offloading() + if not getattr(self, "offload_moe_act", False): + activation_op.disable_cpu_offloading() + fc2_op.disable_cpu_offloading() + # Emulate submodule pre-forward hooks ops.register_forward_pre_hook(self._make_fused_impl_pre_forward_hook()) diff --git a/tests/unit_tests/transformer/moe/test_grouped_mlp.py b/tests/unit_tests/transformer/moe/test_grouped_mlp.py index 21af088e258..bd8645ab2a0 100644 --- a/tests/unit_tests/transformer/moe/test_grouped_mlp.py +++ b/tests/unit_tests/transformer/moe/test_grouped_mlp.py @@ -357,7 +357,11 @@ def register_forward_pre_hook(self, hook): def _make_fake_te_namespace(): """Build a fake TE namespace with the activation classes _make_fused_ops uses.""" - class FakeGroupedLinear(torch.nn.Module): + class FakeCpuOffloadControl: + def disable_cpu_offloading(self, disabled=True): + self.cpu_offloading_disabled = disabled + + class FakeGroupedLinear(FakeCpuOffloadControl, torch.nn.Module): def __init__( self, num_gemms, @@ -383,27 +387,31 @@ def __init__( self.single_grouped_weight = single_grouped_weight self.single_grouped_bias = single_grouped_bias self.delay_wgrad_compute = delay_wgrad_compute + self.cpu_offloading_disabled = False def need_backward_dw(self): return False - class FakeScaledSwiGLU(torch.nn.Module): + class FakeScaledSwiGLU(FakeCpuOffloadControl, torch.nn.Module): def __init__(self, glu_interleave_size, *, activation_recompute_in_mlp=False): super().__init__() self.glu_interleave_size = glu_interleave_size self.activation_recompute_in_mlp = activation_recompute_in_mlp + self.cpu_offloading_disabled = False - class FakeScaledClampedQGeGLU(torch.nn.Module): + class FakeScaledClampedQGeGLU(FakeCpuOffloadControl, torch.nn.Module): def __init__(self, glu_interleave_size, *, activation_recompute_in_mlp=False, limit=None): super().__init__() self.glu_interleave_size = glu_interleave_size self.activation_recompute_in_mlp = activation_recompute_in_mlp self.limit = limit + self.cpu_offloading_disabled = False - class FakeScaledSReLU(torch.nn.Module): + class FakeScaledSReLU(FakeCpuOffloadControl, torch.nn.Module): def __init__(self, *, activation_recompute_in_mlp=False): super().__init__() self.activation_recompute_in_mlp = activation_recompute_in_mlp + self.cpu_offloading_disabled = False class FakeSequential(list): def register_forward_pre_hook(self, hook): @@ -647,6 +655,24 @@ def fake_is_te_min_version(version): assert checked_versions == ["2.17"] +def test_fused_grouped_mlp_activation_offload_requires_te_opt_out_api(monkeypatch): + import megatron.core.extensions.transformer_engine as te_ext + + class FakeGroupedLinearWithoutOptOut: + pass + + fake_te = SimpleNamespace( + pytorch=SimpleNamespace( + ops=SimpleNamespace(GroupedLinear=FakeGroupedLinearWithoutOptOut) + ) + ) + monkeypatch.setattr(te_ext, "te", fake_te) + monkeypatch.setattr(te_ext, "HAVE_TE", True) + monkeypatch.setattr(te_ext, "is_te_min_version", lambda _: True) + + assert te_ext.fused_grouped_mlp_activation_offload_supported() is False + + def test_is_fused_impl_supported_rejects_offload_without_te_217(monkeypatch): fake_te, FakeGroupedLinear = _make_fake_te_namespace() monkeypatch.setattr(experts_module, "te", fake_te) @@ -704,7 +730,17 @@ def test_make_fused_ops_attaches_single_grouped_bias_for_fc1(monkeypatch): ), "bias should not be split into bias{idx} when single_grouped_bias=True" -def test_make_fused_ops_does_not_set_deprecated_offload_opt_out_attrs(monkeypatch): +@pytest.mark.parametrize( + ("offload_expert_fc1", "offload_moe_act", "expected_disabled"), + ( + (True, False, (False, True, True)), + (False, True, (True, False, True)), + (True, True, (False, False, True)), + ), +) +def test_make_fused_ops_configures_te_cpu_offload_opt_out( + monkeypatch, offload_expert_fc1, offload_moe_act, expected_disabled +): fake_te, FakeGroupedLinear = _make_fake_te_namespace() monkeypatch.setattr(experts_module, "te", fake_te) @@ -719,8 +755,8 @@ def test_make_fused_ops_does_not_set_deprecated_offload_opt_out_attrs(monkeypatc ) module.activation_func = F.silu module.activation_recompute = False - module.offload_expert_fc1 = True - module.offload_moe_act = True + module.offload_expert_fc1 = offload_expert_fc1 + module.offload_moe_act = offload_moe_act common = dict( device="cuda", dtype=torch.bfloat16, @@ -736,9 +772,7 @@ def test_make_fused_ops_does_not_set_deprecated_offload_opt_out_attrs(monkeypatc ops = module._make_fused_ops() - assert not hasattr(ops[0], "no_offload_activation") - assert not hasattr(ops[1], "no_offload_activation") - assert not hasattr(ops[2], "no_offload_activation") + assert tuple(op.cpu_offloading_disabled for op in ops) == expected_disabled def test_backward_dw_dispatches_fused_children_in_fc2_then_fc1_order(): From 18164a72ca6260d606b2c1fcf6816ec0dbd69ee5 Mon Sep 17 00:00:00 2001 From: hongbinl Date: Tue, 9 Jun 2026 05:42:37 -0700 Subject: [PATCH 12/26] Rename TE activation offload API usage Signed-off-by: hongbinl --- .../core/extensions/transformer_engine.py | 2 +- megatron/core/transformer/moe/experts.py | 6 ++--- .../transformer/moe/test_grouped_mlp.py | 26 +++++++++---------- 3 files changed, 17 insertions(+), 17 deletions(-) diff --git a/megatron/core/extensions/transformer_engine.py b/megatron/core/extensions/transformer_engine.py index 11f8d0a2a48..a914973a44a 100644 --- a/megatron/core/extensions/transformer_engine.py +++ b/megatron/core/extensions/transformer_engine.py @@ -92,7 +92,7 @@ def fused_grouped_mlp_activation_offload_supported() -> bool: return False grouped_linear_cls = getattr(getattr(te.pytorch, "ops", None), "GroupedLinear", None) return grouped_linear_cls is not None and hasattr( - grouped_linear_cls, "disable_cpu_offloading" + grouped_linear_cls, "disable_activation_offloading" ) diff --git a/megatron/core/transformer/moe/experts.py b/megatron/core/transformer/moe/experts.py index c07b3bfe0d3..9ab17461301 100644 --- a/megatron/core/transformer/moe/experts.py +++ b/megatron/core/transformer/moe/experts.py @@ -531,10 +531,10 @@ def _make_fused_ops(self) -> torch.nn.Module: ) or getattr(self, "offload_moe_act", False) if fine_grained_activation_offloading: if not getattr(self, "offload_expert_fc1", False): - fc1_op.disable_cpu_offloading() + fc1_op.disable_activation_offloading() if not getattr(self, "offload_moe_act", False): - activation_op.disable_cpu_offloading() - fc2_op.disable_cpu_offloading() + activation_op.disable_activation_offloading() + fc2_op.disable_activation_offloading() # Emulate submodule pre-forward hooks ops.register_forward_pre_hook(self._make_fused_impl_pre_forward_hook()) diff --git a/tests/unit_tests/transformer/moe/test_grouped_mlp.py b/tests/unit_tests/transformer/moe/test_grouped_mlp.py index bd8645ab2a0..8162cb14c32 100644 --- a/tests/unit_tests/transformer/moe/test_grouped_mlp.py +++ b/tests/unit_tests/transformer/moe/test_grouped_mlp.py @@ -358,8 +358,8 @@ def _make_fake_te_namespace(): """Build a fake TE namespace with the activation classes _make_fused_ops uses.""" class FakeCpuOffloadControl: - def disable_cpu_offloading(self, disabled=True): - self.cpu_offloading_disabled = disabled + def disable_activation_offloading(self, disabled=True): + self.activation_offloading = not disabled class FakeGroupedLinear(FakeCpuOffloadControl, torch.nn.Module): def __init__( @@ -387,7 +387,7 @@ def __init__( self.single_grouped_weight = single_grouped_weight self.single_grouped_bias = single_grouped_bias self.delay_wgrad_compute = delay_wgrad_compute - self.cpu_offloading_disabled = False + self.activation_offloading = True def need_backward_dw(self): return False @@ -397,7 +397,7 @@ def __init__(self, glu_interleave_size, *, activation_recompute_in_mlp=False): super().__init__() self.glu_interleave_size = glu_interleave_size self.activation_recompute_in_mlp = activation_recompute_in_mlp - self.cpu_offloading_disabled = False + self.activation_offloading = True class FakeScaledClampedQGeGLU(FakeCpuOffloadControl, torch.nn.Module): def __init__(self, glu_interleave_size, *, activation_recompute_in_mlp=False, limit=None): @@ -405,13 +405,13 @@ def __init__(self, glu_interleave_size, *, activation_recompute_in_mlp=False, li self.glu_interleave_size = glu_interleave_size self.activation_recompute_in_mlp = activation_recompute_in_mlp self.limit = limit - self.cpu_offloading_disabled = False + self.activation_offloading = True class FakeScaledSReLU(FakeCpuOffloadControl, torch.nn.Module): def __init__(self, *, activation_recompute_in_mlp=False): super().__init__() self.activation_recompute_in_mlp = activation_recompute_in_mlp - self.cpu_offloading_disabled = False + self.activation_offloading = True class FakeSequential(list): def register_forward_pre_hook(self, hook): @@ -731,15 +731,15 @@ def test_make_fused_ops_attaches_single_grouped_bias_for_fc1(monkeypatch): @pytest.mark.parametrize( - ("offload_expert_fc1", "offload_moe_act", "expected_disabled"), + ("offload_expert_fc1", "offload_moe_act", "expected_activation_offloading"), ( - (True, False, (False, True, True)), - (False, True, (True, False, True)), - (True, True, (False, False, True)), + (True, False, (True, False, False)), + (False, True, (False, True, False)), + (True, True, (True, True, False)), ), ) -def test_make_fused_ops_configures_te_cpu_offload_opt_out( - monkeypatch, offload_expert_fc1, offload_moe_act, expected_disabled +def test_make_fused_ops_configures_te_activation_offload_opt_out( + monkeypatch, offload_expert_fc1, offload_moe_act, expected_activation_offloading ): fake_te, FakeGroupedLinear = _make_fake_te_namespace() monkeypatch.setattr(experts_module, "te", fake_te) @@ -772,7 +772,7 @@ def test_make_fused_ops_configures_te_cpu_offload_opt_out( ops = module._make_fused_ops() - assert tuple(op.cpu_offloading_disabled for op in ops) == expected_disabled + assert tuple(op.activation_offloading for op in ops) == expected_activation_offloading def test_backward_dw_dispatches_fused_children_in_fc2_then_fc1_order(): From 217ffca8fb11aeb689c69c3a52ee64bcbadbc3c7 Mon Sep 17 00:00:00 2001 From: hongbinl Date: Tue, 9 Jun 2026 06:01:27 -0700 Subject: [PATCH 13/26] Use TE activation offload policy setter Signed-off-by: hongbinl --- megatron/core/extensions/transformer_engine.py | 2 +- megatron/core/transformer/moe/experts.py | 6 +++--- tests/unit_tests/transformer/moe/test_grouped_mlp.py | 4 ++-- 3 files changed, 6 insertions(+), 6 deletions(-) diff --git a/megatron/core/extensions/transformer_engine.py b/megatron/core/extensions/transformer_engine.py index a914973a44a..04f3b82852c 100644 --- a/megatron/core/extensions/transformer_engine.py +++ b/megatron/core/extensions/transformer_engine.py @@ -92,7 +92,7 @@ def fused_grouped_mlp_activation_offload_supported() -> bool: return False grouped_linear_cls = getattr(getattr(te.pytorch, "ops", None), "GroupedLinear", None) return grouped_linear_cls is not None and hasattr( - grouped_linear_cls, "disable_activation_offloading" + grouped_linear_cls, "set_activation_offloading" ) diff --git a/megatron/core/transformer/moe/experts.py b/megatron/core/transformer/moe/experts.py index 9ab17461301..01c540e254d 100644 --- a/megatron/core/transformer/moe/experts.py +++ b/megatron/core/transformer/moe/experts.py @@ -531,10 +531,10 @@ def _make_fused_ops(self) -> torch.nn.Module: ) or getattr(self, "offload_moe_act", False) if fine_grained_activation_offloading: if not getattr(self, "offload_expert_fc1", False): - fc1_op.disable_activation_offloading() + fc1_op.set_activation_offloading(False) if not getattr(self, "offload_moe_act", False): - activation_op.disable_activation_offloading() - fc2_op.disable_activation_offloading() + activation_op.set_activation_offloading(False) + fc2_op.set_activation_offloading(False) # Emulate submodule pre-forward hooks ops.register_forward_pre_hook(self._make_fused_impl_pre_forward_hook()) diff --git a/tests/unit_tests/transformer/moe/test_grouped_mlp.py b/tests/unit_tests/transformer/moe/test_grouped_mlp.py index 8162cb14c32..a20b6b34f8c 100644 --- a/tests/unit_tests/transformer/moe/test_grouped_mlp.py +++ b/tests/unit_tests/transformer/moe/test_grouped_mlp.py @@ -358,8 +358,8 @@ def _make_fake_te_namespace(): """Build a fake TE namespace with the activation classes _make_fused_ops uses.""" class FakeCpuOffloadControl: - def disable_activation_offloading(self, disabled=True): - self.activation_offloading = not disabled + def set_activation_offloading(self, enabled): + self.activation_offloading = enabled class FakeGroupedLinear(FakeCpuOffloadControl, torch.nn.Module): def __init__( From 2d553a77751c5ed02a9db0977e189c4fe10dcf7a Mon Sep 17 00:00:00 2001 From: hongbinl Date: Tue, 9 Jun 2026 06:37:13 -0700 Subject: [PATCH 14/26] Skip non-offloadable activation tensors Signed-off-by: hongbinl --- .../fine_grained_activation_offload.py | 21 ++++++-- ...test_fine_grained_activation_offloading.py | 49 +++++++++++++++++++ 2 files changed, 67 insertions(+), 3 deletions(-) diff --git a/megatron/core/pipeline_parallel/fine_grained_activation_offload.py b/megatron/core/pipeline_parallel/fine_grained_activation_offload.py index a3cf97ea1bf..59892f47604 100644 --- a/megatron/core/pipeline_parallel/fine_grained_activation_offload.py +++ b/megatron/core/pipeline_parallel/fine_grained_activation_offload.py @@ -852,8 +852,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 +862,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 +882,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 +899,8 @@ 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 getattr(tensor, "_TE_do_not_offload", False): return False if tensor.numel() < self.min_offloaded_tensor_size: 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..2aa9aaa7ffb 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 @@ -11,6 +11,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, FineGrainedActivationOffloadingInterface as off_interface, ) from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed @@ -32,6 +33,54 @@ 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 offloadability checks." +) +def test_chunk_offload_handler_respects_tensor_offload_opt_out_flags(): + 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, From cfe0d08c2b393a0cc941882d62a86bb3ed6c2dfe Mon Sep 17 00:00:00 2001 From: hongbinl Date: Thu, 11 Jun 2026 23:02:53 -0700 Subject: [PATCH 15/26] Add fused grouped MLP offload module Signed-off-by: hongbinl --- .../fine_grained_activation_offloading.md | 4 +- docs/user-guide/features/paged_stash.md | 2 +- .../core/extensions/transformer_engine.py | 10 -- .../fine_grained_activation_offload.py | 4 +- megatron/core/transformer/moe/README.md | 2 +- megatron/core/transformer/moe/experts.py | 43 ++--- .../core/transformer/transformer_config.py | 28 +++- .../transformer/moe/test_grouped_mlp.py | 150 +++++++----------- 8 files changed, 99 insertions(+), 144 deletions(-) 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/extensions/transformer_engine.py b/megatron/core/extensions/transformer_engine.py index 04f3b82852c..b84565dd1f3 100644 --- a/megatron/core/extensions/transformer_engine.py +++ b/megatron/core/extensions/transformer_engine.py @@ -86,16 +86,6 @@ _TE_CONFIG_TYPE_KEY = "transformer_engine_config_type" -def fused_grouped_mlp_activation_offload_supported() -> bool: - """Return whether TE fused grouped MLP handles fine-grained activation offload.""" - if not HAVE_TE or not is_te_min_version("2.17"): - return False - grouped_linear_cls = getattr(getattr(te.pytorch, "ops", None), "GroupedLinear", None) - return grouped_linear_cls is not None and hasattr( - grouped_linear_cls, "set_activation_offloading" - ) - - class TransformerEngineConfigType(enum.Enum): """Configuration object types in config dictionary""" diff --git a/megatron/core/pipeline_parallel/fine_grained_activation_offload.py b/megatron/core/pipeline_parallel/fine_grained_activation_offload.py index 59892f47604..dc6495885d7 100644 --- a/megatron/core/pipeline_parallel/fine_grained_activation_offload.py +++ b/megatron/core/pipeline_parallel/fine_grained_activation_offload.py @@ -343,9 +343,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 in ("expert_fc1", "moe_act", "expert_fc1_moe_act"): + if name in ("expert_fc1", "moe_act", "expert_fc1_moe_act", "fused_group_mlp"): self.use_cpu_pool = False else: self.use_cpu_pool = True 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 01c540e254d..bde1737fce0 100644 --- a/megatron/core/transformer/moe/experts.py +++ b/megatron/core/transformer/moe/experts.py @@ -18,10 +18,7 @@ from megatron.core.activations import squared_relu from megatron.core.dist_checkpointing.mapping import ShardedStateDict from megatron.core.dist_checkpointing.utils import replace_prefix_for_sharding -from megatron.core.extensions.transformer_engine import ( - HAVE_TE, - fused_grouped_mlp_activation_offload_supported, -) +from megatron.core.extensions.transformer_engine import HAVE_TE from megatron.core.fusions.fused_bias_geglu import quick_gelu, weighted_bias_quick_geglu_impl from megatron.core.fusions.fused_bias_swiglu import weighted_bias_swiglu_impl from megatron.core.fusions.fused_weighted_squared_relu import weighted_squared_relu_impl @@ -253,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 @@ -332,10 +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 ( - getattr(self, "offload_expert_fc1", False) or getattr(self, "offload_moe_act", False) - ) and not fused_grouped_mlp_activation_offload_supported(): - return False # TE fused grouped MLP selective offload support requires a TE opt-out API. + 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 @@ -432,7 +432,6 @@ def _make_fused_ops(self) -> torch.nn.Module: setattr(op, f"bias{idx}", getattr(self.linear_fc1, f"bias{idx}")) if self.linear_fc1.use_bias and fc1_single_grouped_bias: setattr(op, "bias", getattr(self.linear_fc1, "bias")) - fc1_op = op ops.append(op) # Activation and post-multiply probs (SwiGLU, clamped quick-GeGLU, or SReLU) @@ -495,7 +494,6 @@ def _make_fused_ops(self) -> torch.nn.Module: "_make_fused_ops expected SwiGLU, quick_gelu, or weighted squared_relu; " "call _is_fused_impl_supported() before constructing fused ops." ) - activation_op = op ops.append(op) # FC2 @@ -523,19 +521,8 @@ def _make_fused_ops(self) -> torch.nn.Module: setattr(op, f"bias{idx}", getattr(self.linear_fc2, f"bias{idx}")) if self.linear_fc2.use_bias and fc2_single_grouped_bias: setattr(op, "bias", getattr(self.linear_fc2, "bias")) - fc2_op = op ops.append(op) - fine_grained_activation_offloading = getattr( - self, "offload_expert_fc1", False - ) or getattr(self, "offload_moe_act", False) - if fine_grained_activation_offloading: - if not getattr(self, "offload_expert_fc1", False): - fc1_op.set_activation_offloading(False) - if not getattr(self, "offload_moe_act", False): - activation_op.set_activation_offloading(False) - fc2_op.set_activation_offloading(False) - # Emulate submodule pre-forward hooks ops.register_forward_pre_hook(self._make_fused_impl_pre_forward_hook()) @@ -621,18 +608,14 @@ def _fused_forward( ) else: stash_context = nullcontext() - offload_expert_fc1 = getattr(self, "offload_expert_fc1", False) - offload_moe_act = getattr(self, "offload_moe_act", False) - fine_grained_activation_offloading = offload_expert_fc1 or offload_moe_act - offload_name = "_".join( - name - for name, enabled in (("expert_fc1", offload_expert_fc1), ("moe_act", offload_moe_act)) - if enabled - ) + 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 offload_expert_fc1 else [] + forced_released_tensors = ( + [permuted_local_hidden_states] if fine_grained_activation_offloading else [] + ) with stash_context: # Call fused impl output = ops( diff --git a/megatron/core/transformer/transformer_config.py b/megatron/core/transformer/transformer_config.py index 3ba0dcff0fa..bf0adc70461 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,17 @@ 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 +1718,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}" ) @@ -2370,7 +2386,7 @@ def _scope_to_str(s): local_partial_moe_offload = ( self.cuda_graph_impl == "local" and bool(offload_modules) - and offload_modules <= {"expert_fc1", "moe_act"} + and offload_modules <= {"expert_fc1", "moe_act", "fused_group_mlp"} and CudaGraphModule.moe not in self.cuda_graph_modules ) assert ( @@ -2380,8 +2396,8 @@ def _scope_to_str(s): "fine-grained activation offloading is only supported with " "transformer_engine CUDA graph implementation or local CUDA graph " "implementation with full_iteration scope. Local partial CUDA graphs " - "are supported only for expert_fc1/moe_act offload when the full MoE " - "module is not captured." + "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/transformer/moe/test_grouped_mlp.py b/tests/unit_tests/transformer/moe/test_grouped_mlp.py index a20b6b34f8c..7f0448592d8 100644 --- a/tests/unit_tests/transformer/moe/test_grouped_mlp.py +++ b/tests/unit_tests/transformer/moe/test_grouped_mlp.py @@ -357,11 +357,7 @@ def register_forward_pre_hook(self, hook): def _make_fake_te_namespace(): """Build a fake TE namespace with the activation classes _make_fused_ops uses.""" - class FakeCpuOffloadControl: - def set_activation_offloading(self, enabled): - self.activation_offloading = enabled - - class FakeGroupedLinear(FakeCpuOffloadControl, torch.nn.Module): + class FakeGroupedLinear(torch.nn.Module): def __init__( self, num_gemms, @@ -387,31 +383,27 @@ def __init__( self.single_grouped_weight = single_grouped_weight self.single_grouped_bias = single_grouped_bias self.delay_wgrad_compute = delay_wgrad_compute - self.activation_offloading = True def need_backward_dw(self): return False - class FakeScaledSwiGLU(FakeCpuOffloadControl, torch.nn.Module): + class FakeScaledSwiGLU(torch.nn.Module): def __init__(self, glu_interleave_size, *, activation_recompute_in_mlp=False): super().__init__() self.glu_interleave_size = glu_interleave_size self.activation_recompute_in_mlp = activation_recompute_in_mlp - self.activation_offloading = True - class FakeScaledClampedQGeGLU(FakeCpuOffloadControl, torch.nn.Module): + class FakeScaledClampedQGeGLU(torch.nn.Module): def __init__(self, glu_interleave_size, *, activation_recompute_in_mlp=False, limit=None): super().__init__() self.glu_interleave_size = glu_interleave_size self.activation_recompute_in_mlp = activation_recompute_in_mlp self.limit = limit - self.activation_offloading = True - class FakeScaledSReLU(FakeCpuOffloadControl, torch.nn.Module): + class FakeScaledSReLU(torch.nn.Module): def __init__(self, *, activation_recompute_in_mlp=False): super().__init__() self.activation_recompute_in_mlp = activation_recompute_in_mlp - self.activation_offloading = True class FakeSequential(list): def register_forward_pre_hook(self, hook): @@ -574,6 +566,7 @@ def _make_fused_impl_support_module( module.tp_group = SimpleNamespace(size=lambda: 1) module.offload_expert_fc1 = False module.offload_moe_act = False + module.offload_fused_group_mlp = False common = dict( device="cuda", dtype=torch.bfloat16, @@ -639,56 +632,74 @@ def test_is_fused_impl_supported_requires_scaled_srelu_op(monkeypatch): assert module._is_fused_impl_supported() is False -def test_fused_grouped_mlp_activation_offload_requires_te_217(monkeypatch): - import megatron.core.extensions.transformer_engine as te_ext - - checked_versions = [] - - def fake_is_te_min_version(version): - checked_versions.append(version) - return False - - monkeypatch.setattr(te_ext, "HAVE_TE", True) - monkeypatch.setattr(te_ext, "is_te_min_version", fake_is_te_min_version) - - assert te_ext.fused_grouped_mlp_activation_offload_supported() is False - assert checked_versions == ["2.17"] - - -def test_fused_grouped_mlp_activation_offload_requires_te_opt_out_api(monkeypatch): - import megatron.core.extensions.transformer_engine as te_ext - - class FakeGroupedLinearWithoutOptOut: - pass +@pytest.mark.parametrize("offload_attr", ("offload_expert_fc1", "offload_moe_act")) +def test_is_fused_impl_supported_rejects_partial_moe_offload(monkeypatch, offload_attr): + fake_te, FakeGroupedLinear = _make_fake_te_namespace() + monkeypatch.setattr(experts_module, "te", fake_te) + monkeypatch.setattr(experts_module, "HAVE_TE", True) + monkeypatch.setattr(experts_module, "is_te_min_version", lambda _: True) + _install_fake_te_ops_modules(monkeypatch, fake_te) - fake_te = SimpleNamespace( - pytorch=SimpleNamespace( - ops=SimpleNamespace(GroupedLinear=FakeGroupedLinearWithoutOptOut) - ) + module = _make_fused_impl_support_module( + FakeGroupedLinear, activation_func=F.silu, gated_linear_unit=True ) - monkeypatch.setattr(te_ext, "te", fake_te) - monkeypatch.setattr(te_ext, "HAVE_TE", True) - monkeypatch.setattr(te_ext, "is_te_min_version", lambda _: True) + setattr(module, offload_attr, True) - assert te_ext.fused_grouped_mlp_activation_offload_supported() is False + assert module._is_fused_impl_supported() is False -def test_is_fused_impl_supported_rejects_offload_without_te_217(monkeypatch): +def test_is_fused_impl_supported_allows_fused_group_mlp_offload(monkeypatch): fake_te, FakeGroupedLinear = _make_fake_te_namespace() monkeypatch.setattr(experts_module, "te", fake_te) monkeypatch.setattr(experts_module, "HAVE_TE", True) monkeypatch.setattr(experts_module, "is_te_min_version", lambda _: True) - monkeypatch.setattr( - experts_module, "fused_grouped_mlp_activation_offload_supported", lambda: False - ) _install_fake_te_ops_modules(monkeypatch, fake_te) module = _make_fused_impl_support_module( FakeGroupedLinear, activation_func=F.silu, gated_linear_unit=True ) - module.offload_expert_fc1 = True + module.offload_fused_group_mlp = True - assert module._is_fused_impl_supported() is False + assert module._is_fused_impl_supported() is True + + +def test_transformer_config_allows_fused_group_mlp_offload_module(): + config = TransformerConfig( + num_layers=1, + hidden_size=64, + num_attention_heads=4, + fine_grained_activation_offloading=True, + offload_modules=["fused_group_mlp"], + use_transformer_engine_op_fuser=True, + cuda_graph_impl="transformer_engine", + ) + + assert config.offload_modules == ["fused_group_mlp"] + + +def test_transformer_config_rejects_fused_group_mlp_without_op_fuser(): + with pytest.raises(ValueError, match="fused_group_mlp requires"): + TransformerConfig( + num_layers=1, + hidden_size=64, + num_attention_heads=4, + fine_grained_activation_offloading=True, + offload_modules=["fused_group_mlp"], + cuda_graph_impl="transformer_engine", + ) + + +def test_transformer_config_rejects_mixed_fused_group_mlp_and_partial_moe_offload(): + with pytest.raises(ValueError, match="cannot be combined"): + TransformerConfig( + num_layers=1, + hidden_size=64, + num_attention_heads=4, + fine_grained_activation_offloading=True, + offload_modules=["fused_group_mlp", "expert_fc1"], + use_transformer_engine_op_fuser=True, + cuda_graph_impl="transformer_engine", + ) def test_make_fused_ops_attaches_single_grouped_bias_for_fc1(monkeypatch): @@ -730,51 +741,6 @@ def test_make_fused_ops_attaches_single_grouped_bias_for_fc1(monkeypatch): ), "bias should not be split into bias{idx} when single_grouped_bias=True" -@pytest.mark.parametrize( - ("offload_expert_fc1", "offload_moe_act", "expected_activation_offloading"), - ( - (True, False, (True, False, False)), - (False, True, (False, True, False)), - (True, True, (True, True, False)), - ), -) -def test_make_fused_ops_configures_te_activation_offload_opt_out( - monkeypatch, offload_expert_fc1, offload_moe_act, expected_activation_offloading -): - fake_te, FakeGroupedLinear = _make_fake_te_namespace() - monkeypatch.setattr(experts_module, "te", fake_te) - - module = TEGroupedMLP.__new__(TEGroupedMLP) - torch.nn.Module.__init__(module) - module.config = SimpleNamespace( - moe_mlp_glu_interleave_size=2, - delay_wgrad_compute=False, - activation_func_clamp_value=None, - activation_func=F.silu, - gated_linear_unit=True, - ) - module.activation_func = F.silu - module.activation_recompute = False - module.offload_expert_fc1 = offload_expert_fc1 - module.offload_moe_act = offload_moe_act - common = dict( - device="cuda", - dtype=torch.bfloat16, - accumulate_into_main_grad=False, - single_grouped_weight=False, - ) - module.linear_fc1 = FakeGroupedLinear(2, 4, 8, bias=False, **common) - module.linear_fc2 = FakeGroupedLinear(2, 8, 4, bias=False, **common) - module.linear_fc1.weight0 = torch.nn.Parameter(torch.ones(8, 4)) - module.linear_fc1.weight1 = torch.nn.Parameter(torch.ones(8, 4)) - module.linear_fc2.weight0 = torch.nn.Parameter(torch.ones(4, 8)) - module.linear_fc2.weight1 = torch.nn.Parameter(torch.ones(4, 8)) - - ops = module._make_fused_ops() - - assert tuple(op.activation_offloading for op in ops) == expected_activation_offloading - - def test_backward_dw_dispatches_fused_children_in_fc2_then_fc1_order(): """delay_wgrad_compute=True (via wrapper) → backward_dw calls fused [2] then [0], then triggers the wrapper-side wgrad/reduce hooks (PR 4311) so DDP's reduce-scatter From 4c3a50fcf36e4a467be7d3a7c482ce3e0447a6e6 Mon Sep 17 00:00:00 2001 From: hongbinl Date: Thu, 11 Jun 2026 23:45:17 -0700 Subject: [PATCH 16/26] Clean up fused group MLP offload follow-ups Signed-off-by: hongbinl --- .../fine_grained_activation_offload.py | 4 +--- .../dist_checkpointing/test_integrity.py | 24 +++++++------------ ...test_fine_grained_activation_offloading.py | 10 ++------ 3 files changed, 12 insertions(+), 26 deletions(-) diff --git a/megatron/core/pipeline_parallel/fine_grained_activation_offload.py b/megatron/core/pipeline_parallel/fine_grained_activation_offload.py index dc6495885d7..5188c53ac66 100644 --- a/megatron/core/pipeline_parallel/fine_grained_activation_offload.py +++ b/megatron/core/pipeline_parallel/fine_grained_activation_offload.py @@ -345,7 +345,7 @@ def __init__(self, name): # Using memory pool is for the compatibility with cuda graph. # Shapes of tensors for MoE activation offload groups are not known in advance, # so we do not use CPU pool for them. - if name in ("expert_fc1", "moe_act", "expert_fc1_moe_act", "fused_group_mlp"): + if name in ("expert_fc1", "moe_act", "fused_group_mlp"): self.use_cpu_pool = False else: self.use_cpu_pool = True @@ -901,8 +901,6 @@ def tensor_need_offloading_checker(self, tensor): ) if not self._can_manage_tensor_for_offload(tensor): return False - if getattr(tensor, "_TE_do_not_offload", False): - return False if tensor.numel() < self.min_offloaded_tensor_size: return False # Respect tensor's offload preference if specified diff --git a/tests/unit_tests/dist_checkpointing/test_integrity.py b/tests/unit_tests/dist_checkpointing/test_integrity.py index 10aaa4110ae..e87af62af93 100644 --- a/tests/unit_tests/dist_checkpointing/test_integrity.py +++ b/tests/unit_tests/dist_checkpointing/test_integrity.py @@ -64,11 +64,9 @@ def test_save_verify_integrity_manifest_directly(self, init_model_parallel, tmp_ tmp_path_dist_ckpt / 'test_save_integrity_manifest_directly', sync=True ) as ckpt_dir: metadata_file = Path(ckpt_dir / "metadata.json") - if torch.distributed.get_rank() == 0: - with open(metadata_file, "w") as f: - data = {"test_metadata": 1} - json.dump(data, f) - torch.distributed.barrier() + with open(metadata_file, "w") as f: + data = {"test_metadata": 1} + json.dump(data, f) if torch.distributed.get_rank() == 0: save_integrity_manifest(ckpt_dir) @@ -91,21 +89,17 @@ def test_save_verify_integrity_manifest_error(self, init_model_parallel, tmp_pat ) as ckpt_dir: metadata_file = Path(ckpt_dir / "metadata.json") - if torch.distributed.get_rank() == 0: - with open(metadata_file, "w") as f: - data = {"test_metadata": 1} - json.dump(data, f) - torch.distributed.barrier() + with open(metadata_file, "w") as f: + data = {"test_metadata": 1} + json.dump(data, f) if torch.distributed.get_rank() == 0: save_integrity_manifest(ckpt_dir) torch.distributed.barrier() - if torch.distributed.get_rank() == 0: - with open(metadata_file, "w") as f: - data = {"test_metadata": 11} - json.dump(data, f) - torch.distributed.barrier() + with open(metadata_file, "w") as f: + data = {"test_metadata": 11} + json.dump(data, f) # CheckpointingException, hash mismatch with pytest.raises(CheckpointingException): 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 2aa9aaa7ffb..f71d3ad17e1 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 @@ -64,19 +64,13 @@ def test_chunk_offload_handler_skips_non_offloadable_tensor_types(): assert handler.tensor_pop(fake_tensor) is fake_tensor -@pytest.mark.skipif( - not torch.cuda.is_available(), reason="CUDA is required for offloadability checks." -) -def test_chunk_offload_handler_respects_tensor_offload_opt_out_flags(): +@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) From ebf631c6197c802566d59c60873c953d89631c6e Mon Sep 17 00:00:00 2001 From: hongbinl Date: Thu, 11 Jun 2026 23:49:34 -0700 Subject: [PATCH 17/26] Respect TE activation offload opt-out marker Signed-off-by: hongbinl --- .../core/pipeline_parallel/fine_grained_activation_offload.py | 2 ++ .../test_fine_grained_activation_offloading.py | 4 ++++ 2 files changed, 6 insertions(+) diff --git a/megatron/core/pipeline_parallel/fine_grained_activation_offload.py b/megatron/core/pipeline_parallel/fine_grained_activation_offload.py index 5188c53ac66..26609be2745 100644 --- a/megatron/core/pipeline_parallel/fine_grained_activation_offload.py +++ b/megatron/core/pipeline_parallel/fine_grained_activation_offload.py @@ -901,6 +901,8 @@ def tensor_need_offloading_checker(self, tensor): ) if not self._can_manage_tensor_for_offload(tensor): return False + if getattr(tensor, "_TE_do_not_offload", False): + return False if tensor.numel() < self.min_offloaded_tensor_size: return False # Respect tensor's offload preference if specified 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 f71d3ad17e1..7d54122105b 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 @@ -71,6 +71,10 @@ def test_chunk_offload_handler_respects_tensor_offloading_activation_opt_out(): 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) From 2d6f25fb410fc3fb724b88ec53459a900deed9ca Mon Sep 17 00:00:00 2001 From: hongbinl Date: Fri, 12 Jun 2026 00:06:58 -0700 Subject: [PATCH 18/26] Fix transformer config formatting Signed-off-by: hongbinl --- megatron/core/transformer/transformer_config.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/megatron/core/transformer/transformer_config.py b/megatron/core/transformer/transformer_config.py index bf0adc70461..c04b6aba645 100644 --- a/megatron/core/transformer/transformer_config.py +++ b/megatron/core/transformer/transformer_config.py @@ -1701,9 +1701,7 @@ def __post_init__(self): ) 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." - ) + 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( From 8014953c8e480a9c46bcd6ba1a0717c947f67658 Mon Sep 17 00:00:00 2001 From: hongbinl Date: Fri, 12 Jun 2026 00:14:59 -0700 Subject: [PATCH 19/26] Fix activation offload test import order Signed-off-by: hongbinl --- .../test_fine_grained_activation_offloading.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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 7d54122105b..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,8 +10,8 @@ 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 ( - ChunkOffloadHandler, FineGrainedActivationOffloadingInterface as off_interface, ) from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed From b389f3e67a3b82414596af8b9defae7266f5cff0 Mon Sep 17 00:00:00 2001 From: hongbinl Date: Fri, 12 Jun 2026 06:01:48 -0700 Subject: [PATCH 20/26] test: remove grouped mlp offload test changes Signed-off-by: hongbinl --- .../transformer/moe/test_grouped_mlp.py | 71 ------------------- 1 file changed, 71 deletions(-) diff --git a/tests/unit_tests/transformer/moe/test_grouped_mlp.py b/tests/unit_tests/transformer/moe/test_grouped_mlp.py index 7f0448592d8..f6cef08b1b6 100644 --- a/tests/unit_tests/transformer/moe/test_grouped_mlp.py +++ b/tests/unit_tests/transformer/moe/test_grouped_mlp.py @@ -566,7 +566,6 @@ def _make_fused_impl_support_module( module.tp_group = SimpleNamespace(size=lambda: 1) module.offload_expert_fc1 = False module.offload_moe_act = False - module.offload_fused_group_mlp = False common = dict( device="cuda", dtype=torch.bfloat16, @@ -632,76 +631,6 @@ def test_is_fused_impl_supported_requires_scaled_srelu_op(monkeypatch): assert module._is_fused_impl_supported() is False -@pytest.mark.parametrize("offload_attr", ("offload_expert_fc1", "offload_moe_act")) -def test_is_fused_impl_supported_rejects_partial_moe_offload(monkeypatch, offload_attr): - fake_te, FakeGroupedLinear = _make_fake_te_namespace() - monkeypatch.setattr(experts_module, "te", fake_te) - monkeypatch.setattr(experts_module, "HAVE_TE", True) - monkeypatch.setattr(experts_module, "is_te_min_version", lambda _: True) - _install_fake_te_ops_modules(monkeypatch, fake_te) - - module = _make_fused_impl_support_module( - FakeGroupedLinear, activation_func=F.silu, gated_linear_unit=True - ) - setattr(module, offload_attr, True) - - assert module._is_fused_impl_supported() is False - - -def test_is_fused_impl_supported_allows_fused_group_mlp_offload(monkeypatch): - fake_te, FakeGroupedLinear = _make_fake_te_namespace() - monkeypatch.setattr(experts_module, "te", fake_te) - monkeypatch.setattr(experts_module, "HAVE_TE", True) - monkeypatch.setattr(experts_module, "is_te_min_version", lambda _: True) - _install_fake_te_ops_modules(monkeypatch, fake_te) - - module = _make_fused_impl_support_module( - FakeGroupedLinear, activation_func=F.silu, gated_linear_unit=True - ) - module.offload_fused_group_mlp = True - - assert module._is_fused_impl_supported() is True - - -def test_transformer_config_allows_fused_group_mlp_offload_module(): - config = TransformerConfig( - num_layers=1, - hidden_size=64, - num_attention_heads=4, - fine_grained_activation_offloading=True, - offload_modules=["fused_group_mlp"], - use_transformer_engine_op_fuser=True, - cuda_graph_impl="transformer_engine", - ) - - assert config.offload_modules == ["fused_group_mlp"] - - -def test_transformer_config_rejects_fused_group_mlp_without_op_fuser(): - with pytest.raises(ValueError, match="fused_group_mlp requires"): - TransformerConfig( - num_layers=1, - hidden_size=64, - num_attention_heads=4, - fine_grained_activation_offloading=True, - offload_modules=["fused_group_mlp"], - cuda_graph_impl="transformer_engine", - ) - - -def test_transformer_config_rejects_mixed_fused_group_mlp_and_partial_moe_offload(): - with pytest.raises(ValueError, match="cannot be combined"): - TransformerConfig( - num_layers=1, - hidden_size=64, - num_attention_heads=4, - fine_grained_activation_offloading=True, - offload_modules=["fused_group_mlp", "expert_fc1"], - use_transformer_engine_op_fuser=True, - cuda_graph_impl="transformer_engine", - ) - - def test_make_fused_ops_attaches_single_grouped_bias_for_fc1(monkeypatch): """single_grouped_bias=True → bias attached as `bias` (not `bias{idx}`).""" fake_te, FakeGroupedLinear = _make_fake_te_namespace() From 181c93a613f366bc7a279f31a2917fa90f9eb910 Mon Sep 17 00:00:00 2001 From: hongbinl Date: Tue, 16 Jun 2026 00:10:58 -0700 Subject: [PATCH 21/26] fix: override stale streams during full CUDA graph capture Signed-off-by: hongbinl --- megatron/core/full_cuda_graph.py | 35 ++++++++++++++----- .../transformer/test_full_cuda_graph.py | 18 +++++++++- 2 files changed, 43 insertions(+), 10 deletions(-) diff --git a/megatron/core/full_cuda_graph.py b/megatron/core/full_cuda_graph.py index abee2bf811e..03f2c3218db 100644 --- a/megatron/core/full_cuda_graph.py +++ b/megatron/core/full_cuda_graph.py @@ -4,6 +4,7 @@ import gc import logging +from contextlib import contextmanager import torch @@ -52,6 +53,21 @@ def get_graph_pool(use_single_mempool): return torch.cuda.graph_pool_handle() +@contextmanager +def _override_stale_capture_stream(): + """Redirect stale autograd stream refs while capturing full-iteration graphs.""" + set_override = getattr(torch.autograd.graph, "set_override_stale_capture_stream", None) + if set_override is None: + yield + return + + set_override(True) + try: + yield + finally: + set_override(False) + + # The below functions traverse through nested data structures (tuples, lists, dicts) # present in src and creates a deep copy where all PyTorch tensors are cloned, # detached from the computation graph, and moved to CUDA device. Non-tensor objects @@ -214,15 +230,16 @@ def __call__(self, *args, **kwargs): FullCudaGraphWrapper.cuda_graph[training_str].register_generator_state(state) torch.cuda.synchronize() capture_stream = get_shared_capture_stream() - with torch.cuda.graph( - FullCudaGraphWrapper.cuda_graph[training_str], - stream=capture_stream, - pool=get_graph_pool(self.use_single_mempool), - capture_error_mode="thread_local", - ): - FullCudaGraphWrapper.result[training_str] = self.forward_backward_func( - *args, **kwargs - ) + with _override_stale_capture_stream(): + with torch.cuda.graph( + FullCudaGraphWrapper.cuda_graph[training_str], + stream=capture_stream, + pool=get_graph_pool(self.use_single_mempool), + capture_error_mode="thread_local", + ): + FullCudaGraphWrapper.result[training_str] = self.forward_backward_func( + *args, **kwargs + ) torch.cuda.synchronize() torch.distributed.barrier() logger.info(f'CUDA graph capture done for {training_str}!!!') diff --git a/tests/unit_tests/transformer/test_full_cuda_graph.py b/tests/unit_tests/transformer/test_full_cuda_graph.py index 312ae467304..0c8f01f5c25 100644 --- a/tests/unit_tests/transformer/test_full_cuda_graph.py +++ b/tests/unit_tests/transformer/test_full_cuda_graph.py @@ -6,7 +6,7 @@ import megatron.core.pipeline_parallel.schedules as schedule from megatron.core import ModelParallelConfig -from megatron.core.full_cuda_graph import FullCudaGraphWrapper +from megatron.core.full_cuda_graph import FullCudaGraphWrapper, _override_stale_capture_stream from megatron.core.tensor_parallel.random import ( HAVE_TE, initialize_rng_tracker, @@ -18,6 +18,22 @@ rank = Utils.rank +def test_override_stale_capture_stream_toggles_when_available(monkeypatch): + calls = [] + + def set_override(enabled): + calls.append(enabled) + + monkeypatch.setattr( + torch.autograd.graph, "set_override_stale_capture_stream", set_override, raising=False + ) + + with _override_stale_capture_stream(): + assert calls == [True] + + assert calls == [True, False] + + @pytest.mark.skipif( not (HAVE_TE and is_te_min_version("1.5.0")), reason="use_te_rng_tracker requires TransformerEngine version >= 1.5", From d966480ca3fbf6c0b048ae6ddab69499430d8b1f Mon Sep 17 00:00:00 2001 From: hongbinl Date: Tue, 16 Jun 2026 01:49:03 -0700 Subject: [PATCH 22/26] fix: honor TE non-offload marks on tensor wrappers Signed-off-by: hongbinl --- .../fine_grained_activation_offload.py | 18 ++++++++++++- ...test_fine_grained_activation_offloading.py | 25 +++++++++++++++++++ 2 files changed, 42 insertions(+), 1 deletion(-) diff --git a/megatron/core/pipeline_parallel/fine_grained_activation_offload.py b/megatron/core/pipeline_parallel/fine_grained_activation_offload.py index 26609be2745..966c885e603 100644 --- a/megatron/core/pipeline_parallel/fine_grained_activation_offload.py +++ b/megatron/core/pipeline_parallel/fine_grained_activation_offload.py @@ -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. @@ -901,7 +917,7 @@ def tensor_need_offloading_checker(self, tensor): ) if not self._can_manage_tensor_for_offload(tensor): return False - if getattr(tensor, "_TE_do_not_offload", False): + if _te_do_not_offload(tensor): return False if tensor.numel() < self.min_offloaded_tensor_size: return False 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 515f6a01ddf..32a290e8512 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 @@ -79,6 +79,31 @@ def test_chunk_offload_handler_respects_tensor_offloading_activation_opt_out(): assert not handler.tensor_need_offloading_checker(tensor) +def test_chunk_offload_handler_respects_te_do_not_offload_on_wrapper_data_tensors(): + handler = _make_chunk_handler_for_offload_checker() + + class TensorLikeWrapper: + device = torch.device("cuda") + + def __init__(self, data_tensors): + self._data_tensors = data_tensors + + def numel(self): + return 1024 + + def get_data_tensors(self): + return self._data_tensors + + class DataTensor: + pass + + wrapper = TensorLikeWrapper([DataTensor(), DataTensor()]) + assert handler.tensor_need_offloading_checker(wrapper) + + wrapper.get_data_tensors()[1]._TE_do_not_offload = True + assert not handler.tensor_need_offloading_checker(wrapper) + + def _build_gpt_model( *, seed: int, From 04a75da8e21829379a1943913b9e9ba75146196a Mon Sep 17 00:00:00 2001 From: hongbinl Date: Tue, 16 Jun 2026 01:52:57 -0700 Subject: [PATCH 23/26] fix: avoid formatting saved tensors during graph capture Signed-off-by: hongbinl --- .../core/pipeline_parallel/fine_grained_activation_offload.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/megatron/core/pipeline_parallel/fine_grained_activation_offload.py b/megatron/core/pipeline_parallel/fine_grained_activation_offload.py index 966c885e603..a96a81208e9 100644 --- a/megatron/core/pipeline_parallel/fine_grained_activation_offload.py +++ b/megatron/core/pipeline_parallel/fine_grained_activation_offload.py @@ -748,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) From 412f21ad98246a164cc76837b632f9be171143d9 Mon Sep 17 00:00:00 2001 From: hongbinl Date: Tue, 16 Jun 2026 02:05:09 -0700 Subject: [PATCH 24/26] test: remove TE wrapper offload unit test Signed-off-by: hongbinl --- ...test_fine_grained_activation_offloading.py | 25 ------------------- 1 file changed, 25 deletions(-) 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 32a290e8512..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 @@ -79,31 +79,6 @@ def test_chunk_offload_handler_respects_tensor_offloading_activation_opt_out(): assert not handler.tensor_need_offloading_checker(tensor) -def test_chunk_offload_handler_respects_te_do_not_offload_on_wrapper_data_tensors(): - handler = _make_chunk_handler_for_offload_checker() - - class TensorLikeWrapper: - device = torch.device("cuda") - - def __init__(self, data_tensors): - self._data_tensors = data_tensors - - def numel(self): - return 1024 - - def get_data_tensors(self): - return self._data_tensors - - class DataTensor: - pass - - wrapper = TensorLikeWrapper([DataTensor(), DataTensor()]) - assert handler.tensor_need_offloading_checker(wrapper) - - wrapper.get_data_tensors()[1]._TE_do_not_offload = True - assert not handler.tensor_need_offloading_checker(wrapper) - - def _build_gpt_model( *, seed: int, From 6a96a564a6b17362f32191342ccac4957251ea81 Mon Sep 17 00:00:00 2001 From: hongbinl Date: Tue, 16 Jun 2026 02:12:38 -0700 Subject: [PATCH 25/26] chore: update copyright headers Signed-off-by: hongbinl --- megatron/core/full_cuda_graph.py | 2 +- .../core/pipeline_parallel/fine_grained_activation_offload.py | 2 +- tests/unit_tests/transformer/test_full_cuda_graph.py | 2 +- 3 files changed, 3 insertions(+), 3 deletions(-) diff --git a/megatron/core/full_cuda_graph.py b/megatron/core/full_cuda_graph.py index 03f2c3218db..0db2e68957b 100644 --- a/megatron/core/full_cuda_graph.py +++ b/megatron/core/full_cuda_graph.py @@ -1,4 +1,4 @@ -# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved. +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. """Full iteration CUDA graph for training.""" diff --git a/megatron/core/pipeline_parallel/fine_grained_activation_offload.py b/megatron/core/pipeline_parallel/fine_grained_activation_offload.py index a96a81208e9..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 diff --git a/tests/unit_tests/transformer/test_full_cuda_graph.py b/tests/unit_tests/transformer/test_full_cuda_graph.py index 0c8f01f5c25..367cdd51d1a 100644 --- a/tests/unit_tests/transformer/test_full_cuda_graph.py +++ b/tests/unit_tests/transformer/test_full_cuda_graph.py @@ -1,4 +1,4 @@ -# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. import pytest import torch From c6dd28511dc075b678b04478c59d24fcc643836e Mon Sep 17 00:00:00 2001 From: hongbinl Date: Tue, 16 Jun 2026 02:20:55 -0700 Subject: [PATCH 26/26] revert: drop full CUDA graph changes Signed-off-by: hongbinl --- megatron/core/full_cuda_graph.py | 37 +++++-------------- .../transformer/test_full_cuda_graph.py | 20 +--------- 2 files changed, 12 insertions(+), 45 deletions(-) diff --git a/megatron/core/full_cuda_graph.py b/megatron/core/full_cuda_graph.py index 0db2e68957b..abee2bf811e 100644 --- a/megatron/core/full_cuda_graph.py +++ b/megatron/core/full_cuda_graph.py @@ -1,10 +1,9 @@ -# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved. """Full iteration CUDA graph for training.""" import gc import logging -from contextlib import contextmanager import torch @@ -53,21 +52,6 @@ def get_graph_pool(use_single_mempool): return torch.cuda.graph_pool_handle() -@contextmanager -def _override_stale_capture_stream(): - """Redirect stale autograd stream refs while capturing full-iteration graphs.""" - set_override = getattr(torch.autograd.graph, "set_override_stale_capture_stream", None) - if set_override is None: - yield - return - - set_override(True) - try: - yield - finally: - set_override(False) - - # The below functions traverse through nested data structures (tuples, lists, dicts) # present in src and creates a deep copy where all PyTorch tensors are cloned, # detached from the computation graph, and moved to CUDA device. Non-tensor objects @@ -230,16 +214,15 @@ def __call__(self, *args, **kwargs): FullCudaGraphWrapper.cuda_graph[training_str].register_generator_state(state) torch.cuda.synchronize() capture_stream = get_shared_capture_stream() - with _override_stale_capture_stream(): - with torch.cuda.graph( - FullCudaGraphWrapper.cuda_graph[training_str], - stream=capture_stream, - pool=get_graph_pool(self.use_single_mempool), - capture_error_mode="thread_local", - ): - FullCudaGraphWrapper.result[training_str] = self.forward_backward_func( - *args, **kwargs - ) + with torch.cuda.graph( + FullCudaGraphWrapper.cuda_graph[training_str], + stream=capture_stream, + pool=get_graph_pool(self.use_single_mempool), + capture_error_mode="thread_local", + ): + FullCudaGraphWrapper.result[training_str] = self.forward_backward_func( + *args, **kwargs + ) torch.cuda.synchronize() torch.distributed.barrier() logger.info(f'CUDA graph capture done for {training_str}!!!') diff --git a/tests/unit_tests/transformer/test_full_cuda_graph.py b/tests/unit_tests/transformer/test_full_cuda_graph.py index 367cdd51d1a..312ae467304 100644 --- a/tests/unit_tests/transformer/test_full_cuda_graph.py +++ b/tests/unit_tests/transformer/test_full_cuda_graph.py @@ -1,4 +1,4 @@ -# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. import pytest import torch @@ -6,7 +6,7 @@ import megatron.core.pipeline_parallel.schedules as schedule from megatron.core import ModelParallelConfig -from megatron.core.full_cuda_graph import FullCudaGraphWrapper, _override_stale_capture_stream +from megatron.core.full_cuda_graph import FullCudaGraphWrapper from megatron.core.tensor_parallel.random import ( HAVE_TE, initialize_rng_tracker, @@ -18,22 +18,6 @@ rank = Utils.rank -def test_override_stale_capture_stream_toggles_when_available(monkeypatch): - calls = [] - - def set_override(enabled): - calls.append(enabled) - - monkeypatch.setattr( - torch.autograd.graph, "set_override_stale_capture_stream", set_override, raising=False - ) - - with _override_stale_capture_stream(): - assert calls == [True] - - assert calls == [True, False] - - @pytest.mark.skipif( not (HAVE_TE and is_te_min_version("1.5.0")), reason="use_te_rng_tracker requires TransformerEngine version >= 1.5",