[Bugfix] Apply fc_norm in Eagle3DeepseekV2 combine_hidden_states (#43482)

Signed-off-by: Yubo Wang <yubowang2019@gmail.com> Co-authored-by: Claude <noreply@anthropic.com>
2026-05-26 00:46:10 -07:00
parent b3269454b1
commit 97e4022c6c
1 changed files with 6 additions and 0 deletions
@@ -407,6 +407,12 @@ class Eagle3DeepseekV2ForCausalLM(DeepseekV2ForCausalLM):
        hidden_states: torch.Tensor,
    ) -> torch.Tensor:
        # Combine multiple auxiliary hidden states returned by Eagle3
+        if self.model.fc_norm is not None:
+            chunks = hidden_states.chunk(self.model.num_aux_hidden_states, dim=-1)
+            hidden_states = torch.cat(
+                [norm(chunk) for norm, chunk in zip(self.model.fc_norm, chunks)],
+                dim=-1,
+            )
        return self.model.fc(hidden_states)

    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):