From af02072f8db475ab9beeaaaa0bd7354c77ed40ae Mon Sep 17 00:00:00 2001 From: qiyuw Date: Mon, 11 May 2026 17:49:01 -0700 Subject: [PATCH 1/7] fix mxfp8 without overlap Signed-off-by: qiyuw --- megatron/core/optimizer/distrib_optimizer.py | 10 ++-- megatron/core/optimizer/optimizer.py | 55 +++++++++++++++++--- 2 files changed, 55 insertions(+), 10 deletions(-) diff --git a/megatron/core/optimizer/distrib_optimizer.py b/megatron/core/optimizer/distrib_optimizer.py index 95e8a0c407b..75b6129110b 100644 --- a/megatron/core/optimizer/distrib_optimizer.py +++ b/megatron/core/optimizer/distrib_optimizer.py @@ -2832,10 +2832,12 @@ def step_with_ready_grads(self) -> bool: """ update_successful = super().step_with_ready_grads() + should_sync_params = self.ddp_config.use_megatron_fsdp or ( + not self.ddp_config.overlap_param_gather and not getattr(self, '_defer_param_sync', False) + ) timers = self.config.timers - if timers is not None: + if timers is not None and should_sync_params: timers('params-all-gather', log_level=1).start(barrier=self.config.barrier_with_L1_time) - if self.ddp_config.use_megatron_fsdp: # Optionally all-gather Megatron-FSDP sharded main weights # early in preparation for the subsequent forward pass. @@ -2846,10 +2848,10 @@ def step_with_ready_grads(self) -> bool: # communication calls here. If overlapping all-gather for parameters, the following # the first all-gather is launched asynchronously in the next optimizer.zero_grad() # call and subsequent all-gathers are launched in the forward pre-hook. - if not self.ddp_config.overlap_param_gather: + if should_sync_params: for model_chunk in self.model_chunks: model_chunk.start_param_sync() - if timers is not None: + if timers is not None and should_sync_params: timers('params-all-gather').stop() return update_successful diff --git a/megatron/core/optimizer/optimizer.py b/megatron/core/optimizer/optimizer.py index 9ae23bb4b7f..705bc1cab3e 100644 --- a/megatron/core/optimizer/optimizer.py +++ b/megatron/core/optimizer/optimizer.py @@ -1284,12 +1284,55 @@ def prepare_grads(self) -> bool: def step_with_ready_grads(self) -> bool: """Step the optimizer with ready gradients, return successful.""" success = True - for optimizer_idx, optimizer in enumerate(self.chained_optimizers): - success &= optimizer.step_with_ready_grads() - if self.config.overlap_param_gather_with_optimizer_step and optimizer_idx == 0: - assert success - assert len(optimizer.model_chunks) == 1 - optimizer.model_chunks[0].start_param_sync(force_dispatch=True) + # With MXFP8 grad-buffer reuse and non-overlap param gather, each DistOpt stages + # its own updated main-param shards into its param buffers during step. However, + # param sync is a DDP model-chunk operation: model_chunk.start_param_sync() gathers + # both dense and expert bucket groups, copies gathered values into model weights, + # and zeros the shared MXFP8 param/grad buffers. For MoE, dense and expert DistOpts + # may share the same model chunk, so defer param sync until all chained optimizers + # have staged their params, then sync each model chunk once. + defer_param_sync = ( + self.config.reuse_grad_buf_for_mxfp8_param_ag + and not self.config.overlap_param_gather + ) + deferred_model_chunks = [] + deferred_model_chunk_ids = set() + + if defer_param_sync: + from .distrib_optimizer import DistributedOptimizer + + for optimizer in self.chained_optimizers: + if isinstance(optimizer, DistributedOptimizer): + optimizer._defer_param_sync = True + for model_chunk in optimizer.model_chunks: + model_chunk_id = id(model_chunk) + if model_chunk_id not in deferred_model_chunk_ids: + deferred_model_chunk_ids.add(model_chunk_id) + deferred_model_chunks.append(model_chunk) + + try: + for optimizer_idx, optimizer in enumerate(self.chained_optimizers): + success &= optimizer.step_with_ready_grads() + if self.config.overlap_param_gather_with_optimizer_step and optimizer_idx == 0: + assert success + assert len(optimizer.model_chunks) == 1 + optimizer.model_chunks[0].start_param_sync(force_dispatch=True) + finally: + if defer_param_sync: + for optimizer in self.chained_optimizers: + if hasattr(optimizer, '_defer_param_sync'): + optimizer._defer_param_sync = False + + if defer_param_sync and success: + timers = self.config.timers + if timers is not None: + timers('params-all-gather', log_level=1).start( + barrier=self.config.barrier_with_L1_time + ) + for model_chunk in deferred_model_chunks: + model_chunk.start_param_sync() + if timers is not None: + timers('params-all-gather').stop() return success From 79d0bc87a8076acb2bb9076aab0cb4f9d9dde3da Mon Sep 17 00:00:00 2001 From: qiyuw Date: Mon, 11 May 2026 17:53:12 -0700 Subject: [PATCH 2/7] minor Signed-off-by: qiyuw --- megatron/core/optimizer/distrib_optimizer.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/megatron/core/optimizer/distrib_optimizer.py b/megatron/core/optimizer/distrib_optimizer.py index 75b6129110b..6d6104a1e55 100644 --- a/megatron/core/optimizer/distrib_optimizer.py +++ b/megatron/core/optimizer/distrib_optimizer.py @@ -2832,11 +2832,11 @@ def step_with_ready_grads(self) -> bool: """ update_successful = super().step_with_ready_grads() - should_sync_params = self.ddp_config.use_megatron_fsdp or ( - not self.ddp_config.overlap_param_gather and not getattr(self, '_defer_param_sync', False) + should_sync_params = not self.ddp_config.overlap_param_gather and not getattr( + self, '_defer_param_sync', False ) timers = self.config.timers - if timers is not None and should_sync_params: + if timers is not None and (self.ddp_config.use_megatron_fsdp or should_sync_params): timers('params-all-gather', log_level=1).start(barrier=self.config.barrier_with_L1_time) if self.ddp_config.use_megatron_fsdp: # Optionally all-gather Megatron-FSDP sharded main weights @@ -2851,7 +2851,7 @@ def step_with_ready_grads(self) -> bool: if should_sync_params: for model_chunk in self.model_chunks: model_chunk.start_param_sync() - if timers is not None and should_sync_params: + if timers is not None and (self.ddp_config.use_megatron_fsdp or should_sync_params): timers('params-all-gather').stop() return update_successful From 097f85f6c9d5e8aaf3e73cadd74127da5477bb57 Mon Sep 17 00:00:00 2001 From: qiyuw Date: Tue, 12 May 2026 17:13:37 -0700 Subject: [PATCH 3/7] unit test for mxfp8 on moe model Signed-off-by: qiyuw --- tests/unit_tests/test_fp8_param.py | 59 ++++++++++++++++++++++++++---- 1 file changed, 51 insertions(+), 8 deletions(-) diff --git a/tests/unit_tests/test_fp8_param.py b/tests/unit_tests/test_fp8_param.py index 34b504e21de..83f3675e70f 100644 --- a/tests/unit_tests/test_fp8_param.py +++ b/tests/unit_tests/test_fp8_param.py @@ -90,7 +90,10 @@ def model_provider( model_parallel_cuda_manual_seed(_SEED) args = get_args() config = core_transformer_config_from_args(args) - transformer_layer_spec = layer_spec_fn() + transformer_layer_spec = layer_spec_fn( + num_experts=args.num_experts, + moe_grouped_gemm=args.moe_grouped_gemm, + ) return GPTModel( config=config, transformer_layer_spec=transformer_layer_spec, @@ -180,7 +183,7 @@ def _run_test_helper( use_cuda_graph: bool = False, **kwargs, ): - """Test fp8_param with gpt_model.""" + """Test fp8_param with a small GPT model.""" args = self.create_test_args( tp_size, recipe, @@ -199,7 +202,10 @@ def _run_test_helper( set_args(args) torch.manual_seed(_SEED) - Utils.initialize_model_parallel(tensor_model_parallel_size=tp_size) + Utils.initialize_model_parallel( + tensor_model_parallel_size=tp_size, + expert_model_parallel_size=args.expert_model_parallel_size, + ) input_ids, labels, position_ids, attention_mask, loss_mask = self.get_batch( self.seq_length, self.micro_batch_size @@ -237,14 +243,21 @@ def _run_test_helper( if is_float8tensor(param): num_fp8_params += 1 - # Verify the number of fp8 params. fp8_layers = args.num_layers if kwargs.get("first_last_layers_bf16", False): fp8_layers -= kwargs["num_layers_at_start_in_bf16"] fp8_layers -= kwargs["num_layers_at_end_in_bf16"] - # Each layer has 4 GEMM weights: qkv, proj, fc1, fc2. - if fp8_param_gather: - assert num_fp8_params == 4 * fp8_layers + if fp8_param_gather and fp8_layers > 0: + if args.num_experts is None: + # Each dense layer has 4 GEMM weights: qkv, proj, fc1, fc2. + assert num_fp8_params == 4 * fp8_layers + else: + assert num_fp8_params > 0 + assert any( + not getattr(param, 'allreduce', True) for param in gpt_model[0].parameters() + ) + if not inference: + assert len(optimizer.chained_optimizers) >= 2 # Verify that bf16 params (embedding, LN, etc.) in the MXFP8 model are mapped # to the param buffer (shared with grad buffer) rather than allocated separately. @@ -332,7 +345,7 @@ def _run_test_helper( return torch.tensor(loss_list) def run_test(self, tp_size, recipe, inference: bool = False, **kwargs): - """Test fp8_param with gpt_model.""" + """Test fp8_param with a small GPT model.""" if inference: with torch.inference_mode(): self._run_test_helper(tp_size, recipe, inference=True, **kwargs) @@ -442,6 +455,36 @@ def test_mxfp8(self, tp_size, dp_overlap): kwargs = {"overlap_param_gather": dp_overlap[0], "overlap_grad_reduce": dp_overlap[1]} self.run_test(tp_size=tp_size, recipe="mxfp8", **kwargs) + @pytest.mark.skipif( + get_device_arch_version() < 10, reason="MXFP8 is supported since Blackwell architecture" + ) + @pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8) + @pytest.mark.skipif(not is_te_min_version("2.3.0.dev0"), reason="TE 2.3.0.dev0 is required") + @pytest.mark.parametrize("tp_size", [1]) + @pytest.mark.parametrize("dp_overlap", [(False, False), (False, True), (True, True)]) + def test_mxfp8_moe(self, tp_size, dp_overlap): + """ + dp_overlap: (overlap_param_gather, overlap_grad_reduce) + """ + kwargs = { + "overlap_param_gather": dp_overlap[0], + "overlap_grad_reduce": dp_overlap[1], + "num_layers": 4, + "vocal_size": 128800, + "hidden_size": 128, + "num_attention_heads": 8, + "expert_model_parallel_size": 2, + "num_experts": 2, + "moe_grouped_gemm": True, + "moe_token_dispatcher_type": "alltoall", + "moe_router_topk": 1, + "moe_router_pre_softmax": True, + "moe_router_load_balancing_type": "none", + "moe_aux_loss_coeff": 0.0, + "moe_ffn_hidden_size": 128, + } + self.run_test(tp_size=tp_size, recipe="mxfp8", **kwargs) + @pytest.mark.skipif( get_device_arch_version() < 10, reason="MXFP8 is supported since Blackwell architecture" ) From fc3471a69b6b7f732b3b8eacb6d8e47666039354 Mon Sep 17 00:00:00 2001 From: qiyuw Date: Tue, 12 May 2026 17:19:26 -0700 Subject: [PATCH 4/7] lint Signed-off-by: qiyuw --- megatron/core/optimizer/optimizer.py | 3 +-- tests/unit_tests/test_fp8_param.py | 3 +-- 2 files changed, 2 insertions(+), 4 deletions(-) diff --git a/megatron/core/optimizer/optimizer.py b/megatron/core/optimizer/optimizer.py index 705bc1cab3e..eb163bc5622 100644 --- a/megatron/core/optimizer/optimizer.py +++ b/megatron/core/optimizer/optimizer.py @@ -1292,8 +1292,7 @@ def step_with_ready_grads(self) -> bool: # may share the same model chunk, so defer param sync until all chained optimizers # have staged their params, then sync each model chunk once. defer_param_sync = ( - self.config.reuse_grad_buf_for_mxfp8_param_ag - and not self.config.overlap_param_gather + self.config.reuse_grad_buf_for_mxfp8_param_ag and not self.config.overlap_param_gather ) deferred_model_chunks = [] deferred_model_chunk_ids = set() diff --git a/tests/unit_tests/test_fp8_param.py b/tests/unit_tests/test_fp8_param.py index 83f3675e70f..60fb0a8c91c 100644 --- a/tests/unit_tests/test_fp8_param.py +++ b/tests/unit_tests/test_fp8_param.py @@ -91,8 +91,7 @@ def model_provider( args = get_args() config = core_transformer_config_from_args(args) transformer_layer_spec = layer_spec_fn( - num_experts=args.num_experts, - moe_grouped_gemm=args.moe_grouped_gemm, + num_experts=args.num_experts, moe_grouped_gemm=args.moe_grouped_gemm ) return GPTModel( config=config, From 0479d84eac6c1692b73328a0e9dc813ab600a420 Mon Sep 17 00:00:00 2001 From: Philip Petrakian Date: Fri, 22 May 2026 15:49:40 +0000 Subject: [PATCH 5/7] Refactor chained optimizer param sync handling --- megatron/core/optimizer/optimizer.py | 124 +++++++++++++++++---------- 1 file changed, 78 insertions(+), 46 deletions(-) diff --git a/megatron/core/optimizer/optimizer.py b/megatron/core/optimizer/optimizer.py index daaf05449ed..3d1d94b5e6d 100644 --- a/megatron/core/optimizer/optimizer.py +++ b/megatron/core/optimizer/optimizer.py @@ -1280,70 +1280,102 @@ def prepare_grads(self) -> bool: return found_inf_flag - @torch.no_grad() - def step_with_ready_grads(self) -> bool: - """Step the optimizer with ready gradients, return successful.""" + def _step_chained_optimizers(self) -> bool: + """Step all optimizers in this chain.""" success = True + for optimizer_idx, optimizer in enumerate(self.chained_optimizers): + success &= optimizer.step_with_ready_grads() + if self.config.overlap_param_gather_with_optimizer_step and optimizer_idx == 0: + assert success + assert len(optimizer.model_chunks) == 1 + optimizer.model_chunks[0].start_param_sync(force_dispatch=True) + return success + + def _should_defer_mxfp8_param_sync(self) -> bool: + """Return whether MXFP8 param sync should be deferred until chained steps finish.""" + return ( + self.config.reuse_grad_buf_for_mxfp8_param_ag and not self.config.overlap_param_gather + ) + + def _enable_deferred_mxfp8_param_sync(self) -> List[Tuple[Any, Any]]: + """Enable deferred DistOpt param sync and collect bucket groups to sync later.""" + from .distrib_optimizer import DistributedOptimizer + from .layer_wise_optimizer import _bucket_is_managed_by_layer_wise_optimizer + # With MXFP8 grad-buffer reuse and non-overlap param gather, each DistOpt stages # its own updated main-param shards into its param buffers during step. However, # param sync is a DDP bucket-group operation that copies gathered values into model # weights and zeros the shared MXFP8 param/grad buffers. For MoE, dense and expert # DistOpts may share the same model chunk, so defer param sync until all chained # optimizers have staged their params, then sync each DistOpt-managed bucket group once. - defer_param_sync = ( - self.config.reuse_grad_buf_for_mxfp8_param_ag and not self.config.overlap_param_gather - ) deferred_bucket_groups = [] deferred_bucket_group_ids = set() - if defer_param_sync: - from .distrib_optimizer import DistributedOptimizer - from .layer_wise_optimizer import _bucket_is_managed_by_layer_wise_optimizer + for optimizer in self.chained_optimizers: + if not isinstance(optimizer, DistributedOptimizer): + continue - for optimizer in self.chained_optimizers: - if isinstance(optimizer, DistributedOptimizer): - optimizer._defer_param_sync = True - for model_chunk in optimizer.model_chunks: - for bucket_group in ( - model_chunk.bucket_groups + model_chunk.expert_parallel_bucket_groups - ): - if not bucket_group.buckets: - continue - if _bucket_is_managed_by_layer_wise_optimizer( - bucket_group.buckets[0], default_for_untagged=False - ): - continue - bucket_group_id = id(bucket_group) - if bucket_group_id not in deferred_bucket_group_ids: - deferred_bucket_group_ids.add(bucket_group_id) - deferred_bucket_groups.append((model_chunk, bucket_group)) + optimizer._defer_param_sync = True + for model_chunk in optimizer.model_chunks: + for bucket_group in ( + model_chunk.bucket_groups + model_chunk.expert_parallel_bucket_groups + ): + if not bucket_group.buckets: + continue + if _bucket_is_managed_by_layer_wise_optimizer( + bucket_group.buckets[0], default_for_untagged=False + ): + continue + + bucket_group_id = id(bucket_group) + if bucket_group_id in deferred_bucket_group_ids: + continue + + deferred_bucket_group_ids.add(bucket_group_id) + deferred_bucket_groups.append((model_chunk, bucket_group)) + return deferred_bucket_groups + + def _disable_deferred_mxfp8_param_sync(self) -> None: + """Disable deferred DistOpt param sync.""" + for optimizer in self.chained_optimizers: + if hasattr(optimizer, '_defer_param_sync'): + optimizer._defer_param_sync = False + + def _start_deferred_mxfp8_param_sync( + self, deferred_bucket_groups: List[Tuple[Any, Any]] + ) -> None: + """Start param sync for deferred bucket groups.""" + timers = self.config.timers + if timers is not None: + timers('params-all-gather', log_level=1).start( + barrier=self.config.barrier_with_L1_time + ) + for model_chunk, bucket_group in deferred_bucket_groups: + model_chunk._start_bucket_group_param_sync(bucket_group, force_sync=False) + if timers is not None: + timers('params-all-gather').stop() + + def _step_with_deferred_mxfp8_param_sync(self) -> bool: + """Step optimizers with MXFP8 param sync deferred until all steps finish.""" + deferred_bucket_groups = self._enable_deferred_mxfp8_param_sync() try: - for optimizer_idx, optimizer in enumerate(self.chained_optimizers): - success &= optimizer.step_with_ready_grads() - if self.config.overlap_param_gather_with_optimizer_step and optimizer_idx == 0: - assert success - assert len(optimizer.model_chunks) == 1 - optimizer.model_chunks[0].start_param_sync(force_dispatch=True) + success = self._step_chained_optimizers() finally: - if defer_param_sync: - for optimizer in self.chained_optimizers: - if hasattr(optimizer, '_defer_param_sync'): - optimizer._defer_param_sync = False + self._disable_deferred_mxfp8_param_sync() - if defer_param_sync and success and deferred_bucket_groups: - timers = self.config.timers - if timers is not None: - timers('params-all-gather', log_level=1).start( - barrier=self.config.barrier_with_L1_time - ) - for model_chunk, bucket_group in deferred_bucket_groups: - model_chunk._start_bucket_group_param_sync(bucket_group, force_sync=False) - if timers is not None: - timers('params-all-gather').stop() + if success and deferred_bucket_groups: + self._start_deferred_mxfp8_param_sync(deferred_bucket_groups) return success + @torch.no_grad() + def step_with_ready_grads(self) -> bool: + """Step the optimizer with ready gradients, return successful.""" + if self._should_defer_mxfp8_param_sync(): + return self._step_with_deferred_mxfp8_param_sync() + return self._step_chained_optimizers() + def grads_states_parallel_group_is_shared(self): """Check if all optimizers share the same gradient statistics parallel group.""" reference_group = self.chained_optimizers[0].get_grad_stats_parallel_group() From c2b4f6d8639aebf02db15f96d1354f743f9c549a Mon Sep 17 00:00:00 2001 From: Philip Petrakian Date: Fri, 22 May 2026 15:53:15 +0000 Subject: [PATCH 6/7] Format optimizer refactor --- megatron/core/optimizer/optimizer.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/megatron/core/optimizer/optimizer.py b/megatron/core/optimizer/optimizer.py index 3d1d94b5e6d..0bd788409b5 100644 --- a/megatron/core/optimizer/optimizer.py +++ b/megatron/core/optimizer/optimizer.py @@ -1348,9 +1348,7 @@ def _start_deferred_mxfp8_param_sync( """Start param sync for deferred bucket groups.""" timers = self.config.timers if timers is not None: - timers('params-all-gather', log_level=1).start( - barrier=self.config.barrier_with_L1_time - ) + timers('params-all-gather', log_level=1).start(barrier=self.config.barrier_with_L1_time) for model_chunk, bucket_group in deferred_bucket_groups: model_chunk._start_bucket_group_param_sync(bucket_group, force_sync=False) if timers is not None: From 752d246175c3ca6073ee864d7afd68067dbafdb6 Mon Sep 17 00:00:00 2001 From: Philip Petrakian Date: Fri, 22 May 2026 16:25:34 +0000 Subject: [PATCH 7/7] Rename chained optimizer step helper --- megatron/core/optimizer/optimizer.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/megatron/core/optimizer/optimizer.py b/megatron/core/optimizer/optimizer.py index 0bd788409b5..ddc3dd8620e 100644 --- a/megatron/core/optimizer/optimizer.py +++ b/megatron/core/optimizer/optimizer.py @@ -1280,7 +1280,7 @@ def prepare_grads(self) -> bool: return found_inf_flag - def _step_chained_optimizers(self) -> bool: + def _step(self) -> bool: """Step all optimizers in this chain.""" success = True for optimizer_idx, optimizer in enumerate(self.chained_optimizers): @@ -1358,7 +1358,7 @@ def _step_with_deferred_mxfp8_param_sync(self) -> bool: """Step optimizers with MXFP8 param sync deferred until all steps finish.""" deferred_bucket_groups = self._enable_deferred_mxfp8_param_sync() try: - success = self._step_chained_optimizers() + success = self._step() finally: self._disable_deferred_mxfp8_param_sync() @@ -1372,7 +1372,7 @@ def step_with_ready_grads(self) -> bool: """Step the optimizer with ready gradients, return successful.""" if self._should_defer_mxfp8_param_sync(): return self._step_with_deferred_mxfp8_param_sync() - return self._step_chained_optimizers() + return self._step() def grads_states_parallel_group_is_shared(self): """Check if all optimizers share the same gradient statistics parallel group."""