From 11e1aff8825da721b3789103d67a709157d23107 Mon Sep 17 00:00:00 2001 From: Robin Zhang Date: Tue, 4 Nov 2025 21:38:23 -0800 Subject: [PATCH 01/19] [DEV] feat(MoE): Refactor cuda_graph_scope (#1917) Signed-off-by: Robin Zhang --- .../text_generation_controller.py | 2 +- .../common/language_module/language_module.py | 4 +- .../core/models/gpt/fine_grained_callables.py | 5 +- megatron/core/models/gpt/gpt_model.py | 4 +- megatron/core/pipeline_parallel/schedules.py | 8 +- megatron/core/ssm/mamba_block.py | 4 +- megatron/core/transformer/attention.py | 4 +- megatron/core/transformer/cuda_graphs.py | 86 ++- megatron/core/transformer/moe/fused_a2a.py | 2 + megatron/core/transformer/moe/moe_layer.py | 58 +- megatron/core/transformer/moe/moe_utils.py | 251 ++++++- .../core/transformer/moe/token_dispatcher.py | 46 +- .../core/transformer/transformer_block.py | 4 +- .../core/transformer/transformer_config.py | 146 +++- .../core/transformer/transformer_layer.py | 184 ++++- megatron/training/arguments.py | 36 +- megatron/training/training.py | 18 +- .../golden_values_dev_dgx_h100.json | 644 ++++++++++++++++++ .../model_config.yaml | 96 +++ tests/test_utils/recipes/moe.yaml | 5 + .../inference/engines/test_dynamic_engine.py | 7 +- .../transformer/test_cuda_graphs.py | 282 +++++++- tools/checkpoint/checkpoint_inspector.py | 2 + 23 files changed, 1763 insertions(+), 135 deletions(-) create mode 100644 tests/functional_tests/test_cases/moe/gpt3_moe_mcore_te_tp4_ep2_etp2_pp2_scoped_cudagraph/golden_values_dev_dgx_h100.json create mode 100644 tests/functional_tests/test_cases/moe/gpt3_moe_mcore_te_tp4_ep2_etp2_pp2_scoped_cudagraph/model_config.yaml diff --git a/megatron/core/inference/text_generation_controllers/text_generation_controller.py b/megatron/core/inference/text_generation_controllers/text_generation_controller.py index 2b44b418749..0aed3df079e 100644 --- a/megatron/core/inference/text_generation_controllers/text_generation_controller.py +++ b/megatron/core/inference/text_generation_controllers/text_generation_controller.py @@ -922,7 +922,7 @@ def generate_all_output_tokens_static_batch( # Check whether CUDA graphs are enabled enable_cuda_graph = ( model_config.cuda_graph_impl == "local" - and model_config.cuda_graph_scope != "full_iteration" + and "full_iteration" not in model_config.cuda_graph_scope ) # Pad batch tokens if necessary diff --git a/megatron/core/models/common/language_module/language_module.py b/megatron/core/models/common/language_module/language_module.py index d855322c2df..8f90fb3ba47 100644 --- a/megatron/core/models/common/language_module/language_module.py +++ b/megatron/core/models/common/language_module/language_module.py @@ -1,4 +1,4 @@ -# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved. +# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. import logging import os from typing import Optional, Tuple @@ -136,7 +136,7 @@ def compute_language_model_loss(self, labels: Tensor, logits: Tensor) -> Tensor: # Use is_cg_capturable=True for full iteration CUDA graphs to avoid torch.equal checks is_cg_capturable = ( hasattr(self.config, 'cuda_graph_scope') - and self.config.cuda_graph_scope == 'full_iteration' + and 'full_iteration' in self.config.cuda_graph_scope ) if is_cg_capturable and not is_te_min_version("2.7.0"): from megatron.core.utils import get_te_version diff --git a/megatron/core/models/gpt/fine_grained_callables.py b/megatron/core/models/gpt/fine_grained_callables.py index b125ee11255..1f3e5988a73 100644 --- a/megatron/core/models/gpt/fine_grained_callables.py +++ b/megatron/core/models/gpt/fine_grained_callables.py @@ -1,4 +1,4 @@ -# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. import weakref from contextlib import nullcontext @@ -355,7 +355,8 @@ def submodule_post_attn_forward(node: ScheduleNode, hidden_states: torch.Tensor) else: pre_mlp_layernorm_output = layer.pre_mlp_layernorm(hidden_states) - local_tokens, probs, _ = layer.mlp.router_and_preprocess(pre_mlp_layernorm_output) + probs, routing_map = layer.mlp.route(pre_mlp_layernorm_output) + local_tokens, probs, _ = layer.mlp.preprocess(pre_mlp_layernorm_output, probs, routing_map) # Detach here for mlp_bda residual connection node.layer_state.residual = node.detach(hidden_states) diff --git a/megatron/core/models/gpt/gpt_model.py b/megatron/core/models/gpt/gpt_model.py index 25546d36629..77e87917911 100644 --- a/megatron/core/models/gpt/gpt_model.py +++ b/megatron/core/models/gpt/gpt_model.py @@ -1,4 +1,4 @@ -# Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved. +# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. from collections import OrderedDict from typing import Dict, Literal, Optional @@ -371,7 +371,7 @@ def _preprocess( and ( ( self.config.cuda_graph_impl == "local" - and self.config.cuda_graph_scope != "full_iteration" + and "full_iteration" not in self.config.cuda_graph_scope ) or self.config.flash_decode ) diff --git a/megatron/core/pipeline_parallel/schedules.py b/megatron/core/pipeline_parallel/schedules.py index e83f8d90635..db670bbeaf1 100644 --- a/megatron/core/pipeline_parallel/schedules.py +++ b/megatron/core/pipeline_parallel/schedules.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved. +# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. import contextlib from functools import partial @@ -648,7 +648,7 @@ def forward_backward_no_pipelining( if ( hasattr(config, 'cuda_graph_impl') and config.cuda_graph_impl == "local" - and config.cuda_graph_scope != "full_iteration" + and "full_iteration" not in config.cuda_graph_scope ): create_cudagraphs() @@ -1912,7 +1912,7 @@ def pp_post_backward(input_tensor_grad, vp_stage=None): if ( hasattr(config, 'cuda_graph_impl') and config.cuda_graph_impl == "local" - and config.cuda_graph_scope != "full_iteration" + and "full_iteration" not in config.cuda_graph_scope ): create_cudagraphs() nvtx_range_pop(suffix="misc") @@ -2296,7 +2296,7 @@ def enable_grad_sync(): if ( hasattr(config, 'cuda_graph_impl') and config.cuda_graph_impl == "local" - and config.cuda_graph_scope != "full_iteration" + and "full_iteration" not in config.cuda_graph_scope ): create_cudagraphs() diff --git a/megatron/core/ssm/mamba_block.py b/megatron/core/ssm/mamba_block.py index 7d8ca74c8f2..d4a2981178e 100644 --- a/megatron/core/ssm/mamba_block.py +++ b/megatron/core/ssm/mamba_block.py @@ -1,4 +1,4 @@ -# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved. +# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # Copyright (c) 2024, Tri Dao, Albert Gu. # Some of this code was adopted from https://github.com/state-spaces/mamba/ @@ -292,7 +292,7 @@ def forward( ( ( self.config.cuda_graph_impl == "local" - and self.config.cuda_graph_scope != "full_iteration" + and "full_iteration" not in self.config.cuda_graph_scope ) or self.config.flash_decode ) diff --git a/megatron/core/transformer/attention.py b/megatron/core/transformer/attention.py index e221c0ea00d..1bbd38ed368 100644 --- a/megatron/core/transformer/attention.py +++ b/megatron/core/transformer/attention.py @@ -1,4 +1,4 @@ -# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. from abc import ABC, abstractmethod from dataclasses import dataclass @@ -791,7 +791,7 @@ def forward( if ( in_decode_mode and self.config.cuda_graph_impl == "local" - and self.config.cuda_graph_scope != "full_iteration" + and "full_iteration" not in self.config.cuda_graph_scope and inference_context.is_static_batching() ): raise ValueError(f"CUDA graphs must use flash decode with static batching!") diff --git a/megatron/core/transformer/cuda_graphs.py b/megatron/core/transformer/cuda_graphs.py index 9d50e34a2cc..12f15ee980a 100644 --- a/megatron/core/transformer/cuda_graphs.py +++ b/megatron/core/transformer/cuda_graphs.py @@ -1,4 +1,4 @@ -# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved. +# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. import gc import inspect @@ -22,7 +22,7 @@ get_cuda_rng_tracker, ) from megatron.core.transformer.identity_op import IdentityOp -from megatron.core.transformer.module import MegatronModule +from megatron.core.transformer.module import GraphableMegatronModule, MegatronModule from megatron.core.transformer.transformer_config import TransformerConfig from megatron.core.utils import ( get_attr_wrapped_model, @@ -1070,9 +1070,12 @@ def __init__( ), "RNG tracker does not support cudagraphs!" assert config.cuda_graph_impl == "local", "Option cuda_graph_impl=local not enabled." - assert "expandable_segments:True" not in os.getenv("PYTORCH_CUDA_ALLOC_CONF", ""), ( - "expandable_segments:True may not be safe when using CUDA Graphs, and may result in" - "a crash due to illegal memory access or other undefined behaviour." + assert ( + "expandable_segments:True" not in os.getenv("PYTORCH_CUDA_ALLOC_CONF", "") + or os.getenv("NCCL_GRAPH_REGISTER", "") == "0" + ), ( + "Setting NCCL_GRAPH_REGISTER=0 to avoid illegal memory access when using " + "CUDA Graph with PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True." ) self.cudagraph_runners = [] @@ -1326,23 +1329,40 @@ def _layer_is_graphable(layer, config): Check if a layer is graphable. """ + # Only GraphableMegatronModule can be graphed. + if not isinstance(layer, GraphableMegatronModule): + return False + + # If cuda_graph_scope is not set, every layer is graphed. + if not config.cuda_graph_scope: + return True + # import modules here to avoid a circular import from megatron.core.ssm.mamba_layer import MambaLayer from megatron.core.transformer.identity_op import IdentityOp + from megatron.core.transformer.mlp import MLP + from megatron.core.transformer.moe.moe_layer import MoELayer from megatron.core.transformer.transformer_layer import TransformerLayer - if isinstance(layer, MambaLayer) and config.cuda_graph_scope == "full": + if isinstance(layer, MambaLayer) and 'mamba' in config.cuda_graph_scope: # mamba layer. return True if isinstance(layer, TransformerLayer): - if config.cuda_graph_scope == 'attn': - if not ( - isinstance(layer.self_attention, IdentityOp) - and isinstance(layer.cross_attention, IdentityOp) - ): - # attn layer. - return True - else: + if 'attn' in config.cuda_graph_scope and not ( + isinstance(layer.self_attention, IdentityOp) + and isinstance(layer.cross_attention, IdentityOp) + ): + # attn layer. + return True + if ( + 'moe' in config.cuda_graph_scope + or 'moe_router' in config.cuda_graph_scope + or 'moe_preprocess' in config.cuda_graph_scope + ) and isinstance(layer.mlp, MoELayer): + # moe layer. + return True + if 'mlp' in config.cuda_graph_scope and isinstance(layer.mlp, MLP): + # mlp layer. return True return False @@ -1361,18 +1381,17 @@ def __init__(self, model, config, seq_length, micro_batch_size, optimizers=[]): assert ( config.cuda_graph_impl == "transformer_engine" ), "Option cuda_graph_impl=transformer_engine not enabled." - assert "expandable_segments:True" not in os.getenv("PYTORCH_CUDA_ALLOC_CONF", ""), ( - "expandable_segments:True may not be safe when using CUDA Graphs, and may result in" - "a crash due to illegal memory access or other undefined behaviour." + assert ( + "expandable_segments:True" not in os.getenv("PYTORCH_CUDA_ALLOC_CONF", "") + or os.getenv("NCCL_GRAPH_REGISTER", "") == "0" + ), ( + "Setting NCCL_GRAPH_REGISTER=0 to avoid illegal memory access when using " + "CUDA Graph with PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True." ) - assert config.cuda_graph_scope != "full_iteration", ( + assert "full_iteration" not in config.cuda_graph_scope, ( "full_iteration cuda graph is not supported for cuda_graph_impl=transformer_engine. " "Please use cuda_graph_impl=local instead." ) - assert config.cuda_graph_scope in [ - 'full', - 'attn', - ], f"--cuda-graph-scope should be full or attn, got {config.cuda_graph_scope}." self.model = model self.config = config @@ -1455,6 +1474,16 @@ def __init__(self, model, config, seq_length, micro_batch_size, optimizers=[]): f'{len(self.flattened_callables)} graphable layers.', ) + # One helper object can only capture CUDA Graphs once. Use this flag to check if the graphs + # have been created. + self._graphs_created = False + + def graphs_created(self): + """ + Returns whether the CUDA Graphs have been created. + """ + return self._graphs_created + def _get_cuda_graph_input_data(self): """ Create the CUDA Graph capturing input data. @@ -1495,8 +1524,13 @@ def get_rotary_pos_emb(transformer_module, transformer_input): from megatron.core.transformer.identity_op import IdentityOp from megatron.core.transformer.transformer_layer import TransformerLayer - contains_self_attn = isinstance(layer, TransformerLayer) and not isinstance( - layer.self_attention, IdentityOp + contains_self_attn = ( + isinstance(layer, TransformerLayer) + and not isinstance(layer.self_attention, IdentityOp) + and ( + not self.config.cuda_graph_scope + or 'attn' in self.config.cuda_graph_scope + ) ) if is_te_min_version("1.10.0"): # te.make_graphed_callables() accepts keyword arguments since 1.10.0. @@ -1605,6 +1639,8 @@ def _start_capturing(self): """ Start capturing CUDA Graphs. """ + assert not self._graphs_created, "CUDA Graphs have already been created." + torch.distributed.barrier() gc.collect() torch.cuda.empty_cache() @@ -1638,6 +1674,8 @@ def _finish_capturing(self, start_time): gc.collect() torch.cuda.empty_cache() + self._graphs_created = True + def create_cudagraphs(self): """ Capture CUDA Graphs per TransformerLayer per microbatch. diff --git a/megatron/core/transformer/moe/fused_a2a.py b/megatron/core/transformer/moe/fused_a2a.py index 00a840f2b7f..26baca5d5b3 100644 --- a/megatron/core/transformer/moe/fused_a2a.py +++ b/megatron/core/transformer/moe/fused_a2a.py @@ -12,6 +12,8 @@ except ImportError: HAVE_DEEP_EP = False +HAVE_HYBRIDEP = False + import torch _buffer = None diff --git a/megatron/core/transformer/moe/moe_layer.py b/megatron/core/transformer/moe/moe_layer.py index d5a6be9224c..245eada51ce 100644 --- a/megatron/core/transformer/moe/moe_layer.py +++ b/megatron/core/transformer/moe/moe_layer.py @@ -1,4 +1,4 @@ -# Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved. +# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. from abc import ABC, abstractmethod from dataclasses import dataclass @@ -9,7 +9,12 @@ from megatron.core import parallel_state, tensor_parallel, utils from megatron.core.process_groups_config import ProcessGroupCollection from megatron.core.transformer.module import MegatronModule -from megatron.core.transformer.moe.moe_utils import get_default_pg_collection +from megatron.core.transformer.moe.moe_utils import ( + MoECudaGraphPartialCaptureSignal, + MoECudaGraphTensorStore, + get_default_pg_collection, + maybe_skip_or_early_return_by_cudagraph, +) from megatron.core.transformer.moe.router import TopKRouter from megatron.core.transformer.moe.token_dispatcher import ( MoEAllGatherTokenDispatcher, @@ -166,16 +171,29 @@ def __init__( if self.shared_expert_overlap: self.token_dispatcher.set_shared_experts(self.shared_experts) - def router_and_preprocess(self, hidden_states: torch.Tensor): - """Compute and preprocess token routing for dispatch. + # Cudagraph tensor store for resuming the forward pass from the end of the cudagraph. + self.cudagraph_tensor_store = MoECudaGraphTensorStore() + + @maybe_skip_or_early_return_by_cudagraph("route") + def route(self, hidden_states: torch.Tensor): + """Compute token routing for preprocessing. This method uses the router to determine which experts to send each token to, - producing routing probabilities and a mapping. It then preprocesses the - hidden states and probabilities for the token dispatcher. The original - hidden states are returned as a residual connection. + producing routing probabilities and a mapping. """ - residual = hidden_states probs, routing_map = self.router(hidden_states) + return probs, routing_map + + @maybe_skip_or_early_return_by_cudagraph("preprocess") + def preprocess( + self, hidden_states: torch.Tensor, probs: torch.Tensor, routing_map: torch.Tensor + ): + """Preprocess token routing for dispatch. + + This method preprocesses the hidden states and routing probabilities for the token + dispatcher. The original hidden states are returned as a residual connection. + """ + residual = hidden_states hidden_states, probs = self.token_dispatcher.dispatch_preprocess( hidden_states, routing_map, probs ) @@ -183,12 +201,14 @@ def router_and_preprocess(self, hidden_states: torch.Tensor): def dispatch(self, hidden_states: torch.Tensor, probs: torch.Tensor): """Dispatches tokens to assigned expert ranks via communication. + This method performs the actual communication (e.g., All-to-All) to distribute tokens and their associated probabilities to the devices hosting their assigned experts. """ return self.token_dispatcher.token_dispatch(hidden_states, probs) + @maybe_skip_or_early_return_by_cudagraph("shared_experts_compute") def shared_experts_compute(self, hidden_states: torch.Tensor): """Computes the output of the shared experts. @@ -270,8 +290,18 @@ def forward(self, hidden_states: torch.Tensor): # MoE forward: route -> dispatch -> compute -> combine def custom_forward(hidden_states): - shared_expert_output = self.shared_experts_compute(hidden_states) - hidden_states, probs, residual = self.router_and_preprocess(hidden_states) + try: + shared_expert_output = self.shared_experts_compute(hidden_states) + probs, routing_map = self.route(hidden_states) + hidden_states, probs, residual = self.preprocess(hidden_states, probs, routing_map) + except MoECudaGraphPartialCaptureSignal as e: + # This signal is raised from the maybe_skip_or_early_return_by_cudagraph decorator. + # It means we should early-return from the MoE layer forward pass. + # This happens when we are partially capturing the CUDA graph of the MoE layer, + # like cuda_graph_scope=["moe_router", "moe_preprocess"]. + # We need to return the intermediate tensors as CUDA graph outputs. + return e.get_early_return_outputs(hidden_states, shared_expert_output) + dispatched_input, probs = self.dispatch(hidden_states, probs) output, mlp_bias = self.routed_experts_compute(dispatched_input, probs, residual) output = self.combine(output, shared_expert_output) @@ -279,7 +309,7 @@ def custom_forward(hidden_states): if self.moe_layer_recompute: if self.config.fp8: - output, mlp_bias = te_checkpoint( + outputs = te_checkpoint( custom_forward, False, tensor_parallel.random.get_cuda_rng_tracker, @@ -287,11 +317,11 @@ def custom_forward(hidden_states): hidden_states, ) else: - output, mlp_bias = tensor_parallel.checkpoint(custom_forward, False, hidden_states) + outputs = tensor_parallel.checkpoint(custom_forward, False, hidden_states) else: - output, mlp_bias = custom_forward(hidden_states) + outputs = custom_forward(hidden_states) - return output, mlp_bias + return outputs def backward_dw(self): """Compute weight gradients for experts and shared experts.""" diff --git a/megatron/core/transformer/moe/moe_utils.py b/megatron/core/transformer/moe/moe_utils.py index dc857129834..5a0793ef5b9 100644 --- a/megatron/core/transformer/moe/moe_utils.py +++ b/megatron/core/transformer/moe/moe_utils.py @@ -1,12 +1,14 @@ -# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved. +# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. import math +from dataclasses import dataclass from typing import List, Optional, Union import torch from megatron.core import parallel_state from megatron.core.process_groups_config import ProcessGroupCollection +from megatron.core.transformer.cuda_graphs import is_graph_capturing try: import transformer_engine as te # pylint: disable=unused-import @@ -905,12 +907,16 @@ class RandomSTE(torch.autograd.Function): """ generator = None + random_logits = None @staticmethod def forward(ctx, logits): """ Forward pass returns random logits with rank-specific seed. """ + if is_graph_capturing() and RandomSTE.random_logits is not None: + return RandomSTE.random_logits + if RandomSTE.generator is None: global_rank = torch.distributed.get_rank() base_seed = 42 @@ -918,8 +924,8 @@ def forward(ctx, logits): RandomSTE.generator = torch.Generator(device=logits.device) RandomSTE.generator.manual_seed(seed) - random_logits = logits.clone().normal_(generator=RandomSTE.generator) - return random_logits + RandomSTE.random_logits = logits.clone().normal_(generator=RandomSTE.generator) + return RandomSTE.random_logits @staticmethod def backward(ctx, grad_output): @@ -1028,3 +1034,242 @@ def get_default_pg_collection(): with_context_parallel=True ) return pg_collection + + +class MoECudaGraphPartialCaptureSignal(Exception): + """ + Used to early-return from a MoE layer forward pass in CUDA graph capture. + This signal is raised when we are partially capturing the CUDA graph of the MoE layer, + and the related intermediate tensors are recorded in self.kwargs. + Call self.get_early_return_outputs() to collect the CUDA graph outputs. + """ + + def __init__(self, moe_layer, return_step: str, **kwargs): + self.moe_layer = moe_layer + self.return_step = return_step + self.kwargs = kwargs + + def get_early_return_outputs( + self, hidden_states: torch.Tensor, shared_expert_output: torch.Tensor + ): + """ + Get the CUDA graph early return outputs for the MoE layer, including the intermediate + tensors and the intermediate attributes of the token dispatcher. + """ + if self.return_step == "route": + # Capturing the router step returns three intermediate tensors: + # hidden states, routing probabilities, and routing map. + outputs = [hidden_states, self.kwargs['probs'], self.kwargs['routing_map']] + elif self.return_step == "preprocess": + # Capturing the preprocess step returns three intermediate tensors: + # hidden states, routing probabilities, and residual connection. + # It also returns the intermediate attributes of the token dispatcher, recorded in + # "token_dispatcher.cudagraph_attrs". + outputs = [self.kwargs['hidden_states'], self.kwargs['probs'], self.kwargs['residual']] + valid_cudagraph_attrs = [] + for attr_name in self.moe_layer.token_dispatcher.cudagraph_attrs: + hier_attr_name = attr_name.split('.') + attr = self.moe_layer.token_dispatcher + for name in hier_attr_name: + attr = getattr(attr, name, None) + if attr is None: + break + if isinstance(attr, torch.Tensor): + outputs.append(attr) + valid_cudagraph_attrs.append(attr_name) + if self.moe_layer.token_dispatcher.valid_cudagraph_attrs is None: + self.moe_layer.token_dispatcher.valid_cudagraph_attrs = valid_cudagraph_attrs + else: + assert ( + self.moe_layer.token_dispatcher.valid_cudagraph_attrs == valid_cudagraph_attrs + ), ( + "valid_cudagraph_attrs mismatch: " + f"{self.moe_layer.token_dispatcher.valid_cudagraph_attrs} != " + f"{valid_cudagraph_attrs}" + ) + # Also return the shared expert output, if it is not None. + if shared_expert_output is not None: + outputs.append(shared_expert_output) + return outputs + + +@dataclass +class MoECudaGraphTensorStore: + """Storage for tensors used in CUDA graph replay for MoE layers. + + This dataclass stores intermediate tensors computed during CUDA graph replay + that need to be resumed from the end of the CUDA graph scope to skip redundant computations. + + Attributes: + hidden_states (Optional[torch.Tensor]): The hidden states output from the CUDA graph replay. + probs (Optional[torch.Tensor]): The routing probabilities for each token-expert pair. + routing_map (Optional[torch.Tensor]): The sparse mapping indicating which experts + were selected for each token. Used to skip the normal router step. + residual (Optional[torch.Tensor]): The residual connection tensor before routing. + Used to skip the normal preprocess step. + shared_expert_output (Optional[torch.Tensor]): The output from shared experts + computation. Used to skip the normal shared expert computation step. + """ + + hidden_states: Optional[torch.Tensor] = None + probs: Optional[torch.Tensor] = None + routing_map: Optional[torch.Tensor] = None + residual: Optional[torch.Tensor] = None + shared_expert_output: Optional[torch.Tensor] = None + + def is_empty(self) -> bool: + """Check if the store has any non-None tensors. + + Returns: + bool: True if all fields are None, False otherwise. + """ + return all( + getattr(self, field_name) is None + for field_name in [ + 'hidden_states', + 'probs', + 'routing_map', + 'residual', + 'shared_expert_output', + ] + ) + + def set(self, **kwargs): + """Set the tensors in the store from keyword arguments.""" + for field_name, value in kwargs.items(): + assert field_name in [ + 'hidden_states', + 'probs', + 'routing_map', + 'residual', + 'shared_expert_output', + ], f"Invalid field name: {field_name}" + if value is not None: + assert isinstance( + value, torch.Tensor + ), f"Value must be a torch.Tensor, got {type(value)} for field {field_name}" + setattr(self, field_name, value) + + def clear(self): + """Reset all stored tensors to None.""" + for field_name in [ + 'hidden_states', + 'probs', + 'routing_map', + 'residual', + 'shared_expert_output', + ]: + setattr(self, field_name, None) + + +def maybe_skip_or_early_return_by_cudagraph(step_condition): + """ + Decorator to skip certain codepaths in the MoE layer forward pass in CUDA graph replay, + or early return from the MoE layer forward pass in CUDA graph capture. + + Args: + step_condition: The step condition to check. Can be "shared_experts_compute", "route", + or "preprocess". If "shared_experts_compute", the shared experts computation will be + skipped in replay if it is in the CUDA graph scope. If "route" or "preprocess", the + router or preprocess will be skipped in replay if it is in the CUDA graph scope, or + early return from the MoE layer forward pass if it is in CUDA graph capturing mode. + + Returns: + A decorator function that wraps the MoE layer forward pass. + """ + + def maybe_raise_signal(moe_layer, **kwargs): + """ + Check if the MoE layer should early return for CUDA graph capture. + If so, raise a MoECudaGraphPartialCaptureSignal. + """ + if ( + moe_layer.config.cuda_graph_impl == "transformer_engine" + and moe_layer.training + and is_graph_capturing() + ): + if ( + step_condition == "route" + and 'moe_router' in moe_layer.config.cuda_graph_scope + and 'moe_preprocess' not in moe_layer.config.cuda_graph_scope + ): + raise MoECudaGraphPartialCaptureSignal(moe_layer, "route", **kwargs) + elif ( + step_condition == "preprocess" + and 'moe_preprocess' in moe_layer.config.cuda_graph_scope + ): + raise MoECudaGraphPartialCaptureSignal(moe_layer, "preprocess", **kwargs) + + def decorator(func): + def wrapped_func(moe_layer, *args, **kwargs): + """ + Check if we should skip executing the original function based on the current + step condition and the tensor store status. If the tensor can be found in the store, + it indicates that it is already computed by the CUDA graph replay, so we can skip it. + Otherwise, we execute the original function and check if we should raise a signal to + early return in CUDA graph capture. + """ + # The non-cudagraph codepath just calls the original function. + if not is_graph_capturing() and moe_layer.cudagraph_tensor_store.is_empty(): + return func(moe_layer, *args, **kwargs) + + assert ( + not is_graph_capturing() or moe_layer.cudagraph_tensor_store.is_empty() + ), "cudagraph_tensor_store cannot be used when it is capturing cuda graph." + if step_condition == "shared_experts_compute": + if moe_layer.cudagraph_tensor_store.shared_expert_output is None: + # Don't skip the shared expert computation. + shared_expert_output = func(moe_layer, *args, **kwargs) + else: + # Skip the shared expert computation and get value from store. + shared_expert_output = moe_layer.cudagraph_tensor_store.shared_expert_output + return shared_expert_output + elif step_condition == "route": + if moe_layer.cudagraph_tensor_store.probs is None: + # Don't skip the router. + assert ( + moe_layer.cudagraph_tensor_store.routing_map is None + and moe_layer.cudagraph_tensor_store.residual is None + ), "both routing_map and residual must be None if probs is None" + probs, routing_map = func(moe_layer, *args, **kwargs) + + # Maybe early return after the router. + maybe_raise_signal(moe_layer, probs=probs, routing_map=routing_map) + else: + # Skip the router and get value from store. + assert ( + moe_layer.cudagraph_tensor_store.routing_map is not None + or moe_layer.cudagraph_tensor_store.residual is not None + ), "either routing_map or residual must be given if probs is given" + probs, routing_map = ( + moe_layer.cudagraph_tensor_store.probs, + moe_layer.cudagraph_tensor_store.routing_map, + ) + return probs, routing_map + elif step_condition == "preprocess": + if moe_layer.cudagraph_tensor_store.residual is None: + # Don't skip the preprocess. + hidden_states, probs, residual = func(moe_layer, *args, **kwargs) + + # Maybe early return after the preprocess. + maybe_raise_signal( + moe_layer, hidden_states=hidden_states, probs=probs, residual=residual + ) + else: + # Skip the preprocess and get value from store. + assert ( + moe_layer.cudagraph_tensor_store.probs is not None + ), "probs must not be None if residual is not None" + assert ( + moe_layer.cudagraph_tensor_store.routing_map is None + ), "routing_map must be None if residual is not None" + hidden_states, probs, residual = ( + moe_layer.cudagraph_tensor_store.hidden_states, + moe_layer.cudagraph_tensor_store.probs, + moe_layer.cudagraph_tensor_store.residual, + ) + return hidden_states, probs, residual + + return wrapped_func + + return decorator diff --git a/megatron/core/transformer/moe/token_dispatcher.py b/megatron/core/transformer/moe/token_dispatcher.py index 82fb7b00583..7e3a831bbf3 100644 --- a/megatron/core/transformer/moe/token_dispatcher.py +++ b/megatron/core/transformer/moe/token_dispatcher.py @@ -1,4 +1,4 @@ -# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved. +# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. import logging from abc import ABC, abstractmethod @@ -33,6 +33,8 @@ from megatron.core.transformer.moe.shared_experts import SharedExpertMLP from megatron.core.transformer.transformer_config import TransformerConfig +logger = logging.getLogger(__name__) + """ We use the following notation throughout this file: H: hidden size B: micro batch size @@ -71,6 +73,11 @@ def __init__( self.tp_rank = utils.get_pg_rank(self.tp_group) self.ep_size = utils.get_pg_size(self.ep_group) + # Attributes that need to be captured in cudagraph. These attributes are returned + # as cudagraph outputs when the cuda_graph_scope contains moe_preprocess. + self.cudagraph_attrs = [] + self.valid_cudagraph_attrs = None + @abstractmethod def dispatch_preprocess( self, tokens: torch.Tensor, routing_map: torch.Tensor, probs: torch.Tensor @@ -228,6 +235,10 @@ def __init__( # device token permutation is enabled and **AllGahter** is performed. self.global_local_map = None + # Attributes that need to be captured in cudagraph. These attributes are returned + # as cudagraph outputs when the cuda_graph_scope contains moe_preprocess. + self.cudagraph_attrs = ['routing_map'] + def dispatch_preprocess( self, hidden_states: torch.Tensor, routing_map: torch.Tensor, probs: torch.Tensor ): @@ -420,12 +431,38 @@ def __init__( "before_finish": 3, "no_sync": 4, } - self.cuda_dtoh_point = "before_permutation_1" + if ( + config.cuda_graph_impl == "transformer_engine" + and 'moe_preprocess' in config.cuda_graph_scope + ): + self.cuda_dtoh_point = "before_ep_alltoall" + else: + self.cuda_dtoh_point = "before_permutation_1" if MoEAlltoAllTokenDispatcher.cuda_dtoh_stream is None: MoEAlltoAllTokenDispatcher.cuda_dtoh_stream = torch.cuda.Stream() + # Attributes that need to be captured in cudagraph. These attributes are returned + # as cudagraph outputs when the cuda_graph_scope contains moe_preprocess. + self.cudagraph_attrs = [ + 'tokens_per_expert', + 'input_splits', + 'output_splits', + 'output_splits_tp', + 'num_out_tokens', + 'num_global_tokens_per_local_expert', + 'reversed_local_input_permutation_mapping', + 'routing_map', + ] + self.shared_experts = None + def set_shared_experts(self, shared_experts): + """Set shared expert to the dispatcher.""" + super().set_shared_experts(shared_experts) + if shared_experts.use_shared_expert_gate: + self.cudagraph_attrs.append('shared_experts.gate_score') + self.cudagraph_attrs.append('shared_experts.cached_fc1_input') + def preprocess(self, routing_map: torch.Tensor) -> torch.Tensor: """ Preprocesses the token routing map for All-to-All communication and token permutation. @@ -989,7 +1026,9 @@ def dispatch( # DeepEP only supports float32 probs if self.token_probs.dtype != torch.float32: if self.token_probs.dtype in [torch.bfloat16, torch.float16]: - print("DeepEP only supports float32 probs, please set --moe-router-dtype=fp32") + logger.info( + "DeepEP only supports float32 probs, please set --moe-router-dtype=fp32" + ) self.token_probs = self.token_probs.float() # downcast or upcast hidden_states, dispatched_indices, dispatched_probs, num_tokens_per_expert, handle = ( fused_dispatch( @@ -1175,6 +1214,7 @@ def __init__( num_experts=self.tp_size * self.config.num_moe_experts, config=self.config, ) + self.cudagraph_attrs = ['_comm_manager.token_probs', '_comm_manager.token_indices'] def set_shared_experts(self, shared_experts): raise NotImplementedError( diff --git a/megatron/core/transformer/transformer_block.py b/megatron/core/transformer/transformer_block.py index aead6133f22..b61124f04a0 100755 --- a/megatron/core/transformer/transformer_block.py +++ b/megatron/core/transformer/transformer_block.py @@ -1,4 +1,4 @@ -# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved. +# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. import logging from contextlib import nullcontext from dataclasses import dataclass @@ -522,7 +522,7 @@ def _should_call_local_cudagraph(self, *args, **kwargs): kwargs.get('inference_context') is not None or kwargs.get('inference_params') is not None ) - and self.config.cuda_graph_scope == 'full_iteration' + and 'full_iteration' in self.config.cuda_graph_scope ): if kwargs['inference_context'].is_static_batching(): using_cuda_graph = kwargs['inference_context'].is_decode_only() diff --git a/megatron/core/transformer/transformer_config.py b/megatron/core/transformer/transformer_config.py index 7dba2c35017..638c8645642 100644 --- a/megatron/core/transformer/transformer_config.py +++ b/megatron/core/transformer/transformer_config.py @@ -651,11 +651,10 @@ class TransformerConfig(ModelParallelConfig): excluding optimizer) is enabled. "transformer_engine": capture the CUDA graph using TE make_graphed_callables().""" - cuda_graph_scope: str = "full" + cuda_graph_scope: Optional[List[str]] = None """Determines the CUDA graphs capturing scope. - When cuda_graph_impl is set to "transformer_engine", valid values are "full" and "attn". - "Full" scope captures a whole Transformer layer. "Attn" scope only captures operations in - TransformerLayer._forward_attention(). + When cuda_graph_impl is set to "transformer_engine", valid values are "attn", "mlp", "moe", + "moe_router", "moe_preprocess", "mamba". None means ["attn", "mlp"]. When cuda_graph_impl is set to "local", "full_iteration" can be specified as cuda_graph_scope to enable whole iteration CUDA graph. All other values enable layerwise CUDA graph.""" @@ -1421,24 +1420,133 @@ def __post_init__(self): ], f"Invalid cuda graph implementation: {self.cuda_graph_impl}" if self.cpu_offloading: raise ValueError("CUDA graphs not supported with CPU offloading.") - if self.recompute_granularity: - if ( - self.recompute_granularity != "selective" - or self.cuda_graph_impl != "transformer_engine" - or self.cuda_graph_scope != "attn" - ): - raise ValueError("CUDA graphs not supported with activation recomputation.") + + if self.cuda_graph_scope is None: + self.cuda_graph_scope = [] + elif not isinstance(self.cuda_graph_scope, list): + assert isinstance(self.cuda_graph_scope, str), ( + "cuda_graph_scope must be a string or a list of strings, " + f"got {self.cuda_graph_scope}." + ) + self.cuda_graph_scope = [self.cuda_graph_scope] + + if self.cuda_graph_impl == "local": + assert not self.cuda_graph_scope or self.cuda_graph_scope == ["full_iteration"], ( + "For local cuda graph implementation, the only valid value " + "for cuda_graph_scope is full_iteration. " + "To use other scopes, use cuda_graph_impl=transformer_engine." + ) + + if self.cuda_graph_impl == "transformer_engine": + assert "full_iteration" not in self.cuda_graph_scope, ( + "To use full iteration cuda graph, please use " + "cuda_graph_impl=transformer_engine instead of cuda_graph_impl=local." + ) + for scope in self.cuda_graph_scope: + assert scope in [ + 'attn', + 'mlp', + 'moe', + 'moe_router', + 'moe_preprocess', + 'mamba', + ], ( + "--cuda-graph-scope should be attn, mlp, moe, moe_router, moe_preprocess, " + f"or mamba, got {self.cuda_graph_scope}." + ) + + assert ( + 'moe' not in self.cuda_graph_scope or 'moe_router' not in self.cuda_graph_scope + ), 'cuda_graph_scope must not contain both moe and moe_router.' + if 'moe_preprocess' in self.cuda_graph_scope: + assert ( + 'moe_router' in self.cuda_graph_scope + ), 'moe_preprocess cuda graph is only supported with moe_router cuda graph.' + if self.num_moe_experts is None or self.num_moe_experts <= 1: + assert ( + 'moe' not in self.cuda_graph_scope + and 'moe_router' not in self.cuda_graph_scope + ), 'moe cuda graph is only supported for MoE.' else: - for module in self.recompute_modules: - if module in ['core_attn', 'mla_up_proj']: - raise ValueError( - f'attn cuda graph is not supported with {module} recompute.' + if self.moe_layer_freq == 1 or ( + isinstance(self.moe_layer_freq, list) and 0 not in self.moe_layer_freq + ): + assert 'mlp' not in self.cuda_graph_scope, ( + 'mlp cuda graph is only supported for dense layers, ' + 'but not found in the model.' + ) + if ( + self.moe_expert_capacity_factor is None + or not self.moe_pad_expert_input_to_capacity + ): + assert ( + 'moe' not in self.cuda_graph_scope + ), 'moe cuda graph is only supported with drop-padding MoE.' + if self.moe_token_dispatcher_type == 'alltoall' and ( + self.moe_expert_capacity_factor is not None + or self.moe_router_padding_for_fp8 + ): + assert 'moe_preprocess' not in self.cuda_graph_scope, ( + 'moe_preprocess cuda graph is not supported when there are ' + 'DtoH copies and synchronizations in the preprocess step.' ) + + if self.recompute_granularity: + if self.recompute_granularity != "selective" or not self.cuda_graph_scope: + raise ValueError( + "Full-layer CUDA graphs not supported with activation recomputation." + ) + elif self.cuda_graph_scope != ['full_iteration']: + # For scoped CUDA graphs, only the non-graphed parts of the layer can be + # recomputed. So check if there are overlaps between the recomputed parts + # and the graphed parts. + if "attn" in self.cuda_graph_scope: + for module in self.recompute_modules: + if module in ['core_attn', 'mla_up_proj']: + raise ValueError( + f'attn cuda graph is not supported with {module} recompute.' + ) + if "mlp" in self.cuda_graph_scope and "mlp" in self.recompute_modules: + raise ValueError(f'mlp cuda graph is not supported with mlp recompute.') + if "moe" in self.cuda_graph_scope: + for module in self.recompute_modules: + if module in ['moe_act', 'moe', 'shared_experts']: + raise ValueError( + f'moe cuda graph is not supported with {module} recompute.' + ) + if "moe_router" in self.cuda_graph_scope: + for module in self.recompute_modules: + if module in ['moe', 'shared_experts']: + raise ValueError( + f'moe_router cuda graph is not supported with {module} ' + 'recompute.' + ) if "layernorm" in self.recompute_modules: - warnings.warn( - "input_layernorm recompute is not supported with attention " - "cudagraph. Will only recompute the pre_mlp_layernorm." - ) + if ( + "attn" in self.cuda_graph_scope + and "mlp" in self.cuda_graph_scope + and ( + "moe" in self.cuda_graph_scope + or "moe_router" in self.cuda_graph_scope + ) + ): + raise ValueError( + 'cuda graph is not supported with layernorm recompute.' + ) + if "attn" in self.cuda_graph_scope: + warnings.warn( + "input_layernorm recompute is not supported with attention " + "cudagraph. Will only recompute the pre_mlp_layernorm." + ) + if ( + "mlp" in self.cuda_graph_scope + or "moe" in self.cuda_graph_scope + or "moe_router" in self.cuda_graph_scope + ): + warnings.warn( + "pre_mlp_layernorm recompute is not supported with mlp/moe " + "cudagraph. Will only recompute the input_layernorm." + ) if self.moe_token_dispatcher_type in ["allgather"]: if self.variable_seq_lengths is True: diff --git a/megatron/core/transformer/transformer_layer.py b/megatron/core/transformer/transformer_layer.py index a5babece9d0..14e2dfe36a5 100644 --- a/megatron/core/transformer/transformer_layer.py +++ b/megatron/core/transformer/transformer_layer.py @@ -1,4 +1,4 @@ -# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. import logging import warnings @@ -15,6 +15,7 @@ from megatron.core.dist_checkpointing.utils import apply_prefix_mapping from megatron.core.packed_seq_params import PackedSeqParams from megatron.core.process_groups_config import ProcessGroupCollection +from megatron.core.transformer.cuda_graphs import is_graph_capturing from megatron.core.transformer.enums import LayerType from megatron.core.transformer.identity_op import IdentityFuncOp, IdentityOp from megatron.core.transformer.mlp import MLP @@ -371,19 +372,29 @@ def __init__( # [Module 9: BiasDropoutFusion] self.mlp_bda = build_module(submodules.mlp_bda) + self.is_moe_layer = isinstance(self.mlp, MoELayer) + self.recompute_input_layernorm = False self.recompute_pre_mlp_layernorm = False self.recompute_mlp = False if self.config.recompute_granularity == 'selective': if "layernorm" in self.config.recompute_modules: - if ( - not isinstance(self.input_layernorm, IdentityOp) - and self.config.cuda_graph_impl == "none" + if not isinstance(self.input_layernorm, IdentityOp) and ( + self.config.cuda_graph_impl == "none" + or 'attn' not in self.config.cuda_graph_scope ): self.recompute_input_layernorm = True if self.config.fp8: self.self_attention.set_for_recompute_input_layernorm() - if not isinstance(self.pre_mlp_layernorm, IdentityOp): + if not isinstance(self.pre_mlp_layernorm, IdentityOp) and ( + self.config.cuda_graph_impl == "none" + or (not self.is_moe_layer and 'mlp' not in self.config.cuda_graph_scope) + or ( + self.is_moe_layer + and 'moe' not in self.config.cuda_graph_scope + and 'moe_router' not in self.config.cuda_graph_scope + ) + ): self.recompute_pre_mlp_layernorm = True if self.config.fp8: if isinstance(self.mlp, MoELayer): @@ -395,7 +406,7 @@ def __init__( set_save_original_input(self.mlp.linear_fc1) if "mlp" in self.config.recompute_modules: - if not isinstance(self.mlp, MoELayer): + if not self.is_moe_layer: self.recompute_mlp = True # @jcasper how should we handle nvfuser? @@ -584,7 +595,19 @@ def _forward_mlp(self, hidden_states, inference_context=None): and not isinstance(self.mlp, IdentityOp) ) - if self.recompute_mlp: + if ( + self.is_moe_layer + and self.config.cuda_graph_impl == "transformer_engine" + and self.training + and is_graph_capturing() + and 'moe_router' in self.config.cuda_graph_scope + ): + assert ( + not self.recompute_pre_mlp_layernorm + ), "Recomputation is not supported for CUDA graph." + cudagraph_outputs = self.mlp(pre_mlp_layernorm_output) + return cudagraph_outputs + [residual] + elif self.recompute_mlp: if self.config.fp8: # import here to avoid circular import from megatron.core.extensions.transformer_engine import te_checkpoint @@ -613,7 +636,6 @@ def _forward_mlp(self, hidden_states, inference_context=None): bias_chunks = [bias for _, bias in outputs if bias is not None] bias_output = torch.stack(bias_chunks, dim=0).sum(dim=0) if bias_chunks else None mlp_output_with_bias = (mlp_output, bias_output) - else: mlp_output_with_bias = self.mlp(pre_mlp_layernorm_output) @@ -625,6 +647,20 @@ def _forward_mlp(self, hidden_states, inference_context=None): ) nvtx_range_pop(suffix="mlp") + return self._forward_post_mlp(mlp_output_with_bias, residual) + + def _forward_post_mlp(self, mlp_output_with_bias, residual): + """ + Perform operations after the MLP computation. + + Args: + mlp_output_with_bias (Tensor): Output tensor of the MLP layer with bias. + residual (Tensor): Residual tensor. + + Returns: + output (Tensor): Transformed hidden states of shape [s, b, h]. + """ + # TODO: could we move `bias_dropout_add_exec_handler` itself # inside the module provided in the `bias_dropout_add_spec` module? nvtx_range_push(suffix="mlp_bda") @@ -679,7 +715,9 @@ def get_layer_static_inputs(self, seq_length, micro_batch_size): """ static_inputs = super().get_layer_static_inputs(seq_length, micro_batch_size) - if not isinstance(self.self_attention, IdentityOp): + if not isinstance(self.self_attention, IdentityOp) and ( + not self.config.cuda_graph_scope or 'attn' in self.config.cuda_graph_scope + ): slen_per_cp = seq_length // self.config.context_parallel_size static_inputs["attention_mask"] = ( ~(torch.tril(torch.ones((slen_per_cp, seq_length))).bool()) @@ -693,18 +731,28 @@ def _get_submodules_under_cudagraphs(self): """ Get the submodules that are covered by cudagraphs. """ - if self.config.cuda_graph_scope == 'full': - submodules = [self] - else: - assert ( - self.config.cuda_graph_scope == 'attn' - ), f"Invalid cuda_graph_scope {self.config.cuda_graph_scope}" - submodules = [ + if not self.config.cuda_graph_scope: + return super()._get_submodules_under_cudagraphs() + + submodules = [] + if 'attn' in self.config.cuda_graph_scope: + submodules += [ self.input_layernorm, self.self_attention, self.pre_cross_attn_layernorm, self.cross_attention, ] + if (not self.is_moe_layer and 'mlp' in self.config.cuda_graph_scope) or ( + self.is_moe_layer and 'moe' in self.config.cuda_graph_scope + ): + submodules += [self.pre_mlp_layernorm, self.mlp] + elif self.is_moe_layer and 'moe_router' in self.config.cuda_graph_scope: + submodules += [self.pre_mlp_layernorm, self.mlp.router] + if ( + self.config.moe_shared_expert_intermediate_size is not None + and not self.config.moe_shared_expert_overlap + ): + submodules += [self.mlp.shared_experts] return submodules def _te_cuda_graph_capture(self, *args, **kwargs): @@ -715,12 +763,31 @@ def _te_cuda_graph_capture(self, *args, **kwargs): attribute can be set to control the scope of the CUDA graph. 2. If context is None, it cannot be returned as output. """ - hidden_states, context = self._forward_attention(*args, **kwargs) - - if self.config.cuda_graph_scope == "full": + context = None + if not self.config.cuda_graph_scope or 'attn' in self.config.cuda_graph_scope: + hidden_states, context = self._forward_attention(*args, **kwargs) + else: + if len(args) > 0: + hidden_states = args[0] + else: + hidden_states = kwargs.pop("hidden_states") + + if ( + not self.config.cuda_graph_scope + or (not self.is_moe_layer and 'mlp' in self.config.cuda_graph_scope) + or ( + self.is_moe_layer + and ( + 'moe' in self.config.cuda_graph_scope + or 'moe_router' in self.config.cuda_graph_scope + ) + ) + ): hidden_states = self._forward_mlp(hidden_states) - cuda_graph_outputs = [hidden_states] - + if not isinstance(hidden_states, list) and not isinstance(hidden_states, tuple): + cuda_graph_outputs = [hidden_states] + else: + cuda_graph_outputs = list(hidden_states) if context is not None: cuda_graph_outputs.append(context) return tuple(cuda_graph_outputs) @@ -732,6 +799,11 @@ def _te_cuda_graph_replay(self, *args, **kwargs): However, CUDA graph accepts only Tensor inputs. Hence, `inference_context` and `packed_seq_params` are excluded from input list. """ + context = None + if self.config.cuda_graph_scope and 'attn' not in self.config.cuda_graph_scope: + hidden_states, context = self._forward_attention(*args, **kwargs) + args = (hidden_states,) + kwargs = {} assert (kwargs.get('inference_context') is None) and ( kwargs.get('packed_seq_params') is None @@ -741,19 +813,69 @@ def _te_cuda_graph_replay(self, *args, **kwargs): "For inference cuda graph, please use cuda_graph_impl=local instead." ) - cuda_graph_output = super()._te_cuda_graph_replay(*args, **kwargs) + cuda_graph_output = list(super()._te_cuda_graph_replay(*args, **kwargs)) if kwargs.get('context') is not None: - context = cuda_graph_output[-1] - cuda_graph_output = cuda_graph_output[:-1] + context = cuda_graph_output.pop() + + if ( + not self.config.cuda_graph_scope + or (not self.is_moe_layer and 'mlp' in self.config.cuda_graph_scope) + or (self.is_moe_layer and 'moe' in self.config.cuda_graph_scope) + ): + # CUDA Graph captures the whole MLP/MoE part. CUDA Graph output is the layer output. + assert len(cuda_graph_output) == 1, "CUDA Graph output should be the layer output." + output = cuda_graph_output.pop() + elif self.is_moe_layer and 'moe_router' in self.config.cuda_graph_scope: + # CUDA Graph partially captures the MoE. + # The rest of the layer should go to the normal pass. + shared_expert_output, routing_map, residual = None, None, None + mlp_residual = cuda_graph_output.pop() + if ( + self.config.moe_shared_expert_intermediate_size is not None + and not self.config.moe_shared_expert_overlap + ): + # The shared expert output is the fourth element in the CUDA graph output. + shared_expert_output = cuda_graph_output.pop() + + # Split cudagraph outputs into function outputs and attribute outputs, and + # process them separately. Function outputs should have three tensors. + func_output, attr_outputs = cuda_graph_output[:3], cuda_graph_output[3:] + if 'moe_preprocess' in self.config.cuda_graph_scope: + hidden_states, probs, residual = func_output + valid_cudagraph_attrs = self.mlp.token_dispatcher.valid_cudagraph_attrs + assert len(attr_outputs) == len( + valid_cudagraph_attrs + ), f"attr_outputs: {len(attr_outputs)} != {len(valid_cudagraph_attrs)}" + for i, attr_name in enumerate(valid_cudagraph_attrs): + hier_attr_name = attr_name.split('.') + attr = self.mlp.token_dispatcher + for name in hier_attr_name[:-1]: + attr = getattr(attr, name) + setattr(attr, hier_attr_name[-1], attr_outputs[i]) + else: + hidden_states, probs, routing_map = func_output + assert not attr_outputs, "cuda_graph_attr_outputs should be empty" + + # Resume the MoELayer forward pass from the end of the CUDA graph scope. + # The MoE layer will skip redundant computations when we pass in the calculated values + # through the keyword arguments. See MoELayer.forward docstring for more details. + nvtx_range_push(suffix="mlp") + self.mlp.cudagraph_tensor_store.set( + hidden_states=hidden_states, + probs=probs, + routing_map=routing_map, + residual=residual, + shared_expert_output=shared_expert_output, + ) + mlp_output_with_bias = self.mlp(hidden_states) + self.mlp.cudagraph_tensor_store.clear() + nvtx_range_pop(suffix="mlp") + + output = self._forward_post_mlp(mlp_output_with_bias, mlp_residual) else: - context = None - if self.config.cuda_graph_scope == "attn": - # CUDA Graph only covers the attention layer. Feed-forward - # layer still goes through the normal pass. + # CUDA Graph does not capture the MLP/MoE part at all. output = self._forward_mlp(*cuda_graph_output) - else: - output = cuda_graph_output[0] return output, context def _get_te_cuda_graph_replay_args(self, *args, **kwargs): @@ -826,7 +948,7 @@ def _should_call_local_cudagraph(self, *args, **kwargs): (kwargs.get('inference_context') is not None) or (kwargs.get('inference_params') is not None) ) - and self.config.cuda_graph_scope != 'full_iteration' + and 'full_iteration' not in self.config.cuda_graph_scope ): if kwargs['inference_context'].is_static_batching(): using_cuda_graph = kwargs['inference_context'].is_decode_only() diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py index b5f777a30c6..f3fa79888b2 100644 --- a/megatron/training/arguments.py +++ b/megatron/training/arguments.py @@ -746,7 +746,7 @@ def validate_args(args, defaults={}): if args.rank == 0: print('accumulate and all-reduce gradients in fp32 for ' 'bfloat16 data type.', flush=True) - if args.cuda_graph_impl == "local" and args.cuda_graph_scope=="full_iteration": + if args.cuda_graph_impl == "local" and "full_iteration" in args.cuda_graph_scope: if not args.inference_dynamic_batching: assert not args.check_for_nan_in_loss_and_grad, \ "--no-check-for-nan-in-loss-and-grad should be set with full_iteration CUDA graph" @@ -1197,9 +1197,12 @@ def validate_args(args, defaults={}): if args.transformer_impl == 'transformer_engine' and not args.te_rng_tracker: args.te_rng_tracker = True warn_rank_0("te_rng_tracker is not enabled, enabling it for CUDA graphs.", args.rank) - assert "expandable_segments:True" not in os.getenv("PYTORCH_CUDA_ALLOC_CONF", ""), ( - "expandable_segments:True may not be safe when using CUDA Graphs with some specific parallel settings. " - "The training may crash with illegal memory access." + assert ( + "expandable_segments:True" not in os.getenv("PYTORCH_CUDA_ALLOC_CONF", "") + or os.getenv("NCCL_GRAPH_REGISTER", "") == "0" + ), ( + "Setting NCCL_GRAPH_REGISTER=0 to avoid illegal memory access when using " + "CUDA Graph with PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True." ) assert ( args.recompute_granularity != 'full' @@ -1411,22 +1414,27 @@ def _add_inference_args(parser): help="Number of CUDA graph warmup steps") group.add_argument('--external-cuda-graph', action='store_true', help='Deprecated. Use --cuda-graph-impl=transformer_engine instead. ' - 'Use TE make_graphed_callables() to capture the CUDA graph.') + 'Use TE make_graphed_callables() to capture the CUDA graph. ' + 'Use --cuda-graph-scope=\"attn\", \"mlp\", \"moe\", \"moe_router\", \"moe_preprocess\", \"mamba\" for partial capture. ') group.add_argument('--cuda-graph-impl', type=str, default='none', choices=['none', 'local', 'transformer_engine'], help='Determines the CUDA graph capture implementation. ' '"none": no CUDA graph. ' '"local": capture the CUDA graph using MCore local implementation. --cuda-graph-scope=\"full_iteration\" enables whole iteration CUDA graph. ' '"transformer_engine": capture the CUDA graph using TE make_graphed_callables().') - group.add_argument('--cuda-graph-scope', type=str, default='full', - choices=['full', 'attn', 'full_iteration'], - help='Determines the CUDA graphs capturing scope. Valid values are ' - '\"full\", \"attn\" and \"full_iteration\". \"Full\" scope captures a whole ' - 'Transformer layer. \"Attn\" scope only captures operations in ' - 'TransformerLayer._forward_attention(). \"ful_iteration\" scope captures a ' - 'whole iteration. ' - 'full_iteration scope is only supported with --cuda-graph-impl=local, ' - 'attn scope is only supported with --cuda-graph-impl=transformer_engine.') + group.add_argument('--cuda-graph-scope', nargs='+', type=str, default=[], + help='Determines the CUDA graphs capturing scope. ' + 'choices: "attn", "mlp", "moe", "moe_router", "moe_preprocess", "mamba", "full_iteration". ' + '"attn": captures operations in TransformerLayer._forward_attention(). ' + '"mlp": captures operations in TransformerLayer._forward_mlp() for a dense layer. ' + '"moe": captures operations in TransformerLayer._forward_mlp() for a MoE layer. ' + '"moe_router": captures operations in TransformerLayer._forward_mlp() up to MoELayer.router(), ' + 'including the shared experts if they are not overlapped with EP comm. ' + '"moe_preprocess": captures operations in MoELayer.preprocess(). Must be used together with "moe_router". ' + '"mamba": captures the mamba layer. ' + '"full_iteration": captures a whole iteration. ' + 'full_iteration scope is only supported with --cuda-graph-impl=local, other scopes are only supported with --cuda-graph-impl=transformer_engine. ' + 'If not specified, the default scope is to capture the whole Transformer layer.') group.add_argument('--use-legacy-static-engine', action='store_true', default=False, help='Use legacy static engine. (Current static engine uses dynamic engine under the hood)', dest='use_legacy_static_engine') diff --git a/megatron/training/training.py b/megatron/training/training.py index b162aa87acf..2a608dd69d8 100644 --- a/megatron/training/training.py +++ b/megatron/training/training.py @@ -1,4 +1,4 @@ -# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved. +# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. """Pretrain utilities.""" @@ -2168,7 +2168,7 @@ def train( eval_iterations = 0 # Wrap forward_backward_func for Full iteration CUDA graph forward_backward_func = get_forward_backward_func() - if args.cuda_graph_impl == "local" and args.cuda_graph_scope=="full_iteration": + if args.cuda_graph_impl == "local" and "full_iteration" in args.cuda_graph_scope: forward_backward_func = FullCudaGraphWrapper(forward_backward_func, cuda_graph_warmup_steps=args.cuda_graph_warmup_steps) def get_e2e_base_metrics(): @@ -2301,12 +2301,13 @@ def get_e2e_base_metrics(): # Capture CUDA Graphs. if ( args.cuda_graph_impl == "transformer_engine" - and iteration == args.cuda_graph_warmup_steps + and not cuda_graph_helper.graphs_created() + and iteration - start_iteration == args.cuda_graph_warmup_steps ): - if iteration > start_iteration and should_disable_forward_pre_hook(args): + if args.cuda_graph_warmup_steps > 0 and should_disable_forward_pre_hook(args): disable_forward_pre_hook(model, param_sync=False) cuda_graph_helper.create_cudagraphs() - if iteration > start_iteration and should_disable_forward_pre_hook(args): + if args.cuda_graph_warmup_steps > 0 and should_disable_forward_pre_hook(args): enable_forward_pre_hook(model) cuda_graph_helper.cuda_graph_set_manual_hooks() @@ -2381,8 +2382,11 @@ def get_e2e_base_metrics(): # Set the manual hooks here since it's not set right after the capturing. if ( args.cuda_graph_impl == "transformer_engine" - and iteration == args.cuda_graph_warmup_steps + and args.cuda_graph_warmup_steps == 0 ): + assert ( + cuda_graph_helper.graphs_created() + ), "CUDA Graphs should have been created." cuda_graph_helper.cuda_graph_set_manual_hooks() iteration += 1 @@ -2579,7 +2583,7 @@ def evaluate( eval_batch_size = args.global_batch_size eval_num_microbatches = eval_batch_size // (args.micro_batch_size * args.data_parallel_size) forward_backward_func = get_forward_backward_func() - if args.cuda_graph_impl == "local" and args.cuda_graph_scope=="full_iteration": + if args.cuda_graph_impl == "local" and "full_iteration" in args.cuda_graph_scope: forward_backward_func = FullCudaGraphWrapper(forward_backward_func, cuda_graph_warmup_steps=args.cuda_graph_warmup_steps) if eval_iters is None: diff --git a/tests/functional_tests/test_cases/moe/gpt3_moe_mcore_te_tp4_ep2_etp2_pp2_scoped_cudagraph/golden_values_dev_dgx_h100.json b/tests/functional_tests/test_cases/moe/gpt3_moe_mcore_te_tp4_ep2_etp2_pp2_scoped_cudagraph/golden_values_dev_dgx_h100.json new file mode 100644 index 00000000000..309b2533461 --- /dev/null +++ b/tests/functional_tests/test_cases/moe/gpt3_moe_mcore_te_tp4_ep2_etp2_pp2_scoped_cudagraph/golden_values_dev_dgx_h100.json @@ -0,0 +1,644 @@ +{ + "lm loss": { + "start_step": 1, + "end_step": 100, + "step_interval": 1, + "values": { + "1": 10.93663, + "2": 10.9327, + "3": 10.94263, + "4": 10.94969, + "5": 10.95052, + "6": 10.94157, + "7": 10.94484, + "8": 10.93674, + "9": 10.94996, + "10": 10.93686, + "11": 10.94102, + "12": 10.93763, + "13": 10.9235, + "14": 10.93428, + "15": 10.88791, + "16": 10.87434, + "17": 10.86896, + "18": 10.86065, + "19": 10.86311, + "20": 10.78063, + "21": 10.73125, + "22": 10.60283, + "23": 10.73278, + "24": 10.61888, + "25": 10.55212, + "26": 10.62704, + "27": 10.6391, + "28": 10.5908, + "29": 10.59809, + "30": 10.37777, + "31": 10.1201, + "32": 10.46078, + "33": 10.45538, + "34": 10.20107, + "35": 10.25779, + "36": 10.20889, + "37": 10.33688, + "38": 10.16827, + "39": 10.40875, + "40": 10.05239, + "41": 10.09432, + "42": 10.17894, + "43": 9.74205, + "44": 9.8904, + "45": 9.74009, + "46": 9.72707, + "47": 10.09139, + "48": 9.75298, + "49": 9.40106, + "50": 9.83667, + "51": 9.77071, + "52": 9.65705, + "53": 10.03051, + "54": 9.87899, + "55": 9.79604, + "56": 9.52924, + "57": 9.36583, + "58": 9.75331, + "59": 9.48065, + "60": 9.40785, + "61": 9.60145, + "62": 9.90753, + "63": 9.2583, + "64": 9.68397, + "65": 8.80003, + "66": 9.60779, + "67": 9.25408, + "68": 9.71438, + "69": 9.71682, + "70": 9.6617, + "71": 9.52466, + "72": 9.47116, + "73": 9.38822, + "74": 8.80223, + "75": 9.33966, + "76": 8.93574, + "77": 9.99333, + "78": 9.64731, + "79": 9.28114, + "80": 9.29588, + "81": 9.39589, + "82": 9.60893, + "83": 9.21629, + "84": 9.33891, + "85": 9.52979, + "86": 8.95817, + "87": 9.51641, + "88": 9.68228, + "89": 9.50664, + "90": 9.75348, + "91": 9.23465, + "92": 9.25972, + "93": 8.94517, + "94": 8.69188, + "95": 9.44591, + "96": 9.4101, + "97": 9.20087, + "98": 9.58175, + "99": 8.75818, + "100": 9.29466 + } + }, + "num-zeros": { + "start_step": 1, + "end_step": 100, + "step_interval": 1, + "values": { + "1": 22750260.0, + "2": 22953110.0, + "3": 22604450.0, + "4": 23266322.0, + "5": 22735560.0, + "6": 23061920.0, + "7": 22793342.0, + "8": 22960820.0, + "9": 22865664.0, + "10": 22950364.0, + "11": 22499674.0, + "12": 22456088.0, + "13": 22948060.0, + "14": 22384512.0, + "15": 22846272.0, + "16": 22856858.0, + "17": 22836412.0, + "18": 22590058.0, + "19": 22627048.0, + "20": 22712308.0, + "21": 22762624.0, + "22": 22816888.0, + "23": 22545124.0, + "24": 22794440.0, + "25": 22841936.0, + "26": 22549680.0, + "27": 22464820.0, + "28": 22453684.0, + "29": 22534640.0, + "30": 22636152.0, + "31": 22989488.0, + "32": 22594070.0, + "33": 22566010.0, + "34": 22855504.0, + "35": 22813688.0, + "36": 22595396.0, + "37": 22499360.0, + "38": 22926126.0, + "39": 22825392.0, + "40": 22675666.0, + "41": 22671586.0, + "42": 22682140.0, + "43": 23013940.0, + "44": 22764458.0, + "45": 22678992.0, + "46": 22915276.0, + "47": 22642868.0, + "48": 22954190.0, + "49": 23786668.0, + "50": 22934008.0, + "51": 23866222.0, + "52": 23807290.0, + "53": 24007532.0, + "54": 22871610.0, + "55": 23571284.0, + "56": 23954310.0, + "57": 24211632.0, + "58": 23914404.0, + "59": 23771838.0, + "60": 23813560.0, + "61": 23797288.0, + "62": 23739984.0, + "63": 23916692.0, + "64": 23895952.0, + "65": 24150562.0, + "66": 23796504.0, + "67": 25032232.0, + "68": 23673188.0, + "69": 23648580.0, + "70": 23903504.0, + "71": 24864636.0, + "72": 24767108.0, + "73": 24850612.0, + "74": 24132990.0, + "75": 24146528.0, + "76": 25025540.0, + "77": 24358472.0, + "78": 24910064.0, + "79": 23810516.0, + "80": 24821440.0, + "81": 25020512.0, + "82": 23851244.0, + "83": 24961024.0, + "84": 25144020.0, + "85": 24823608.0, + "86": 23153096.0, + "87": 24850204.0, + "88": 24749150.0, + "89": 22505554.0, + "90": 24059620.0, + "91": 23839038.0, + "92": 23874568.0, + "93": 24769548.0, + "94": 23992452.0, + "95": 25189838.0, + "96": 23909262.0, + "97": 24713068.0, + "98": 23832506.0, + "99": 23983474.0, + "100": 24101108.0 + } + }, + "mem-allocated-bytes": { + "start_step": 1, + "end_step": 100, + "step_interval": 1, + "values": { + "1": 763142656.0, + "2": 778734592.0, + "3": 772525056.0, + "4": 803593216.0, + "5": 803593216.0, + "6": 803593216.0, + "7": 801299456.0, + "8": 803593216.0, + "9": 801840128.0, + "10": 803593216.0, + "11": 802987008.0, + "12": 803593216.0, + "13": 802987008.0, + "14": 801299456.0, + "15": 803593216.0, + "16": 801840128.0, + "17": 803593216.0, + "18": 802987008.0, + "19": 801299456.0, + "20": 803593216.0, + "21": 801299456.0, + "22": 803593216.0, + "23": 801299456.0, + "24": 803593216.0, + "25": 801299456.0, + "26": 803593216.0, + "27": 801299456.0, + "28": 803593216.0, + "29": 801299456.0, + "30": 803593216.0, + "31": 801299456.0, + "32": 803593216.0, + "33": 801840128.0, + "34": 803593216.0, + "35": 801840128.0, + "36": 803593216.0, + "37": 802987008.0, + "38": 801299456.0, + "39": 803593216.0, + "40": 801299456.0, + "41": 803593216.0, + "42": 801840128.0, + "43": 803593216.0, + "44": 801840128.0, + "45": 803593216.0, + "46": 801840128.0, + "47": 803593216.0, + "48": 801840128.0, + "49": 803593216.0, + "50": 801840128.0, + "51": 801299456.0, + "52": 803593216.0, + "53": 801299456.0, + "54": 803593216.0, + "55": 801840128.0, + "56": 803593216.0, + "57": 801840128.0, + "58": 803593216.0, + "59": 801840128.0, + "60": 803593216.0, + "61": 801299456.0, + "62": 803593216.0, + "63": 801299456.0, + "64": 802987008.0, + "65": 803593216.0, + "66": 801299456.0, + "67": 803593216.0, + "68": 801299456.0, + "69": 803593216.0, + "70": 801840128.0, + "71": 803593216.0, + "72": 801299456.0, + "73": 803593216.0, + "74": 803593216.0, + "75": 802987008.0, + "76": 803593216.0, + "77": 801840128.0, + "78": 803593216.0, + "79": 801299456.0, + "80": 802987008.0, + "81": 803593216.0, + "82": 801840128.0, + "83": 803593216.0, + "84": 801299456.0, + "85": 802987008.0, + "86": 803593216.0, + "87": 801840128.0, + "88": 803593216.0, + "89": 801299456.0, + "90": 802987008.0, + "91": 803593216.0, + "92": 801299456.0, + "93": 803593216.0, + "94": 801299456.0, + "95": 803593216.0, + "96": 801299456.0, + "97": 803593216.0, + "98": 801299456.0, + "99": 802987008.0, + "100": 803593216.0 + } + }, + "mem-max-allocated-bytes": { + "start_step": 1, + "end_step": 100, + "step_interval": 1, + "values": { + "1": 993582592.0, + "2": 1210942464.0, + "3": 1210942464.0, + "4": 1210942464.0, + "5": 1210942464.0, + "6": 1210942464.0, + "7": 1210942464.0, + "8": 1210942464.0, + "9": 1210942464.0, + "10": 1210942464.0, + "11": 1210942464.0, + "12": 1210942464.0, + "13": 1210942464.0, + "14": 1210942464.0, + "15": 1210942464.0, + "16": 1210942464.0, + "17": 1210942464.0, + "18": 1210942464.0, + "19": 1210942464.0, + "20": 1210942464.0, + "21": 1210942464.0, + "22": 1210942464.0, + "23": 1210942464.0, + "24": 1210942464.0, + "25": 1210942464.0, + "26": 1210942464.0, + "27": 1210942464.0, + "28": 1210942464.0, + "29": 1210942464.0, + "30": 1210942464.0, + "31": 1210942464.0, + "32": 1210942464.0, + "33": 1210942464.0, + "34": 1210942464.0, + "35": 1210942464.0, + "36": 1210942464.0, + "37": 1210942464.0, + "38": 1210942464.0, + "39": 1210942464.0, + "40": 1210942464.0, + "41": 1210942464.0, + "42": 1210942464.0, + "43": 1210942464.0, + "44": 1210942464.0, + "45": 1210942464.0, + "46": 1210942464.0, + "47": 1210942464.0, + "48": 1210942464.0, + "49": 1210942464.0, + "50": 1210942464.0, + "51": 1210942464.0, + "52": 1210942464.0, + "53": 1210942464.0, + "54": 1210942464.0, + "55": 1210942464.0, + "56": 1210942464.0, + "57": 1210942464.0, + "58": 1210942464.0, + "59": 1210942464.0, + "60": 1210942464.0, + "61": 1210942464.0, + "62": 1210942464.0, + "63": 1210942464.0, + "64": 1210942464.0, + "65": 1210942464.0, + "66": 1210942464.0, + "67": 1210942464.0, + "68": 1210942464.0, + "69": 1210942464.0, + "70": 1210942464.0, + "71": 1210942464.0, + "72": 1210942464.0, + "73": 1210942464.0, + "74": 1210942464.0, + "75": 1210942464.0, + "76": 1210942464.0, + "77": 1210942464.0, + "78": 1210942464.0, + "79": 1210942464.0, + "80": 1210942464.0, + "81": 1210942464.0, + "82": 1210942464.0, + "83": 1210942464.0, + "84": 1210942464.0, + "85": 1210942464.0, + "86": 1210942464.0, + "87": 1210942464.0, + "88": 1210942464.0, + "89": 1210942464.0, + "90": 1210942464.0, + "91": 1210942464.0, + "92": 1210942464.0, + "93": 1210942464.0, + "94": 1210942464.0, + "95": 1210942464.0, + "96": 1210942464.0, + "97": 1210942464.0, + "98": 1210942464.0, + "99": 1210942464.0, + "100": 1210942464.0 + } + }, + "mtp_1 loss": { + "start_step": 1, + "end_step": 100, + "step_interval": 1, + "values": { + "1": 10.88689, + "2": 10.90485, + "3": 10.90869, + "4": 10.86903, + "5": 10.91601, + "6": 10.906, + "7": 10.90268, + "8": 10.88984, + "9": 10.90425, + "10": 10.89144, + "11": 10.93384, + "12": 10.91647, + "13": 10.91108, + "14": 10.91974, + "15": 10.88488, + "16": 10.9077, + "17": 10.87571, + "18": 10.91379, + "19": 10.9092, + "20": 10.87837, + "21": 10.87896, + "22": 10.85583, + "23": 10.88007, + "24": 10.87245, + "25": 10.85859, + "26": 10.8696, + "27": 10.87702, + "28": 10.88641, + "29": 10.88866, + "30": 10.85422, + "31": 10.79713, + "32": 10.86631, + "33": 10.8781, + "34": 10.83982, + "35": 10.84165, + "36": 10.85012, + "37": 10.85556, + "38": 10.83674, + "39": 10.86355, + "40": 10.82887, + "41": 10.8341, + "42": 10.84469, + "43": 10.78828, + "44": 10.82123, + "45": 10.78831, + "46": 10.7823, + "47": 10.82898, + "48": 10.78985, + "49": 10.71269, + "50": 10.77382, + "51": 10.76639, + "52": 10.7397, + "53": 10.80285, + "54": 10.77365, + "55": 10.76066, + "56": 10.71068, + "57": 10.66686, + "58": 10.74378, + "59": 10.69209, + "60": 10.66474, + "61": 10.7073, + "62": 10.77206, + "63": 10.61812, + "64": 10.7178, + "65": 10.49439, + "66": 10.67106, + "67": 10.57534, + "68": 10.6873, + "69": 10.6816, + "70": 10.66836, + "71": 10.64586, + "72": 10.60925, + "73": 10.56508, + "74": 10.37144, + "75": 10.51183, + "76": 10.39914, + "77": 10.75182, + "78": 10.6268, + "79": 10.46827, + "80": 10.47524, + "81": 10.51083, + "82": 10.58769, + "83": 10.4381, + "84": 10.45057, + "85": 10.55084, + "86": 10.28076, + "87": 10.51088, + "88": 10.60323, + "89": 10.50794, + "90": 10.60274, + "91": 10.38238, + "92": 10.38703, + "93": 10.23076, + "94": 10.08438, + "95": 10.42616, + "96": 10.44905, + "97": 10.32215, + "98": 10.4966, + "99": 10.04765, + "100": 10.33491 + } + }, + "iteration-time": { + "start_step": 1, + "end_step": 100, + "step_interval": 1, + "values": { + "1": 51.30209, + "2": 1.41746, + "3": 1.28029, + "4": 10.57024, + "5": 0.66643, + "6": 0.67893, + "7": 0.65727, + "8": 0.66196, + "9": 0.66227, + "10": 0.65877, + "11": 0.65828, + "12": 0.65862, + "13": 0.65727, + "14": 0.65896, + "15": 0.65851, + "16": 0.66826, + "17": 0.65878, + "18": 0.65573, + "19": 0.65631, + "20": 0.65579, + "21": 0.65091, + "22": 0.65603, + "23": 0.65158, + "24": 0.65266, + "25": 0.65816, + "26": 0.65194, + "27": 0.6541, + "28": 0.65515, + "29": 0.65439, + "30": 0.65241, + "31": 0.65597, + "32": 0.65551, + "33": 0.65318, + "34": 0.6553, + "35": 0.65725, + "36": 0.65926, + "37": 0.65606, + "38": 0.65571, + "39": 0.65846, + "40": 0.65642, + "41": 0.65509, + "42": 0.66105, + "43": 0.65448, + "44": 0.65534, + "45": 0.65304, + "46": 0.65227, + "47": 0.64871, + "48": 0.65257, + "49": 0.65485, + "50": 0.65054, + "51": 0.67883, + "52": 0.6571, + "53": 0.65671, + "54": 0.65877, + "55": 0.65584, + "56": 0.65072, + "57": 0.64951, + "58": 0.65703, + "59": 0.65106, + "60": 0.64536, + "61": 0.64416, + "62": 0.64816, + "63": 0.64084, + "64": 0.6396, + "65": 0.64182, + "66": 0.64004, + "67": 0.64101, + "68": 0.63928, + "69": 0.65723, + "70": 0.6828, + "71": 0.64052, + "72": 0.64287, + "73": 0.64136, + "74": 0.64252, + "75": 0.64617, + "76": 0.64857, + "77": 0.64304, + "78": 0.64068, + "79": 0.64048, + "80": 0.64091, + "81": 0.64179, + "82": 0.64793, + "83": 0.641, + "84": 0.64077, + "85": 0.64011, + "86": 0.64018, + "87": 0.64132, + "88": 0.63901, + "89": 0.6407, + "90": 0.64277, + "91": 0.64132, + "92": 0.64123, + "93": 0.65051, + "94": 0.65036, + "95": 0.64542, + "96": 0.64561, + "97": 0.6504, + "98": 0.64563, + "99": 0.64524, + "100": 0.65049 + } + } +} \ No newline at end of file diff --git a/tests/functional_tests/test_cases/moe/gpt3_moe_mcore_te_tp4_ep2_etp2_pp2_scoped_cudagraph/model_config.yaml b/tests/functional_tests/test_cases/moe/gpt3_moe_mcore_te_tp4_ep2_etp2_pp2_scoped_cudagraph/model_config.yaml new file mode 100644 index 00000000000..ef2b76069a1 --- /dev/null +++ b/tests/functional_tests/test_cases/moe/gpt3_moe_mcore_te_tp4_ep2_etp2_pp2_scoped_cudagraph/model_config.yaml @@ -0,0 +1,96 @@ +ENV_VARS: + CUDA_DEVICE_MAX_CONNECTIONS: 1 + NVTE_ALLOW_NONDETERMINISTIC_ALGO: 0 + NCCL_ALGO: Ring + CUBLAS_WORKSPACE_CONFIG: :4096:8 +MODEL_ARGS: + --num-layers: 13 + --hidden-size: 512 + --num-attention-heads: 8 + --mtp-num-layers: 1 + --micro-batch-size: 2 + --global-batch-size: 32 + --seq-length: 1024 + --max-position-embeddings: 1024 + --position-embedding-type: rope + --rotary-base: 10000 + --untie-embeddings-and-output-weights: true + --disable-bias-linear: true + --attention-dropout: 0.0 + --hidden-dropout: 0.0 + --train-iters: 100 + --lr-decay-iters: 320000 + --split: 949,50,1 + --distributed-backend: nccl + --lr: 0.00015 + --lr-decay-style: cosine + --min-lr: 1.0e-5 + --weight-decay: 1e-2 + --clip-grad: 1.0 + --lr-warmup-fraction: .01 + --transformer-impl: transformer_engine + --tensor-model-parallel-size: 4 + --pipeline-model-parallel-size: 2 + --expert-model-parallel-size: 2 + --expert-tensor-parallel-size: 2 + --pipeline-model-parallel-layout: Et\\|\\(tt\\|\\)*6mL # Et|(tt|)*6mL + --sequence-parallel: true + --num-experts: 8 + --use-distributed-optimizer: true + --overlap-grad-reduce: true + --overlap-param-gather: true + --moe-token-dispatcher-type: alltoall + --moe-router-load-balancing-type: global_aux_loss + --moe-router-topk: 2 + --moe-router-dtype: fp32 + --moe-router-fusion: true + --moe-router-enable-expert-bias: true + --moe-router-score-function: sigmoid + --moe-router-pre-softmax: true + --moe-ffn-hidden-size: 1024 + --moe-shared-expert-intermediate-size: 512 + --moe-grouped-gemm: true + --moe-layer-freq: ([0]*4+[1]*9) + --moe-permute-fusion: true + --deterministic-mode: true + --no-gradient-accumulation-fusion: true + --attention-softmax-in-fp32: true + --use-checkpoint-opt_param-scheduler: true + --use-mcore-models: true + --bf16: true + --fp8-format: hybrid + --fp8-recipe: blockwise + --first-last-layers-bf16: true + --no-bias-gelu-fusion: true + --recompute-granularity: selective + --recompute-modules: "[moe_act]" + --cuda-graph-impl: transformer_engine + --cuda-graph-scope: "[attn mlp moe_router moe_preprocess]" + --log-memory-to-tensorboard: true + --log-params-norm: true + --log-num-zeros-in-grad: true + --log-validation-ppl-to-tensorboard: true + --log-timers-to-tensorboard: true + --tensorboard-dir: ${TENSORBOARD_PATH} + --log-interval: 1 + --timing-log-level: 0 + --save-interval: 50 + --eval-interval: 1000 + --eval-iters: 10 + --data-path: ${DATA_PATH}/text/the_pile/shard00/my-gpt3_00_text_document + --data-cache-path: ${DATA_CACHE_PATH} + --vocab-file: ${DATA_PATH}/text/the_pile/shard00/bpe/vocab.json + --merge-file: ${DATA_PATH}/text/the_pile/shard00/bpe/merges.txt + --save: ${CHECKPOINT_SAVE_PATH} + --load: ${CHECKPOINT_LOAD_PATH} + --ckpt-fully-parallel-load: true + --ckpt-format: torch_dist + --ckpt-assume-constant-structure: true +TEST_TYPE: ckpt-resume +METRICS: + - "iteration-time" + - "lm loss" + - "num-zeros" + - "mem-allocated-bytes" + - "mem-max-allocated-bytes" + - "mtp_1 loss" diff --git a/tests/test_utils/recipes/moe.yaml b/tests/test_utils/recipes/moe.yaml index 5dc65df3619..bced5fe729f 100644 --- a/tests/test_utils/recipes/moe.yaml +++ b/tests/test_utils/recipes/moe.yaml @@ -180,6 +180,11 @@ products: - environment: [dev] scope: [mr, mr-github] platforms: [dgx_h100] + - test_case: [gpt3_moe_mcore_te_tp4_ep2_etp2_pp2_scoped_cudagraph] + products: + - environment: [dev] + scope: [mr] + platforms: [dgx_h100] ####################################################################### # Super important mr, mr-github tests that run for both DEV and LTS per mr, mr-github # ####################################################################### diff --git a/tests/unit_tests/inference/engines/test_dynamic_engine.py b/tests/unit_tests/inference/engines/test_dynamic_engine.py index b3ba2ed6d5e..97242eff292 100644 --- a/tests/unit_tests/inference/engines/test_dynamic_engine.py +++ b/tests/unit_tests/inference/engines/test_dynamic_engine.py @@ -101,7 +101,7 @@ class DynamicEngineTestConfig: return_log_probs: bool = False materialize_only_last_token_logits: bool = True skip_prompt_log_probs: bool = False - cuda_graph_scope: str = "full_iteration" + cuda_graph_scope: List[str] = None force_build_cuda_graphs: bool = False transformer_impl: str = "local" # If False, do not build cuda graphs in the tests, even if @@ -124,6 +124,9 @@ def __post_init__(self): assert self.num_tokens_total is not None self.max_sequence_length = self.num_tokens_total + if self.cuda_graph_scope is None: + self.cuda_graph_scope = ["full_iteration"] + @dataclass class DynamicEngineTestEnv: @@ -506,7 +509,7 @@ def teardown_method(self, method): ) @pytest.mark.parametrize("model_provider", ["gpt", "mamba"]) @pytest.mark.parametrize("num_cuda_graphs", [None, 1, 4]) - @pytest.mark.parametrize("cuda_graph_scope", ["full", "full_iteration"]) + @pytest.mark.parametrize("cuda_graph_scope", [[], ["full_iteration"]]) def test_simple(self, model_provider, num_cuda_graphs, cuda_graph_scope) -> None: """Simple test that runs without errors, and validates output.""" skip_if_mamba_sequence_packing_not_available(model_provider) diff --git a/tests/unit_tests/transformer/test_cuda_graphs.py b/tests/unit_tests/transformer/test_cuda_graphs.py index b92ff383d82..ee33545ce8f 100644 --- a/tests/unit_tests/transformer/test_cuda_graphs.py +++ b/tests/unit_tests/transformer/test_cuda_graphs.py @@ -1,11 +1,20 @@ # Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +import gc +import os +import sys + import pytest import torch -from megatron.core.models.gpt.gpt_layer_specs import get_gpt_layer_with_transformer_engine_spec +from megatron.core.enums import ModelType +from megatron.core.models.gpt.gpt_layer_specs import ( + get_gpt_layer_with_transformer_engine_spec, + get_gpt_mtp_block_spec, +) from megatron.core.models.gpt.gpt_model import GPTModel from megatron.core.models.mamba.mamba_layer_specs import mamba_stack_spec +from megatron.core.num_microbatches_calculator import destroy_num_microbatches_calculator from megatron.core.pipeline_parallel.schedules import set_current_microbatch from megatron.core.process_groups_config import ProcessGroupCollection from megatron.core.ssm.mamba_block import MambaStack @@ -18,6 +27,14 @@ from megatron.core.transformer.transformer_block import TransformerBlock from megatron.core.transformer.transformer_config import TransformerConfig from megatron.core.utils import is_fa_min_version, is_te_min_version +from megatron.training.arguments import core_transformer_config_from_args, parse_args, validate_args +from megatron.training.global_vars import ( + destroy_global_vars, + get_args, + set_args, + set_global_variables, +) +from megatron.training.training import setup_model_and_optimizer from tests.unit_tests.test_utilities import Utils @@ -497,6 +514,264 @@ def test_gpu_cudagraph(self): del parallel_mamba_block.layers[_].cudagraph_manager.cudagraph_runners[0].fwd_graph +def is_deep_ep_available(): + from megatron.core.transformer.moe.fused_a2a import HAVE_DEEP_EP + + return HAVE_DEEP_EP + + +def is_hybrid_ep_available(): + from megatron.core.transformer.moe.fused_a2a import HAVE_HYBRIDEP + + return HAVE_HYBRIDEP + + +class TestPartialCudaGraph: + """Test that CUDA graph outputs match non-CUDA graph outputs for various scopes.""" + + def setup_method(self, method): + self.seq_length = 512 + self.micro_batch_size = 2 + # Store original environment variable values + self.original_env = { + 'CUDA_DEVICE_MAX_CONNECTIONS': os.environ.get('CUDA_DEVICE_MAX_CONNECTIONS'), + 'NVTE_ALLOW_NONDETERMINISTIC_ALGO': os.environ.get('NVTE_ALLOW_NONDETERMINISTIC_ALGO'), + } + os.environ['CUDA_DEVICE_MAX_CONNECTIONS'] = '1' + os.environ['NVTE_ALLOW_NONDETERMINISTIC_ALGO'] = '0' + + def teardown_method(self, method): + # Restore original environment variable values + for key, value in self.original_env.items(): + if value is None: + os.environ.pop(key, None) + else: + os.environ[key] = value + Utils.destroy_model_parallel() + destroy_global_vars() + destroy_num_microbatches_calculator() + gc.collect() + + def model_provider( + self, + pre_process=True, + post_process=True, + layer_spec_fn=get_gpt_layer_with_transformer_engine_spec, + **config_kwargs, + ): + model_parallel_cuda_manual_seed(123) + args = get_args() + config = core_transformer_config_from_args(args) + transformer_layer_spec = layer_spec_fn() + if args.mtp_num_layers: + mtp_block_spec = get_gpt_mtp_block_spec( + config, transformer_layer_spec, use_transformer_engine=True + ) + else: + mtp_block_spec = None + return GPTModel( + config=config, + transformer_layer_spec=transformer_layer_spec, + vocab_size=args.vocab_size, + max_sequence_length=args.max_position_embeddings, + pre_process=pre_process, + post_process=post_process, + fp16_lm_cross_entropy=args.fp16_lm_cross_entropy, + parallel_output=True, + share_embeddings_and_output_weights=not args.untie_embeddings_and_output_weights, + position_embedding_type=args.position_embedding_type, + rotary_percent=args.rotary_percent, + mtp_block_spec=mtp_block_spec, + ) + + def create_test_args( + self, cuda_graph_impl, cuda_graph_scope, cuda_graph_warmup_steps, ep_size, **kwargs + ): + destroy_global_vars() + destroy_num_microbatches_calculator() + + sys.argv = ['test_cuda_graphs.py'] + args = parse_args() + args.num_layers = 4 + args.mtp_num_layers = 1 + args.vocab_size = 1024 + args.hidden_size = 128 + args.num_attention_heads = 8 + args.max_position_embeddings = 512 + args.global_batch_size = self.micro_batch_size * 8 + args.micro_batch_size = self.micro_batch_size + args.create_attention_mask_in_dataloader = True + args.seq_length = self.seq_length + args.tensor_model_parallel_size = 2 + args.sequence_parallel = True + args.pipeline_model_parallel_size = 1 + args.context_parallel_size = 1 + args.expert_model_parallel_size = ep_size + args.train_iters = 10 + args.lr = 3e-5 + args.bf16 = True + args.add_bias_linear = False + args.swiglu = True + args.use_distributed_optimizer = True + args.position_embedding_type = "rope" + args.rotary_percent = 1.0 + args.hidden_dropout = 0.0 + args.attention_dropout = 0.0 + + # MoE settings + args.num_experts = 4 + args.expert_model_parallel_size = ep_size + args.moe_shared_expert_intermediate_size = 1024 + args.moe_layer_freq = "[0,0,1,1]" + args.moe_permute_fusion = True + args.moe_router_fusion = True + args.moe_router_topk = 2 + + # CUDA graph settings + args.cuda_graph_impl = cuda_graph_impl + args.cuda_graph_scope = cuda_graph_scope + args.cuda_graph_warmup_steps = cuda_graph_warmup_steps + args.use_te_rng_tracker = cuda_graph_impl != "none" + + for key, value in kwargs.items(): + assert hasattr(args, key) + setattr(args, key, value) + + validate_args(args) + set_global_variables(args, False) + return args + + def get_batch(self, seq_length, micro_batch_size): + data = list(range(seq_length)) + input_ids = torch.tensor(data, dtype=torch.int64).repeat((micro_batch_size, 1)).cuda() + labels = 1 + torch.tensor(data, dtype=torch.int64).repeat((micro_batch_size, 1)).cuda() + position_ids = torch.tensor(data, dtype=torch.int64).repeat((micro_batch_size, 1)).cuda() + attention_mask = torch.ones( + (micro_batch_size, 1, seq_length, seq_length), dtype=bool + ).cuda() + loss_mask = torch.ones(seq_length).repeat((micro_batch_size, 1)).cuda() + return input_ids, labels, position_ids, attention_mask, loss_mask + + def _run_test_helper( + self, ep_size, cuda_graph_impl, cuda_graph_scope, cuda_graph_warmup_steps, **kwargs + ): + """Test fp8_param with gpt_model.""" + args = self.create_test_args( + cuda_graph_impl, cuda_graph_scope, cuda_graph_warmup_steps, ep_size, **kwargs + ) + + set_args(args) + torch.manual_seed(123) + Utils.initialize_model_parallel( + tensor_model_parallel_size=2, expert_model_parallel_size=ep_size + ) + + input_ids, labels, position_ids, attention_mask, loss_mask = self.get_batch( + self.seq_length, self.micro_batch_size + ) + + gpt_model, optimizer, _ = setup_model_and_optimizer( + self.model_provider, ModelType.encoder_or_decoder + ) + assert len(gpt_model) == 1 # Assume only one model in the model provider. + + loss_list = [] + + cuda_graph_helper = None + if cuda_graph_impl == "transformer_engine": + from megatron.core.transformer.cuda_graphs import TECudaGraphHelper + + cuda_graph_helper = TECudaGraphHelper( + model=gpt_model, + config=gpt_model[0].config, + seq_length=self.seq_length, + micro_batch_size=self.micro_batch_size, + optimizers=[optimizer], + ) + + for i in range(100): + gpt_model[0].zero_grad_buffer() + optimizer.zero_grad() + + # Capture CUDA graphs after warmup if helper is provided + if cuda_graph_helper is not None and i == cuda_graph_warmup_steps: + cuda_graph_helper.create_cudagraphs() + + output = gpt_model[0].forward( + input_ids=input_ids, + position_ids=position_ids, + attention_mask=attention_mask, + labels=labels, + loss_mask=loss_mask, + ) + + # Check output shapes + assert output.shape[0] == self.micro_batch_size + assert output.shape[1] == self.seq_length + + # Verify gradients + loss = output.mean() + loss.backward() + + for param in gpt_model[0].parameters(): + assert param.main_grad is not None + + update_successful, _, _ = optimizer.step() + assert update_successful + + loss_list.append(loss.item()) + + return torch.tensor(loss_list) + + @pytest.mark.skipif( + not (HAVE_TE and is_te_min_version("1.14.0")), + reason="Partial CUDA graph support requires TransformerEngine version >= 1.14.0", + ) + @pytest.mark.parametrize("ep_size", [1, 4]) + @pytest.mark.parametrize("moe_dropless_dispatcher", [False, True]) + @pytest.mark.parametrize("moe_dispatcher_type", ["alltoall", "deepep", "hybridep"]) + def test_moe_partial_cudagraph(self, ep_size, moe_dropless_dispatcher, moe_dispatcher_type): + extra_kwargs = {} + if moe_dispatcher_type == "deepep": + if not is_deep_ep_available(): + pytest.skip("Deep EP is not available") + extra_kwargs["moe_token_dispatcher_type"] = "flex" + extra_kwargs["moe_enable_deepep"] = True + elif moe_dispatcher_type == "hybridep": + if not is_hybrid_ep_available(): + pytest.skip("Hybrid EP is not available") + extra_kwargs["moe_token_dispatcher_type"] = "flex" + extra_kwargs["moe_flex_dispatcher_backend"] = "hybridep" + else: + extra_kwargs["moe_token_dispatcher_type"] = moe_dispatcher_type + if not moe_dropless_dispatcher: + if moe_dispatcher_type == "deepep": + pytest.skip("Deep EP doesn't support drop&pad MoE") + extra_kwargs["moe_expert_capacity_factor"] = 1.0 + extra_kwargs["moe_pad_expert_input_to_capacity"] = True + + loss_list_ref = self._run_test_helper(ep_size, "none", None, 0, **extra_kwargs) + for cuda_graph_scope in [ + None, + ["attn"], + ["moe"], + ["mlp", "moe_router"], + ["attn", "mlp", "moe_router", "moe_preprocess"], + ]: + if moe_dropless_dispatcher and (cuda_graph_scope is None or "moe" in cuda_graph_scope): + # Dropless MoE doesn't work with "moe" scope cudagraph. Skip. + continue + cuda_graph_warmup_steps = 3 + loss_list = self._run_test_helper( + ep_size, + "transformer_engine", + cuda_graph_scope, + cuda_graph_warmup_steps, + **extra_kwargs, + ) + assert torch.equal(loss_list, loss_list_ref) + + if __name__ == "__main__": test = TestParallelTransformerBlockCudagraphs() @@ -508,3 +783,8 @@ def test_gpu_cudagraph(self): llava_test.setup_method(method=None) llava_test.test_llava_cudagraph_is_last_layer_logic() llava_test.teardown_method(method=None) + + test = TestPartialCudaGraph() + test.setup_method(method=None) + test.test_moe_partial_cudagraph(4, True, "alltoall") + test.teardown_method(method=None) diff --git a/tools/checkpoint/checkpoint_inspector.py b/tools/checkpoint/checkpoint_inspector.py index c62f0ca7417..3d03f4db959 100644 --- a/tools/checkpoint/checkpoint_inspector.py +++ b/tools/checkpoint/checkpoint_inspector.py @@ -1,3 +1,5 @@ +# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + # python checkpoint_inspector.py inspect /path/to/checkpoint # torchrun --nproc_per_node=8 --nnodes=1 checkpoint_inspector.py convert-torch-dist-to-fsdp-dtensor /path/to/input_checkpoint /path/to/output_checkpoint --swiglu import gc From b16a78c568239b8be2c8f4cfe4c8c8ebd64af743 Mon Sep 17 00:00:00 2001 From: Robin Zhang Date: Tue, 4 Nov 2025 21:39:09 -0800 Subject: [PATCH 02/19] main golden Signed-off-by: Robin Zhang --- .../golden_values_dev_dgx_h100.json | 1134 ++++++++--------- 1 file changed, 567 insertions(+), 567 deletions(-) diff --git a/tests/functional_tests/test_cases/moe/gpt3_moe_mcore_te_tp4_ep2_etp2_pp2_scoped_cudagraph/golden_values_dev_dgx_h100.json b/tests/functional_tests/test_cases/moe/gpt3_moe_mcore_te_tp4_ep2_etp2_pp2_scoped_cudagraph/golden_values_dev_dgx_h100.json index 309b2533461..09eb6ef15bf 100644 --- a/tests/functional_tests/test_cases/moe/gpt3_moe_mcore_te_tp4_ep2_etp2_pp2_scoped_cudagraph/golden_values_dev_dgx_h100.json +++ b/tests/functional_tests/test_cases/moe/gpt3_moe_mcore_te_tp4_ep2_etp2_pp2_scoped_cudagraph/golden_values_dev_dgx_h100.json @@ -7,103 +7,103 @@ "1": 10.93663, "2": 10.9327, "3": 10.94263, - "4": 10.94969, - "5": 10.95052, - "6": 10.94157, - "7": 10.94484, - "8": 10.93674, - "9": 10.94996, - "10": 10.93686, - "11": 10.94102, - "12": 10.93763, - "13": 10.9235, - "14": 10.93428, - "15": 10.88791, - "16": 10.87434, - "17": 10.86896, - "18": 10.86065, - "19": 10.86311, - "20": 10.78063, - "21": 10.73125, - "22": 10.60283, - "23": 10.73278, - "24": 10.61888, - "25": 10.55212, - "26": 10.62704, - "27": 10.6391, - "28": 10.5908, - "29": 10.59809, - "30": 10.37777, - "31": 10.1201, - "32": 10.46078, - "33": 10.45538, - "34": 10.20107, - "35": 10.25779, - "36": 10.20889, - "37": 10.33688, - "38": 10.16827, - "39": 10.40875, - "40": 10.05239, - "41": 10.09432, - "42": 10.17894, - "43": 9.74205, - "44": 9.8904, - "45": 9.74009, - "46": 9.72707, - "47": 10.09139, - "48": 9.75298, - "49": 9.40106, - "50": 9.83667, - "51": 9.77071, - "52": 9.65705, - "53": 10.03051, - "54": 9.87899, - "55": 9.79604, - "56": 9.52924, - "57": 9.36583, - "58": 9.75331, - "59": 9.48065, - "60": 9.40785, - "61": 9.60145, - "62": 9.90753, - "63": 9.2583, - "64": 9.68397, - "65": 8.80003, - "66": 9.60779, - "67": 9.25408, - "68": 9.71438, - "69": 9.71682, - "70": 9.6617, - "71": 9.52466, - "72": 9.47116, - "73": 9.38822, - "74": 8.80223, - "75": 9.33966, - "76": 8.93574, - "77": 9.99333, - "78": 9.64731, - "79": 9.28114, - "80": 9.29588, - "81": 9.39589, - "82": 9.60893, - "83": 9.21629, - "84": 9.33891, - "85": 9.52979, - "86": 8.95817, - "87": 9.51641, - "88": 9.68228, - "89": 9.50664, - "90": 9.75348, - "91": 9.23465, - "92": 9.25972, - "93": 8.94517, - "94": 8.69188, - "95": 9.44591, - "96": 9.4101, - "97": 9.20087, - "98": 9.58175, - "99": 8.75818, - "100": 9.29466 + "4": 10.94963, + "5": 10.95058, + "6": 10.94173, + "7": 10.94479, + "8": 10.93683, + "9": 10.94964, + "10": 10.93727, + "11": 10.94087, + "12": 10.93752, + "13": 10.92358, + "14": 10.93403, + "15": 10.88705, + "16": 10.87477, + "17": 10.86854, + "18": 10.86075, + "19": 10.86302, + "20": 10.78056, + "21": 10.73152, + "22": 10.60338, + "23": 10.73311, + "24": 10.61889, + "25": 10.55146, + "26": 10.62716, + "27": 10.63933, + "28": 10.59173, + "29": 10.59786, + "30": 10.37829, + "31": 10.12096, + "32": 10.46077, + "33": 10.45508, + "34": 10.20099, + "35": 10.25824, + "36": 10.20892, + "37": 10.33713, + "38": 10.16915, + "39": 10.40904, + "40": 10.05252, + "41": 10.09429, + "42": 10.17849, + "43": 9.74072, + "44": 9.89045, + "45": 9.73992, + "46": 9.72688, + "47": 10.0918, + "48": 9.75311, + "49": 9.4017, + "50": 9.83702, + "51": 9.77105, + "52": 9.65558, + "53": 10.03094, + "54": 9.87894, + "55": 9.79551, + "56": 9.53279, + "57": 9.36625, + "58": 9.75325, + "59": 9.48161, + "60": 9.40822, + "61": 9.60147, + "62": 9.90763, + "63": 9.25792, + "64": 9.68418, + "65": 8.79865, + "66": 9.60782, + "67": 9.25445, + "68": 9.71388, + "69": 9.71675, + "70": 9.66147, + "71": 9.52492, + "72": 9.47142, + "73": 9.38853, + "74": 8.80276, + "75": 9.33982, + "76": 8.93585, + "77": 9.99344, + "78": 9.64759, + "79": 9.28191, + "80": 9.29645, + "81": 9.39618, + "82": 9.60863, + "83": 9.2168, + "84": 9.33921, + "85": 9.52982, + "86": 8.95634, + "87": 9.51667, + "88": 9.68206, + "89": 9.50613, + "90": 9.75311, + "91": 9.23456, + "92": 9.26029, + "93": 8.94457, + "94": 8.69219, + "95": 9.44595, + "96": 9.40999, + "97": 9.20134, + "98": 9.58177, + "99": 8.75845, + "100": 9.29494 } }, "num-zeros": { @@ -114,103 +114,103 @@ "1": 22750260.0, "2": 22953110.0, "3": 22604450.0, - "4": 23266322.0, - "5": 22735560.0, - "6": 23061920.0, - "7": 22793342.0, - "8": 22960820.0, - "9": 22865664.0, - "10": 22950364.0, - "11": 22499674.0, - "12": 22456088.0, - "13": 22948060.0, - "14": 22384512.0, - "15": 22846272.0, - "16": 22856858.0, - "17": 22836412.0, - "18": 22590058.0, - "19": 22627048.0, - "20": 22712308.0, - "21": 22762624.0, - "22": 22816888.0, - "23": 22545124.0, - "24": 22794440.0, - "25": 22841936.0, - "26": 22549680.0, - "27": 22464820.0, - "28": 22453684.0, - "29": 22534640.0, - "30": 22636152.0, - "31": 22989488.0, - "32": 22594070.0, - "33": 22566010.0, - "34": 22855504.0, - "35": 22813688.0, - "36": 22595396.0, - "37": 22499360.0, - "38": 22926126.0, - "39": 22825392.0, - "40": 22675666.0, - "41": 22671586.0, - "42": 22682140.0, - "43": 23013940.0, - "44": 22764458.0, - "45": 22678992.0, - "46": 22915276.0, - "47": 22642868.0, - "48": 22954190.0, - "49": 23786668.0, - "50": 22934008.0, - "51": 23866222.0, - "52": 23807290.0, - "53": 24007532.0, - "54": 22871610.0, - "55": 23571284.0, - "56": 23954310.0, - "57": 24211632.0, - "58": 23914404.0, - "59": 23771838.0, - "60": 23813560.0, - "61": 23797288.0, - "62": 23739984.0, - "63": 23916692.0, - "64": 23895952.0, - "65": 24150562.0, - "66": 23796504.0, - "67": 25032232.0, - "68": 23673188.0, - "69": 23648580.0, - "70": 23903504.0, - "71": 24864636.0, - "72": 24767108.0, - "73": 24850612.0, - "74": 24132990.0, - "75": 24146528.0, - "76": 25025540.0, - "77": 24358472.0, - "78": 24910064.0, - "79": 23810516.0, - "80": 24821440.0, - "81": 25020512.0, - "82": 23851244.0, - "83": 24961024.0, - "84": 25144020.0, - "85": 24823608.0, - "86": 23153096.0, - "87": 24850204.0, - "88": 24749150.0, - "89": 22505554.0, - "90": 24059620.0, - "91": 23839038.0, - "92": 23874568.0, - "93": 24769548.0, - "94": 23992452.0, - "95": 25189838.0, - "96": 23909262.0, - "97": 24713068.0, - "98": 23832506.0, - "99": 23983474.0, - "100": 24101108.0 + "4": 23266268.0, + "5": 22735536.0, + "6": 23061796.0, + "7": 22793352.0, + "8": 22960904.0, + "9": 22865576.0, + "10": 22950376.0, + "11": 22499640.0, + "12": 22456172.0, + "13": 22948096.0, + "14": 22384552.0, + "15": 22846232.0, + "16": 22856848.0, + "17": 22836380.0, + "18": 22590114.0, + "19": 22626994.0, + "20": 22712320.0, + "21": 22762648.0, + "22": 22816768.0, + "23": 22545256.0, + "24": 22794352.0, + "25": 22841912.0, + "26": 22549662.0, + "27": 22464856.0, + "28": 22453796.0, + "29": 22534624.0, + "30": 22636184.0, + "31": 22989510.0, + "32": 22594012.0, + "33": 22565972.0, + "34": 22855552.0, + "35": 22813580.0, + "36": 22595488.0, + "37": 22499324.0, + "38": 22926252.0, + "39": 22825296.0, + "40": 22675744.0, + "41": 22671458.0, + "42": 22682356.0, + "43": 23014132.0, + "44": 22768892.0, + "45": 22683210.0, + "46": 22915234.0, + "47": 23691900.0, + "48": 22954106.0, + "49": 23786656.0, + "50": 22931628.0, + "51": 23866150.0, + "52": 23807384.0, + "53": 24007520.0, + "54": 22867936.0, + "55": 23571382.0, + "56": 23954146.0, + "57": 24211700.0, + "58": 23914532.0, + "59": 22725000.0, + "60": 23813604.0, + "61": 23810256.0, + "62": 23740378.0, + "63": 23916450.0, + "64": 23899026.0, + "65": 24150662.0, + "66": 23795982.0, + "67": 25032318.0, + "68": 23675500.0, + "69": 23644168.0, + "70": 23903738.0, + "71": 24864580.0, + "72": 24767012.0, + "73": 24850692.0, + "74": 24133088.0, + "75": 24143564.0, + "76": 25025588.0, + "77": 24358278.0, + "78": 24909920.0, + "79": 23808164.0, + "80": 23772264.0, + "81": 25020498.0, + "82": 23851236.0, + "83": 23912290.0, + "84": 25143922.0, + "85": 24823454.0, + "86": 23153236.0, + "87": 24850100.0, + "88": 24749292.0, + "89": 22504736.0, + "90": 24059580.0, + "91": 23838524.0, + "92": 24923932.0, + "93": 24769600.0, + "94": 23992332.0, + "95": 25192816.0, + "96": 23909096.0, + "97": 24713200.0, + "98": 23832510.0, + "99": 23980812.0, + "100": 24101050.0 } }, "mem-allocated-bytes": { @@ -218,105 +218,105 @@ "end_step": 100, "step_interval": 1, "values": { - "1": 763142656.0, - "2": 778734592.0, - "3": 772525056.0, - "4": 803593216.0, - "5": 803593216.0, + "1": 773939712.0, + "2": 781904896.0, + "3": 771107840.0, + "4": 801299456.0, + "5": 801299456.0, "6": 803593216.0, - "7": 801299456.0, - "8": 803593216.0, - "9": 801840128.0, - "10": 803593216.0, - "11": 802987008.0, - "12": 803593216.0, - "13": 802987008.0, - "14": 801299456.0, - "15": 803593216.0, - "16": 801840128.0, + "7": 802446336.0, + "8": 801299456.0, + "9": 803593216.0, + "10": 801840128.0, + "11": 803593216.0, + "12": 802987008.0, + "13": 801299456.0, + "14": 803593216.0, + "15": 801840128.0, + "16": 803593216.0, "17": 803593216.0, - "18": 802987008.0, - "19": 801299456.0, + "18": 801299456.0, + "19": 803593216.0, "20": 803593216.0, "21": 801299456.0, "22": 803593216.0, - "23": 801299456.0, - "24": 803593216.0, - "25": 801299456.0, - "26": 803593216.0, - "27": 801299456.0, - "28": 803593216.0, + "23": 802987008.0, + "24": 801299456.0, + "25": 803593216.0, + "26": 801299456.0, + "27": 803593216.0, + "28": 802987008.0, "29": 801299456.0, "30": 803593216.0, - "31": 801299456.0, - "32": 803593216.0, - "33": 801840128.0, - "34": 803593216.0, - "35": 801840128.0, - "36": 803593216.0, - "37": 802987008.0, - "38": 801299456.0, - "39": 803593216.0, - "40": 801299456.0, - "41": 803593216.0, - "42": 801840128.0, - "43": 803593216.0, - "44": 801840128.0, + "31": 802446336.0, + "32": 801299456.0, + "33": 803593216.0, + "34": 801299456.0, + "35": 803593216.0, + "36": 801840128.0, + "37": 803593216.0, + "38": 802987008.0, + "39": 801299456.0, + "40": 803593216.0, + "41": 801299456.0, + "42": 803593216.0, + "43": 802987008.0, + "44": 801299456.0, "45": 803593216.0, "46": 801840128.0, - "47": 803593216.0, - "48": 801840128.0, - "49": 803593216.0, - "50": 801840128.0, + "47": 801299456.0, + "48": 803593216.0, + "49": 801299456.0, + "50": 803593216.0, "51": 801299456.0, "52": 803593216.0, - "53": 801299456.0, - "54": 803593216.0, - "55": 801840128.0, - "56": 803593216.0, - "57": 801840128.0, + "53": 803593216.0, + "54": 801299456.0, + "55": 803593216.0, + "56": 801840128.0, + "57": 803593216.0, "58": 803593216.0, - "59": 801840128.0, + "59": 801299456.0, "60": 803593216.0, - "61": 801299456.0, - "62": 803593216.0, - "63": 801299456.0, - "64": 802987008.0, + "61": 802446336.0, + "62": 801299456.0, + "63": 801840128.0, + "64": 801840128.0, "65": 803593216.0, - "66": 801299456.0, - "67": 803593216.0, - "68": 801299456.0, - "69": 803593216.0, - "70": 801840128.0, + "66": 803593216.0, + "67": 801299456.0, + "68": 803593216.0, + "69": 802987008.0, + "70": 801299456.0, "71": 803593216.0, - "72": 801299456.0, + "72": 801840128.0, "73": 803593216.0, "74": 803593216.0, - "75": 802987008.0, + "75": 801299456.0, "76": 803593216.0, "77": 801840128.0, "78": 803593216.0, - "79": 801299456.0, - "80": 802987008.0, + "79": 803593216.0, + "80": 801299456.0, "81": 803593216.0, - "82": 801840128.0, - "83": 803593216.0, - "84": 801299456.0, - "85": 802987008.0, + "82": 802987008.0, + "83": 801299456.0, + "84": 803593216.0, + "85": 801840128.0, "86": 803593216.0, - "87": 801840128.0, - "88": 803593216.0, - "89": 801299456.0, - "90": 802987008.0, - "91": 803593216.0, - "92": 801299456.0, - "93": 803593216.0, - "94": 801299456.0, - "95": 803593216.0, + "87": 803593216.0, + "88": 801299456.0, + "89": 803593216.0, + "90": 802446336.0, + "91": 801299456.0, + "92": 803593216.0, + "93": 801299456.0, + "94": 803593216.0, + "95": 802987008.0, "96": 801299456.0, - "97": 803593216.0, - "98": 801299456.0, - "99": 802987008.0, + "97": 801299456.0, + "98": 803593216.0, + "99": 801299456.0, "100": 803593216.0 } }, @@ -325,106 +325,106 @@ "end_step": 100, "step_interval": 1, "values": { - "1": 993582592.0, - "2": 1210942464.0, - "3": 1210942464.0, - "4": 1210942464.0, - "5": 1210942464.0, - "6": 1210942464.0, - "7": 1210942464.0, - "8": 1210942464.0, - "9": 1210942464.0, - "10": 1210942464.0, - "11": 1210942464.0, - "12": 1210942464.0, - "13": 1210942464.0, - "14": 1210942464.0, - "15": 1210942464.0, - "16": 1210942464.0, - "17": 1210942464.0, - "18": 1210942464.0, - "19": 1210942464.0, - "20": 1210942464.0, - "21": 1210942464.0, - "22": 1210942464.0, - "23": 1210942464.0, - "24": 1210942464.0, - "25": 1210942464.0, - "26": 1210942464.0, - "27": 1210942464.0, - "28": 1210942464.0, - "29": 1210942464.0, - "30": 1210942464.0, - "31": 1210942464.0, - "32": 1210942464.0, - "33": 1210942464.0, - "34": 1210942464.0, - "35": 1210942464.0, - "36": 1210942464.0, - "37": 1210942464.0, - "38": 1210942464.0, - "39": 1210942464.0, - "40": 1210942464.0, - "41": 1210942464.0, - "42": 1210942464.0, - "43": 1210942464.0, - "44": 1210942464.0, - "45": 1210942464.0, - "46": 1210942464.0, - "47": 1210942464.0, - "48": 1210942464.0, - "49": 1210942464.0, - "50": 1210942464.0, - "51": 1210942464.0, - "52": 1210942464.0, - "53": 1210942464.0, - "54": 1210942464.0, - "55": 1210942464.0, - "56": 1210942464.0, - "57": 1210942464.0, - "58": 1210942464.0, - "59": 1210942464.0, - "60": 1210942464.0, - "61": 1210942464.0, - "62": 1210942464.0, - "63": 1210942464.0, - "64": 1210942464.0, - "65": 1210942464.0, - "66": 1210942464.0, - "67": 1210942464.0, - "68": 1210942464.0, - "69": 1210942464.0, - "70": 1210942464.0, - "71": 1210942464.0, - "72": 1210942464.0, - "73": 1210942464.0, - "74": 1210942464.0, - "75": 1210942464.0, - "76": 1210942464.0, - "77": 1210942464.0, - "78": 1210942464.0, - "79": 1210942464.0, - "80": 1210942464.0, - "81": 1210942464.0, - "82": 1210942464.0, - "83": 1210942464.0, - "84": 1210942464.0, - "85": 1210942464.0, - "86": 1210942464.0, - "87": 1210942464.0, - "88": 1210942464.0, - "89": 1210942464.0, - "90": 1210942464.0, - "91": 1210942464.0, - "92": 1210942464.0, - "93": 1210942464.0, - "94": 1210942464.0, - "95": 1210942464.0, - "96": 1210942464.0, - "97": 1210942464.0, - "98": 1210942464.0, - "99": 1210942464.0, - "100": 1210942464.0 + "1": 991873024.0, + "2": 1206563328.0, + "3": 1206675456.0, + "4": 1206675456.0, + "5": 1206675456.0, + "6": 1206675456.0, + "7": 1206675456.0, + "8": 1206675456.0, + "9": 1206675456.0, + "10": 1206675456.0, + "11": 1206675456.0, + "12": 1206675456.0, + "13": 1206675456.0, + "14": 1206675456.0, + "15": 1206675456.0, + "16": 1206675456.0, + "17": 1206675456.0, + "18": 1206675456.0, + "19": 1206675456.0, + "20": 1206675456.0, + "21": 1206675456.0, + "22": 1206675456.0, + "23": 1206675456.0, + "24": 1206675456.0, + "25": 1206675456.0, + "26": 1206675456.0, + "27": 1206675456.0, + "28": 1206675456.0, + "29": 1206675456.0, + "30": 1206675456.0, + "31": 1206675456.0, + "32": 1206675456.0, + "33": 1206675456.0, + "34": 1206675456.0, + "35": 1206675456.0, + "36": 1206675456.0, + "37": 1206675456.0, + "38": 1206675456.0, + "39": 1206675456.0, + "40": 1206675456.0, + "41": 1206675456.0, + "42": 1206675456.0, + "43": 1206675456.0, + "44": 1206675456.0, + "45": 1206675456.0, + "46": 1206675456.0, + "47": 1206675456.0, + "48": 1206675456.0, + "49": 1206675456.0, + "50": 1206675456.0, + "51": 1206675456.0, + "52": 1206675456.0, + "53": 1206675456.0, + "54": 1206675456.0, + "55": 1206675456.0, + "56": 1206675456.0, + "57": 1206675456.0, + "58": 1206675456.0, + "59": 1206675456.0, + "60": 1206675456.0, + "61": 1206675456.0, + "62": 1206675456.0, + "63": 1206675456.0, + "64": 1206675456.0, + "65": 1206675456.0, + "66": 1206675456.0, + "67": 1206675456.0, + "68": 1206675456.0, + "69": 1206675456.0, + "70": 1206675456.0, + "71": 1206675456.0, + "72": 1206675456.0, + "73": 1206675456.0, + "74": 1206675456.0, + "75": 1206675456.0, + "76": 1206675456.0, + "77": 1206675456.0, + "78": 1206675456.0, + "79": 1206675456.0, + "80": 1206675456.0, + "81": 1206675456.0, + "82": 1206675456.0, + "83": 1206675456.0, + "84": 1206675456.0, + "85": 1206675456.0, + "86": 1206675456.0, + "87": 1206675456.0, + "88": 1206675456.0, + "89": 1206675456.0, + "90": 1206675456.0, + "91": 1206675456.0, + "92": 1206675456.0, + "93": 1206675456.0, + "94": 1206675456.0, + "95": 1206675456.0, + "96": 1206675456.0, + "97": 1206675456.0, + "98": 1206675456.0, + "99": 1206675456.0, + "100": 1206675456.0 } }, "mtp_1 loss": { @@ -435,103 +435,103 @@ "1": 10.88689, "2": 10.90485, "3": 10.90869, - "4": 10.86903, - "5": 10.91601, - "6": 10.906, - "7": 10.90268, - "8": 10.88984, - "9": 10.90425, - "10": 10.89144, - "11": 10.93384, - "12": 10.91647, - "13": 10.91108, - "14": 10.91974, - "15": 10.88488, - "16": 10.9077, - "17": 10.87571, - "18": 10.91379, - "19": 10.9092, - "20": 10.87837, - "21": 10.87896, - "22": 10.85583, - "23": 10.88007, - "24": 10.87245, - "25": 10.85859, - "26": 10.8696, - "27": 10.87702, - "28": 10.88641, - "29": 10.88866, - "30": 10.85422, - "31": 10.79713, - "32": 10.86631, - "33": 10.8781, - "34": 10.83982, - "35": 10.84165, - "36": 10.85012, - "37": 10.85556, - "38": 10.83674, - "39": 10.86355, - "40": 10.82887, - "41": 10.8341, - "42": 10.84469, - "43": 10.78828, - "44": 10.82123, - "45": 10.78831, - "46": 10.7823, - "47": 10.82898, - "48": 10.78985, - "49": 10.71269, - "50": 10.77382, - "51": 10.76639, - "52": 10.7397, - "53": 10.80285, - "54": 10.77365, - "55": 10.76066, - "56": 10.71068, - "57": 10.66686, - "58": 10.74378, - "59": 10.69209, - "60": 10.66474, - "61": 10.7073, - "62": 10.77206, - "63": 10.61812, - "64": 10.7178, - "65": 10.49439, - "66": 10.67106, - "67": 10.57534, - "68": 10.6873, - "69": 10.6816, - "70": 10.66836, - "71": 10.64586, - "72": 10.60925, - "73": 10.56508, - "74": 10.37144, - "75": 10.51183, - "76": 10.39914, - "77": 10.75182, - "78": 10.6268, - "79": 10.46827, - "80": 10.47524, - "81": 10.51083, - "82": 10.58769, - "83": 10.4381, - "84": 10.45057, - "85": 10.55084, - "86": 10.28076, - "87": 10.51088, - "88": 10.60323, - "89": 10.50794, - "90": 10.60274, - "91": 10.38238, - "92": 10.38703, - "93": 10.23076, - "94": 10.08438, - "95": 10.42616, - "96": 10.44905, - "97": 10.32215, - "98": 10.4966, - "99": 10.04765, - "100": 10.33491 + "4": 10.86909, + "5": 10.91592, + "6": 10.90606, + "7": 10.90233, + "8": 10.89037, + "9": 10.90421, + "10": 10.89128, + "11": 10.93353, + "12": 10.91634, + "13": 10.91128, + "14": 10.92022, + "15": 10.88422, + "16": 10.90792, + "17": 10.87526, + "18": 10.91408, + "19": 10.90945, + "20": 10.87827, + "21": 10.8792, + "22": 10.85485, + "23": 10.8798, + "24": 10.87243, + "25": 10.85788, + "26": 10.87005, + "27": 10.87713, + "28": 10.88659, + "29": 10.88861, + "30": 10.8548, + "31": 10.79738, + "32": 10.86611, + "33": 10.87796, + "34": 10.83937, + "35": 10.84218, + "36": 10.85039, + "37": 10.85607, + "38": 10.83659, + "39": 10.86362, + "40": 10.82843, + "41": 10.83391, + "42": 10.84457, + "43": 10.78795, + "44": 10.82116, + "45": 10.7887, + "46": 10.78286, + "47": 10.82922, + "48": 10.79061, + "49": 10.71287, + "50": 10.77384, + "51": 10.76674, + "52": 10.73969, + "53": 10.80245, + "54": 10.77312, + "55": 10.76014, + "56": 10.70949, + "57": 10.66696, + "58": 10.74347, + "59": 10.69265, + "60": 10.66514, + "61": 10.70862, + "62": 10.77182, + "63": 10.61887, + "64": 10.71829, + "65": 10.49496, + "66": 10.67136, + "67": 10.57536, + "68": 10.68754, + "69": 10.68221, + "70": 10.6686, + "71": 10.64536, + "72": 10.60792, + "73": 10.56485, + "74": 10.37018, + "75": 10.51084, + "76": 10.39856, + "77": 10.75174, + "78": 10.62701, + "79": 10.46673, + "80": 10.4747, + "81": 10.51096, + "82": 10.58798, + "83": 10.43986, + "84": 10.45062, + "85": 10.55152, + "86": 10.28423, + "87": 10.51171, + "88": 10.60343, + "89": 10.50915, + "90": 10.60399, + "91": 10.38243, + "92": 10.3873, + "93": 10.2309, + "94": 10.08356, + "95": 10.42578, + "96": 10.44886, + "97": 10.32148, + "98": 10.4967, + "99": 10.04648, + "100": 10.33494 } }, "iteration-time": { @@ -539,106 +539,106 @@ "end_step": 100, "step_interval": 1, "values": { - "1": 51.30209, - "2": 1.41746, - "3": 1.28029, - "4": 10.57024, - "5": 0.66643, - "6": 0.67893, - "7": 0.65727, - "8": 0.66196, - "9": 0.66227, - "10": 0.65877, - "11": 0.65828, - "12": 0.65862, - "13": 0.65727, - "14": 0.65896, - "15": 0.65851, - "16": 0.66826, - "17": 0.65878, - "18": 0.65573, - "19": 0.65631, - "20": 0.65579, - "21": 0.65091, - "22": 0.65603, - "23": 0.65158, - "24": 0.65266, - "25": 0.65816, - "26": 0.65194, - "27": 0.6541, - "28": 0.65515, - "29": 0.65439, - "30": 0.65241, - "31": 0.65597, - "32": 0.65551, - "33": 0.65318, - "34": 0.6553, - "35": 0.65725, - "36": 0.65926, - "37": 0.65606, - "38": 0.65571, - "39": 0.65846, - "40": 0.65642, - "41": 0.65509, - "42": 0.66105, - "43": 0.65448, - "44": 0.65534, - "45": 0.65304, - "46": 0.65227, - "47": 0.64871, - "48": 0.65257, - "49": 0.65485, - "50": 0.65054, - "51": 0.67883, - "52": 0.6571, - "53": 0.65671, - "54": 0.65877, - "55": 0.65584, - "56": 0.65072, - "57": 0.64951, - "58": 0.65703, - "59": 0.65106, - "60": 0.64536, - "61": 0.64416, - "62": 0.64816, - "63": 0.64084, - "64": 0.6396, - "65": 0.64182, - "66": 0.64004, - "67": 0.64101, - "68": 0.63928, - "69": 0.65723, - "70": 0.6828, - "71": 0.64052, - "72": 0.64287, - "73": 0.64136, - "74": 0.64252, - "75": 0.64617, - "76": 0.64857, - "77": 0.64304, - "78": 0.64068, - "79": 0.64048, - "80": 0.64091, - "81": 0.64179, - "82": 0.64793, - "83": 0.641, - "84": 0.64077, - "85": 0.64011, - "86": 0.64018, - "87": 0.64132, - "88": 0.63901, - "89": 0.6407, - "90": 0.64277, - "91": 0.64132, - "92": 0.64123, - "93": 0.65051, - "94": 0.65036, - "95": 0.64542, - "96": 0.64561, - "97": 0.6504, - "98": 0.64563, - "99": 0.64524, - "100": 0.65049 + "1": 64.34324, + "2": 1.36866, + "3": 1.2499, + "4": 11.21345, + "5": 0.66104, + "6": 0.66604, + "7": 0.66704, + "8": 0.67004, + "9": 0.66612, + "10": 0.65729, + "11": 0.65845, + "12": 0.65975, + "13": 0.66533, + "14": 0.6636, + "15": 0.66469, + "16": 0.66338, + "17": 0.66867, + "18": 0.66738, + "19": 0.66795, + "20": 0.6669, + "21": 0.66551, + "22": 0.66394, + "23": 0.66081, + "24": 0.66215, + "25": 0.66157, + "26": 0.66301, + "27": 0.6607, + "28": 0.6622, + "29": 0.6694, + "30": 0.66325, + "31": 0.66685, + "32": 0.66303, + "33": 0.66009, + "34": 0.65792, + "35": 0.66044, + "36": 0.65963, + "37": 0.65843, + "38": 0.65953, + "39": 0.65999, + "40": 0.66088, + "41": 0.6637, + "42": 0.66204, + "43": 0.66164, + "44": 0.66899, + "45": 0.66336, + "46": 0.66424, + "47": 0.66522, + "48": 0.66191, + "49": 0.65777, + "50": 0.65822, + "51": 0.73732, + "52": 0.66251, + "53": 0.66453, + "54": 0.66439, + "55": 0.66241, + "56": 0.66212, + "57": 0.66118, + "58": 0.66706, + "59": 0.66319, + "60": 0.65811, + "61": 0.66307, + "62": 0.65991, + "63": 0.73521, + "64": 1.79942, + "65": 0.66718, + "66": 0.66282, + "67": 0.66704, + "68": 0.66242, + "69": 0.66434, + "70": 0.66342, + "71": 0.66384, + "72": 0.66529, + "73": 0.6611, + "74": 0.66271, + "75": 0.66187, + "76": 0.66384, + "77": 0.66282, + "78": 0.66306, + "79": 0.6631, + "80": 0.66853, + "81": 0.66405, + "82": 0.66281, + "83": 0.66326, + "84": 0.66301, + "85": 0.6644, + "86": 0.66439, + "87": 0.66159, + "88": 0.66392, + "89": 0.66145, + "90": 0.66064, + "91": 0.66076, + "92": 0.66227, + "93": 0.664, + "94": 0.66102, + "95": 0.66295, + "96": 0.663, + "97": 0.69741, + "98": 0.66084, + "99": 0.66604, + "100": 0.6651 } } } \ No newline at end of file From 3681d7a1de4ad996a569bc267e68871a6b6a4d7b Mon Sep 17 00:00:00 2001 From: Robin Zhang Date: Thu, 13 Nov 2025 04:28:02 -0800 Subject: [PATCH 03/19] fix cudagraph ut Signed-off-by: Robin Zhang --- megatron/core/transformer/cuda_graphs.py | 23 +++++ .../core/transformer/transformer_config.py | 22 ++--- megatron/training/training.py | 4 + .../transformer/test_cuda_graphs.py | 90 +++++++++++++------ 4 files changed, 101 insertions(+), 38 deletions(-) diff --git a/megatron/core/transformer/cuda_graphs.py b/megatron/core/transformer/cuda_graphs.py index 12f15ee980a..dfc7bb4ede3 100644 --- a/megatron/core/transformer/cuda_graphs.py +++ b/megatron/core/transformer/cuda_graphs.py @@ -1712,3 +1712,26 @@ def cuda_graph_set_manual_hooks(self): model_chunk = self.model[chunk_number] for layer in layers: layer.setup_manual_hooks(model_chunk._make_forward_pre_hook) + + def destroy_cudagraphs(self): + """ + Destroy CUDA Graphs. + """ + assert self._graphs_created, "CUDA Graphs have not been created." + graphs_destoryed, graphs_not_destroyed = 0, 0 + for _, layers in enumerate(self.callables_per_chunk): + for layer in layers: + for graph in layer.cuda_graphs: + if is_te_min_version("2.10.0"): + graph.reset() + graphs_destoryed += 1 + else: + graphs_not_destroyed += 1 + layer.cuda_graphs = [] + layer.cuda_graph_manual_hooks = [] + log_single_rank( + logger, + logging.INFO, + f'{graphs_destoryed} graphs destroyed, {graphs_not_destroyed} graphs not destroyed.', + ) + self._graphs_created = False diff --git a/megatron/core/transformer/transformer_config.py b/megatron/core/transformer/transformer_config.py index 638c8645642..91165dfe0de 100644 --- a/megatron/core/transformer/transformer_config.py +++ b/megatron/core/transformer/transformer_config.py @@ -654,7 +654,7 @@ class TransformerConfig(ModelParallelConfig): cuda_graph_scope: Optional[List[str]] = None """Determines the CUDA graphs capturing scope. When cuda_graph_impl is set to "transformer_engine", valid values are "attn", "mlp", "moe", - "moe_router", "moe_preprocess", "mamba". None means ["attn", "mlp"]. + "moe_router", "moe_preprocess", "mamba". None means the full layer. When cuda_graph_impl is set to "local", "full_iteration" can be specified as cuda_graph_scope to enable whole iteration CUDA graph. All other values enable layerwise CUDA graph.""" @@ -1413,23 +1413,25 @@ def __post_init__(self): 'use cuda_graph_impl=transformer_engine instead.' ) self.cuda_graph_impl = "transformer_engine" + + if self.cuda_graph_scope is None: + self.cuda_graph_scope = [] + elif not isinstance(self.cuda_graph_scope, list): + assert isinstance(self.cuda_graph_scope, str), ( + "cuda_graph_scope must be a string or a list of strings, " + f"got {self.cuda_graph_scope}." + ) + self.cuda_graph_scope = [self.cuda_graph_scope] + if self.cuda_graph_impl != "none": assert self.cuda_graph_impl in [ "transformer_engine", "local", ], f"Invalid cuda graph implementation: {self.cuda_graph_impl}" + if self.cpu_offloading: raise ValueError("CUDA graphs not supported with CPU offloading.") - if self.cuda_graph_scope is None: - self.cuda_graph_scope = [] - elif not isinstance(self.cuda_graph_scope, list): - assert isinstance(self.cuda_graph_scope, str), ( - "cuda_graph_scope must be a string or a list of strings, " - f"got {self.cuda_graph_scope}." - ) - self.cuda_graph_scope = [self.cuda_graph_scope] - if self.cuda_graph_impl == "local": assert not self.cuda_graph_scope or self.cuda_graph_scope == ["full_iteration"], ( "For local cuda graph implementation, the only valid value " diff --git a/megatron/training/training.py b/megatron/training/training.py index 2a608dd69d8..193eea747fb 100644 --- a/megatron/training/training.py +++ b/megatron/training/training.py @@ -2510,6 +2510,10 @@ def get_e2e_base_metrics(): if should_exit: break + # Destroy CUDA Graphs. + if args.cuda_graph_impl == "transformer_engine" and cuda_graph_helper.graphs_created(): + cuda_graph_helper.destroy_cudagraphs() + one_logger_utils.track_e2e_metrics() # Flush TensorBoard, WandB writers and one-logger. diff --git a/tests/unit_tests/transformer/test_cuda_graphs.py b/tests/unit_tests/transformer/test_cuda_graphs.py index ee33545ce8f..54e54f9b574 100644 --- a/tests/unit_tests/transformer/test_cuda_graphs.py +++ b/tests/unit_tests/transformer/test_cuda_graphs.py @@ -6,9 +6,11 @@ import pytest import torch +from transformer_engine.pytorch.fp8 import check_fp8_support from megatron.core.enums import ModelType from megatron.core.models.gpt.gpt_layer_specs import ( + get_gpt_decoder_block_spec, get_gpt_layer_with_transformer_engine_spec, get_gpt_mtp_block_spec, ) @@ -37,6 +39,8 @@ from megatron.training.training import setup_model_and_optimizer from tests.unit_tests.test_utilities import Utils +fp8_available, _ = check_fp8_support() + class TestParallelTransformerBlockCudagraphs: def setup_method(self, method): @@ -532,6 +536,9 @@ class TestPartialCudaGraph: def setup_method(self, method): self.seq_length = 512 self.micro_batch_size = 2 + self.tp_size = 2 + self.cp_size = 2 + self.cuda_graph_helper = None # Store original environment variable values self.original_env = { 'CUDA_DEVICE_MAX_CONNECTIONS': os.environ.get('CUDA_DEVICE_MAX_CONNECTIONS'), @@ -547,22 +554,28 @@ def teardown_method(self, method): os.environ.pop(key, None) else: os.environ[key] = value - Utils.destroy_model_parallel() destroy_global_vars() destroy_num_microbatches_calculator() + if self.cuda_graph_helper is not None and self.cuda_graph_helper.graphs_created(): + self.cuda_graph_helper.destroy_cudagraphs() + self.cuda_graph_helper = None gc.collect() def model_provider( self, pre_process=True, post_process=True, - layer_spec_fn=get_gpt_layer_with_transformer_engine_spec, + layer_spec_fn=get_gpt_decoder_block_spec, **config_kwargs, ): - model_parallel_cuda_manual_seed(123) args = get_args() config = core_transformer_config_from_args(args) - transformer_layer_spec = layer_spec_fn() + transformer_layer_spec = layer_spec_fn( + config, + use_transformer_engine=True, + normalization=args.normalization, + qk_l2_norm=args.qk_l2_norm, + ) if args.mtp_num_layers: mtp_block_spec = get_gpt_mtp_block_spec( config, transformer_layer_spec, use_transformer_engine=True @@ -598,15 +611,14 @@ def create_test_args( args.hidden_size = 128 args.num_attention_heads = 8 args.max_position_embeddings = 512 - args.global_batch_size = self.micro_batch_size * 8 + args.global_batch_size = self.micro_batch_size * 8 // self.tp_size // self.cp_size args.micro_batch_size = self.micro_batch_size args.create_attention_mask_in_dataloader = True args.seq_length = self.seq_length - args.tensor_model_parallel_size = 2 - args.sequence_parallel = True + args.tensor_model_parallel_size = self.tp_size + args.sequence_parallel = True if self.tp_size > 1 else False args.pipeline_model_parallel_size = 1 - args.context_parallel_size = 1 - args.expert_model_parallel_size = ep_size + args.context_parallel_size = self.cp_size args.train_iters = 10 args.lr = 3e-5 args.bf16 = True @@ -621,17 +633,26 @@ def create_test_args( # MoE settings args.num_experts = 4 args.expert_model_parallel_size = ep_size + args.expert_tensor_parallel_size = 1 if ep_size > 1 else self.tp_size args.moe_shared_expert_intermediate_size = 1024 - args.moe_layer_freq = "[0,0,1,1]" + args.moe_layer_freq = [0, 0, 1, 1] args.moe_permute_fusion = True args.moe_router_fusion = True args.moe_router_topk = 2 + args.moe_router_dtype = "fp32" # CUDA graph settings args.cuda_graph_impl = cuda_graph_impl args.cuda_graph_scope = cuda_graph_scope args.cuda_graph_warmup_steps = cuda_graph_warmup_steps - args.use_te_rng_tracker = cuda_graph_impl != "none" + + # fp8 settings + if fp8_available: + args.fp8 = "e4m3" + args.fp8_recipe = "tensorwise" + args.first_last_layers_bf16 = True + args.num_layers_at_start_in_bf16 = 1 + args.num_layers_at_end_in_bf16 = 1 for key, value in kwargs.items(): assert hasattr(args, key) @@ -641,15 +662,15 @@ def create_test_args( set_global_variables(args, False) return args - def get_batch(self, seq_length, micro_batch_size): - data = list(range(seq_length)) + def get_batch(self, seq_length, micro_batch_size, cp_size): + data = list(range(seq_length // cp_size)) input_ids = torch.tensor(data, dtype=torch.int64).repeat((micro_batch_size, 1)).cuda() labels = 1 + torch.tensor(data, dtype=torch.int64).repeat((micro_batch_size, 1)).cuda() position_ids = torch.tensor(data, dtype=torch.int64).repeat((micro_batch_size, 1)).cuda() attention_mask = torch.ones( - (micro_batch_size, 1, seq_length, seq_length), dtype=bool + (micro_batch_size, 1, seq_length // cp_size, seq_length), dtype=bool ).cuda() - loss_mask = torch.ones(seq_length).repeat((micro_batch_size, 1)).cuda() + loss_mask = torch.ones(seq_length // cp_size).repeat((micro_batch_size, 1)).cuda() return input_ids, labels, position_ids, attention_mask, loss_mask def _run_test_helper( @@ -662,12 +683,10 @@ def _run_test_helper( set_args(args) torch.manual_seed(123) - Utils.initialize_model_parallel( - tensor_model_parallel_size=2, expert_model_parallel_size=ep_size - ) + model_parallel_cuda_manual_seed(123) input_ids, labels, position_ids, attention_mask, loss_mask = self.get_batch( - self.seq_length, self.micro_batch_size + self.seq_length, self.micro_batch_size, self.cp_size ) gpt_model, optimizer, _ = setup_model_and_optimizer( @@ -675,13 +694,10 @@ def _run_test_helper( ) assert len(gpt_model) == 1 # Assume only one model in the model provider. - loss_list = [] - - cuda_graph_helper = None if cuda_graph_impl == "transformer_engine": from megatron.core.transformer.cuda_graphs import TECudaGraphHelper - cuda_graph_helper = TECudaGraphHelper( + self.cuda_graph_helper = TECudaGraphHelper( model=gpt_model, config=gpt_model[0].config, seq_length=self.seq_length, @@ -689,14 +705,17 @@ def _run_test_helper( optimizers=[optimizer], ) + loss_list = [] + for i in range(100): gpt_model[0].zero_grad_buffer() optimizer.zero_grad() # Capture CUDA graphs after warmup if helper is provided - if cuda_graph_helper is not None and i == cuda_graph_warmup_steps: - cuda_graph_helper.create_cudagraphs() + if self.cuda_graph_helper is not None and i == cuda_graph_warmup_steps: + self.cuda_graph_helper.create_cudagraphs() + gpt_model[0].set_is_first_microbatch() output = gpt_model[0].forward( input_ids=input_ids, position_ids=position_ids, @@ -707,7 +726,7 @@ def _run_test_helper( # Check output shapes assert output.shape[0] == self.micro_batch_size - assert output.shape[1] == self.seq_length + assert output.shape[1] == self.seq_length // self.cp_size # Verify gradients loss = output.mean() @@ -721,16 +740,29 @@ def _run_test_helper( loss_list.append(loss.item()) + if self.cuda_graph_helper is not None and self.cuda_graph_helper.graphs_created(): + self.cuda_graph_helper.destroy_cudagraphs() + self.cuda_graph_helper = None + return torch.tensor(loss_list) @pytest.mark.skipif( - not (HAVE_TE and is_te_min_version("1.14.0")), - reason="Partial CUDA graph support requires TransformerEngine version >= 1.14.0", + not (HAVE_TE and is_te_min_version("2.10.0")), + reason="Partial CUDA graph UT support requires TransformerEngine version >= 2.10.0", ) @pytest.mark.parametrize("ep_size", [1, 4]) @pytest.mark.parametrize("moe_dropless_dispatcher", [False, True]) @pytest.mark.parametrize("moe_dispatcher_type", ["alltoall", "deepep", "hybridep"]) def test_moe_partial_cudagraph(self, ep_size, moe_dropless_dispatcher, moe_dispatcher_type): + initialize_rng_tracker(use_te_rng_tracker=True, force_reset=True) + Utils.initialize_model_parallel( + tensor_model_parallel_size=self.tp_size, + context_parallel_size=self.cp_size, + pipeline_model_parallel_size=1, + expert_tensor_parallel_size=1 if ep_size > 1 else self.tp_size, + expert_model_parallel_size=ep_size, + ) + extra_kwargs = {} if moe_dispatcher_type == "deepep": if not is_deep_ep_available(): @@ -771,6 +803,8 @@ def test_moe_partial_cudagraph(self, ep_size, moe_dropless_dispatcher, moe_dispa ) assert torch.equal(loss_list, loss_list_ref) + Utils.destroy_model_parallel() + if __name__ == "__main__": From 6f139f103cbaa4aaa3fd049f976a3f1c415c81e8 Mon Sep 17 00:00:00 2001 From: Robin Zhang Date: Thu, 13 Nov 2025 04:30:21 -0800 Subject: [PATCH 04/19] enum CudaGraphScope Signed-off-by: Robin Zhang --- .../text_generation_controller.py | 3 +- .../common/language_module/language_module.py | 5 +- megatron/core/models/gpt/gpt_model.py | 4 +- megatron/core/pipeline_parallel/schedules.py | 7 +- megatron/core/ssm/mamba_block.py | 3 +- megatron/core/transformer/attention.py | 3 +- megatron/core/transformer/cuda_graphs.py | 17 ++-- megatron/core/transformer/enums.py | 12 +++ megatron/core/transformer/moe/moe_utils.py | 7 +- .../core/transformer/moe/token_dispatcher.py | 3 +- .../core/transformer/transformer_block.py | 4 +- .../core/transformer/transformer_config.py | 98 +++++++++++-------- .../core/transformer/transformer_layer.py | 45 +++++---- megatron/training/arguments.py | 18 +++- megatron/training/training.py | 5 +- .../inference/engines/test_dynamic_engine.py | 7 +- .../transformer/test_cuda_graphs.py | 18 +++- 17 files changed, 157 insertions(+), 102 deletions(-) diff --git a/megatron/core/inference/text_generation_controllers/text_generation_controller.py b/megatron/core/inference/text_generation_controllers/text_generation_controller.py index 0aed3df079e..4cb1bdd3806 100644 --- a/megatron/core/inference/text_generation_controllers/text_generation_controller.py +++ b/megatron/core/inference/text_generation_controllers/text_generation_controller.py @@ -33,6 +33,7 @@ ) from megatron.core.inference.sampling_params import SamplingParams from megatron.core.inference.utils import get_attention_mask, set_decode_expert_padding +from megatron.core.transformer.enums import CudaGraphScope from megatron.core.transformer.moe.moe_layer import BaseMoELayer from megatron.core.transformer.utils import set_model_to_sequence_parallel from megatron.core.utils import get_asyncio_loop, get_model_config, unwrap_model @@ -922,7 +923,7 @@ def generate_all_output_tokens_static_batch( # Check whether CUDA graphs are enabled enable_cuda_graph = ( model_config.cuda_graph_impl == "local" - and "full_iteration" not in model_config.cuda_graph_scope + and CudaGraphScope.full_iteration not in model_config.cuda_graph_scope ) # Pad batch tokens if necessary diff --git a/megatron/core/models/common/language_module/language_module.py b/megatron/core/models/common/language_module/language_module.py index 8f90fb3ba47..d511bf814de 100644 --- a/megatron/core/models/common/language_module/language_module.py +++ b/megatron/core/models/common/language_module/language_module.py @@ -21,7 +21,7 @@ is_vp_last_stage, ) from megatron.core.process_groups_config import ProcessGroupCollection -from megatron.core.transformer.enums import AttnBackend +from megatron.core.transformer.enums import AttnBackend, CudaGraphScope from megatron.core.transformer.module import MegatronModule from megatron.core.transformer.transformer_config import TransformerConfig from megatron.core.utils import is_te_min_version, make_tp_sharded_tensor_for_checkpoint @@ -136,7 +136,8 @@ def compute_language_model_loss(self, labels: Tensor, logits: Tensor) -> Tensor: # Use is_cg_capturable=True for full iteration CUDA graphs to avoid torch.equal checks is_cg_capturable = ( hasattr(self.config, 'cuda_graph_scope') - and 'full_iteration' in self.config.cuda_graph_scope + and self.config.cuda_graph_scope + and CudaGraphScope.full_iteration in self.config.cuda_graph_scope ) if is_cg_capturable and not is_te_min_version("2.7.0"): from megatron.core.utils import get_te_version diff --git a/megatron/core/models/gpt/gpt_model.py b/megatron/core/models/gpt/gpt_model.py index 77e87917911..9d71234b93a 100644 --- a/megatron/core/models/gpt/gpt_model.py +++ b/megatron/core/models/gpt/gpt_model.py @@ -21,7 +21,7 @@ from megatron.core.process_groups_config import ProcessGroupCollection from megatron.core.quantization.utils import get_quant_config_or_none from megatron.core.tensor_parallel import gather_from_sequence_parallel_region -from megatron.core.transformer.enums import ModelType +from megatron.core.transformer.enums import CudaGraphScope, ModelType from megatron.core.transformer.multi_token_prediction import ( MTPLossAutoScaler, MTPLossLoggingHelper, @@ -371,7 +371,7 @@ def _preprocess( and ( ( self.config.cuda_graph_impl == "local" - and "full_iteration" not in self.config.cuda_graph_scope + and CudaGraphScope.full_iteration not in self.config.cuda_graph_scope ) or self.config.flash_decode ) diff --git a/megatron/core/pipeline_parallel/schedules.py b/megatron/core/pipeline_parallel/schedules.py index db670bbeaf1..0181cf6c9d1 100644 --- a/megatron/core/pipeline_parallel/schedules.py +++ b/megatron/core/pipeline_parallel/schedules.py @@ -18,6 +18,7 @@ ) from megatron.core.process_groups_config import ProcessGroupCollection from megatron.core.transformer.cuda_graphs import create_cudagraphs +from megatron.core.transformer.enums import CudaGraphScope from megatron.core.transformer.moe.router import MoEAuxLossAutoScaler from megatron.core.utils import ( drain_embedding_wgrad_compute, @@ -648,7 +649,7 @@ def forward_backward_no_pipelining( if ( hasattr(config, 'cuda_graph_impl') and config.cuda_graph_impl == "local" - and "full_iteration" not in config.cuda_graph_scope + and CudaGraphScope.full_iteration not in config.cuda_graph_scope ): create_cudagraphs() @@ -1912,7 +1913,7 @@ def pp_post_backward(input_tensor_grad, vp_stage=None): if ( hasattr(config, 'cuda_graph_impl') and config.cuda_graph_impl == "local" - and "full_iteration" not in config.cuda_graph_scope + and CudaGraphScope.full_iteration not in config.cuda_graph_scope ): create_cudagraphs() nvtx_range_pop(suffix="misc") @@ -2296,7 +2297,7 @@ def enable_grad_sync(): if ( hasattr(config, 'cuda_graph_impl') and config.cuda_graph_impl == "local" - and "full_iteration" not in config.cuda_graph_scope + and CudaGraphScope.full_iteration not in config.cuda_graph_scope ): create_cudagraphs() diff --git a/megatron/core/ssm/mamba_block.py b/megatron/core/ssm/mamba_block.py index d4a2981178e..937dc595cab 100644 --- a/megatron/core/ssm/mamba_block.py +++ b/megatron/core/ssm/mamba_block.py @@ -25,6 +25,7 @@ from megatron.core.ssm.mamba_hybrid_layer_allocation import allocate_layers from megatron.core.tensor_parallel import get_cuda_rng_tracker from megatron.core.transformer import TransformerConfig +from megatron.core.transformer.enums import CudaGraphScope from megatron.core.transformer.identity_op import IdentityOp from megatron.core.transformer.module import MegatronModule from megatron.core.transformer.spec_utils import ModuleSpec, build_module @@ -292,7 +293,7 @@ def forward( ( ( self.config.cuda_graph_impl == "local" - and "full_iteration" not in self.config.cuda_graph_scope + and CudaGraphScope.full_iteration not in self.config.cuda_graph_scope ) or self.config.flash_decode ) diff --git a/megatron/core/transformer/attention.py b/megatron/core/transformer/attention.py index 1bbd38ed368..a3b05fa35fa 100644 --- a/megatron/core/transformer/attention.py +++ b/megatron/core/transformer/attention.py @@ -23,6 +23,7 @@ get_tensor_model_parallel_world_size, ) from megatron.core.process_groups_config import ProcessGroupCollection +from megatron.core.transformer.enums import CudaGraphScope from megatron.core.transformer.identity_op import IdentityOp from megatron.core.transformer.module import MegatronModule from megatron.core.transformer.spec_utils import ModuleSpec, build_module @@ -791,7 +792,7 @@ def forward( if ( in_decode_mode and self.config.cuda_graph_impl == "local" - and "full_iteration" not in self.config.cuda_graph_scope + and CudaGraphScope.full_iteration not in self.config.cuda_graph_scope and inference_context.is_static_batching() ): raise ValueError(f"CUDA graphs must use flash decode with static batching!") diff --git a/megatron/core/transformer/cuda_graphs.py b/megatron/core/transformer/cuda_graphs.py index dfc7bb4ede3..bdea58cb976 100644 --- a/megatron/core/transformer/cuda_graphs.py +++ b/megatron/core/transformer/cuda_graphs.py @@ -21,6 +21,7 @@ get_all_rng_states, get_cuda_rng_tracker, ) +from megatron.core.transformer.enums import CudaGraphScope from megatron.core.transformer.identity_op import IdentityOp from megatron.core.transformer.module import GraphableMegatronModule, MegatronModule from megatron.core.transformer.transformer_config import TransformerConfig @@ -1344,24 +1345,24 @@ def _layer_is_graphable(layer, config): from megatron.core.transformer.moe.moe_layer import MoELayer from megatron.core.transformer.transformer_layer import TransformerLayer - if isinstance(layer, MambaLayer) and 'mamba' in config.cuda_graph_scope: + if isinstance(layer, MambaLayer) and CudaGraphScope.mamba in config.cuda_graph_scope: # mamba layer. return True if isinstance(layer, TransformerLayer): - if 'attn' in config.cuda_graph_scope and not ( + if CudaGraphScope.attn in config.cuda_graph_scope and not ( isinstance(layer.self_attention, IdentityOp) and isinstance(layer.cross_attention, IdentityOp) ): # attn layer. return True if ( - 'moe' in config.cuda_graph_scope - or 'moe_router' in config.cuda_graph_scope - or 'moe_preprocess' in config.cuda_graph_scope + CudaGraphScope.moe in config.cuda_graph_scope + or CudaGraphScope.moe_router in config.cuda_graph_scope + or CudaGraphScope.moe_preprocess in config.cuda_graph_scope ) and isinstance(layer.mlp, MoELayer): # moe layer. return True - if 'mlp' in config.cuda_graph_scope and isinstance(layer.mlp, MLP): + if CudaGraphScope.mlp in config.cuda_graph_scope and isinstance(layer.mlp, MLP): # mlp layer. return True return False @@ -1388,7 +1389,7 @@ def __init__(self, model, config, seq_length, micro_batch_size, optimizers=[]): "Setting NCCL_GRAPH_REGISTER=0 to avoid illegal memory access when using " "CUDA Graph with PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True." ) - assert "full_iteration" not in config.cuda_graph_scope, ( + assert CudaGraphScope.full_iteration not in config.cuda_graph_scope, ( "full_iteration cuda graph is not supported for cuda_graph_impl=transformer_engine. " "Please use cuda_graph_impl=local instead." ) @@ -1529,7 +1530,7 @@ def get_rotary_pos_emb(transformer_module, transformer_input): and not isinstance(layer.self_attention, IdentityOp) and ( not self.config.cuda_graph_scope - or 'attn' in self.config.cuda_graph_scope + or CudaGraphScope.attn in self.config.cuda_graph_scope ) ) if is_te_min_version("1.10.0"): diff --git a/megatron/core/transformer/enums.py b/megatron/core/transformer/enums.py index 52b82029f90..d7b37dd8a03 100644 --- a/megatron/core/transformer/enums.py +++ b/megatron/core/transformer/enums.py @@ -65,3 +65,15 @@ class AttnBackend(enum.Enum): unfused = 3 local = 4 auto = 5 + + +class CudaGraphScope(enum.Enum): + """Cuda Graph Scope""" + + full_iteration = 1 + attn = 2 + mlp = 3 + moe = 4 # only used for MoeLayer + moe_router = 5 # only used for MoeLayer + moe_preprocess = 6 # only used for MoeLayer + mamba = 7 # only used for MambaLayer diff --git a/megatron/core/transformer/moe/moe_utils.py b/megatron/core/transformer/moe/moe_utils.py index 5a0793ef5b9..433e6deb3d6 100644 --- a/megatron/core/transformer/moe/moe_utils.py +++ b/megatron/core/transformer/moe/moe_utils.py @@ -9,6 +9,7 @@ from megatron.core import parallel_state from megatron.core.process_groups_config import ProcessGroupCollection from megatron.core.transformer.cuda_graphs import is_graph_capturing +from megatron.core.transformer.enums import CudaGraphScope try: import transformer_engine as te # pylint: disable=unused-import @@ -1190,13 +1191,13 @@ def maybe_raise_signal(moe_layer, **kwargs): ): if ( step_condition == "route" - and 'moe_router' in moe_layer.config.cuda_graph_scope - and 'moe_preprocess' not in moe_layer.config.cuda_graph_scope + and CudaGraphScope.moe_router in moe_layer.config.cuda_graph_scope + and CudaGraphScope.moe_preprocess not in moe_layer.config.cuda_graph_scope ): raise MoECudaGraphPartialCaptureSignal(moe_layer, "route", **kwargs) elif ( step_condition == "preprocess" - and 'moe_preprocess' in moe_layer.config.cuda_graph_scope + and CudaGraphScope.moe_preprocess in moe_layer.config.cuda_graph_scope ): raise MoECudaGraphPartialCaptureSignal(moe_layer, "preprocess", **kwargs) diff --git a/megatron/core/transformer/moe/token_dispatcher.py b/megatron/core/transformer/moe/token_dispatcher.py index 7e3a831bbf3..7f8e3ae2428 100644 --- a/megatron/core/transformer/moe/token_dispatcher.py +++ b/megatron/core/transformer/moe/token_dispatcher.py @@ -16,6 +16,7 @@ gather_from_sequence_parallel_region, reduce_scatter_to_sequence_parallel_region, ) +from megatron.core.transformer.enums import CudaGraphScope from megatron.core.transformer.moe.fused_a2a import ( fused_combine, fused_dispatch, @@ -433,7 +434,7 @@ def __init__( } if ( config.cuda_graph_impl == "transformer_engine" - and 'moe_preprocess' in config.cuda_graph_scope + and CudaGraphScope.moe_preprocess in config.cuda_graph_scope ): self.cuda_dtoh_point = "before_ep_alltoall" else: diff --git a/megatron/core/transformer/transformer_block.py b/megatron/core/transformer/transformer_block.py index b61124f04a0..1581febce97 100755 --- a/megatron/core/transformer/transformer_block.py +++ b/megatron/core/transformer/transformer_block.py @@ -18,7 +18,7 @@ from megatron.core.packed_seq_params import PackedSeqParams from megatron.core.pipeline_parallel.utils import is_vp_first_stage, is_vp_last_stage from megatron.core.process_groups_config import ProcessGroupCollection -from megatron.core.transformer.enums import LayerType +from megatron.core.transformer.enums import CudaGraphScope, LayerType from megatron.core.transformer.module import GraphableMegatronModule, MegatronModule from megatron.core.transformer.spec_utils import ModuleSpec, build_module from megatron.core.transformer.transformer_config import TransformerConfig @@ -522,7 +522,7 @@ def _should_call_local_cudagraph(self, *args, **kwargs): kwargs.get('inference_context') is not None or kwargs.get('inference_params') is not None ) - and 'full_iteration' in self.config.cuda_graph_scope + and CudaGraphScope.full_iteration in self.config.cuda_graph_scope ): if kwargs['inference_context'].is_static_batching(): using_cuda_graph = kwargs['inference_context'].is_decode_only() diff --git a/megatron/core/transformer/transformer_config.py b/megatron/core/transformer/transformer_config.py index 91165dfe0de..4049b9fdd38 100644 --- a/megatron/core/transformer/transformer_config.py +++ b/megatron/core/transformer/transformer_config.py @@ -9,7 +9,7 @@ from megatron.core.enums import Fp4Recipe, Fp8Recipe from megatron.core.quantization.quant_config import RecipeConfig -from megatron.core.transformer.enums import AttnBackend +from megatron.core.transformer.enums import AttnBackend, CudaGraphScope from megatron.core.transformer.pipeline_parallel_layer_layout import PipelineParallelLayerLayout from ..fusions.fused_bias_geglu import quick_gelu @@ -651,7 +651,7 @@ class TransformerConfig(ModelParallelConfig): excluding optimizer) is enabled. "transformer_engine": capture the CUDA graph using TE make_graphed_callables().""" - cuda_graph_scope: Optional[List[str]] = None + cuda_graph_scope: Optional[List[CudaGraphScope]] = None """Determines the CUDA graphs capturing scope. When cuda_graph_impl is set to "transformer_engine", valid values are "attn", "mlp", "moe", "moe_router", "moe_preprocess", "mamba". None means the full layer. @@ -1417,11 +1417,30 @@ def __post_init__(self): if self.cuda_graph_scope is None: self.cuda_graph_scope = [] elif not isinstance(self.cuda_graph_scope, list): - assert isinstance(self.cuda_graph_scope, str), ( - "cuda_graph_scope must be a string or a list of strings, " - f"got {self.cuda_graph_scope}." - ) - self.cuda_graph_scope = [self.cuda_graph_scope] + if isinstance(self.cuda_graph_scope, CudaGraphScope): + self.cuda_graph_scope = [self.cuda_graph_scope] + else: + assert isinstance(self.cuda_graph_scope, str), ( + "cuda_graph_scope must be a string that can be converted to a list of " + f"CudaGraphScope, got {self.cuda_graph_scope}." + ) + self.cuda_graph_scope = self.cuda_graph_scope.split(',') + if all(isinstance(scope, str) for scope in self.cuda_graph_scope): + # Backward compatibility for "full" scope. Now we use an empty list instead. + if "full" in self.cuda_graph_scope: + assert self.cuda_graph_scope == [ + "full" + ], "full scope cannot be used with other scopes." + warnings.warn( + "full scope is deprecated. " + "Use empty cuda_graph_scope to capture the whole layer." + ) + self.cuda_graph_scope = [] + else: + self.cuda_graph_scope = [CudaGraphScope[scope] for scope in self.cuda_graph_scope] + assert all( + isinstance(scope, CudaGraphScope) for scope in self.cuda_graph_scope + ), f"cuda_graph_scope must be a list of CudaGraphScope, got {self.cuda_graph_scope}." if self.cuda_graph_impl != "none": assert self.cuda_graph_impl in [ @@ -1433,47 +1452,37 @@ def __post_init__(self): raise ValueError("CUDA graphs not supported with CPU offloading.") if self.cuda_graph_impl == "local": - assert not self.cuda_graph_scope or self.cuda_graph_scope == ["full_iteration"], ( + assert not self.cuda_graph_scope or self.cuda_graph_scope == [ + CudaGraphScope.full_iteration + ], ( "For local cuda graph implementation, the only valid value " "for cuda_graph_scope is full_iteration. " "To use other scopes, use cuda_graph_impl=transformer_engine." ) if self.cuda_graph_impl == "transformer_engine": - assert "full_iteration" not in self.cuda_graph_scope, ( + assert CudaGraphScope.full_iteration not in self.cuda_graph_scope, ( "To use full iteration cuda graph, please use " "cuda_graph_impl=transformer_engine instead of cuda_graph_impl=local." ) - for scope in self.cuda_graph_scope: - assert scope in [ - 'attn', - 'mlp', - 'moe', - 'moe_router', - 'moe_preprocess', - 'mamba', - ], ( - "--cuda-graph-scope should be attn, mlp, moe, moe_router, moe_preprocess, " - f"or mamba, got {self.cuda_graph_scope}." - ) - assert ( - 'moe' not in self.cuda_graph_scope or 'moe_router' not in self.cuda_graph_scope + CudaGraphScope.moe not in self.cuda_graph_scope + or CudaGraphScope.moe_router not in self.cuda_graph_scope ), 'cuda_graph_scope must not contain both moe and moe_router.' - if 'moe_preprocess' in self.cuda_graph_scope: + if CudaGraphScope.moe_preprocess in self.cuda_graph_scope: assert ( - 'moe_router' in self.cuda_graph_scope + CudaGraphScope.moe_router in self.cuda_graph_scope ), 'moe_preprocess cuda graph is only supported with moe_router cuda graph.' if self.num_moe_experts is None or self.num_moe_experts <= 1: assert ( - 'moe' not in self.cuda_graph_scope - and 'moe_router' not in self.cuda_graph_scope + CudaGraphScope.moe not in self.cuda_graph_scope + and CudaGraphScope.moe_router not in self.cuda_graph_scope ), 'moe cuda graph is only supported for MoE.' else: if self.moe_layer_freq == 1 or ( isinstance(self.moe_layer_freq, list) and 0 not in self.moe_layer_freq ): - assert 'mlp' not in self.cuda_graph_scope, ( + assert CudaGraphScope.mlp not in self.cuda_graph_scope, ( 'mlp cuda graph is only supported for dense layers, ' 'but not found in the model.' ) @@ -1482,13 +1491,13 @@ def __post_init__(self): or not self.moe_pad_expert_input_to_capacity ): assert ( - 'moe' not in self.cuda_graph_scope + CudaGraphScope.moe not in self.cuda_graph_scope ), 'moe cuda graph is only supported with drop-padding MoE.' if self.moe_token_dispatcher_type == 'alltoall' and ( self.moe_expert_capacity_factor is not None or self.moe_router_padding_for_fp8 ): - assert 'moe_preprocess' not in self.cuda_graph_scope, ( + assert CudaGraphScope.moe_preprocess not in self.cuda_graph_scope, ( 'moe_preprocess cuda graph is not supported when there are ' 'DtoH copies and synchronizations in the preprocess step.' ) @@ -1498,25 +1507,28 @@ def __post_init__(self): raise ValueError( "Full-layer CUDA graphs not supported with activation recomputation." ) - elif self.cuda_graph_scope != ['full_iteration']: + elif self.cuda_graph_scope != [CudaGraphScope.full_iteration]: # For scoped CUDA graphs, only the non-graphed parts of the layer can be # recomputed. So check if there are overlaps between the recomputed parts # and the graphed parts. - if "attn" in self.cuda_graph_scope: + if CudaGraphScope.attn in self.cuda_graph_scope: for module in self.recompute_modules: if module in ['core_attn', 'mla_up_proj']: raise ValueError( f'attn cuda graph is not supported with {module} recompute.' ) - if "mlp" in self.cuda_graph_scope and "mlp" in self.recompute_modules: + if ( + CudaGraphScope.mlp in self.cuda_graph_scope + and "mlp" in self.recompute_modules + ): raise ValueError(f'mlp cuda graph is not supported with mlp recompute.') - if "moe" in self.cuda_graph_scope: + if CudaGraphScope.moe in self.cuda_graph_scope: for module in self.recompute_modules: if module in ['moe_act', 'moe', 'shared_experts']: raise ValueError( f'moe cuda graph is not supported with {module} recompute.' ) - if "moe_router" in self.cuda_graph_scope: + if CudaGraphScope.moe_router in self.cuda_graph_scope: for module in self.recompute_modules: if module in ['moe', 'shared_experts']: raise ValueError( @@ -1525,25 +1537,25 @@ def __post_init__(self): ) if "layernorm" in self.recompute_modules: if ( - "attn" in self.cuda_graph_scope - and "mlp" in self.cuda_graph_scope + CudaGraphScope.attn in self.cuda_graph_scope + and CudaGraphScope.mlp in self.cuda_graph_scope and ( - "moe" in self.cuda_graph_scope - or "moe_router" in self.cuda_graph_scope + CudaGraphScope.moe in self.cuda_graph_scope + or CudaGraphScope.moe_router in self.cuda_graph_scope ) ): raise ValueError( 'cuda graph is not supported with layernorm recompute.' ) - if "attn" in self.cuda_graph_scope: + if CudaGraphScope.attn in self.cuda_graph_scope: warnings.warn( "input_layernorm recompute is not supported with attention " "cudagraph. Will only recompute the pre_mlp_layernorm." ) if ( - "mlp" in self.cuda_graph_scope - or "moe" in self.cuda_graph_scope - or "moe_router" in self.cuda_graph_scope + CudaGraphScope.mlp in self.cuda_graph_scope + or CudaGraphScope.moe in self.cuda_graph_scope + or CudaGraphScope.moe_router in self.cuda_graph_scope ): warnings.warn( "pre_mlp_layernorm recompute is not supported with mlp/moe " diff --git a/megatron/core/transformer/transformer_layer.py b/megatron/core/transformer/transformer_layer.py index 14e2dfe36a5..55babb42f24 100644 --- a/megatron/core/transformer/transformer_layer.py +++ b/megatron/core/transformer/transformer_layer.py @@ -16,7 +16,7 @@ from megatron.core.packed_seq_params import PackedSeqParams from megatron.core.process_groups_config import ProcessGroupCollection from megatron.core.transformer.cuda_graphs import is_graph_capturing -from megatron.core.transformer.enums import LayerType +from megatron.core.transformer.enums import CudaGraphScope, LayerType from megatron.core.transformer.identity_op import IdentityFuncOp, IdentityOp from megatron.core.transformer.mlp import MLP from megatron.core.transformer.module import GraphableMegatronModule @@ -381,18 +381,21 @@ def __init__( if "layernorm" in self.config.recompute_modules: if not isinstance(self.input_layernorm, IdentityOp) and ( self.config.cuda_graph_impl == "none" - or 'attn' not in self.config.cuda_graph_scope + or CudaGraphScope.attn not in self.config.cuda_graph_scope ): self.recompute_input_layernorm = True if self.config.fp8: self.self_attention.set_for_recompute_input_layernorm() if not isinstance(self.pre_mlp_layernorm, IdentityOp) and ( self.config.cuda_graph_impl == "none" - or (not self.is_moe_layer and 'mlp' not in self.config.cuda_graph_scope) + or ( + not self.is_moe_layer + and CudaGraphScope.mlp not in self.config.cuda_graph_scope + ) or ( self.is_moe_layer - and 'moe' not in self.config.cuda_graph_scope - and 'moe_router' not in self.config.cuda_graph_scope + and CudaGraphScope.moe not in self.config.cuda_graph_scope + and CudaGraphScope.moe_router not in self.config.cuda_graph_scope ) ): self.recompute_pre_mlp_layernorm = True @@ -600,7 +603,7 @@ def _forward_mlp(self, hidden_states, inference_context=None): and self.config.cuda_graph_impl == "transformer_engine" and self.training and is_graph_capturing() - and 'moe_router' in self.config.cuda_graph_scope + and CudaGraphScope.moe_router in self.config.cuda_graph_scope ): assert ( not self.recompute_pre_mlp_layernorm @@ -716,7 +719,7 @@ def get_layer_static_inputs(self, seq_length, micro_batch_size): static_inputs = super().get_layer_static_inputs(seq_length, micro_batch_size) if not isinstance(self.self_attention, IdentityOp) and ( - not self.config.cuda_graph_scope or 'attn' in self.config.cuda_graph_scope + not self.config.cuda_graph_scope or CudaGraphScope.attn in self.config.cuda_graph_scope ): slen_per_cp = seq_length // self.config.context_parallel_size static_inputs["attention_mask"] = ( @@ -735,18 +738,18 @@ def _get_submodules_under_cudagraphs(self): return super()._get_submodules_under_cudagraphs() submodules = [] - if 'attn' in self.config.cuda_graph_scope: + if CudaGraphScope.attn in self.config.cuda_graph_scope: submodules += [ self.input_layernorm, self.self_attention, self.pre_cross_attn_layernorm, self.cross_attention, ] - if (not self.is_moe_layer and 'mlp' in self.config.cuda_graph_scope) or ( - self.is_moe_layer and 'moe' in self.config.cuda_graph_scope + if (not self.is_moe_layer and CudaGraphScope.mlp in self.config.cuda_graph_scope) or ( + self.is_moe_layer and CudaGraphScope.moe in self.config.cuda_graph_scope ): submodules += [self.pre_mlp_layernorm, self.mlp] - elif self.is_moe_layer and 'moe_router' in self.config.cuda_graph_scope: + elif self.is_moe_layer and CudaGraphScope.moe_router in self.config.cuda_graph_scope: submodules += [self.pre_mlp_layernorm, self.mlp.router] if ( self.config.moe_shared_expert_intermediate_size is not None @@ -764,7 +767,7 @@ def _te_cuda_graph_capture(self, *args, **kwargs): 2. If context is None, it cannot be returned as output. """ context = None - if not self.config.cuda_graph_scope or 'attn' in self.config.cuda_graph_scope: + if not self.config.cuda_graph_scope or CudaGraphScope.attn in self.config.cuda_graph_scope: hidden_states, context = self._forward_attention(*args, **kwargs) else: if len(args) > 0: @@ -774,12 +777,12 @@ def _te_cuda_graph_capture(self, *args, **kwargs): if ( not self.config.cuda_graph_scope - or (not self.is_moe_layer and 'mlp' in self.config.cuda_graph_scope) + or (not self.is_moe_layer and CudaGraphScope.mlp in self.config.cuda_graph_scope) or ( self.is_moe_layer and ( - 'moe' in self.config.cuda_graph_scope - or 'moe_router' in self.config.cuda_graph_scope + CudaGraphScope.moe in self.config.cuda_graph_scope + or CudaGraphScope.moe_router in self.config.cuda_graph_scope ) ) ): @@ -800,7 +803,7 @@ def _te_cuda_graph_replay(self, *args, **kwargs): Hence, `inference_context` and `packed_seq_params` are excluded from input list. """ context = None - if self.config.cuda_graph_scope and 'attn' not in self.config.cuda_graph_scope: + if self.config.cuda_graph_scope and CudaGraphScope.attn not in self.config.cuda_graph_scope: hidden_states, context = self._forward_attention(*args, **kwargs) args = (hidden_states,) kwargs = {} @@ -820,13 +823,13 @@ def _te_cuda_graph_replay(self, *args, **kwargs): if ( not self.config.cuda_graph_scope - or (not self.is_moe_layer and 'mlp' in self.config.cuda_graph_scope) - or (self.is_moe_layer and 'moe' in self.config.cuda_graph_scope) + or (not self.is_moe_layer and CudaGraphScope.mlp in self.config.cuda_graph_scope) + or (self.is_moe_layer and CudaGraphScope.moe in self.config.cuda_graph_scope) ): # CUDA Graph captures the whole MLP/MoE part. CUDA Graph output is the layer output. assert len(cuda_graph_output) == 1, "CUDA Graph output should be the layer output." output = cuda_graph_output.pop() - elif self.is_moe_layer and 'moe_router' in self.config.cuda_graph_scope: + elif self.is_moe_layer and CudaGraphScope.moe_router in self.config.cuda_graph_scope: # CUDA Graph partially captures the MoE. # The rest of the layer should go to the normal pass. shared_expert_output, routing_map, residual = None, None, None @@ -841,7 +844,7 @@ def _te_cuda_graph_replay(self, *args, **kwargs): # Split cudagraph outputs into function outputs and attribute outputs, and # process them separately. Function outputs should have three tensors. func_output, attr_outputs = cuda_graph_output[:3], cuda_graph_output[3:] - if 'moe_preprocess' in self.config.cuda_graph_scope: + if CudaGraphScope.moe_preprocess in self.config.cuda_graph_scope: hidden_states, probs, residual = func_output valid_cudagraph_attrs = self.mlp.token_dispatcher.valid_cudagraph_attrs assert len(attr_outputs) == len( @@ -948,7 +951,7 @@ def _should_call_local_cudagraph(self, *args, **kwargs): (kwargs.get('inference_context') is not None) or (kwargs.get('inference_params') is not None) ) - and 'full_iteration' not in self.config.cuda_graph_scope + and CudaGraphScope.full_iteration not in self.config.cuda_graph_scope ): if kwargs['inference_context'].is_static_batching(): using_cuda_graph = kwargs['inference_context'].is_decode_only() diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py index f3fa79888b2..cbf7e75e603 100644 --- a/megatron/training/arguments.py +++ b/megatron/training/arguments.py @@ -22,7 +22,7 @@ from megatron.core.rerun_state_machine import RerunStateMachine from megatron.core.transformer import MLATransformerConfig, TransformerConfig from megatron.core.transformer.pipeline_parallel_layer_layout import PipelineParallelLayerLayout -from megatron.core.transformer.enums import AttnBackend +from megatron.core.transformer.enums import AttnBackend, CudaGraphScope from megatron.core.transformer.heterogeneous.heterogeneous_config import ( HeterogeneousTransformerConfig, MLPConfig, @@ -746,7 +746,7 @@ def validate_args(args, defaults={}): if args.rank == 0: print('accumulate and all-reduce gradients in fp32 for ' 'bfloat16 data type.', flush=True) - if args.cuda_graph_impl == "local" and "full_iteration" in args.cuda_graph_scope: + if args.cuda_graph_impl == "local" and CudaGraphScope.full_iteration in args.cuda_graph_scope: if not args.inference_dynamic_batching: assert not args.check_for_nan_in_loss_and_grad, \ "--no-check-for-nan-in-loss-and-grad should be set with full_iteration CUDA graph" @@ -1207,6 +1207,15 @@ def validate_args(args, defaults={}): assert ( args.recompute_granularity != 'full' ), 'recompute_granularity must not be full when CUDA Graphs are enabled.' + if args.cuda_graph_scope == "full" or ( + isinstance(args.cuda_graph_scope, list) and "full" in args.cuda_graph_scope + ): + if isinstance(args.cuda_graph_scope, list): + assert args.cuda_graph_scope == ["full"], "full scope cannot be used with other scopes." + args.cuda_graph_scope = [] + warn_rank_0( + 'full scope is deprecated. Use empty cuda_graph_scope to capture the whole layer.' + ) # Print arguments. _print_args("arguments", args) @@ -1422,7 +1431,7 @@ def _add_inference_args(parser): '"none": no CUDA graph. ' '"local": capture the CUDA graph using MCore local implementation. --cuda-graph-scope=\"full_iteration\" enables whole iteration CUDA graph. ' '"transformer_engine": capture the CUDA graph using TE make_graphed_callables().') - group.add_argument('--cuda-graph-scope', nargs='+', type=str, default=[], + group.add_argument('--cuda-graph-scope', nargs='+', type=lambda scope: CudaGraphScope[scope] if isinstance(scope, str) and scope != "full" else scope, default=[], help='Determines the CUDA graphs capturing scope. ' 'choices: "attn", "mlp", "moe", "moe_router", "moe_preprocess", "mamba", "full_iteration". ' '"attn": captures operations in TransformerLayer._forward_attention(). ' @@ -1434,7 +1443,8 @@ def _add_inference_args(parser): '"mamba": captures the mamba layer. ' '"full_iteration": captures a whole iteration. ' 'full_iteration scope is only supported with --cuda-graph-impl=local, other scopes are only supported with --cuda-graph-impl=transformer_engine. ' - 'If not specified, the default scope is to capture the whole Transformer layer.') + 'If not specified, the default scope is to capture the whole Transformer layer. ' + 'For backward compatibility, we still allow passing "full" to specify capturing the whole layer, and convert it to an empty list.') group.add_argument('--use-legacy-static-engine', action='store_true', default=False, help='Use legacy static engine. (Current static engine uses dynamic engine under the hood)', dest='use_legacy_static_engine') diff --git a/megatron/training/training.py b/megatron/training/training.py index 193eea747fb..0ad2325feec 100644 --- a/megatron/training/training.py +++ b/megatron/training/training.py @@ -60,6 +60,7 @@ from megatron.training.checkpointing import checkpoint_exists from megatron.core.full_cuda_graph import FullCudaGraphWrapper from megatron.core.transformer.cuda_graphs import TECudaGraphHelper +from megatron.core.transformer.enums import CudaGraphScope from megatron.core.transformer.module import Float16Module from megatron.core.distributed import DistributedDataParallelConfig, TorchFullyShardedDataParallelConfig from megatron.core.distributed import DistributedDataParallel as DDP @@ -2168,7 +2169,7 @@ def train( eval_iterations = 0 # Wrap forward_backward_func for Full iteration CUDA graph forward_backward_func = get_forward_backward_func() - if args.cuda_graph_impl == "local" and "full_iteration" in args.cuda_graph_scope: + if args.cuda_graph_impl == "local" and CudaGraphScope.full_iteration in args.cuda_graph_scope: forward_backward_func = FullCudaGraphWrapper(forward_backward_func, cuda_graph_warmup_steps=args.cuda_graph_warmup_steps) def get_e2e_base_metrics(): @@ -2587,7 +2588,7 @@ def evaluate( eval_batch_size = args.global_batch_size eval_num_microbatches = eval_batch_size // (args.micro_batch_size * args.data_parallel_size) forward_backward_func = get_forward_backward_func() - if args.cuda_graph_impl == "local" and "full_iteration" in args.cuda_graph_scope: + if args.cuda_graph_impl == "local" and CudaGraphScope.full_iteration in args.cuda_graph_scope: forward_backward_func = FullCudaGraphWrapper(forward_backward_func, cuda_graph_warmup_steps=args.cuda_graph_warmup_steps) if eval_iters is None: diff --git a/tests/unit_tests/inference/engines/test_dynamic_engine.py b/tests/unit_tests/inference/engines/test_dynamic_engine.py index 97242eff292..8b0c09b57a5 100644 --- a/tests/unit_tests/inference/engines/test_dynamic_engine.py +++ b/tests/unit_tests/inference/engines/test_dynamic_engine.py @@ -42,6 +42,7 @@ from megatron.core.models.mamba.mamba_model import MambaModel from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed from megatron.core.transformer.cuda_graphs import CudaGraphManager, _CudagraphGlobalRecord +from megatron.core.transformer.enums import CudaGraphScope from megatron.core.transformer.transformer_config import TransformerConfig from megatron.core.utils import ( check_mamba_sequence_packing_support, @@ -101,7 +102,7 @@ class DynamicEngineTestConfig: return_log_probs: bool = False materialize_only_last_token_logits: bool = True skip_prompt_log_probs: bool = False - cuda_graph_scope: List[str] = None + cuda_graph_scope: List[CudaGraphScope] = None force_build_cuda_graphs: bool = False transformer_impl: str = "local" # If False, do not build cuda graphs in the tests, even if @@ -125,7 +126,7 @@ def __post_init__(self): self.max_sequence_length = self.num_tokens_total if self.cuda_graph_scope is None: - self.cuda_graph_scope = ["full_iteration"] + self.cuda_graph_scope = [CudaGraphScope.full_iteration] @dataclass @@ -509,7 +510,7 @@ def teardown_method(self, method): ) @pytest.mark.parametrize("model_provider", ["gpt", "mamba"]) @pytest.mark.parametrize("num_cuda_graphs", [None, 1, 4]) - @pytest.mark.parametrize("cuda_graph_scope", [[], ["full_iteration"]]) + @pytest.mark.parametrize("cuda_graph_scope", [[], [CudaGraphScope.full_iteration]]) def test_simple(self, model_provider, num_cuda_graphs, cuda_graph_scope) -> None: """Simple test that runs without errors, and validates output.""" skip_if_mamba_sequence_packing_not_available(model_provider) diff --git a/tests/unit_tests/transformer/test_cuda_graphs.py b/tests/unit_tests/transformer/test_cuda_graphs.py index 54e54f9b574..edfc184cd51 100644 --- a/tests/unit_tests/transformer/test_cuda_graphs.py +++ b/tests/unit_tests/transformer/test_cuda_graphs.py @@ -26,6 +26,7 @@ model_parallel_cuda_manual_seed, ) from megatron.core.transformer.cuda_graphs import CudaGraphManager, _CudagraphGlobalRecord +from megatron.core.transformer.enums import CudaGraphScope from megatron.core.transformer.transformer_block import TransformerBlock from megatron.core.transformer.transformer_config import TransformerConfig from megatron.core.utils import is_fa_min_version, is_te_min_version @@ -785,12 +786,19 @@ def test_moe_partial_cudagraph(self, ep_size, moe_dropless_dispatcher, moe_dispa loss_list_ref = self._run_test_helper(ep_size, "none", None, 0, **extra_kwargs) for cuda_graph_scope in [ None, - ["attn"], - ["moe"], - ["mlp", "moe_router"], - ["attn", "mlp", "moe_router", "moe_preprocess"], + [CudaGraphScope.attn], + [CudaGraphScope.moe], + [CudaGraphScope.mlp, CudaGraphScope.moe_router], + [ + CudaGraphScope.attn, + CudaGraphScope.mlp, + CudaGraphScope.moe_router, + CudaGraphScope.moe_preprocess, + ], ]: - if moe_dropless_dispatcher and (cuda_graph_scope is None or "moe" in cuda_graph_scope): + if moe_dropless_dispatcher and ( + cuda_graph_scope is None or CudaGraphScope.moe in cuda_graph_scope + ): # Dropless MoE doesn't work with "moe" scope cudagraph. Skip. continue cuda_graph_warmup_steps = 3 From a35489f3ae2670f5b10dda306ead45b4cf814488 Mon Sep 17 00:00:00 2001 From: Robin Zhang Date: Fri, 14 Nov 2025 00:56:09 -0800 Subject: [PATCH 05/19] minor fixes Signed-off-by: Robin Zhang --- megatron/core/transformer/cuda_graphs.py | 6 +++--- megatron/core/transformer/transformer_config.py | 8 ++++---- 2 files changed, 7 insertions(+), 7 deletions(-) diff --git a/megatron/core/transformer/cuda_graphs.py b/megatron/core/transformer/cuda_graphs.py index bdea58cb976..30680a8ea32 100644 --- a/megatron/core/transformer/cuda_graphs.py +++ b/megatron/core/transformer/cuda_graphs.py @@ -1719,13 +1719,13 @@ def destroy_cudagraphs(self): Destroy CUDA Graphs. """ assert self._graphs_created, "CUDA Graphs have not been created." - graphs_destoryed, graphs_not_destroyed = 0, 0 + graphs_destroyed, graphs_not_destroyed = 0, 0 for _, layers in enumerate(self.callables_per_chunk): for layer in layers: for graph in layer.cuda_graphs: if is_te_min_version("2.10.0"): graph.reset() - graphs_destoryed += 1 + graphs_destroyed += 1 else: graphs_not_destroyed += 1 layer.cuda_graphs = [] @@ -1733,6 +1733,6 @@ def destroy_cudagraphs(self): log_single_rank( logger, logging.INFO, - f'{graphs_destoryed} graphs destroyed, {graphs_not_destroyed} graphs not destroyed.', + f'{graphs_destroyed} graphs destroyed, {graphs_not_destroyed} graphs not destroyed.', ) self._graphs_created = False diff --git a/megatron/core/transformer/transformer_config.py b/megatron/core/transformer/transformer_config.py index 4049b9fdd38..29434726e06 100644 --- a/megatron/core/transformer/transformer_config.py +++ b/megatron/core/transformer/transformer_config.py @@ -1455,15 +1455,15 @@ def __post_init__(self): assert not self.cuda_graph_scope or self.cuda_graph_scope == [ CudaGraphScope.full_iteration ], ( - "For local cuda graph implementation, the only valid value " - "for cuda_graph_scope is full_iteration. " - "To use other scopes, use cuda_graph_impl=transformer_engine." + "For local cuda graph implementation, the only valid value for " + "cuda_graph_scope is full_iteration, or an empty list to denote layerwise " + "graphs. To use other scopes, use cuda_graph_impl=transformer_engine." ) if self.cuda_graph_impl == "transformer_engine": assert CudaGraphScope.full_iteration not in self.cuda_graph_scope, ( "To use full iteration cuda graph, please use " - "cuda_graph_impl=transformer_engine instead of cuda_graph_impl=local." + "cuda_graph_impl=local instead of cuda_graph_impl=transformer_engine." ) assert ( CudaGraphScope.moe not in self.cuda_graph_scope From 0337f2053bafd3affb50a56a0d01c8650141b98d Mon Sep 17 00:00:00 2001 From: Robin Zhang Date: Wed, 19 Nov 2025 21:35:21 -0800 Subject: [PATCH 06/19] minor updates Signed-off-by: Robin Zhang --- megatron/core/transformer/cuda_graphs.py | 10 +++++----- megatron/core/transformer/moe/fused_a2a.py | 2 -- megatron/core/transformer/moe/token_dispatcher.py | 3 +-- megatron/core/transformer/transformer_config.py | 4 ++-- megatron/training/training.py | 2 +- .../inference/engines/test_dynamic_engine.py | 7 +++---- tests/unit_tests/transformer/test_cuda_graphs.py | 4 ++-- 7 files changed, 14 insertions(+), 18 deletions(-) diff --git a/megatron/core/transformer/cuda_graphs.py b/megatron/core/transformer/cuda_graphs.py index 02f9e218292..aa20e053304 100644 --- a/megatron/core/transformer/cuda_graphs.py +++ b/megatron/core/transformer/cuda_graphs.py @@ -1731,25 +1731,25 @@ def cuda_graph_set_manual_hooks(self): for layer in layers: layer.setup_manual_hooks(model_chunk._make_forward_pre_hook) - def destroy_cudagraphs(self): + def delete_cuda_graphs(self): """ Destroy CUDA Graphs. """ assert self._graphs_created, "CUDA Graphs have not been created." - graphs_destroyed, graphs_not_destroyed = 0, 0 + graphs_deleted, graphs_not_deleted = 0, 0 for _, layers in enumerate(self.callables_per_chunk): for layer in layers: for graph in layer.cuda_graphs: if is_te_min_version("2.10.0"): graph.reset() - graphs_destroyed += 1 + graphs_deleted += 1 else: - graphs_not_destroyed += 1 + graphs_not_deleted += 1 layer.cuda_graphs = [] layer.cuda_graph_manual_hooks = [] log_single_rank( logger, logging.INFO, - f'{graphs_destroyed} graphs destroyed, {graphs_not_destroyed} graphs not destroyed.', + f'{graphs_deleted} graphs deleted, {graphs_not_deleted} graphs not deleted.', ) self._graphs_created = False diff --git a/megatron/core/transformer/moe/fused_a2a.py b/megatron/core/transformer/moe/fused_a2a.py index 6cff70dcbe9..60b0b11a32c 100644 --- a/megatron/core/transformer/moe/fused_a2a.py +++ b/megatron/core/transformer/moe/fused_a2a.py @@ -12,8 +12,6 @@ except ImportError: HAVE_DEEP_EP = False -HAVE_HYBRIDEP = False - import torch _buffer = None diff --git a/megatron/core/transformer/moe/token_dispatcher.py b/megatron/core/transformer/moe/token_dispatcher.py index ad61cd5ca12..3d143a3ffa1 100644 --- a/megatron/core/transformer/moe/token_dispatcher.py +++ b/megatron/core/transformer/moe/token_dispatcher.py @@ -37,8 +37,6 @@ from megatron.core.transformer.moe.shared_experts import SharedExpertMLP from megatron.core.transformer.transformer_config import TransformerConfig -logger = logging.getLogger(__name__) - """ We use the following notation throughout this file: H: hidden size B: micro batch size @@ -1369,6 +1367,7 @@ def __init__( num_experts=self.tp_size * self.config.num_moe_experts, config=self.config, ) + self.cudagraph_attrs = ['_comm_manager.token_probs', '_comm_manager.routing_map'] else: raise ValueError( f"Invalid backend: {self.config.moe_flex_dispatcher_backend}" diff --git a/megatron/core/transformer/transformer_config.py b/megatron/core/transformer/transformer_config.py index e63e221a67c..c162624126d 100644 --- a/megatron/core/transformer/transformer_config.py +++ b/megatron/core/transformer/transformer_config.py @@ -1,7 +1,7 @@ # Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. import warnings -from dataclasses import dataclass +from dataclasses import dataclass, field from typing import Callable, List, Literal, Optional, Tuple, Union import torch @@ -664,7 +664,7 @@ class TransformerConfig(ModelParallelConfig): excluding optimizer) is enabled. "transformer_engine": capture the CUDA graph using TE make_graphed_callables().""" - cuda_graph_scope: Optional[List[CudaGraphScope]] = None + cuda_graph_scope: List[CudaGraphScope] = field(default_factory=list) """Determines the CUDA graphs capturing scope. When cuda_graph_impl is set to "transformer_engine", valid values are "attn", "mlp", "moe", "moe_router", "moe_preprocess", "mamba". None means the full layer. diff --git a/megatron/training/training.py b/megatron/training/training.py index 0ad2325feec..531c7d03b49 100644 --- a/megatron/training/training.py +++ b/megatron/training/training.py @@ -2513,7 +2513,7 @@ def get_e2e_base_metrics(): # Destroy CUDA Graphs. if args.cuda_graph_impl == "transformer_engine" and cuda_graph_helper.graphs_created(): - cuda_graph_helper.destroy_cudagraphs() + cuda_graph_helper.delete_cuda_graphs() one_logger_utils.track_e2e_metrics() diff --git a/tests/unit_tests/inference/engines/test_dynamic_engine.py b/tests/unit_tests/inference/engines/test_dynamic_engine.py index 9e5e4b10360..b25a4e2f531 100644 --- a/tests/unit_tests/inference/engines/test_dynamic_engine.py +++ b/tests/unit_tests/inference/engines/test_dynamic_engine.py @@ -106,7 +106,9 @@ class DynamicEngineTestConfig: return_log_probs: bool = False materialize_only_last_token_logits: bool = True skip_prompt_log_probs: bool = False - cuda_graph_scope: List[CudaGraphScope] = None + cuda_graph_scope: List[CudaGraphScope] = field( + default_factory=lambda: [CudaGraphScope.full_iteration] + ) force_build_cuda_graphs: bool = False transformer_impl: str = "local" # If False, do not build cuda graphs in the tests, even if @@ -130,9 +132,6 @@ def __post_init__(self): assert self.num_tokens_total is not None self.max_sequence_length = self.num_tokens_total - if self.cuda_graph_scope is None: - self.cuda_graph_scope = [CudaGraphScope.full_iteration] - @dataclass class DynamicEngineTestEnv: diff --git a/tests/unit_tests/transformer/test_cuda_graphs.py b/tests/unit_tests/transformer/test_cuda_graphs.py index edfc184cd51..fd0e6c45d00 100644 --- a/tests/unit_tests/transformer/test_cuda_graphs.py +++ b/tests/unit_tests/transformer/test_cuda_graphs.py @@ -558,7 +558,7 @@ def teardown_method(self, method): destroy_global_vars() destroy_num_microbatches_calculator() if self.cuda_graph_helper is not None and self.cuda_graph_helper.graphs_created(): - self.cuda_graph_helper.destroy_cudagraphs() + self.cuda_graph_helper.delete_cuda_graphs() self.cuda_graph_helper = None gc.collect() @@ -742,7 +742,7 @@ def _run_test_helper( loss_list.append(loss.item()) if self.cuda_graph_helper is not None and self.cuda_graph_helper.graphs_created(): - self.cuda_graph_helper.destroy_cudagraphs() + self.cuda_graph_helper.delete_cuda_graphs() self.cuda_graph_helper = None return torch.tensor(loss_list) From 20cb0133ffc3819433449668c6453ed1575b8866 Mon Sep 17 00:00:00 2001 From: Robin Zhang Date: Thu, 20 Nov 2025 06:05:19 -0800 Subject: [PATCH 07/19] update Signed-off-by: Robin Zhang --- megatron/core/transformer/attention.py | 3 +-- megatron/core/transformer/cuda_graphs.py | 2 +- tests/unit_tests/transformer/test_cuda_graphs.py | 2 +- 3 files changed, 3 insertions(+), 4 deletions(-) diff --git a/megatron/core/transformer/attention.py b/megatron/core/transformer/attention.py index 49e2ec00f28..e29b986938c 100644 --- a/megatron/core/transformer/attention.py +++ b/megatron/core/transformer/attention.py @@ -23,7 +23,6 @@ get_tensor_model_parallel_world_size, ) from megatron.core.process_groups_config import ProcessGroupCollection -from megatron.core.transformer.enums import CudaGraphScope from megatron.core.transformer.identity_op import IdentityOp from megatron.core.transformer.module import MegatronModule from megatron.core.transformer.spec_utils import ModuleSpec, build_module @@ -41,7 +40,7 @@ from ..models.common.embeddings.yarn_rotary_pos_embedding import ( _yarn_get_concentration_factor_from_config, ) -from .enums import AttnMaskType +from .enums import AttnMaskType, CudaGraphScope from .transformer_config import TransformerConfig try: diff --git a/megatron/core/transformer/cuda_graphs.py b/megatron/core/transformer/cuda_graphs.py index aa20e053304..31296cd57d1 100644 --- a/megatron/core/transformer/cuda_graphs.py +++ b/megatron/core/transformer/cuda_graphs.py @@ -1733,7 +1733,7 @@ def cuda_graph_set_manual_hooks(self): def delete_cuda_graphs(self): """ - Destroy CUDA Graphs. + Delete all CUDA graphs. """ assert self._graphs_created, "CUDA Graphs have not been created." graphs_deleted, graphs_not_deleted = 0, 0 diff --git a/tests/unit_tests/transformer/test_cuda_graphs.py b/tests/unit_tests/transformer/test_cuda_graphs.py index fd0e6c45d00..4fc5e9611ad 100644 --- a/tests/unit_tests/transformer/test_cuda_graphs.py +++ b/tests/unit_tests/transformer/test_cuda_graphs.py @@ -769,7 +769,7 @@ def test_moe_partial_cudagraph(self, ep_size, moe_dropless_dispatcher, moe_dispa if not is_deep_ep_available(): pytest.skip("Deep EP is not available") extra_kwargs["moe_token_dispatcher_type"] = "flex" - extra_kwargs["moe_enable_deepep"] = True + extra_kwargs["moe_flex_dispatcher_backend"] = "deepep" elif moe_dispatcher_type == "hybridep": if not is_hybrid_ep_available(): pytest.skip("Hybrid EP is not available") From 63a958ff898b25fd21bb98f62be4a8b337efc740 Mon Sep 17 00:00:00 2001 From: Robin Zhang Date: Thu, 20 Nov 2025 16:53:28 -0800 Subject: [PATCH 08/19] remove None check in language_module Signed-off-by: Robin Zhang --- megatron/core/models/common/language_module/language_module.py | 1 - 1 file changed, 1 deletion(-) diff --git a/megatron/core/models/common/language_module/language_module.py b/megatron/core/models/common/language_module/language_module.py index d511bf814de..fc7e7d26637 100644 --- a/megatron/core/models/common/language_module/language_module.py +++ b/megatron/core/models/common/language_module/language_module.py @@ -136,7 +136,6 @@ def compute_language_model_loss(self, labels: Tensor, logits: Tensor) -> Tensor: # Use is_cg_capturable=True for full iteration CUDA graphs to avoid torch.equal checks is_cg_capturable = ( hasattr(self.config, 'cuda_graph_scope') - and self.config.cuda_graph_scope and CudaGraphScope.full_iteration in self.config.cuda_graph_scope ) if is_cg_capturable and not is_te_min_version("2.7.0"): From a3607afc583b6a5ceb96dd698d504b091f824486 Mon Sep 17 00:00:00 2001 From: Robin Zhang Date: Fri, 21 Nov 2025 04:01:27 -0800 Subject: [PATCH 09/19] update hybridep cudagraph ut Signed-off-by: Robin Zhang --- megatron/core/transformer/moe/fused_a2a.py | 8 ++++++++ megatron/core/transformer/moe/token_dispatcher.py | 5 +++-- tests/unit_tests/transformer/test_cuda_graphs.py | 11 ++++++++--- 3 files changed, 19 insertions(+), 5 deletions(-) diff --git a/megatron/core/transformer/moe/fused_a2a.py b/megatron/core/transformer/moe/fused_a2a.py index 60b0b11a32c..045a93039b3 100644 --- a/megatron/core/transformer/moe/fused_a2a.py +++ b/megatron/core/transformer/moe/fused_a2a.py @@ -320,6 +320,14 @@ def init_hybrid_ep_buffer( ) +def reset_hybrid_ep_buffer(): + ''' + Reset the HybridEP buffer + ''' + global _hybrid_ep_buffer + _hybrid_ep_buffer = None + + class HybridEPDispatch(torch.autograd.Function): ''' Fused dispatch operation for permute + dispatch a2a + permute using the HybridEP backend diff --git a/megatron/core/transformer/moe/token_dispatcher.py b/megatron/core/transformer/moe/token_dispatcher.py index 3d143a3ffa1..a30c315bd91 100644 --- a/megatron/core/transformer/moe/token_dispatcher.py +++ b/megatron/core/transformer/moe/token_dispatcher.py @@ -1078,8 +1078,9 @@ def combine( ) # Release the used handle/num_permuted_tokens which could change in each iteration self.handle = None - self.num_permuted_tokens = None - self.num_dispatched_tokens = None + if not self.drop_and_pad: + self.num_permuted_tokens = None + self.num_dispatched_tokens = None return hidden_states def get_permuted_hidden_states_by_experts(self, hidden_states: torch.Tensor) -> torch.Tensor: diff --git a/tests/unit_tests/transformer/test_cuda_graphs.py b/tests/unit_tests/transformer/test_cuda_graphs.py index 4fc5e9611ad..1f34cc0ca83 100644 --- a/tests/unit_tests/transformer/test_cuda_graphs.py +++ b/tests/unit_tests/transformer/test_cuda_graphs.py @@ -27,6 +27,7 @@ ) from megatron.core.transformer.cuda_graphs import CudaGraphManager, _CudagraphGlobalRecord from megatron.core.transformer.enums import CudaGraphScope +from megatron.core.transformer.moe.fused_a2a import reset_hybrid_ep_buffer from megatron.core.transformer.transformer_block import TransformerBlock from megatron.core.transformer.transformer_config import TransformerConfig from megatron.core.utils import is_fa_min_version, is_te_min_version @@ -609,7 +610,7 @@ def create_test_args( args.num_layers = 4 args.mtp_num_layers = 1 args.vocab_size = 1024 - args.hidden_size = 128 + args.hidden_size = 512 args.num_attention_heads = 8 args.max_position_embeddings = 512 args.global_batch_size = self.micro_batch_size * 8 // self.tp_size // self.cp_size @@ -780,6 +781,8 @@ def test_moe_partial_cudagraph(self, ep_size, moe_dropless_dispatcher, moe_dispa if not moe_dropless_dispatcher: if moe_dispatcher_type == "deepep": pytest.skip("Deep EP doesn't support drop&pad MoE") + if moe_dispatcher_type == "hybridep" and ep_size == 1: + pytest.skip("Hybrid EP doesn't support drop&pad MoE with ep_size == 1") extra_kwargs["moe_expert_capacity_factor"] = 1.0 extra_kwargs["moe_pad_expert_input_to_capacity"] = True @@ -796,10 +799,10 @@ def test_moe_partial_cudagraph(self, ep_size, moe_dropless_dispatcher, moe_dispa CudaGraphScope.moe_preprocess, ], ]: - if moe_dropless_dispatcher and ( + if (moe_dropless_dispatcher or moe_dispatcher_type == "hybridep") and ( cuda_graph_scope is None or CudaGraphScope.moe in cuda_graph_scope ): - # Dropless MoE doesn't work with "moe" scope cudagraph. Skip. + # Dropless MoE or Hybrid EP doesn't work with "moe" scope cudagraph. Skip. continue cuda_graph_warmup_steps = 3 loss_list = self._run_test_helper( @@ -811,6 +814,8 @@ def test_moe_partial_cudagraph(self, ep_size, moe_dropless_dispatcher, moe_dispa ) assert torch.equal(loss_list, loss_list_ref) + if moe_dispatcher_type == "hybridep": + reset_hybrid_ep_buffer() Utils.destroy_model_parallel() From be5c46237b2234fded8b19c63e870c73dacb035e Mon Sep 17 00:00:00 2001 From: Robin Zhang Date: Tue, 25 Nov 2025 22:45:11 -0800 Subject: [PATCH 10/19] add nvtx_range_pop Signed-off-by: Robin Zhang --- megatron/core/transformer/transformer_layer.py | 1 + 1 file changed, 1 insertion(+) diff --git a/megatron/core/transformer/transformer_layer.py b/megatron/core/transformer/transformer_layer.py index f044f9bf44e..c8db4597bde 100644 --- a/megatron/core/transformer/transformer_layer.py +++ b/megatron/core/transformer/transformer_layer.py @@ -609,6 +609,7 @@ def _forward_mlp(self, hidden_states, inference_context=None): not self.recompute_pre_mlp_layernorm ), "Recomputation is not supported for CUDA graph." cudagraph_outputs = self.mlp(pre_mlp_layernorm_output) + nvtx_range_pop(suffix="mlp") return cudagraph_outputs + [residual] elif self.recompute_mlp: if self.config.fp8 or self.config.fp4: From 97f73e9b24433784e2fadc39e94a904fddad1a28 Mon Sep 17 00:00:00 2001 From: Robin Zhang Date: Tue, 25 Nov 2025 23:29:37 -0800 Subject: [PATCH 11/19] revert breaking API change Signed-off-by: Robin Zhang --- megatron/core/transformer/transformer_config.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/megatron/core/transformer/transformer_config.py b/megatron/core/transformer/transformer_config.py index c162624126d..f34ff0010b6 100644 --- a/megatron/core/transformer/transformer_config.py +++ b/megatron/core/transformer/transformer_config.py @@ -1,7 +1,7 @@ # Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. import warnings -from dataclasses import dataclass, field +from dataclasses import dataclass from typing import Callable, List, Literal, Optional, Tuple, Union import torch @@ -664,10 +664,12 @@ class TransformerConfig(ModelParallelConfig): excluding optimizer) is enabled. "transformer_engine": capture the CUDA graph using TE make_graphed_callables().""" - cuda_graph_scope: List[CudaGraphScope] = field(default_factory=list) + cuda_graph_scope: Union[str, CudaGraphScope, List[str], List[CudaGraphScope]] = "full" """Determines the CUDA graphs capturing scope. When cuda_graph_impl is set to "transformer_engine", valid values are "attn", "mlp", "moe", - "moe_router", "moe_preprocess", "mamba". None means the full layer. + "moe_router", "moe_preprocess", "mamba". "full" or an empty list means the full layer. "full" + is actually deprecated, but for backward compatibility, we still use "full" as the default + value. It will be transformed to an empty list in __post_init__. When cuda_graph_impl is set to "local", "full_iteration" can be specified as cuda_graph_scope to enable whole iteration CUDA graph. All other values enable layerwise CUDA graph.""" From 40a89ce0869a5a274e58261a22faaf690d09d19a Mon Sep 17 00:00:00 2001 From: Robin Zhang Date: Wed, 26 Nov 2025 00:31:34 -0800 Subject: [PATCH 12/19] revert breaking API change Signed-off-by: Robin Zhang --- megatron/core/transformer/moe/moe_layer.py | 7 +++++++ megatron/core/transformer/moe/token_dispatcher.py | 3 +-- 2 files changed, 8 insertions(+), 2 deletions(-) diff --git a/megatron/core/transformer/moe/moe_layer.py b/megatron/core/transformer/moe/moe_layer.py index eadf3d32ac6..a4a17a0fa3b 100644 --- a/megatron/core/transformer/moe/moe_layer.py +++ b/megatron/core/transformer/moe/moe_layer.py @@ -267,6 +267,13 @@ def combine(self, output: torch.Tensor, shared_expert_output: Optional[torch.Ten output = output + shared_expert_output return output + def router_and_preprocess(self, hidden_states: torch.Tensor): + """This method is a combined method of route and preprocess. Deprecated.""" + + probs, routing_map = self.route(hidden_states) + hidden_states, probs, residual = self.preprocess(hidden_states, probs, routing_map) + return hidden_states, probs, residual + def forward(self, hidden_states: torch.Tensor): """Forward pass for the MoE layer. diff --git a/megatron/core/transformer/moe/token_dispatcher.py b/megatron/core/transformer/moe/token_dispatcher.py index a30c315bd91..fcf9ab5cea4 100644 --- a/megatron/core/transformer/moe/token_dispatcher.py +++ b/megatron/core/transformer/moe/token_dispatcher.py @@ -435,13 +435,12 @@ def __init__( "before_finish": 3, "no_sync": 4, } + self.cuda_dtoh_point = "before_permutation_1" if ( config.cuda_graph_impl == "transformer_engine" and CudaGraphScope.moe_preprocess in config.cuda_graph_scope ): self.cuda_dtoh_point = "before_ep_alltoall" - else: - self.cuda_dtoh_point = "before_permutation_1" if MoEAlltoAllTokenDispatcher.cuda_dtoh_stream is None: MoEAlltoAllTokenDispatcher.cuda_dtoh_stream = torch.cuda.Stream() From 9639ab97ce63782790636afe88b83ac39cd3891d Mon Sep 17 00:00:00 2001 From: Robin Zhang Date: Mon, 1 Dec 2025 04:29:45 -0800 Subject: [PATCH 13/19] cherry-pick dev PR 2353 Signed-off-by: Robin Zhang --- megatron/core/transformer/cuda_graphs.py | 25 ++++++++++++------- megatron/core/transformer/enums.py | 18 ++++++------- .../core/transformer/moe/token_dispatcher.py | 4 ++- megatron/training/arguments.py | 2 +- 4 files changed, 29 insertions(+), 20 deletions(-) diff --git a/megatron/core/transformer/cuda_graphs.py b/megatron/core/transformer/cuda_graphs.py index aabfc4e48a8..bbd7c3c2d35 100644 --- a/megatron/core/transformer/cuda_graphs.py +++ b/megatron/core/transformer/cuda_graphs.py @@ -1740,20 +1740,27 @@ def delete_cuda_graphs(self): Delete all CUDA graphs. """ assert self._graphs_created, "CUDA Graphs have not been created." - graphs_deleted, graphs_not_deleted = 0, 0 - for _, layers in enumerate(self.callables_per_chunk): + + graph_resettable = is_te_min_version("2.10.0") + graphs_reset, graphs_not_reset = 0, 0 + for layers in self.callables_per_chunk: for layer in layers: for graph in layer.cuda_graphs: - if is_te_min_version("2.10.0"): + if graph_resettable: graph.reset() - graphs_deleted += 1 + graphs_reset += 1 else: - graphs_not_deleted += 1 + graphs_not_reset += 1 layer.cuda_graphs = [] layer.cuda_graph_manual_hooks = [] - log_single_rank( - logger, - logging.INFO, - f'{graphs_deleted} graphs deleted, {graphs_not_deleted} graphs not deleted.', + + log_on_each_pipeline_stage( + logger=logger, + tp_group=None, + dp_cp_group=None, + level=logging.INFO, + msg=f'Rank {torch.distributed.get_rank()}: ' + f'{graphs_reset} graphs deleted with explicit reset, ' + f'{graphs_not_reset} graphs deleted without explicit reset.', ) self._graphs_created = False diff --git a/megatron/core/transformer/enums.py b/megatron/core/transformer/enums.py index d7b37dd8a03..d06d58d65f2 100644 --- a/megatron/core/transformer/enums.py +++ b/megatron/core/transformer/enums.py @@ -68,12 +68,12 @@ class AttnBackend(enum.Enum): class CudaGraphScope(enum.Enum): - """Cuda Graph Scope""" - - full_iteration = 1 - attn = 2 - mlp = 3 - moe = 4 # only used for MoeLayer - moe_router = 5 # only used for MoeLayer - moe_preprocess = 6 # only used for MoeLayer - mamba = 7 # only used for MambaLayer + """Cuda Graph Scope - defines which parts of the model to capture.""" + + full_iteration = 1 # Captures the entire training/inference iteration + attn = 2 # Captures attention layers + mlp = 3 # Captures MLP layers (dense layers only) + moe = 4 # Captures MoE layers (drop-and-pad MoE layers only) + moe_router = 5 # Captures MoE router part + moe_preprocess = 6 # Captures MoE preprocessing part (requires moe_router) + mamba = 7 # Captures Mamba layers diff --git a/megatron/core/transformer/moe/token_dispatcher.py b/megatron/core/transformer/moe/token_dispatcher.py index fcf9ab5cea4..5a7644bf3c9 100644 --- a/megatron/core/transformer/moe/token_dispatcher.py +++ b/megatron/core/transformer/moe/token_dispatcher.py @@ -1075,7 +1075,9 @@ def combine( num_permuted_tokens=self.num_permuted_tokens, pad_multiple=self.pad_multiple, ) - # Release the used handle/num_permuted_tokens which could change in each iteration + # Release the used handle/num_permuted_tokens which could change in each iteration. + # For drop_and_pad mode, we don't need to reset the num_permuted_tokens and + # num_dispatched_tokens, because their values never change. self.handle = None if not self.drop_and_pad: self.num_permuted_tokens = None diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py index 30064dee38d..a0a4822f4de 100644 --- a/megatron/training/arguments.py +++ b/megatron/training/arguments.py @@ -1443,7 +1443,7 @@ def _add_inference_args(parser): '"none": no CUDA graph. ' '"local": capture the CUDA graph using MCore local implementation. --cuda-graph-scope=\"full_iteration\" enables whole iteration CUDA graph. ' '"transformer_engine": capture the CUDA graph using TE make_graphed_callables().') - group.add_argument('--cuda-graph-scope', nargs='+', type=lambda scope: CudaGraphScope[scope] if isinstance(scope, str) and scope != "full" else scope, default=[], + group.add_argument('--cuda-graph-scope', nargs='+', type=lambda scope: CudaGraphScope[scope] if scope != "full" else scope, default=[], help='Determines the CUDA graphs capturing scope. ' 'choices: "attn", "mlp", "moe", "moe_router", "moe_preprocess", "mamba", "full_iteration". ' '"attn": captures operations in TransformerLayer._forward_attention(). ' From 8ce611c7642b376508299056f614c50fccedb087 Mon Sep 17 00:00:00 2001 From: Robin Zhang Date: Tue, 2 Dec 2025 11:15:30 +0800 Subject: [PATCH 14/19] Replay "[Dev] feat(MoE): Refactor cuda_graph_scope - part2 (#2353)" (#2447) --- megatron/core/safe_globals.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/megatron/core/safe_globals.py b/megatron/core/safe_globals.py index cc5eb8809e8..ddb1dd25399 100755 --- a/megatron/core/safe_globals.py +++ b/megatron/core/safe_globals.py @@ -13,7 +13,7 @@ from megatron.core.enums import ModelType from megatron.core.optimizer import OptimizerConfig from megatron.core.rerun_state_machine import RerunDiagnostic, RerunMode, RerunState -from megatron.core.transformer.enums import AttnBackend +from megatron.core.transformer.enums import AttnBackend, CudaGraphScope SAFE_GLOBALS = [ SimpleNamespace, @@ -24,6 +24,7 @@ UInt32DType, Namespace, AttnBackend, + CudaGraphScope, ModelType, OptimizerConfig, RerunDiagnostic, From f9857d27449fd66e1f0cded580387c378733cfd2 Mon Sep 17 00:00:00 2001 From: Robin Zhang Date: Wed, 3 Dec 2025 19:47:50 -0800 Subject: [PATCH 15/19] Add functools.wraps Signed-off-by: Robin Zhang --- megatron/core/transformer/moe/moe_utils.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/megatron/core/transformer/moe/moe_utils.py b/megatron/core/transformer/moe/moe_utils.py index 220134b687d..d4ec0cde77a 100644 --- a/megatron/core/transformer/moe/moe_utils.py +++ b/megatron/core/transformer/moe/moe_utils.py @@ -1,5 +1,6 @@ # Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +import functools import math from dataclasses import dataclass from typing import List, Optional, Union @@ -1222,6 +1223,8 @@ def maybe_raise_signal(moe_layer, **kwargs): raise MoECudaGraphPartialCaptureSignal(moe_layer, "preprocess", **kwargs) def decorator(func): + + @functools.wraps(func) def wrapped_func(moe_layer, *args, **kwargs): """ Check if we should skip executing the original function based on the current From 625e5b7e2eac7fe97bd8d2392b1341747289f2f8 Mon Sep 17 00:00:00 2001 From: Robin Zhang Date: Thu, 4 Dec 2025 17:34:45 -0800 Subject: [PATCH 16/19] update test_fp8_param cudagraph ut Signed-off-by: Robin Zhang --- tests/unit_tests/test_fp8_param.py | 24 +++++++++++++++++------- 1 file changed, 17 insertions(+), 7 deletions(-) diff --git a/tests/unit_tests/test_fp8_param.py b/tests/unit_tests/test_fp8_param.py index 0b8d41769ec..361698f7127 100644 --- a/tests/unit_tests/test_fp8_param.py +++ b/tests/unit_tests/test_fp8_param.py @@ -1,4 +1,4 @@ -# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. import contextlib import gc @@ -36,7 +36,10 @@ try: from transformer_engine.pytorch.tensor.utils import post_all_gather_processing - cuda_graph_supported = True + if is_te_min_version("2.10.0"): + cuda_graph_supported = True + else: + reason_for_no_cuda_graph = "Need newer TransformerEngine" except ImportError: reason_for_no_cuda_graph = "Need newer TransformerEngine" @@ -65,12 +68,16 @@ class TestFP8Param: def setup_method(self, method): self.seq_length = 512 self.micro_batch_size = 2 + self.cuda_graph_helper = None os.environ['CUDA_DEVICE_MAX_CONNECTIONS'] = '1' def teardown_method(self, method): Utils.destroy_model_parallel() destroy_global_vars() destroy_num_microbatches_calculator() + if self.cuda_graph_helper is not None and self.cuda_graph_helper.graphs_created(): + self.cuda_graph_helper.delete_cuda_graphs() + self.cuda_graph_helper = None gc.collect() def model_provider( @@ -209,13 +216,12 @@ def _run_test_helper( ) assert len(gpt_model) == 1 # Assume only one model in the model provider. - cuda_graph_helper = None # Hard coded to use cuda_graph_impl="transformer_engine" cuda_graph_impl = "transformer_engine" if use_cuda_graph and cuda_graph_impl == "transformer_engine": from megatron.core.transformer.cuda_graphs import TECudaGraphHelper - cuda_graph_helper = TECudaGraphHelper( + self.cuda_graph_helper = TECudaGraphHelper( model=gpt_model, config=gpt_model[0].config, seq_length=self.seq_length, @@ -250,13 +256,13 @@ def _run_test_helper( # Capture CUDA graphs after warmup if helper is provided. # Hard coded cuda_graph_warmup_steps = 0. cuda_graph_warmup_steps = 0 - if cuda_graph_helper is not None and i == cuda_graph_warmup_steps: + if self.cuda_graph_helper is not None and i == cuda_graph_warmup_steps: if should_disable_forward_pre_hook(args): disable_forward_pre_hook(gpt_model, param_sync=False) - cuda_graph_helper.create_cudagraphs() + self.cuda_graph_helper.create_cudagraphs() if should_disable_forward_pre_hook(args): enable_forward_pre_hook(gpt_model) - cuda_graph_helper.cuda_graph_set_manual_hooks() + self.cuda_graph_helper.cuda_graph_set_manual_hooks() # For the mxfp8_param with reuse_grad_buf_for_mxfp8_param_ag and dp_ag_overlap, # we need to call the _copy_main_params_to_param_buffer() after the grad buffer @@ -297,6 +303,10 @@ def _run_test_helper( loss_list.append(loss.item()) + if self.cuda_graph_helper is not None and self.cuda_graph_helper.graphs_created(): + self.cuda_graph_helper.delete_cuda_graphs() + self.cuda_graph_helper = None + return torch.tensor(loss_list) def run_test(self, tp_size, recipe, inference: bool = False, **kwargs): From 0ccb4250914973c47c7614fd672adc6ca0bde07f Mon Sep 17 00:00:00 2001 From: Robin Zhang Date: Wed, 17 Dec 2025 04:32:40 -0800 Subject: [PATCH 17/19] improve recompute checks Signed-off-by: Robin Zhang --- megatron/core/tensor_parallel/random.py | 5 + megatron/core/transformer/cuda_graphs.py | 7 +- megatron/core/transformer/moe/moe_layer.py | 13 +- megatron/core/transformer/moe/moe_utils.py | 64 ++++------ .../core/transformer/transformer_config.py | 104 +++++++--------- .../core/transformer/transformer_layer.py | 114 ++++++++++++------ megatron/training/arguments.py | 3 - 7 files changed, 158 insertions(+), 152 deletions(-) diff --git a/megatron/core/tensor_parallel/random.py b/megatron/core/tensor_parallel/random.py index 54cac0e41e3..d86aa3e010b 100644 --- a/megatron/core/tensor_parallel/random.py +++ b/megatron/core/tensor_parallel/random.py @@ -555,6 +555,11 @@ def checkpoint(self, run_function, *args): def _recompute(self, _): """Used as a hook to recompute the output.""" + + if self.ctx is None: + # The recomputation has been triggered already. Just return. + return + if not torch.autograd._is_checkpoint_valid(): raise RuntimeError( "Checkpointing is not compatible with .grad(), " diff --git a/megatron/core/transformer/cuda_graphs.py b/megatron/core/transformer/cuda_graphs.py index bbd7c3c2d35..923135c0817 100644 --- a/megatron/core/transformer/cuda_graphs.py +++ b/megatron/core/transformer/cuda_graphs.py @@ -1603,7 +1603,12 @@ def get_rotary_pos_emb(transformer_module, transformer_input): ) def get_make_graphed_callables_kwargs(): - kwargs = {'num_warmup_iters': 11, 'allow_unused_input': True, '_order': order} + kwargs = { + 'num_warmup_iters': 11, + 'allow_unused_input': True, + '_order': order, + 'retain_graph_in_backward': self.config.cuda_graph_retain_backward_graph, + } if is_te_min_version("2.6.0"): # Starting from TE 2.6.0, make_graphed_callables() accepts different number diff --git a/megatron/core/transformer/moe/moe_layer.py b/megatron/core/transformer/moe/moe_layer.py index 1e9b0158d42..9e8bcfb7a56 100644 --- a/megatron/core/transformer/moe/moe_layer.py +++ b/megatron/core/transformer/moe/moe_layer.py @@ -217,9 +217,8 @@ def preprocess( """Preprocess token routing for dispatch. This method preprocesses the hidden states and routing probabilities for the token - dispatcher. The original hidden states are returned as a residual connection. + dispatcher. """ - residual = hidden_states # Project the hidden_states from hidden dimension down to latent dimenion. if self.config.moe_latent_size: assert ( @@ -229,7 +228,7 @@ def preprocess( hidden_states, probs = self.token_dispatcher.dispatch_preprocess( hidden_states, routing_map, probs ) - return hidden_states, probs, residual + return hidden_states, probs def dispatch(self, hidden_states: torch.Tensor, probs: torch.Tensor): """Dispatches tokens to assigned expert ranks via communication. @@ -268,9 +267,7 @@ def shared_experts_compute(self, hidden_states: torch.Tensor): return shared_expert_output - def routed_experts_compute( - self, hidden_states: torch.Tensor, probs: torch.Tensor, residual: torch.Tensor - ): + def routed_experts_compute(self, hidden_states: torch.Tensor, probs: torch.Tensor): """Computes the output of the routed experts on the dispatched tokens. This method first post-processes the dispatched input to get permuted tokens @@ -336,7 +333,7 @@ def custom_forward(hidden_states): try: shared_expert_output = self.shared_experts_compute(hidden_states) probs, routing_map = self.route(hidden_states) - hidden_states, probs, residual = self.preprocess(hidden_states, probs, routing_map) + hidden_states, probs = self.preprocess(hidden_states, probs, routing_map) except MoECudaGraphPartialCaptureSignal as e: # This signal is raised from the maybe_skip_or_early_return_by_cudagraph decorator. # It means we should early-return from the MoE layer forward pass. @@ -346,7 +343,7 @@ def custom_forward(hidden_states): return e.get_early_return_outputs(hidden_states, shared_expert_output) dispatched_input, probs = self.dispatch(hidden_states, probs) - output, mlp_bias = self.routed_experts_compute(dispatched_input, probs, residual) + output, mlp_bias = self.routed_experts_compute(dispatched_input, probs) assert mlp_bias is None, f"mlp_bias is not supported for {type(self.token_dispatcher)}" output = self.combine(output, shared_expert_output) diff --git a/megatron/core/transformer/moe/moe_utils.py b/megatron/core/transformer/moe/moe_utils.py index 3e3730af9c2..76923b530d2 100644 --- a/megatron/core/transformer/moe/moe_utils.py +++ b/megatron/core/transformer/moe/moe_utils.py @@ -1083,17 +1083,24 @@ def get_early_return_outputs( """ Get the CUDA graph early return outputs for the MoE layer, including the intermediate tensors and the intermediate attributes of the token dispatcher. + + The returned output tensors are in the order of: + - routed experts path outputs + - hidden states, probs, and routing map for capturing router + - hidden states and probs for capturing router and preprocess + - intermediate attributes of the token dispatcher (if capturing the preprocess step) + - shared expert path output (if exists) """ if self.return_step == "route": # Capturing the router step returns three intermediate tensors: # hidden states, routing probabilities, and routing map. outputs = [hidden_states, self.kwargs['probs'], self.kwargs['routing_map']] elif self.return_step == "preprocess": - # Capturing the preprocess step returns three intermediate tensors: - # hidden states, routing probabilities, and residual connection. + # Capturing the preprocess step returns two intermediate tensors: + # hidden states and routing probabilities. # It also returns the intermediate attributes of the token dispatcher, recorded in # "token_dispatcher.cudagraph_attrs". - outputs = [self.kwargs['hidden_states'], self.kwargs['probs'], self.kwargs['residual']] + outputs = [self.kwargs['hidden_states'], self.kwargs['probs']] valid_cudagraph_attrs = [] for attr_name in self.moe_layer.token_dispatcher.cudagraph_attrs: hier_attr_name = attr_name.split('.') @@ -1133,8 +1140,6 @@ class MoECudaGraphTensorStore: probs (Optional[torch.Tensor]): The routing probabilities for each token-expert pair. routing_map (Optional[torch.Tensor]): The sparse mapping indicating which experts were selected for each token. Used to skip the normal router step. - residual (Optional[torch.Tensor]): The residual connection tensor before routing. - Used to skip the normal preprocess step. shared_expert_output (Optional[torch.Tensor]): The output from shared experts computation. Used to skip the normal shared expert computation step. """ @@ -1142,7 +1147,6 @@ class MoECudaGraphTensorStore: hidden_states: Optional[torch.Tensor] = None probs: Optional[torch.Tensor] = None routing_map: Optional[torch.Tensor] = None - residual: Optional[torch.Tensor] = None shared_expert_output: Optional[torch.Tensor] = None def is_empty(self) -> bool: @@ -1153,13 +1157,7 @@ def is_empty(self) -> bool: """ return all( getattr(self, field_name) is None - for field_name in [ - 'hidden_states', - 'probs', - 'routing_map', - 'residual', - 'shared_expert_output', - ] + for field_name in ['hidden_states', 'probs', 'routing_map', 'shared_expert_output'] ) def set(self, **kwargs): @@ -1169,7 +1167,6 @@ def set(self, **kwargs): 'hidden_states', 'probs', 'routing_map', - 'residual', 'shared_expert_output', ], f"Invalid field name: {field_name}" if value is not None: @@ -1180,13 +1177,7 @@ def set(self, **kwargs): def clear(self): """Reset all stored tensors to None.""" - for field_name in [ - 'hidden_states', - 'probs', - 'routing_map', - 'residual', - 'shared_expert_output', - ]: + for field_name in ['hidden_states', 'probs', 'routing_map', 'shared_expert_output']: setattr(self, field_name, None) @@ -1259,46 +1250,39 @@ def wrapped_func(moe_layer, *args, **kwargs): # Don't skip the router. assert ( moe_layer.cudagraph_tensor_store.routing_map is None - and moe_layer.cudagraph_tensor_store.residual is None - ), "both routing_map and residual must be None if probs is None" + ), "routing_map must be None if probs is None" probs, routing_map = func(moe_layer, *args, **kwargs) # Maybe early return after the router. maybe_raise_signal(moe_layer, probs=probs, routing_map=routing_map) else: # Skip the router and get value from store. - assert ( - moe_layer.cudagraph_tensor_store.routing_map is not None - or moe_layer.cudagraph_tensor_store.residual is not None - ), "either routing_map or residual must be given if probs is given" probs, routing_map = ( moe_layer.cudagraph_tensor_store.probs, moe_layer.cudagraph_tensor_store.routing_map, ) return probs, routing_map elif step_condition == "preprocess": - if moe_layer.cudagraph_tensor_store.residual is None: + if ( + moe_layer.cudagraph_tensor_store.is_empty() + or moe_layer.cudagraph_tensor_store.routing_map is not None + ): # Don't skip the preprocess. - hidden_states, probs, residual = func(moe_layer, *args, **kwargs) + hidden_states, probs = func(moe_layer, *args, **kwargs) # Maybe early return after the preprocess. - maybe_raise_signal( - moe_layer, hidden_states=hidden_states, probs=probs, residual=residual - ) + maybe_raise_signal(moe_layer, hidden_states=hidden_states, probs=probs) else: # Skip the preprocess and get value from store. assert ( - moe_layer.cudagraph_tensor_store.probs is not None - ), "probs must not be None if residual is not None" - assert ( - moe_layer.cudagraph_tensor_store.routing_map is None - ), "routing_map must be None if residual is not None" - hidden_states, probs, residual = ( + moe_layer.cudagraph_tensor_store.hidden_states is not None + and moe_layer.cudagraph_tensor_store.probs is not None + ), "hidden_states and probs must be given in moe_preprocess cudagraph replay" + hidden_states, probs = ( moe_layer.cudagraph_tensor_store.hidden_states, moe_layer.cudagraph_tensor_store.probs, - moe_layer.cudagraph_tensor_store.residual, ) - return hidden_states, probs, residual + return hidden_states, probs return wrapped_func diff --git a/megatron/core/transformer/transformer_config.py b/megatron/core/transformer/transformer_config.py index e6e960f68d4..4ec14e1a8d5 100644 --- a/megatron/core/transformer/transformer_config.py +++ b/megatron/core/transformer/transformer_config.py @@ -654,11 +654,11 @@ class TransformerConfig(ModelParallelConfig): determines the scope of graph capture.""" cuda_graph_use_single_mempool: bool = False - """When set to true, cudagraphs will be captured inside a single mempool, in which all - cudagraphs may only be used once per step. If false, cudagraphs may be reused across - microbatches. Enabling may reduce cudagraph memory overheads due to memory fragmentation, - however may greatly increase the number of cudagraphs created when the number of microbatches - is high.""" + """[For `local` implementation only] When set to true, cudagraphs will be captured inside a + single mempool, in which all cudagraphs may only be used once per step. If false, cudagraphs may + be reused across microbatches. Enabling may reduce cudagraph memory overheads due to memory + fragmentation, however may greatly increase the number of cudagraphs created when the number of + microbatches is high.""" cuda_graph_retain_backward_graph: bool = False """When set to true, cudagraph backward passes will be graph captured with 'retain_grad=True' @@ -1564,64 +1564,46 @@ def __post_init__(self): ) if self.recompute_granularity: - if self.recompute_granularity != "selective" or not self.cuda_graph_scope: - raise ValueError( - "Full-layer CUDA graphs not supported with activation recomputation." - ) - elif self.cuda_graph_scope != [CudaGraphScope.full_iteration]: - # For scoped CUDA graphs, only the non-graphed parts of the layer can be - # recomputed. So check if there are overlaps between the recomputed parts - # and the graphed parts. - if CudaGraphScope.attn in self.cuda_graph_scope: - for module in self.recompute_modules: - if module in ['core_attn', 'mla_up_proj']: - raise ValueError( - f'attn cuda graph is not supported with {module} recompute.' - ) + if self.recompute_granularity != "selective": + assert self.cuda_graph_scope == [ + CudaGraphScope.full_iteration + ], "full recompute is only supported with full iteration CUDA graph." + else: + # The recompute module should be inside or outside of the graph scope. + # Recompute module coverring graph scope is not allowed. + if "moe" in self.recompute_modules: + assert ( + CudaGraphScope.moe_router not in self.cuda_graph_scope + ), "moe recompute is not supported with moe_router CUDA graph." + # Graphed recompute module doesn't accept random number. if ( - CudaGraphScope.mlp in self.cuda_graph_scope - and "mlp" in self.recompute_modules + not self.cuda_graph_scope + or CudaGraphScope.full_iteration in self.cuda_graph_scope ): - raise ValueError(f'mlp cuda graph is not supported with mlp recompute.') - if CudaGraphScope.moe in self.cuda_graph_scope: - for module in self.recompute_modules: - if module in ['moe_act', 'moe', 'shared_experts']: - raise ValueError( - f'moe cuda graph is not supported with {module} recompute.' - ) - if CudaGraphScope.moe_router in self.cuda_graph_scope: - for module in self.recompute_modules: - if module in ['moe', 'shared_experts']: - raise ValueError( - f'moe_router cuda graph is not supported with {module} ' - 'recompute.' - ) - if "layernorm" in self.recompute_modules: - if ( - CudaGraphScope.attn in self.cuda_graph_scope - and CudaGraphScope.mlp in self.cuda_graph_scope - and ( - CudaGraphScope.moe in self.cuda_graph_scope - or CudaGraphScope.moe_router in self.cuda_graph_scope - ) - ): - raise ValueError( - 'cuda graph is not supported with layernorm recompute.' - ) - if CudaGraphScope.attn in self.cuda_graph_scope: - warnings.warn( - "input_layernorm recompute is not supported with attention " - "cudagraph. Will only recompute the pre_mlp_layernorm." - ) - if ( - CudaGraphScope.mlp in self.cuda_graph_scope - or CudaGraphScope.moe in self.cuda_graph_scope - or CudaGraphScope.moe_router in self.cuda_graph_scope - ): - warnings.warn( - "pre_mlp_layernorm recompute is not supported with mlp/moe " - "cudagraph. Will only recompute the input_layernorm." - ) + full_cudagraph = True + else: + full_cudagraph = False + if self.attention_dropout != 0.0: + assert ( + not full_cudagraph and CudaGraphScope.attn not in self.cuda_graph_scope + ) or "core_attn" not in self.recompute_modules, ( + "attention dropout is not supported with graphed attention " + "recomputation." + ) + if self.hidden_dropout != 0.0: + assert ( + (not full_cudagraph and CudaGraphScope.mlp not in self.cuda_graph_scope) + or "mlp" not in self.recompute_modules + ) and ( + (not full_cudagraph and CudaGraphScope.moe not in self.cuda_graph_scope) + or "moe" not in self.recompute_modules + ), "hidden dropout is not supported with graphed MLP/MoE recomputation." + if self.moe_input_jitter_eps is not None: + assert ( + not full_cudagraph and CudaGraphScope.moe not in self.cuda_graph_scope + ) or "moe" not in self.recompute_modules, ( + "moe_input_jitter_eps is not supported with graphed moe recomputation." + ) if self.moe_token_dispatcher_type in ["allgather"]: if self.variable_seq_lengths is True: diff --git a/megatron/core/transformer/transformer_layer.py b/megatron/core/transformer/transformer_layer.py index c8db4597bde..472f7ce2f0b 100644 --- a/megatron/core/transformer/transformer_layer.py +++ b/megatron/core/transformer/transformer_layer.py @@ -379,24 +379,55 @@ def __init__( self.recompute_mlp = False if self.config.recompute_granularity == 'selective': if "layernorm" in self.config.recompute_modules: - if not isinstance(self.input_layernorm, IdentityOp) and ( - self.config.cuda_graph_impl == "none" - or CudaGraphScope.attn not in self.config.cuda_graph_scope - ): + if not isinstance(self.input_layernorm, IdentityOp): self.recompute_input_layernorm = True if self.config.fp8 or self.config.fp4: self.self_attention.set_for_recompute_input_layernorm() - if not isinstance(self.pre_mlp_layernorm, IdentityOp) and ( - self.config.cuda_graph_impl == "none" - or ( + + def can_recompute_pre_mlp_layernorm_for_cudagraph(): + if ( not self.is_moe_layer - and CudaGraphScope.mlp not in self.config.cuda_graph_scope - ) - or ( - self.is_moe_layer - and CudaGraphScope.moe not in self.config.cuda_graph_scope - and CudaGraphScope.moe_router not in self.config.cuda_graph_scope + or CudaGraphScope.moe_router not in self.config.cuda_graph_scope + ): + # Not a MoE layer, or not capturing the router part. + return True + if ( + self.config.moe_shared_expert_intermediate_size is not None + and self.config.moe_shared_expert_overlap + ): + # If shared expert overlap is used, we cannot make the pre-mlp layernorm + # recomputation, because the shared expert takes the layernorm output as + # input, and it is outside of the CUDA graph scope. + log_single_rank( + logger, + logging.WARNING, + "pre_mlp_layernorm recompute is not supported with moe router " + "cudagraph + shared expert overlap. Disabling pre_mlp_layernorm " + "recompute.", + ) + return False + if CudaGraphScope.moe_preprocess in self.config.cuda_graph_scope and ( + self.config.moe_token_dispatcher_type == "alltoall" + or self.config.moe_latent_size + ): + # Only when capturing the preprocess part and using alltoall token + # dispatcher or latent MoE can we make the pre-mlp layernorm recomputation. + # Because in other cases the layernorm output returns directly as one of the + # outputs of the cudagraph, which will be allocated a static buffer, thus + # not able to be released. + return True + log_single_rank( + logger, + logging.WARNING, + "pre_mlp_layernorm recompute is only supported with moe router + " + "preprocess cudagraph will alltoall token dispatcher or latent MoE. " + "Disabling pre_mlp_layernorm recompute.", ) + return False + + if ( + not isinstance(self.pre_mlp_layernorm, IdentityOp) + and can_recompute_pre_mlp_layernorm_for_cudagraph() ): self.recompute_pre_mlp_layernorm = True if self.config.fp8 or self.config.fp4: @@ -598,20 +629,7 @@ def _forward_mlp(self, hidden_states, inference_context=None): and not isinstance(self.mlp, IdentityOp) ) - if ( - self.is_moe_layer - and self.config.cuda_graph_impl == "transformer_engine" - and self.training - and is_graph_capturing() - and CudaGraphScope.moe_router in self.config.cuda_graph_scope - ): - assert ( - not self.recompute_pre_mlp_layernorm - ), "Recomputation is not supported for CUDA graph." - cudagraph_outputs = self.mlp(pre_mlp_layernorm_output) - nvtx_range_pop(suffix="mlp") - return cudagraph_outputs + [residual] - elif self.recompute_mlp: + if self.recompute_mlp: if self.config.fp8 or self.config.fp4: # import here to avoid circular import from megatron.core.extensions.transformer_engine import te_checkpoint @@ -651,7 +669,23 @@ def _forward_mlp(self, hidden_states, inference_context=None): ) nvtx_range_pop(suffix="mlp") - return self._forward_post_mlp(mlp_output_with_bias, residual) + if ( + self.is_moe_layer + and self.config.cuda_graph_impl == "transformer_engine" + and self.training + and is_graph_capturing() + and CudaGraphScope.moe_router in self.config.cuda_graph_scope + ): + if self.recompute_pre_mlp_layernorm: + # Register the recompute hooks to all the cudagraph output tensors, because some + # tensors are in parallel execution paths and they all need pre_mlp_layernorm to be + # recomputed in backward pass. For example, the router path and the shared expert + # path. So only register in one path is risky. + for tensor in mlp_output_with_bias[1:]: + self.pre_mlp_norm_checkpoint.discard_output_and_register_recompute(tensor) + return list(mlp_output_with_bias) + [residual] + else: + return self._forward_post_mlp(mlp_output_with_bias, residual) def _forward_post_mlp(self, mlp_output_with_bias, residual): """ @@ -833,20 +867,19 @@ def _te_cuda_graph_replay(self, *args, **kwargs): elif self.is_moe_layer and CudaGraphScope.moe_router in self.config.cuda_graph_scope: # CUDA Graph partially captures the MoE. # The rest of the layer should go to the normal pass. - shared_expert_output, routing_map, residual = None, None, None - mlp_residual = cuda_graph_output.pop() + shared_expert_output, routing_map = None, None + # residual is the last element in the CUDA graph output. + residual = cuda_graph_output.pop() if ( self.config.moe_shared_expert_intermediate_size is not None and not self.config.moe_shared_expert_overlap ): - # The shared expert output is the fourth element in the CUDA graph output. + # The shared expert output is the last second element in the CUDA graph output. shared_expert_output = cuda_graph_output.pop() - # Split cudagraph outputs into function outputs and attribute outputs, and - # process them separately. Function outputs should have three tensors. - func_output, attr_outputs = cuda_graph_output[:3], cuda_graph_output[3:] if CudaGraphScope.moe_preprocess in self.config.cuda_graph_scope: - hidden_states, probs, residual = func_output + # CUDA graph output is [hidden_states, probs] + attributes outputs. + (hidden_states, probs), attr_outputs = cuda_graph_output[:2], cuda_graph_output[2:] valid_cudagraph_attrs = self.mlp.token_dispatcher.valid_cudagraph_attrs assert len(attr_outputs) == len( valid_cudagraph_attrs @@ -858,8 +891,12 @@ def _te_cuda_graph_replay(self, *args, **kwargs): attr = getattr(attr, name) setattr(attr, hier_attr_name[-1], attr_outputs[i]) else: - hidden_states, probs, routing_map = func_output - assert not attr_outputs, "cuda_graph_attr_outputs should be empty" + # CUDA graph output is [hidden_states, probs, routing_map]. + assert len(cuda_graph_output) == 3, ( + "CUDA graph output should be [hidden_states, probs, routing_map], " + f"but got {len(cuda_graph_output)} elements" + ) + hidden_states, probs, routing_map = cuda_graph_output # Resume the MoELayer forward pass from the end of the CUDA graph scope. # The MoE layer will skip redundant computations when we pass in the calculated values @@ -869,14 +906,13 @@ def _te_cuda_graph_replay(self, *args, **kwargs): hidden_states=hidden_states, probs=probs, routing_map=routing_map, - residual=residual, shared_expert_output=shared_expert_output, ) mlp_output_with_bias = self.mlp(hidden_states) self.mlp.cudagraph_tensor_store.clear() nvtx_range_pop(suffix="mlp") - output = self._forward_post_mlp(mlp_output_with_bias, mlp_residual) + output = self._forward_post_mlp(mlp_output_with_bias, residual) else: # CUDA Graph does not capture the MLP/MoE part at all. output = self._forward_mlp(*cuda_graph_output) diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py index 3ddd441667b..64bca4fa075 100644 --- a/megatron/training/arguments.py +++ b/megatron/training/arguments.py @@ -1256,9 +1256,6 @@ def validate_args(args, defaults={}): "Setting NCCL_GRAPH_REGISTER=0 to avoid illegal memory access when using " "CUDA Graph with PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True." ) - assert ( - args.recompute_granularity != 'full' - ), 'recompute_granularity must not be full when CUDA Graphs are enabled.' if args.cuda_graph_scope == "full" or ( isinstance(args.cuda_graph_scope, list) and "full" in args.cuda_graph_scope ): From 85583d0ada082292d00153a2e64ff31c5dcbbe35 Mon Sep 17 00:00:00 2001 From: Robin Zhang Date: Thu, 18 Dec 2025 00:33:09 -0800 Subject: [PATCH 18/19] fix backward compatibility Signed-off-by: Robin Zhang --- megatron/core/models/gpt/fine_grained_callables.py | 4 ++-- megatron/core/transformer/moe/moe_layer.py | 2 ++ megatron/core/transformer/moe/moe_utils.py | 2 ++ 3 files changed, 6 insertions(+), 2 deletions(-) diff --git a/megatron/core/models/gpt/fine_grained_callables.py b/megatron/core/models/gpt/fine_grained_callables.py index 5d35512290c..47f773fc156 100644 --- a/megatron/core/models/gpt/fine_grained_callables.py +++ b/megatron/core/models/gpt/fine_grained_callables.py @@ -359,7 +359,7 @@ def submodule_post_attn_forward(node: ScheduleNode, hidden_states: torch.Tensor) pre_mlp_layernorm_output = layer.pre_mlp_layernorm(hidden_states) probs, routing_map = layer.mlp.route(pre_mlp_layernorm_output) - local_tokens, probs, _ = layer.mlp.preprocess(pre_mlp_layernorm_output, probs, routing_map) + local_tokens, probs = layer.mlp.preprocess(pre_mlp_layernorm_output, probs, routing_map) # Detach here for mlp_bda residual connection node.layer_state.residual = node.detach(hidden_states) @@ -401,7 +401,7 @@ def submodule_moe_forward(node: ScheduleNode, dispatched_tokens: torch.Tensor): pre_mlp_layernorm_output = getattr(node.layer_state, 'pre_mlp_layernorm_output', None) shared_expert_output = layer.mlp.shared_experts_compute(pre_mlp_layernorm_output) expert_output, mlp_bias = layer.mlp.routed_experts_compute( - dispatched_tokens, dispatched_probs, pre_mlp_layernorm_output + dispatched_tokens, dispatched_probs ) if layer.recompute_pre_mlp_layernorm: diff --git a/megatron/core/transformer/moe/moe_layer.py b/megatron/core/transformer/moe/moe_layer.py index 4ce0da57d65..8f46f6886ed 100644 --- a/megatron/core/transformer/moe/moe_layer.py +++ b/megatron/core/transformer/moe/moe_layer.py @@ -24,6 +24,7 @@ ) from megatron.core.transformer.spec_utils import ModuleSpec, build_module from megatron.core.transformer.transformer_config import TransformerConfig +from megatron.core.utils import internal_api try: import transformer_engine as te # pylint: disable=unused-import @@ -269,6 +270,7 @@ def shared_experts_compute(self, hidden_states: torch.Tensor): return shared_expert_output + @internal_api def routed_experts_compute(self, hidden_states: torch.Tensor, probs: torch.Tensor): """Computes the output of the routed experts on the dispatched tokens. diff --git a/megatron/core/transformer/moe/moe_utils.py b/megatron/core/transformer/moe/moe_utils.py index 76923b530d2..05525b5daba 100644 --- a/megatron/core/transformer/moe/moe_utils.py +++ b/megatron/core/transformer/moe/moe_utils.py @@ -14,6 +14,7 @@ from megatron.core.transformer.cuda_graphs import is_graph_capturing from megatron.core.transformer.enums import CudaGraphScope from megatron.core.transformer.transformer_config import TransformerConfig +from megatron.core.utils import internal_api try: import transformer_engine as te # pylint: disable=unused-import @@ -1128,6 +1129,7 @@ def get_early_return_outputs( return outputs +@internal_api @dataclass class MoECudaGraphTensorStore: """Storage for tensors used in CUDA graph replay for MoE layers. From 0698a5913377788533ee3ab3b7e741fe660ac7ae Mon Sep 17 00:00:00 2001 From: Robin Zhang Date: Thu, 18 Dec 2025 02:54:07 -0800 Subject: [PATCH 19/19] disable hybridep ut Signed-off-by: Robin Zhang --- tests/unit_tests/transformer/test_cuda_graphs.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/tests/unit_tests/transformer/test_cuda_graphs.py b/tests/unit_tests/transformer/test_cuda_graphs.py index fa03406b07d..346fd393a2f 100644 --- a/tests/unit_tests/transformer/test_cuda_graphs.py +++ b/tests/unit_tests/transformer/test_cuda_graphs.py @@ -1029,6 +1029,10 @@ def test_moe_partial_cudagraph(self, ep_size, moe_dropless_dispatcher, moe_dispa extra_kwargs["moe_token_dispatcher_type"] = "flex" extra_kwargs["moe_flex_dispatcher_backend"] = "deepep" elif moe_dispatcher_type == "hybridep": + pytest.skip( + "Currently, the Hybrid EP is broken. " + "Temporarily skip the test and wait for the fix." + ) if not is_hybrid_ep_available(): pytest.skip("Hybrid EP is not available") extra_kwargs["moe_token_dispatcher_type"] = "flex"