diff --git a/megatron/core/models/gpt/gpt_model.py b/megatron/core/models/gpt/gpt_model.py index ad5f46c1f97..b0e6db63079 100644 --- a/megatron/core/models/gpt/gpt_model.py +++ b/megatron/core/models/gpt/gpt_model.py @@ -29,6 +29,7 @@ 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 InferenceCudaGraphScope, ModelType +from megatron.core.transformer.module import GraphableMegatronModule from megatron.core.transformer.moe.paged_stash import paged_stash_init_chunk_handler from megatron.core.transformer.multi_token_prediction import ( MultiTokenPredictionBlock, @@ -48,7 +49,7 @@ logger = logging.getLogger(__name__) -class GPTModel(LanguageModule): +class GPTModel(LanguageModule, GraphableMegatronModule): """GPT Transformer language model. Args: @@ -237,6 +238,8 @@ def __init__( pg_collection=self.pg_collection, vp_stage=vp_stage, ) + if hasattr(self, 'cudagraph_manager') and hasattr(self.decoder, 'cudagraph_manager'): + del self.decoder.cudagraph_manager if self.mtp_process: self.mtp = MultiTokenPredictionBlock( @@ -522,6 +525,42 @@ def preprocess_for_paged_stash(self): vp_size=self.config.virtual_pipeline_model_parallel_size, vp_stage=self.vp_stage ) + def _should_call_local_cudagraph(self, *args, **kwargs): + """ + Check if we should call the local cudagraph path. + """ + if ( + InferenceMode.is_active() + and hasattr(self, 'cudagraph_manager') + and ( + kwargs.get('inference_context') is not None + or kwargs.get('inference_params') is not None + ) + and self.config.inference_cuda_graph_scope == InferenceCudaGraphScope.block + ): + if kwargs['inference_context'].is_static_batching(): + using_cuda_graph = kwargs['inference_context'].is_decode_only() + else: + using_cuda_graph = kwargs['inference_context'].using_cuda_graph_this_step() + + if using_cuda_graph: + return True + return False + + def __call__(self, *args, **kwargs): + if self._should_call_local_cudagraph(*args, **kwargs): + return super().__call__(*args, **kwargs)[0] + return super().__call__(*args, **kwargs) + + def create_mcore_cudagraph_manager(self, config): + """ + Create the cudagraph manager for the full iteration inference scope + """ + if config.inference_cuda_graph_scope == InferenceCudaGraphScope.block: + from megatron.core.transformer.cuda_graphs import CudaGraphManager + + self.cudagraph_manager = CudaGraphManager(config) + def forward( self, input_ids: Tensor, diff --git a/tests/unit_tests/inference/engines/test_dynamic_engine.py b/tests/unit_tests/inference/engines/test_dynamic_engine.py index af91409ef7b..ac7cd1bb57c 100644 --- a/tests/unit_tests/inference/engines/test_dynamic_engine.py +++ b/tests/unit_tests/inference/engines/test_dynamic_engine.py @@ -213,6 +213,7 @@ def _assert_inference_cuda_graphs_disabled(env) -> None: assert env.engine.cuda_graph_impl == "full_iteration" assert env.engine.inference_cuda_graph_scope == InferenceCudaGraphScope.none assert env.engine.capture_stats is None + assert not hasattr(model, 'cudagraph_manager') assert not hasattr(model.decoder, 'cudagraph_manager') for layer in model.decoder.layers: assert not hasattr(layer, 'cudagraph_manager') @@ -845,11 +846,10 @@ def test_simple(self, model_provider, num_cuda_graphs, inference_cuda_graph_scop assert env.engine.context.cuda_graph_batch_dimensions_list model = env.engine.controller.inference_wrapped_model.model if inference_cuda_graph_scope == InferenceCudaGraphScope.block: - # hybrid models attach cudagraph_manager to the model; others attach to the decoder - if model_provider == "hybrid": - assert model.cudagraph_manager.cudagraph_runners - else: - assert model.decoder.cudagraph_manager.cudagraph_runners + # GPT and hybrid models both own the block-scope graph at model level; + # GPT removes the decoder's fallback manager at construction. + assert model.cudagraph_manager.cudagraph_runners + assert not hasattr(model.decoder, 'cudagraph_manager') else: # check if cudagraph runners are created at the layer level for layer in model.decoder.layers: @@ -959,7 +959,8 @@ def test_deprecated_full_iteration_inference_scope_matches_new_flag_runtime_beha model = env.engine.controller.inference_wrapped_model.model assert model.config.inference_cuda_graph_scope == InferenceCudaGraphScope.block assert model.config.cuda_graph_modules == [] - assert model.decoder.cudagraph_manager.cudagraph_runners + assert model.cudagraph_manager.cudagraph_runners + assert not hasattr(model.decoder, 'cudagraph_manager') for layer in model.decoder.layers: assert not hasattr(layer, 'cudagraph_manager')