cleanup

Signed-off-by: Jieming Zhang <jiemingz@nvidia.com>

Author: jiemingz

megatron/core/transformer/cuda_graphs.py

@@ -131,27 +131,6 @@ def _check_supported_type(meta):
     ), f"Cudagraphs recieved an arg of type {meta.type} which is not supported."
 
 
-def _determine_if_transformer_decoder_layer(base_module):
-    """Determine if the given module is a transformer decoder layer."""
-    # import modules here to avoid a circular import
-    from megatron.core.ssm.mamba_layer import MambaLayer
-    from megatron.core.transformer.transformer_layer import BaseTransformerLayer, TransformerLayer
-
-    is_potential_decoder_layer = isinstance(
-        base_module, (TransformerLayer, BaseTransformerLayer, MambaLayer)
-    )
-    if not is_potential_decoder_layer:
-        return False
-    if isinstance(base_module, TransformerLayer) and not isinstance(
-        base_module.cross_attention, IdentityOp
-    ):
-        # If the layer has a cross attention, it is not a decoder layer
-        return False
-    else:
-        # Otherwise it is a decoder layer
-        return True
-
-
 def _determine_if_first_last_layer_of_this_vp_chunk(base_module):
     """Determine if the given module is the first/last layer of the PP+VPP chunk it belongs to.
     Returns a tuple of two booleans indicating if the module is the first/last layer of the chunk.
@@ -242,10 +221,6 @@ def create_cudagraphs(cls):
         gc.collect()
         torch.cuda.empty_cache()
 
-        _set_capture_start()
-        if has_te_modules:
-            te_set_capture_start()
-
         def format_mem_bytes(mem_bytes):
             for power, suffix in [(4, "tb"), (3, "gb"), (2, "mb"), (1, "kb"), (0, "bytes")]:
                 suffix_bytes = 1024**power
@@ -279,9 +254,8 @@ def format_mem_bytes(mem_bytes):
                 runner.create_bwd_graph(global_tensor_pool)
 
         global bwd_buffer_reuse_ref_count, fwd_buffer_reuse_ref_count
-        # assert bwd_buffer_reuse_ref_count == 0
-        # assert fwd_buffer_reuse_ref_count == 0
-
+        assert bwd_buffer_reuse_ref_count == 0
+        assert fwd_buffer_reuse_ref_count == 0
 
         # Memory usage.
         time_end = time.time()
@@ -317,11 +291,6 @@ def format_mem_bytes(mem_bytes):
         cls.cudagraph_created = True
         cls.cudagraph_record = []
 
-        # Finished capturing.
-        _set_capture_end()
-        if has_te_modules:
-            te_set_capture_end()
-
         # Return capture time and memory usage.
         return capture_stats
 
@@ -547,8 +516,8 @@ def __init__(
         self.fp8_enabled = False
         self.fp4_enabled = False
         self.deallocate_pipeline_outputs = False
+        self.num_warmup_steps = 1
 
-        self.is_transformer_decoder_layer = _determine_if_transformer_decoder_layer(base_module)
         self.grad_enabled = need_backward and torch.is_grad_enabled()
         self.func = super(MegatronModule, self.base_module).__call__ if func is None else func
         self.is_first_layer, self.is_last_layer = (
@@ -571,14 +540,17 @@ def __init__(
             self.fp8_runtime_enabled = None
             self.fp4_runtime_enabled = None
 
-            if self.fp8_enabled:
-                self.fp8_recipe = FP8GlobalStateManager.get_fp8_recipe()
-                FP8GlobalStateManager.set_skip_fp8_weight_update_tensor(False)
+            if HAVE_TE_GRAPHS:
+                self.has_te_modules = any(
+                    [isinstance(m, TransformerEngineBaseModule) for m in self.base_module.modules()]
+                )
 
-            if self.fp4_enabled:
-                from megatron.core.fp4_utils import get_fp4_recipe  # to avoid circular import
+                if self.fp8_enabled:
+                    self.fp8_recipe = FP8GlobalStateManager.get_fp8_recipe()
+                if self.fp4_enabled:
+                    from megatron.core.fp4_utils import get_fp4_recipe  # to avoid circular import
+                    self.fp4_recipe = get_fp4_recipe(self.base_module.config)
 
-                self.fp4_recipe = get_fp4_recipe(self.base_module.config)
                 FP8GlobalStateManager.set_skip_fp8_weight_update_tensor(False)
 
     def __str__(self):
@@ -669,7 +641,7 @@ def get_fwd_input_buffer(ten):
 
         # cache the moe aux loss if needed, this is needed because the moe aux loss is accumulated inside
         # the transformer layer forward pass:
-        is_moe = self.is_transformer_decoder_layer and hasattr(self.base_module, "is_moe_layer") and self.base_module.is_moe_layer
+        is_moe = hasattr(self.base_module, "is_moe_layer") and self.base_module.is_moe_layer
         if is_moe:
             from megatron.core.transformer.moe.moe_utils import get_moe_layer_wise_logging_tracker
             tracker = get_moe_layer_wise_logging_tracker()
@@ -710,7 +682,7 @@ def get_fwd_input_buffer(ten):
 
         with ctx:
             # warmup again as case graph capture mode may execute a different codepath
-            for _ in range(1):
+            for _ in range(num_warmup_steps):
                 with self.get_quantization_context():
                     def clone_ten(ten):
                         clone = torch.zeros_like(ten)
@@ -744,6 +716,8 @@ def clone_ten(ten):
 
             with self.get_quantization_context():
                 _set_capture_start()
+                if self.has_te_modules: te_set_capture_start()
+
                 # Freeze GC, to speed up capture time ~15-20x.
                 if FREEZE_GC:
                     gc.freeze()
@@ -763,7 +737,9 @@ def clone_ten(ten):
                     # per-device to avoid slowing down graph creation.
                     if self.is_last_layer:
                         gc.collect()
+
                 _set_capture_end()
+                if self.has_te_modules: te_set_capture_end()
 
         # save cudagraph output buffer
         if isinstance(fwd_graph_outputs, torch.Tensor):
@@ -843,6 +819,9 @@ def create_bwd_graph(self, global_tensor_pool):
         input_tensors = self.get_tensors(self.fwd_graph_input_args, self.fwd_graph_input_kwargs)
         fwd_input_surface = input_tensors + tuple(self.params_to_backprop)
 
+        _set_capture_start()
+        if self.has_te_modules: te_set_capture_start()
+
         # Freeze GC, to speed up capture time ~15-20x.
         if FREEZE_GC:
             gc.freeze()
@@ -861,6 +840,9 @@ def create_bwd_graph(self, global_tensor_pool):
         if FREEZE_GC:
             gc.unfreeze()
 
+        _set_capture_end()
+        if self.has_te_modules: te_set_capture_end()
+
         grad_inputs = list(grad_inputs)
 
         self.static_grad_outputs = static_grad_outputs
@@ -966,19 +948,17 @@ def record_graph_capture(self, args, kwargs):
             o.is_cudagraph_output = True
 
         if not self.fwd_graph_recorded:
-            if HAVE_TE_GRAPHS:
-                if FP8GlobalStateManager.is_fp8_enabled():
-                    # check if the low precision recipe is either fp4 or fp8
-                    if is_te_min_version("2.7.0.dev0"):
-                        from transformer_engine.common.recipe import NVFP4BlockScaling
-                        recipe = FP8GlobalStateManager.get_fp8_recipe()
-                        if isinstance(recipe, NVFP4BlockScaling):
-                            self.fp4_runtime_enabled = True
-                        else:
-                            self.fp8_runtime_enabled = True
-                    else:        
+            if self.fp8_enabled or self.fp4_enabled:
+                # check if any low precision recipe is enabled
+                if is_te_min_version("2.7.0.dev0"):
+                    from transformer_engine.common.recipe import NVFP4BlockScaling
+                    recipe = FP8GlobalStateManager.get_fp8_recipe()
+                    if isinstance(recipe, NVFP4BlockScaling):
+                        self.fp4_runtime_enabled = True
+                    else:
                         self.fp8_runtime_enabled = True
-
+                else:        
+                    self.fp8_runtime_enabled = True
 
             logger.debug(f"Recording forward graph creation...")
             m_args, m_kwargs = self.replace_tensors_with_weak_refs(args, kwargs, cache_refs=True)
@@ -1233,9 +1213,8 @@ def wrapped_func(*args, **kwargs):
         self.inference_cudagraphs_lookup_table = defaultdict(lambda: None)
         self.is_first_microbatch = False
 
-        # Without pipeline parallelism, microbatches execute one at a time.
-        # Therefore modules will always execute in the same order, so cudagraphs
-        # can both be reused and share a single mempool.
+        # Without pipeline parallelism, modules execute one at a time in the same order, so cudagraphs
+        # may be reused across microbatches
         self.reuse_cudagraphs = parallel_state.get_pipeline_model_parallel_world_size() == 1
         if CudaGraphManager.global_mempool is None:
             CudaGraphManager.global_mempool = torch.cuda.graph_pool_handle()