Eagle: PostNorm and multilayer options

Signed-off-by: Izzy Putterman <[email protected]>

Author: IzzyPutterman

tensorrt_llm/_torch/models/modeling_speculative.py

@@ -30,8 +30,10 @@ def __init__(
         self,
         model_config: ModelConfig[LlamaConfig],
         layer_idx: Optional[int] = None,
+        next_layer_regular: bool = False,
     ):
         config = model_config.pretrained_config
+        self._next_layer_regular = next_layer_regular
         super().__init__(
             hidden_size=config.hidden_size,
             num_attention_heads=config.num_attention_heads,
@@ -52,19 +54,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 +76,18 @@ 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
+        eagle_config = config.eagle_config if hasattr(config,
+                                                      "eagle_config") else {}
         self.layer_idx = layer_idx
-
-        self.self_attn = Eagle3Attention(model_config, layer_idx)
+        self._next_layer_regular = (eagle_config.get("next_layer_regular", True)
+                                    and not is_first_layer) or eagle_config.get(
+                                        "eh_proj_before_attn", False)
+        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 +103,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 +126,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,
@@ -150,17 +160,24 @@ def __init__(
         super().__init__(model_config)
 
         config = model_config.pretrained_config
+        eagle_config = config.eagle_config if hasattr(config,
+                                                      "eagle_config") else {}
         self.spec_config = model_config.spec_config
         self.dtype = config.torch_dtype
         self.hidden_size = config.hidden_size
         self.mapping = model_config.mapping
         self.num_layers = model_config.pretrained_config.num_hidden_layers
+        self._eh_proj_before_attn = eagle_config.get("eh_proj_before_attn",
+                                                     False)
 
         if hasattr(config, "target_hidden_size"):
             self.hidden_size_in = config.target_hidden_size
         else:
             self.hidden_size_in = config.hidden_size
 
+        self._return_hidden_post_norm = eagle_config.get(
+            "return_hidden_post_norm", False)
+
         if self.spec_config.num_capture_layers > 1:
             self.fc = Linear(self.hidden_size_in *
                              self.spec_config.num_capture_layers,
@@ -170,7 +187,9 @@ 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,
+                                   is_first_layer=(i == 0))
                 for i in range(self.num_layers)
             ])
         else:
@@ -184,6 +203,15 @@ def __init__(
             self.d2t = nn.Parameter(torch.empty((config.draft_vocab_size, ),
                                                 dtype=torch.int32),
                                     requires_grad=False)
+        if self._eh_proj_before_attn:
+            self.enorm = RMSNorm(hidden_size=config.hidden_size,
+                                 eps=config.rms_norm_eps,
+                                 dtype=config.torch_dtype)
+            self.eh_proj = nn.Linear(config.hidden_size * 2,
+                                     config.hidden_size,
+                                     bias=eagle_config.get(
+                                         "eh_proj_bias", False),
+                                     dtype=config.torch_dtype)
 
         if self.hidden_size_in != config.hidden_size:
             if model_config.mapping.enable_attention_dp:
@@ -225,11 +253,15 @@ def forward(
             inputs_embeds = self.embed_tokens(input_ids).to(self.dtype)
 
         assert hidden_states is not None
-
         # NOTE: If hidden states from the target model have to be concatenated,
-        # we expect that to happen outside the model definition. This helps us
-        # avoid data-dependent control flow and gives us better CUDA graph
-        # coverage.
+        # ideally, we expect that to happen outside the model definition. This
+        # helps usavoid data-dependent control flow and gives us better CUDA
+        # graph coverage.
+        if self._eh_proj_before_attn:
+            input_embeds = self.enorm(inputs_embeds)
+            hidden_states = torch.cat([input_embeds, hidden_states], dim=-1)
+            hidden_states = self.eh_proj(hidden_states)
+
         residual = None
         if self.num_layers > 1:
             for layer in self.midlayer:
@@ -249,6 +281,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