diff --git a/examples/multimodal/layer_scaling.py b/examples/multimodal/layer_scaling.py index a82afa7cc5e..b3d2881b315 100644 --- a/examples/multimodal/layer_scaling.py +++ b/examples/multimodal/layer_scaling.py @@ -4,6 +4,8 @@ import torch from megatron.core.transformer.transformer_layer import TransformerLayer +from megatron.core.typed_torch import copy_signature + def _bias_dropout_add_func_layer_scaling(ls, x_with_bias, residual, prob, training): x, bias = x_with_bias # unpack @@ -36,6 +38,7 @@ def get_bias_dropout_add_layer_scaling(ls, training, fused): # Add LayerScaling to our default TransformerLayer. class LayerScalingTransformerLayer(TransformerLayer): + @copy_signature(TransformerLayer.__init__) def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) self.ls1 = torch.nn.Parameter(torch.ones(self.config.hidden_size)) diff --git a/examples/multimodal/nvlm/internvit.py b/examples/multimodal/nvlm/internvit.py index cb95129c02a..0018bb5ccb9 100644 --- a/examples/multimodal/nvlm/internvit.py +++ b/examples/multimodal/nvlm/internvit.py @@ -36,9 +36,9 @@ from megatron.core.transformer.module import MegatronModule from megatron.core.transformer.spec_utils import ModuleSpec from megatron.core.transformer.transformer_config import TransformerConfig -from megatron.core.transformer.transformer_layer import TransformerLayer, TransformerLayerSubmodules +from megatron.core.transformer.transformer_layer import TransformerLayerSubmodules from megatron.core.transformer.utils import make_sharded_tensors_for_checkpoint -from megatron.core.typed_torch import not_none +from megatron.core.typed_torch import copy_signature, not_none from megatron.core.utils import divide try: @@ -173,6 +173,7 @@ def get_mlp_module_spec(use_te: bool = True) -> ModuleSpec: # Override a few things that are special in InternViT and not supported by the SelfAttention class. class InternViTSelfAttention(SelfAttention): + @copy_signature(SelfAttention.__init__) def __init__( self, config: TransformerConfig, submodules: SelfAttentionSubmodules, *args, **kwargs ): @@ -213,6 +214,7 @@ def __init__( class InternViTTEDotProductAttention(TEDotProductAttention): """Adjusted Attention for InternViT""" + @copy_signature(TEDotProductAttention.forward) def forward(self, *args, **kwargs): """Regular TEDotProductAttention + zero-out dummy attention heads.""" out = super().forward(*args, **kwargs) diff --git a/megatron/core/extensions/transformer_engine.py b/megatron/core/extensions/transformer_engine.py index 996330f5674..bb913d97446 100644 --- a/megatron/core/extensions/transformer_engine.py +++ b/megatron/core/extensions/transformer_engine.py @@ -51,6 +51,7 @@ is_layer_window_attention, make_sharded_tensors_for_checkpoint, ) +from megatron.core.typed_torch import copy_signature from megatron.core.utils import ( get_pg_rank, get_pg_size, @@ -1925,6 +1926,7 @@ def sharded_state_dict(self, prefix="", sharded_offsets=(), metadata=None): class TEFusedMLP(MLP): """MLP wrapper using Transformer Engine's operation-based API.""" + @copy_signature(MLP.__init__) def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) diff --git a/megatron/core/models/gpt/fine_grained_callables.py b/megatron/core/models/gpt/fine_grained_callables.py index 7d09773d4eb..6658b6363ea 100644 --- a/megatron/core/models/gpt/fine_grained_callables.py +++ b/megatron/core/models/gpt/fine_grained_callables.py @@ -22,7 +22,7 @@ get_mtp_layer_offset, ) from megatron.core.transformer.transformer_layer import TransformerLayer, make_viewless_tensor -from megatron.core.typed_torch import apply_module +from megatron.core.typed_torch import apply_module, copy_signature from megatron.core.utils import internal_api @@ -613,6 +613,7 @@ def submodule_combine_forward(node: ScheduleNode, output: torch.Tensor): output = make_viewless_tensor(inp=output, requires_grad=True, keep_graph=True) return output + @copy_signature(layer._forward_mlp, handle_first_dst_param='preserve') def mlp_wrapper(node: ScheduleNode, *args, **kwargs): """Wrapper for Dense forward.""" return layer._forward_mlp(*args, **kwargs) diff --git a/megatron/core/post_training/modelopt/layers.py b/megatron/core/post_training/modelopt/layers.py index 45aabf3db66..4a82048d255 100644 --- a/megatron/core/post_training/modelopt/layers.py +++ b/megatron/core/post_training/modelopt/layers.py @@ -12,6 +12,7 @@ from megatron.core.transformer.transformer_config import TransformerConfig from megatron.core.transformer.transformer_layer import TransformerLayer from megatron.core.transformer.utils import make_sharded_tensors_for_checkpoint +from megatron.core.typed_torch import copy_signature logger = logging.getLogger(__name__) @@ -202,6 +203,7 @@ class RealQuantTransformerLayer(TransformerLayer): verbose: bool = False real_quant_cfg: str = "None" + @copy_signature(TransformerLayer.__init__) def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) diff --git a/megatron/core/rerun_state_machine.py b/megatron/core/rerun_state_machine.py index 8fce2beaa85..1136a493e72 100644 --- a/megatron/core/rerun_state_machine.py +++ b/megatron/core/rerun_state_machine.py @@ -16,6 +16,7 @@ from megatron.core._rank_utils import log_single_rank, safe_get_rank from megatron.core.dist_checkpointing.mapping import ShardedObject +from megatron.core.typed_torch import copy_signature """DISCLAIMER: THIS IS AN EXPERIMENTAL FEATURE. @@ -1338,13 +1339,14 @@ def load_state_dict(self, state_dict: SerializableStateType) -> None: self.injected_error_type = state_dict["injected_error_type"] -def initialize_rerun_state_machine(**kwargs) -> None: +@copy_signature(RerunStateMachine.__init__, handle_first_src_param='skip') +def initialize_rerun_state_machine(*args, **kwargs) -> None: """Helper function to initialize the rerun machine instance. Check the RerunStateMachine class for the details. """ - rerun_state_machine: RerunStateMachine = RerunStateMachine(**kwargs) + rerun_state_machine: RerunStateMachine = RerunStateMachine(*args, **kwargs) _set_rerun_state_machine(rerun_state_machine) diff --git a/megatron/core/transformer/transformer_layer.py b/megatron/core/transformer/transformer_layer.py index c9cf57a4eb0..d3dc3bfeec1 100644 --- a/megatron/core/transformer/transformer_layer.py +++ b/megatron/core/transformer/transformer_layer.py @@ -25,7 +25,7 @@ from megatron.core.transformer.spec_utils import ModuleSpec, build_module from megatron.core.transformer.torch_norm import LayerNormBuilder from megatron.core.transformer.transformer_config import TransformerConfig -from megatron.core.typed_torch import apply_module +from megatron.core.typed_torch import apply_module, copy_signature from megatron.core.utils import ( deprecate_inference_params, get_pg_rank, @@ -515,21 +515,6 @@ def _get_layer_offset(config: TransformerConfig): ) return get_transformer_layer_offset(config) - def forward(self, *args, **kwargs): - """ - Perform a forward pass through the transformer layer. - - This method calls the core computation of a transformer layer, including - self-attention, cross-attention (if applicable), and feed-forward operations. - """ - hidden_states, context = self._forward_attention(*args, **kwargs) - output = self._forward_mlp( - hidden_states, - kwargs.get("inference_context", None), - padding_mask=kwargs.get("padding_mask", None), - ) - return output, context - def _forward_attention( self, hidden_states: Tensor, @@ -675,6 +660,22 @@ def _forward_attention( return hidden_states, context + @copy_signature(_forward_attention) + def forward(self, *args, **kwargs): + """ + Perform a forward pass through the transformer layer. + + This method calls the core computation of a transformer layer, including + self-attention, cross-attention (if applicable), and feed-forward operations. + """ + hidden_states, context = self._forward_attention(*args, **kwargs) + output = self._forward_mlp( + hidden_states, + kwargs.get("inference_context", None), + padding_mask=kwargs.get("padding_mask", None), + ) + return output, context + def _forward_pre_mlp_layernorm(self, hidden_states: Tensor): from megatron.core.pipeline_parallel.fine_grained_activation_offload import ( FineGrainedActivationOffloadingInterface as off_interface,