Use copy_signature to preserve typing of pass-through methods

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

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)

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)
 

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)

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)
 

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)
 
 

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,