From 7f67ed2d75ec393813a5f473042961082705594c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?oliver=20k=C3=B6nig?= Date: Tue, 26 May 2026 07:10:33 +0000 Subject: [PATCH] fix(optimizer): gate ChainedOptimizer MXFP8 defer-sync on DDP-level overlap_param_gather MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit PR #4800 introduced the deferred MXFP8 param-sync path in ChainedOptimizer to fix a numerical race when DP param-gather is not overlapped. The gate read self.config.overlap_param_gather, which resolves to the OptimizerConfig field. However, the actual overlap behavior is determined by DistributedDataParallelConfig.overlap_param_gather (a separate field that can diverge). When the DDP-level gather IS overlapped but OptimizerConfig says it is not, the deferred path activates unnecessarily and adds per-step bucket-group all-gather work via _start_bucket_group_param_sync, causing a ~3-5%% throughput regression on hybrid mamba + MoE + MXFP8 pretraining (e.g. nemotron_3_super_64gpu_b300_fp8_mx). Probe the underlying DistributedOptimizer instances and only defer when at least one chained DistOpt actually has its DDP-level param-gather disabled -- the only case where the race fix is needed. Verified empirically by bisect (11 probes localizing the regression to this commit) and by V1/V2 control pipelines on bia/B300: V1 (mcore parent fa7a23bad3): GPU util -0.94%% vs golden (healthy) V2 (mcore 5f791187ee): GPU util -5.00%% vs golden (regression) Signed-off-by: oliver könig --- megatron/core/optimizer/optimizer.py | 23 +++++++++++++++++++---- 1 file changed, 19 insertions(+), 4 deletions(-) diff --git a/megatron/core/optimizer/optimizer.py b/megatron/core/optimizer/optimizer.py index ddc3dd8620e..1206c92e144 100644 --- a/megatron/core/optimizer/optimizer.py +++ b/megatron/core/optimizer/optimizer.py @@ -1292,10 +1292,25 @@ def _step(self) -> bool: 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 - ) + """Return whether MXFP8 param sync should be deferred until chained steps finish. + + The deferred-sync path is only needed when MXFP8 grad/param buffer reuse is active + AND the DDP-level param gather is not overlapped (i.e. the race fixed by PR #4800 + can occur). The OptimizerConfig.overlap_param_gather field is unreliable as a proxy + for the DDP-level setting -- the two configs can diverge -- so probe the underlying + DistOpts directly. + """ + if not self.config.reuse_grad_buf_for_mxfp8_param_ag: + return False + + from .distrib_optimizer import DistributedOptimizer + + for optimizer in self.chained_optimizers: + if not isinstance(optimizer, DistributedOptimizer): + continue + if not optimizer.ddp_config.overlap_param_gather: + return True + return False def _enable_deferred_mxfp8_param_sync(self) -> List[Tuple[Any, Any]]: """Enable deferred DistOpt param sync and collect bucket groups to sync later."""