From 8b36432e2e6d99c7f1727686574f252c35a9d312 Mon Sep 17 00:00:00 2001 From: Pingtian Li Date: Tue, 10 Mar 2026 18:56:00 -0700 Subject: [PATCH 1/6] Support Delay wgrad gemm --- .../fsdp/src/megatron_fsdp/megatron_fsdp.py | 50 ++++++++++- .../core/extensions/transformer_engine.py | 10 ++- megatron/core/model_parallel_config.py | 9 ++ megatron/core/transformer/moe/moe_layer.py | 88 +++++++++++++++++++ .../core/transformer/transformer_config.py | 13 +++ tests/unit_tests/a2a_overlap/utils.py | 11 +++ 6 files changed, 177 insertions(+), 4 deletions(-) diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/megatron_fsdp.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/megatron_fsdp.py index 6987729ba8f..4fcdb830506 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/megatron_fsdp.py +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/megatron_fsdp.py @@ -50,6 +50,7 @@ from megatron.core.distributed.distributed_data_parallel_config import ( DistributedDataParallelConfig, ) + from megatron.core.transformer import TransformerLayer from megatron.core.utils import is_submodule except ImportError: # Megatron-LM is not installed, use Megatron-FSDP as a standalone module. @@ -73,6 +74,48 @@ class TrainingState(Enum): IDLE = auto() +def _maybe_setup_delayed_wgrad_for_experts(module, process_post_backward_gradients_fn): + """Configure delayed wgrad gradient processing for MoE expert parameters. + + When ``delay_wgrad_compute_for_te_grouped_gemm`` is enabled on a TransformerLayer, + this function: + 1. Marks expert parameters so the normal post-accumulate-grad hook is skipped. + 2. Registers a callback on the MoE layer that invokes FSDP's gradient + reduce-scatter after the delayed wgrad computation completes. + + Args: + module: The module being processed in the forward pre-hook. Only + ``TransformerLayer`` instances with the delayed wgrad config flag + enabled are affected; all other modules are no-ops. + process_post_backward_gradients_fn: The FSDP gradient processing function + (``_process_post_backward_gradients``) to be called after the delayed + wgrad computation finishes. + """ + try: + if not (isinstance(module, TransformerLayer) and module.is_moe_layer): + return + except NameError: + return + + if not getattr(module.config, 'delay_wgrad_compute_for_te_grouped_gemm', False): + return + + expert_params = list(module.mlp.experts.parameters()) + for p in expert_params: + p._fsdp_delay_grad_reduce = True + + def _make_process_expert_grads(mlp_module): + def _process_expert_grads(): + params = list(mlp_module.experts.parameters()) + process_post_backward_gradients_fn(params) + + return _process_expert_grads + + if module.mlp._process_expert_grads_fn is not None: + return + module.mlp.register_process_expert_grads_fn(_make_process_expert_grads(module.mlp)) + + class MegatronFSDP(torch.nn.Module): """Fully Sharded Data Parallel training. @@ -719,6 +762,7 @@ def _pre_forward_param_unshard( prefetch=fsdp_forward_prefetch, prefetch_order=PrefetchOrder.FORWARD_PASS_ORDER, ) + _maybe_setup_delayed_wgrad_for_experts(module, _process_post_backward_gradients) return args, kwargs @torch.compiler.disable @@ -1022,7 +1066,11 @@ def _register_pre_backward_param_unshard_hook(module): for param in grad_acc_param_list: self.grad_acc_hooks[f"grad_acc and reduce for {self.param_to_name[param]}"] = ( param.register_post_accumulate_grad_hook( - lambda p: _process_post_backward_gradients([p]) + lambda p: ( + None + if getattr(p, '_fsdp_delay_grad_reduce', False) + else _process_post_backward_gradients([p]) + ) ) ) diff --git a/megatron/core/extensions/transformer_engine.py b/megatron/core/extensions/transformer_engine.py index e901e40597a..c43f4a2e673 100644 --- a/megatron/core/extensions/transformer_engine.py +++ b/megatron/core/extensions/transformer_engine.py @@ -1701,10 +1701,14 @@ def __init__( self.disable_parameter_transpose_cache = self.config.disable_parameter_transpose_cache extra_kwargs = _get_extra_te_kwargs(config) + self.delay_wgrad_compute = ( + self.config.delay_wgrad_compute + or self.config.delay_wgrad_compute_for_te_grouped_gemm + ) - if self.config.delay_wgrad_compute: + if self.delay_wgrad_compute: if is_te_min_version("2.3.0"): - extra_kwargs["delay_wgrad_compute"] = self.config.delay_wgrad_compute + extra_kwargs["delay_wgrad_compute"] = True else: raise RuntimeError( "Only TE with version >=2.3.0 supports delay_wgrad_compute now." @@ -2012,7 +2016,7 @@ def backward_dw(self): Compute weight gradients during the backward pass if delay_wgrad_compute is enabled. """ - if self.config.delay_wgrad_compute: + if self.delay_wgrad_compute: super().backward_dw() class TEColumnParallelGroupedLinear(TEGroupedLinear): diff --git a/megatron/core/model_parallel_config.py b/megatron/core/model_parallel_config.py index 5da94b802ce..15493531691 100644 --- a/megatron/core/model_parallel_config.py +++ b/megatron/core/model_parallel_config.py @@ -261,6 +261,15 @@ class ModelParallelConfig: delay_wgrad_compute: bool = False """Delay the weight gradient computation to improve batch-level communication overlapping""" + delay_wgrad_compute_for_te_grouped_gemm: bool = False + """Delay the weight gradient computation for TE Grouped GEMM MoE experts. + When enabled with FSDP, the expert weight gradients are computed on a separate + CUDA stream after the data gradients finish, allowing overlap of wgrad compute + with the backward pass of earlier layers. The FSDP gradient reduce-scatter for + expert parameters is deferred until the delayed wgrad computation completes. + This requires transformer_engine with GroupedLinear support (TE >= 2.3.0). + """ + ep_overlap_early_attn_memory_release: bool = False """Enable early memory release of attention activations during EP overlap. EP overlap can increase peak memory usage when the overlapped forward module allocates diff --git a/megatron/core/transformer/moe/moe_layer.py b/megatron/core/transformer/moe/moe_layer.py index 8277486b03b..e9c36451c0d 100644 --- a/megatron/core/transformer/moe/moe_layer.py +++ b/megatron/core/transformer/moe/moe_layer.py @@ -283,6 +283,14 @@ def __init__( self.cudagraph_tensor_store = MoECudaGraphTensorStore() self.fwd_execution_map = ["route", "expert_compute", "postprocess"] + # Delay wgrad computation for TE grouped GEMM + self._delayed_wgrad_event: Optional[torch.cuda.Event] = None + self._delayed_wgrad_stream: Optional[torch.cuda.Stream] = None + self._process_expert_grads_fn = None + if self.config.delay_wgrad_compute_for_te_grouped_gemm: + self._delayed_wgrad_event = torch.cuda.Event() + self._delayed_wgrad_stream = torch.cuda.Stream(device="cuda") + def _setup_inference_mode(self, pg_collection): """Set up inference-optimized token dispatcher and state. @@ -373,6 +381,8 @@ def dispatch(self, hidden_states: torch.Tensor, probs: torch.Tensor): tokens and their associated probabilities to the devices hosting their assigned experts. """ + if self.config.delay_wgrad_compute_for_te_grouped_gemm: + hidden_states = _RegisterDelayedWgradForExperts.apply(self, hidden_states) return self.token_dispatcher.token_dispatch(hidden_states, probs) @maybe_skip_or_early_return_by_cudagraph("shared_experts_compute") @@ -411,6 +421,10 @@ def routed_experts_compute(self, hidden_states: torch.Tensor, probs: torch.Tenso for each expert. It then passes the tokens through the local experts. The output from the experts is preprocessed for the combine step. """ + if self.config.delay_wgrad_compute_for_te_grouped_gemm: + hidden_states = _RecordExpertDgradCompletion.apply( + self._delayed_wgrad_event, hidden_states + ) dispatched_input, tokens_per_expert, permuted_probs = ( self.token_dispatcher.dispatch_postprocess(hidden_states, probs) ) @@ -584,3 +598,77 @@ def set_for_recompute_pre_mlp_layernorm(self): from megatron.core.extensions.transformer_engine import set_save_original_input set_save_original_input(self.shared_experts.linear_fc1) + + def register_process_expert_grads_fn(self, fn): + """Register a callback to process expert gradients after delayed wgrad computation. + + This is used by FSDP to defer the reduce-scatter of expert parameter + gradients until the delayed wgrad computation has completed. + + Args: + fn: A callable that processes expert gradients (e.g., triggers + FSDP reduce-scatter for expert parameters). + """ + self._process_expert_grads_fn = fn + + +class _RecordExpertDgradCompletion(torch.autograd.Function): + """Autograd function that records a CUDA event when expert data gradients finish. + + Placed in the forward graph just before the expert computation so that during + the backward pass, when the expert dgrad completes, we record an event. The + subsequent ``_RegisterDelayedWgradForExperts`` waits on this event before + launching the delayed wgrad computation on a separate CUDA stream. + """ + + @staticmethod + def forward(ctx, event: torch.cuda.Event, *inputs): + """Forward pass that stores the event and passes through inputs unchanged.""" + ctx.event = event + return inputs[0] if len(inputs) == 1 else inputs + + @staticmethod + def backward(ctx, *grad_outputs): + """Backward pass that records the event when expert dgrad completes.""" + ctx.event.record(torch.cuda.current_stream()) + ctx.event = None + return (None,) + grad_outputs + + +class _RegisterDelayedWgradForExperts(torch.autograd.Function): + """Autograd function that orchestrates delayed wgrad computation for MoE experts. + + Placed in the forward graph at the dispatch boundary. During the backward pass, + this function: + 1. Records an event on the current (backward) stream to signal the dgrad is done. + 2. Executes the delayed wgrad computation on a dedicated CUDA stream. + 3. Waits for the wgrad computation to complete. + 4. Invokes the registered gradient processing callback (e.g., FSDP reduce-scatter). + """ + + @staticmethod + def forward(ctx, module: MoELayer, *inputs): + """Forward pass that stores the MoE module and passes through inputs unchanged.""" + ctx.module = module + return inputs[0] if len(inputs) == 1 else inputs + + @staticmethod + def backward(ctx, *grad_outputs): + """Backward pass that executes delayed wgrad computation on a separate stream.""" + module = ctx.module + event = module._delayed_wgrad_event + wgrad_stream = module._delayed_wgrad_stream + + wgrad_stream.wait_event(event) + with torch.cuda.stream(wgrad_stream): + with torch.cuda.nvtx.range("delayed_expert_wgrad"): + module.backward_dw(routed_experts=True, shared_experts=True) + event.record(wgrad_stream) + + torch.cuda.current_stream().wait_event(event) + + if module._process_expert_grads_fn is not None: + module._process_expert_grads_fn() + + ctx.module = None + return (None,) + grad_outputs diff --git a/megatron/core/transformer/transformer_config.py b/megatron/core/transformer/transformer_config.py index fd4025de9f7..5411a7be260 100644 --- a/megatron/core/transformer/transformer_config.py +++ b/megatron/core/transformer/transformer_config.py @@ -2101,6 +2101,19 @@ def __post_init__(self): 'partial cuda graph' ) + if self.delay_wgrad_compute_for_te_grouped_gemm: + assert not self.overlap_moe_expert_parallel_comm, ( + 'overlap_moe_expert_parallel_comm must be disabled when enabling ' + 'delay_wgrad_compute_for_te_grouped_gemm.' + ) + assert is_te_min_version( + "2.3.0" + ), 'TE version >= 2.3.0 is required for delay_wgrad_compute_for_te_grouped_gemm' + assert not self.delay_wgrad_compute, ( + 'delay_wgrad_compute and delay_wgrad_compute_for_te_grouped_gemm ' + 'are mutually exclusive; use only one' + ) + if self.ep_overlap_early_attn_memory_release: assert self.overlap_moe_expert_parallel_comm, ( 'overlap_moe_expert_parallel_comm must be enabled when enabling ' diff --git a/tests/unit_tests/a2a_overlap/utils.py b/tests/unit_tests/a2a_overlap/utils.py index a52843956df..528f1c25de4 100644 --- a/tests/unit_tests/a2a_overlap/utils.py +++ b/tests/unit_tests/a2a_overlap/utils.py @@ -231,6 +231,17 @@ def get_valid_token_dispatcher_types(): return ["alltoall"] +def get_valid_flex_dispatcher_backends(): + from megatron.core.transformer.moe.fused_a2a import HAVE_DEEP_EP, HAVE_HYBRIDEP + + if HAVE_HYBRIDEP: + return ["hybridep"] + elif HAVE_DEEP_EP: + return ["deepep"] + else: + return [None] + + def get_valid_fp8_flags(): from megatron.core.enums import Fp8Recipe From 71e07e7d652903d5e91d26977706d1dec368196c Mon Sep 17 00:00:00 2001 From: Pingtian Li Date: Tue, 10 Mar 2026 19:09:05 -0700 Subject: [PATCH 2/6] add test --- .../a2a_overlap/test_delay_wgrad_compute.py | 211 ++++++++++++++++++ 1 file changed, 211 insertions(+) create mode 100644 tests/unit_tests/a2a_overlap/test_delay_wgrad_compute.py diff --git a/tests/unit_tests/a2a_overlap/test_delay_wgrad_compute.py b/tests/unit_tests/a2a_overlap/test_delay_wgrad_compute.py new file mode 100644 index 00000000000..63e9369ff1f --- /dev/null +++ b/tests/unit_tests/a2a_overlap/test_delay_wgrad_compute.py @@ -0,0 +1,211 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +import gc + +import pytest +import torch + +from megatron.core.models.gpt.gpt_layer_specs import get_gpt_decoder_block_spec +from megatron.core.models.gpt.gpt_model import GPTModel +from megatron.core.transformer import TransformerLayer +from megatron.core.transformer.module import float16_to_fp32 +from megatron.core.utils import is_te_min_version +from tests.unit_tests.a2a_overlap.utils import ( + deterministic_mode, + get_test_config, + get_valid_flex_dispatcher_backends, + get_valid_fp8_flags, + get_valid_token_dispatcher_types, + reset_model, +) +from tests.unit_tests.test_utilities import Utils + +NUM_STEPS = 3 +SEQ_LEN = 128 +VOCAB_SIZE = 512 +LR = 0.01 + + +def _build_gpt_model(config): + """Build and return a GPTModel on CUDA from the given config.""" + layer_spec = get_gpt_decoder_block_spec(config=config, use_transformer_engine=True) + model = GPTModel( + config=config, + transformer_layer_spec=layer_spec, + vocab_size=VOCAB_SIZE, + pre_process=True, + post_process=True, + max_sequence_length=300, + ) + model.cuda() + return model + + +def _build_input_data(): + """Build fixed input data for the model.""" + return { + "input_ids": torch.randint(0, VOCAB_SIZE, (1, SEQ_LEN), dtype=torch.int64).cuda(), + "labels": torch.randint(0, VOCAB_SIZE, (1, SEQ_LEN), dtype=torch.int64).cuda(), + "position_ids": torch.arange(SEQ_LEN, dtype=torch.int64).unsqueeze(0).cuda(), + "attention_mask": torch.ones((1, 1, SEQ_LEN, SEQ_LEN), dtype=bool).cuda(), + } + + +def _train_step(model, optimizer, data): + """Run one forward-backward-optimizer step. Return the detached loss.""" + optimizer.zero_grad() + loss = model.forward(**data) + loss = float16_to_fp32(loss) + loss.backward(torch.ones_like(loss)) + optimizer.step() + return loss.detach().clone() + + +def _assert_models_equal(ref_model, test_model): + """Assert that all parameters of two models are bit-identical.""" + rank = torch.distributed.get_rank() + for (name_r, param_r), (_, param_t) in zip( + ref_model.named_parameters(), test_model.named_parameters() + ): + assert torch.equal( + param_r.data, param_t.data + ), f"[rank {rank}] Parameter mismatch after training: {name_r}" + + +class TestDelayWgradCompute: + """Verify that delay_wgrad_compute_for_te_grouped_gemm produces identical + training behaviour (per-step loss and final weights) as the non-delayed baseline + across multiple forward-backward-optimizer steps on the full GPTModel. + """ + + def setup_method(self, method): + Utils.initialize_model_parallel( + tensor_model_parallel_size=1, + pipeline_model_parallel_size=1, + expert_model_parallel_size=4, + ) + + def teardown_method(self, method): + Utils.destroy_model_parallel() + + @pytest.mark.skipif(not is_te_min_version("2.3.0"), reason="Requires TE >= 2.3.0") + @pytest.mark.parametrize("num_layers", [2, 3]) + @pytest.mark.parametrize("shared_expert_intermediate_size", [None, 512]) + @pytest.mark.parametrize("dispatcher_type", get_valid_token_dispatcher_types()) + @pytest.mark.parametrize("flex_dispatcher_backend", get_valid_flex_dispatcher_backends()) + @pytest.mark.parametrize("fp8_flag", get_valid_fp8_flags()) + def test_delay_wgrad_compute_for_te_grouped_gemm( + self, + num_layers, + shared_expert_intermediate_size, + dispatcher_type, + flex_dispatcher_backend, + fp8_flag, + ): + """Verify that delay_wgrad_compute_for_te_grouped_gemm produces identical + per-step loss and final weights as the non-delayed baseline across multiple + forward-backward-optimizer steps on the full GPTModel. + + Covers single/multi-layer, with/without shared experts, dispatcher types, + and FP8 modes. + """ + extra_kwargs = {"moe_token_dispatcher_type": dispatcher_type} + if dispatcher_type == "flex": + extra_kwargs["moe_flex_dispatcher_backend"] = flex_dispatcher_backend + extra_kwargs["moe_router_dtype"] = "fp32" + if fp8_flag is not None: + extra_kwargs["fp8"] = fp8_flag[0] + extra_kwargs["fp8_recipe"] = fp8_flag[1] + if shared_expert_intermediate_size is not None: + extra_kwargs["moe_shared_expert_intermediate_size"] = shared_expert_intermediate_size + + with deterministic_mode(): + ref_config = get_test_config(num_layers=num_layers, extra_kwargs=extra_kwargs) + ref_model = _build_gpt_model(ref_config) + init_params = reset_model(ref_model) + + delay_kwargs = {**extra_kwargs, "delay_wgrad_compute_for_te_grouped_gemm": True} + test_config = get_test_config(num_layers=num_layers, extra_kwargs=delay_kwargs) + test_model = _build_gpt_model(test_config) + reset_model(test_model, init_params) + + data = _build_input_data() + ref_opt = torch.optim.SGD(ref_model.parameters(), lr=LR) + test_opt = torch.optim.SGD(test_model.parameters(), lr=LR) + + rank = torch.distributed.get_rank() + for step in range(NUM_STEPS): + ref_loss = _train_step(ref_model, ref_opt, data) + test_loss = _train_step(test_model, test_opt, data) + assert torch.equal(ref_loss, test_loss), ( + f"[rank {rank}] Loss mismatch at step {step}: " + f"ref={ref_loss.item()}, test={test_loss.item()}" + ) + + _assert_models_equal(ref_model, test_model) + + del ref_model, test_model + gc.collect() + torch.cuda.empty_cache() + + @pytest.mark.skipif(not is_te_min_version("2.3.0"), reason="Requires TE >= 2.3.0") + @pytest.mark.parametrize("num_layers", [2, 3]) + @pytest.mark.parametrize("shared_expert_intermediate_size", [None, 512]) + @pytest.mark.parametrize("dispatcher_type", get_valid_token_dispatcher_types()) + @pytest.mark.parametrize("flex_dispatcher_backend", get_valid_flex_dispatcher_backends()) + def test_delay_wgrad_compute_for_te_grouped_gemm_with_fsdp( + self, num_layers, shared_expert_intermediate_size, dispatcher_type, flex_dispatcher_backend + ): + """Verify delayed wgrad with MegatronFSDP wrapping. + + The delayed wgrad path defers the FSDP reduce-scatter for expert + parameters until the wgrad computation completes on a separate stream. + This test checks that the deferred reduce-scatter produces identical + per-step loss and final weights as the non-delayed FSDP baseline. + """ + from megatron.core.distributed.fsdp.src.megatron_fsdp.fully_shard import ( + fully_shard_model, + fully_shard_optimizer, + ) + + extra_kwargs = {"moe_token_dispatcher_type": dispatcher_type} + if dispatcher_type == "flex": + extra_kwargs["moe_flex_dispatcher_backend"] = flex_dispatcher_backend + extra_kwargs["moe_router_dtype"] = "fp32" + if shared_expert_intermediate_size is not None: + extra_kwargs["moe_shared_expert_intermediate_size"] = shared_expert_intermediate_size + + with deterministic_mode(): + # Build reference model (no delay) and wrap with FSDP + ref_config = get_test_config(num_layers=num_layers, extra_kwargs=extra_kwargs) + ref_model = _build_gpt_model(ref_config) + init_params = reset_model(ref_model) + + ref_fsdp = fully_shard_model(module=ref_model, fsdp_unit_modules=[TransformerLayer]) + ref_opt = torch.optim.SGD(ref_fsdp.parameters(), lr=LR) + ref_opt = fully_shard_optimizer(optimizer=ref_opt) + + # Build test model (with delay) and wrap with FSDP + delay_kwargs = {**extra_kwargs, "delay_wgrad_compute_for_te_grouped_gemm": True} + test_config = get_test_config(num_layers=num_layers, extra_kwargs=delay_kwargs) + test_model = _build_gpt_model(test_config) + reset_model(test_model, init_params) + + test_fsdp = fully_shard_model(module=test_model, fsdp_unit_modules=[TransformerLayer]) + test_opt = torch.optim.SGD(test_fsdp.parameters(), lr=LR) + test_opt = fully_shard_optimizer(optimizer=test_opt) + + data = _build_input_data() + rank = torch.distributed.get_rank() + for step in range(NUM_STEPS): + ref_loss = _train_step(ref_fsdp, ref_opt, data) + test_loss = _train_step(test_fsdp, test_opt, data) + assert torch.equal(ref_loss, test_loss), ( + f"[rank {rank}] Loss mismatch at step {step}: " + f"ref={ref_loss.item()}, test={test_loss.item()}" + ) + + _assert_models_equal(ref_fsdp, test_fsdp) + + del ref_fsdp, test_fsdp, ref_opt, test_opt + gc.collect() + torch.cuda.empty_cache() From f70e899ce5c20ce4ca80d4408a55dd1c1eed6801 Mon Sep 17 00:00:00 2001 From: Pingtian Li <158665726+Wohox@users.noreply.github.com> Date: Wed, 11 Mar 2026 10:09:43 +0800 Subject: [PATCH 3/6] Update megatron/core/transformer/moe/moe_layer.py Co-authored-by: claude[bot] <209825114+claude[bot]@users.noreply.github.com> --- megatron/core/transformer/moe/moe_layer.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/megatron/core/transformer/moe/moe_layer.py b/megatron/core/transformer/moe/moe_layer.py index e9c36451c0d..da0a4493f6e 100644 --- a/megatron/core/transformer/moe/moe_layer.py +++ b/megatron/core/transformer/moe/moe_layer.py @@ -662,7 +662,7 @@ def backward(ctx, *grad_outputs): wgrad_stream.wait_event(event) with torch.cuda.stream(wgrad_stream): with torch.cuda.nvtx.range("delayed_expert_wgrad"): - module.backward_dw(routed_experts=True, shared_experts=True) + module.backward_dw(routed_experts=True, shared_experts=False) event.record(wgrad_stream) torch.cuda.current_stream().wait_event(event) From b6e1b5d92b7ed96170a1555cb8659952878474a7 Mon Sep 17 00:00:00 2001 From: Pingtian Li Date: Tue, 10 Mar 2026 19:43:02 -0700 Subject: [PATCH 4/6] fix ut and repeated status setting --- .../fsdp/src/megatron_fsdp/megatron_fsdp.py | 15 +++++---- .../a2a_overlap/test_delay_wgrad_compute.py | 23 ++++--------- tests/unit_tests/a2a_overlap/utils.py | 32 ++++++++++++------- 3 files changed, 36 insertions(+), 34 deletions(-) diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/megatron_fsdp.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/megatron_fsdp.py index 4fcdb830506..6b6ba57e1de 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/megatron_fsdp.py +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/megatron_fsdp.py @@ -100,10 +100,6 @@ def _maybe_setup_delayed_wgrad_for_experts(module, process_post_backward_gradien if not getattr(module.config, 'delay_wgrad_compute_for_te_grouped_gemm', False): return - expert_params = list(module.mlp.experts.parameters()) - for p in expert_params: - p._fsdp_delay_grad_reduce = True - def _make_process_expert_grads(mlp_module): def _process_expert_grads(): params = list(mlp_module.experts.parameters()) @@ -111,8 +107,10 @@ def _process_expert_grads(): return _process_expert_grads - if module.mlp._process_expert_grads_fn is not None: - return + expert_params = list(module.mlp.experts.parameters()) + for p in expert_params: + p._fsdp_delay_grad_reduce = True + module.mlp.register_process_expert_grads_fn(_make_process_expert_grads(module.mlp)) @@ -762,7 +760,10 @@ def _pre_forward_param_unshard( prefetch=fsdp_forward_prefetch, prefetch_order=PrefetchOrder.FORWARD_PASS_ORDER, ) - _maybe_setup_delayed_wgrad_for_experts(module, _process_post_backward_gradients) + + # Set post backward hook for TE grouped gemm if enabled comm overlap + if getattr(module.mlp, '_process_expert_grads_fn') is None: + _maybe_setup_delayed_wgrad_for_experts(module, _process_post_backward_gradients) return args, kwargs @torch.compiler.disable diff --git a/tests/unit_tests/a2a_overlap/test_delay_wgrad_compute.py b/tests/unit_tests/a2a_overlap/test_delay_wgrad_compute.py index 63e9369ff1f..95bbab94ad4 100644 --- a/tests/unit_tests/a2a_overlap/test_delay_wgrad_compute.py +++ b/tests/unit_tests/a2a_overlap/test_delay_wgrad_compute.py @@ -12,7 +12,7 @@ from tests.unit_tests.a2a_overlap.utils import ( deterministic_mode, get_test_config, - get_valid_flex_dispatcher_backends, + get_valid_flex_dispatcher_backend, get_valid_fp8_flags, get_valid_token_dispatcher_types, reset_model, @@ -88,18 +88,11 @@ def teardown_method(self, method): Utils.destroy_model_parallel() @pytest.mark.skipif(not is_te_min_version("2.3.0"), reason="Requires TE >= 2.3.0") - @pytest.mark.parametrize("num_layers", [2, 3]) @pytest.mark.parametrize("shared_expert_intermediate_size", [None, 512]) @pytest.mark.parametrize("dispatcher_type", get_valid_token_dispatcher_types()) - @pytest.mark.parametrize("flex_dispatcher_backend", get_valid_flex_dispatcher_backends()) @pytest.mark.parametrize("fp8_flag", get_valid_fp8_flags()) def test_delay_wgrad_compute_for_te_grouped_gemm( - self, - num_layers, - shared_expert_intermediate_size, - dispatcher_type, - flex_dispatcher_backend, - fp8_flag, + self, shared_expert_intermediate_size, dispatcher_type, fp8_flag ): """Verify that delay_wgrad_compute_for_te_grouped_gemm produces identical per-step loss and final weights as the non-delayed baseline across multiple @@ -108,10 +101,10 @@ def test_delay_wgrad_compute_for_te_grouped_gemm( Covers single/multi-layer, with/without shared experts, dispatcher types, and FP8 modes. """ + num_layers = 4 extra_kwargs = {"moe_token_dispatcher_type": dispatcher_type} if dispatcher_type == "flex": - extra_kwargs["moe_flex_dispatcher_backend"] = flex_dispatcher_backend - extra_kwargs["moe_router_dtype"] = "fp32" + extra_kwargs["moe_flex_dispatcher_backend"] = get_valid_flex_dispatcher_backend() if fp8_flag is not None: extra_kwargs["fp8"] = fp8_flag[0] extra_kwargs["fp8_recipe"] = fp8_flag[1] @@ -148,12 +141,10 @@ def test_delay_wgrad_compute_for_te_grouped_gemm( torch.cuda.empty_cache() @pytest.mark.skipif(not is_te_min_version("2.3.0"), reason="Requires TE >= 2.3.0") - @pytest.mark.parametrize("num_layers", [2, 3]) @pytest.mark.parametrize("shared_expert_intermediate_size", [None, 512]) @pytest.mark.parametrize("dispatcher_type", get_valid_token_dispatcher_types()) - @pytest.mark.parametrize("flex_dispatcher_backend", get_valid_flex_dispatcher_backends()) def test_delay_wgrad_compute_for_te_grouped_gemm_with_fsdp( - self, num_layers, shared_expert_intermediate_size, dispatcher_type, flex_dispatcher_backend + self, shared_expert_intermediate_size, dispatcher_type ): """Verify delayed wgrad with MegatronFSDP wrapping. @@ -167,10 +158,10 @@ def test_delay_wgrad_compute_for_te_grouped_gemm_with_fsdp( fully_shard_optimizer, ) + num_layers = 4 extra_kwargs = {"moe_token_dispatcher_type": dispatcher_type} if dispatcher_type == "flex": - extra_kwargs["moe_flex_dispatcher_backend"] = flex_dispatcher_backend - extra_kwargs["moe_router_dtype"] = "fp32" + extra_kwargs["moe_flex_dispatcher_backend"] = get_valid_flex_dispatcher_backend() if shared_expert_intermediate_size is not None: extra_kwargs["moe_shared_expert_intermediate_size"] = shared_expert_intermediate_size diff --git a/tests/unit_tests/a2a_overlap/utils.py b/tests/unit_tests/a2a_overlap/utils.py index 528f1c25de4..9a644ee8cc8 100644 --- a/tests/unit_tests/a2a_overlap/utils.py +++ b/tests/unit_tests/a2a_overlap/utils.py @@ -216,44 +216,54 @@ def get_test_config(num_layers=1, num_moe_experts=8, extra_kwargs={}, moe_groupe multi_latent_attention=True, num_moe_experts=num_moe_experts, moe_grouped_gemm=moe_grouped_gemm, + moe_router_dtype="fp32", **extra_kwargs, ) return config def get_valid_token_dispatcher_types(): - try: - from deep_ep import Buffer - from deep_ep.utils import EventHandle, EventOverlap + from megatron.core.transformer.moe.fused_a2a import HAVE_DEEP_EP, HAVE_HYBRIDEP + if HAVE_HYBRIDEP or HAVE_DEEP_EP: return ["alltoall", "flex"] - except ImportError: + else: return ["alltoall"] -def get_valid_flex_dispatcher_backends(): +def get_valid_flex_dispatcher_backend(): from megatron.core.transformer.moe.fused_a2a import HAVE_DEEP_EP, HAVE_HYBRIDEP if HAVE_HYBRIDEP: - return ["hybridep"] + return "hybridep" elif HAVE_DEEP_EP: - return ["deepep"] + return "deepep" else: - return [None] + return None def get_valid_fp8_flags(): from megatron.core.enums import Fp8Recipe + from megatron.training.utils import get_device_arch_version fp8_types = ["e4m3", "hybrid"] recipes = [] - valid_flags = [] + arch = get_device_arch_version() + if is_te_min_version("2.3.0.dev0"): - recipes.append(Fp8Recipe.blockwise) - recipes.append(Fp8Recipe.tensorwise) + recipes.append(Fp8Recipe.tensorwise) # Hopper + Blackwell + + if is_te_min_version("2.4.0.dev0") and arch == 9: + recipes.append(Fp8Recipe.blockwise) # Hopper only + if is_te_min_version("2.3.0.dev0") and arch >= 10: + recipes.append(Fp8Recipe.mxfp8) # Blackwell only + + valid_flags = [] for fp8_type in fp8_types: for recipe in recipes: + if fp8_type == "hybrid" and recipe == Fp8Recipe.mxfp8: + continue valid_flags.append((fp8_type, recipe)) valid_flags.append(None) From 95c62e0f63c0ffd708cb03a231724c6aef1f6946 Mon Sep 17 00:00:00 2001 From: Pingtian Li Date: Tue, 10 Mar 2026 20:05:28 -0700 Subject: [PATCH 5/6] fix --- .../fsdp/src/megatron_fsdp/megatron_fsdp.py | 11 +++++------ 1 file changed, 5 insertions(+), 6 deletions(-) diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/megatron_fsdp.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/megatron_fsdp.py index 6b6ba57e1de..cbb8e7bd209 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/megatron_fsdp.py +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/megatron_fsdp.py @@ -91,15 +91,15 @@ def _maybe_setup_delayed_wgrad_for_experts(module, process_post_backward_gradien (``_process_post_backward_gradients``) to be called after the delayed wgrad computation finishes. """ - try: - if not (isinstance(module, TransformerLayer) and module.is_moe_layer): - return - except NameError: + if not (isinstance(module, TransformerLayer) and module.is_moe_layer): return if not getattr(module.config, 'delay_wgrad_compute_for_te_grouped_gemm', False): return + if getattr(module.mlp, '_process_expert_grads_fn', None) is not None: + return + def _make_process_expert_grads(mlp_module): def _process_expert_grads(): params = list(mlp_module.experts.parameters()) @@ -762,8 +762,7 @@ def _pre_forward_param_unshard( ) # Set post backward hook for TE grouped gemm if enabled comm overlap - if getattr(module.mlp, '_process_expert_grads_fn') is None: - _maybe_setup_delayed_wgrad_for_experts(module, _process_post_backward_gradients) + _maybe_setup_delayed_wgrad_for_experts(module, _process_post_backward_gradients) return args, kwargs @torch.compiler.disable From a8861740df5e73babc38dfd9273a5c55277682d4 Mon Sep 17 00:00:00 2001 From: Pingtian Li Date: Tue, 10 Mar 2026 20:21:20 -0700 Subject: [PATCH 6/6] fix comment --- megatron/core/model_parallel_config.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/megatron/core/model_parallel_config.py b/megatron/core/model_parallel_config.py index 15493531691..15007cc5810 100644 --- a/megatron/core/model_parallel_config.py +++ b/megatron/core/model_parallel_config.py @@ -265,7 +265,7 @@ class ModelParallelConfig: """Delay the weight gradient computation for TE Grouped GEMM MoE experts. When enabled with FSDP, the expert weight gradients are computed on a separate CUDA stream after the data gradients finish, allowing overlap of wgrad compute - with the backward pass of earlier layers. The FSDP gradient reduce-scatter for + with EP A2A communication. The FSDP gradient reduce-scatter for expert parameters is deferred until the delayed wgrad computation completes. This requires transformer_engine with GroupedLinear support (TE >= 2.3.0). """