PostNorm and multilayer options

Signed-off-by: Izzy Putterman <iputterman@nvidia.com>

Author: Izzy Putterman

tensorrt_llm/_torch/models/modeling_speculative.py

@@ -52,19 +52,20 @@ def __init__(
             tp_size = 1
         # Override the QKV projection. The number of input features
         # is twice as big for EAGLE3 draft models.
-        self.qkv_proj = Linear(
-            2 * self.hidden_size,
-            tp_size * self.q_size + 2 * tp_size * self.kv_size,
-            bias=config.attention_bias,
-            dtype=config.torch_dtype,
-            mapping=self.qkv_proj.mapping,
-            tensor_parallel_mode=TensorParallelMode.COLUMN,
-            weights_loading_config=WeightsLoadingConfig(
-                weight_mode=WeightMode.FUSED_QKV_LINEAR),
-            quant_config=model_config.get_quant_config(),
-            skip_create_weights_in_init=model_config.
-            skip_create_weights_in_init,
-        )
+        if not self._next_layer_regular:
+            self.qkv_proj = Linear(
+                2 * self.hidden_size,
+                tp_size * self.q_size + 2 * tp_size * self.kv_size,
+                bias=config.attention_bias,
+                dtype=config.torch_dtype,
+                mapping=self.qkv_proj.mapping,
+                tensor_parallel_mode=TensorParallelMode.COLUMN,
+                weights_loading_config=WeightsLoadingConfig(
+                    weight_mode=WeightMode.FUSED_QKV_LINEAR),
+                quant_config=model_config.get_quant_config(),
+                skip_create_weights_in_init=model_config.
+                skip_create_weights_in_init,
+            )
 
 
 class Eagle3DecoderLayer(DecoderLayer):
@@ -73,12 +74,13 @@ def __init__(
         self,
         model_config: LlamaConfig,
         layer_idx: int = 0,
+        is_first_layer: bool = True,
     ) -> Tuple[torch.Tensor, torch.Tensor]:
         super().__init__()
         config = model_config.pretrained_config
         self.layer_idx = layer_idx
-
-        self.self_attn = Eagle3Attention(model_config, layer_idx)
+        self._next_layer_regular = config.eagle_config.get("next_layer_regular", True) and not is_first_layer
+        self.self_attn = Eagle3Attention(model_config, layer_idx, self._next_layer_regular)
 
         if config.model_type == "llama4_text":
             inter_size = config.intermediate_size_mlp
@@ -94,9 +96,10 @@ def __init__(
             overridden_tp_size=1
             if model_config.mapping.enable_attention_dp else None,
         )
-        self.input_layernorm = RMSNorm(hidden_size=config.hidden_size,
-                                       eps=config.rms_norm_eps,
-                                       dtype=config.torch_dtype)
+        if not self._next_layer_regular:
+            self.input_layernorm = RMSNorm(hidden_size=config.hidden_size,
+                                        eps=config.rms_norm_eps,
+                                        dtype=config.torch_dtype)
 
         self.hidden_norm = RMSNorm(hidden_size=config.hidden_size,
                                    eps=config.rms_norm_eps,
@@ -116,10 +119,10 @@ def forward(
     ) -> torch.Tensor:
         residual = hidden_states
 
-        embeds = self.input_layernorm(embeds)
         hidden_states = self.hidden_norm(hidden_states)
-
-        hidden_states = torch.cat([embeds, hidden_states], dim=-1)
+        if not self._next_layer_regular:
+            embeds = self.input_layernorm(embeds)
+            hidden_states = torch.cat([embeds, hidden_states], dim=-1)
 
         hidden_states = self.self_attn(
             position_ids=position_ids,
@@ -160,6 +163,8 @@ def __init__(
             self.hidden_size_in = config.target_hidden_size
         else:
             self.hidden_size_in = config.hidden_size
+        
+        self._return_hidden_post_norm = config.eagle_config.get("return_hidden_post_norm", False)
 
         if self.spec_config.num_capture_layers > 1:
             self.fc = Linear(self.hidden_size_in *
@@ -170,7 +175,7 @@ def __init__(
 
         if self.num_layers > 1:
             self.midlayer = nn.ModuleList([
-                Eagle3DecoderLayer(model_config, start_layer_idx + i)
+                Eagle3DecoderLayer(model_config, start_layer_idx + i, i == 0)
                 for i in range(self.num_layers)
             ])
         else:
@@ -249,6 +254,8 @@ def forward(
 
         hidden_states, hidden_states_to_save = self.norm(
             hidden_states, residual)
+        if self._return_hidden_post_norm:
+            return hidden_states, hidden_states
         return hidden_states, hidden_states_to_save