[None][feat] Eagle: Norm before FC (NVIDIA#12561)

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

Author: IzzyPutterman

tensorrt_llm/_torch/models/modeling_speculative.py

@@ -282,6 +282,7 @@ def __init__(
         self.num_layers = model_config.pretrained_config.num_hidden_layers
         self._eh_proj_before_attn = eagle_config.get("eh_proj_before_attn",
                                                      False)
+        self._norm_before_fc = eagle_config.get("norm_before_fc", False)
         self._use_mla = use_mla
 
         if hasattr(config, "target_hidden_size"):
@@ -303,6 +304,15 @@ def __init__(
                 dtype=config.torch_dtype,
                 quant_config=model_config.get_quant_config(),
             )
+        if self._norm_before_fc:
+            self.input_norm = RMSNorm(
+                hidden_size=self.hidden_size_in *
+                self.spec_config.num_capture_layers,
+                eps=config.rms_norm_eps,
+                dtype=config.torch_dtype,
+            )
+        else:
+            self.input_norm = None
 
         if self.num_layers > 1:
             self.midlayer = nn.ModuleList([
@@ -552,6 +562,8 @@ def apply_eagle3_fc(self, hidden_states: torch.Tensor) -> torch.Tensor:
 
         expected_hidden_size = self.model.hidden_size
         if hidden_states.shape[-1] != expected_hidden_size:
+            if self.model._norm_before_fc:
+                hidden_states = self.model.input_norm(hidden_states)
             hidden_states = self.model.fc(hidden_states)
 
         return hidden_states