From edf80e8b40d3136a5356c44fd736dcb0c001fa19 Mon Sep 17 00:00:00 2001 From: root Date: Fri, 24 Apr 2026 07:24:05 -0700 Subject: [PATCH 1/2] fix cg convergence Signed-off-by: root --- megatron/core/transformer/cuda_graphs.py | 26 ++++++++---------------- 1 file changed, 8 insertions(+), 18 deletions(-) diff --git a/megatron/core/transformer/cuda_graphs.py b/megatron/core/transformer/cuda_graphs.py index d1850ff9bd5..877044deae6 100644 --- a/megatron/core/transformer/cuda_graphs.py +++ b/megatron/core/transformer/cuda_graphs.py @@ -797,6 +797,10 @@ def create_fwd_graph(self, args, kwargs, outputs=None, clone_inputs=True): self.outputs = outputs # save grads and other variables that may be affected by graph warmup + buffer_backup = [] + for buf in self.base_module.buffers(): + buffer_backup.append(buf.clone()) + if self.training and torch.is_grad_enabled(): grad_backup = [] for param in self.base_module.parameters(): @@ -842,7 +846,6 @@ def create_fwd_graph(self, args, kwargs, outputs=None, clone_inputs=True): def _resolve_input_buffer(ten): if not isinstance(ten, ArgMetadata): return ten - # the input tensor is resued from another cudagraph's input or output if ( hasattr(ten, "cg_buffer_metadata") @@ -913,7 +916,7 @@ def _resolve_input_buffer(ten): def clone_ten(ten): if not torch.is_tensor(ten): return ten - return torch.zeros_like(ten).requires_grad_(ten.requires_grad) + return torch.clone(ten).detach().requires_grad_(ten.requires_grad) warmup_args = tree_map(clone_ten, self.fwd_graph_input_args) warmup_kwargs = tree_map(clone_ten, self.fwd_graph_input_kwargs) @@ -986,17 +989,6 @@ def clone_ten(ten): o.cg_buffer_metadata.fwd_cudagraph_buffer = fwd_graph_out fwd_buffer_reuse_ref_count += 1 - # if an input buffer requires a copy, and does not have metadata attached to it at this - # point, it will not be reused after this forward pass, so return it to the pool - for buf in self.fwd_graph_input_surface: - if ( - hasattr(buf, "can_skip_replay_copy") - and not buf.can_skip_replay_copy - and not hasattr(buf, "cg_buffer_metadata") - ): - assert _CudagraphGlobalRecord.tensor_reuse_pool.owns(buf) - _CudagraphGlobalRecord.tensor_reuse_pool.insert(buf) - if self.training and torch.is_grad_enabled(): assert ( len(self.fwd_graph_output_surface) > 0 @@ -1011,7 +1003,9 @@ def clone_ten(ten): if self.fp8_enabled: restore_fp8_tensors([self.base_module], saved_fp8_tensors) - # restore cached grads + # restore cached grads and buffers + for buf_copy, buf in zip(buffer_backup, self.base_module.buffers()): + buf.copy_(buf_copy) for main_grad_copy, param in zip(grad_backup, self.base_module.parameters()): if main_grad_copy is not None: param.main_grad.copy_(main_grad_copy) @@ -1625,10 +1619,6 @@ def __call__(self, megatron_module, args, kwargs): runner = self.get_cudagraph_runner( megatron_module, args, kwargs, self.reuse_cudagraphs ) - # check if a layer is frozen during training. - if not torch.is_grad_enabled(): - # If the layer is frozen, we need to set the runner to eval mode. - runner.eval() out = runner.record_graph_capture(args, kwargs) else: # No cudagraphs were found in training mode with grad disabled, so fallback to From db9e078e8c5553e85345c56714f802003f8ca287 Mon Sep 17 00:00:00 2001 From: root Date: Fri, 24 Apr 2026 08:49:47 -0700 Subject: [PATCH 2/2] Address comments Signed-off-by: root --- megatron/core/transformer/cuda_graphs.py | 13 +++++++++---- 1 file changed, 9 insertions(+), 4 deletions(-) diff --git a/megatron/core/transformer/cuda_graphs.py b/megatron/core/transformer/cuda_graphs.py index c007ce0aa6e..36e5cd7cc60 100644 --- a/megatron/core/transformer/cuda_graphs.py +++ b/megatron/core/transformer/cuda_graphs.py @@ -796,7 +796,10 @@ def create_fwd_graph(self, args, kwargs, outputs=None, clone_inputs=True): self.kwargs = kwargs self.outputs = outputs - # save grads and other variables that may be affected by graph warmup + # Save buffers, grads, and other variables that may be affected by graph warmup. + # For example, megatron/core/transformer/moe/router.py's expert_bias is a persistent + # buffer updated each forward pass by '_apply_expert_bias()'. So we need to ensure + # graph capture's forward passes do not corrupt its value. buffer_backup = [] for buf in self.base_module.buffers(): buffer_backup.append(buf.clone()) @@ -1003,13 +1006,15 @@ def clone_ten(ten): if self.fp8_enabled: restore_fp8_tensors([self.base_module], saved_fp8_tensors) - # restore cached grads and buffers - for buf_copy, buf in zip(buffer_backup, self.base_module.buffers()): - buf.copy_(buf_copy) + # restore cached grads for main_grad_copy, param in zip(grad_backup, self.base_module.parameters()): if main_grad_copy is not None: param.main_grad.copy_(main_grad_copy) + # restore cached buffers + for buf_copy, buf in zip(buffer_backup, self.base_module.buffers()): + buf.copy_(buf_copy) + if is_moe: for name in tracker: tracker[name]["values"].copy_(cached_aux_losses[name])