minor updates

Signed-off-by: Robin Zhang <[email protected]>

Author: buptzyb

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

megatron/core/transformer/moe/fused_a2a.py

@@ -12,8 +12,6 @@
 except ImportError:
     HAVE_DEEP_EP = False
 
-HAVE_HYBRIDEP = False
-
 import torch
 
 _buffer = None

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}"

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.

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()
 

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:

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)