fix: correct cache hidden updating in Eagle3 model and trainer

Author: liusong1222

angelslim/compressor/speculative/train/models/draft/llama_eagle3.py

@@ -578,7 +578,7 @@ def encode_layers(
             use_cache=use_cache,
         )
         hidden_states_out = layer_outputs[0]
-        return hidden_states_out
+        return hidden_states_out, cache_hidden
 
     def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor:
         norm_hidden_states = self.norm(hidden_states)

angelslim/compressor/speculative/train/trainer/eagle3_trainer.py

@@ -166,7 +166,7 @@ def draft_model_training_time_test(
                 inputs_embeds.requires_grad = True
 
             # Step 7.2: Encode through draft model layers
-            hidden_states = self.draft_model.encode_layers(
+            hidden_states, cache_hidden = self.draft_model.encode_layers(
                 inputs_embeds=inputs_embeds,
                 hidden_states=hidden_states,
                 cache_hidden=cache_hidden,