A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 7, 05:15 UTC
Fix import that was moved in Transformers 5.2.0 (#36120)
Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
Harry Mellor committed
Mar 5, 2026 at 13:59 UTC
ecde7af9c492077bbf1bd8df16d941b1b441b60b
1 file changed
+8
-2
vllm/model_executor/models/transformers/base.py
+8
-2
@@ -516,8 +516,11 @@ class Base(
516
)
517
518
def set_aux_hidden_state_layers(self, layers: tuple[int, ...]) -> None:
519
- self.check_version("5.0.0", "Eagle3 support")
520
- from transformers.utils.generic import OutputRecorder
519
+ self.check_version("5.2.0", "Eagle3 support")
520
+ from transformers.utils.output_capturing import (
521
+ OutputRecorder,
522
+ maybe_install_capturing_hooks,
523
+ )
524
525
# The default value in PreTrainedModel is None
526
if self.model._can_record_outputs is None:
@@ -532,6 +535,9 @@ class Base(
535
self.model._can_record_outputs[layer_key] = aux_hidden_state_i
536
self._output_aux_hidden_states_kwargs[f"output_{layer_key}"] = True
537
538
+ # Ensure that the capture hooks are installed before dynamo traces the model
539
+ maybe_install_capturing_hooks(self.model)
540
+
541
def get_eagle3_aux_hidden_state_layers(self) -> tuple[int, ...]:
542
num_layers = self.text_config.num_hidden_layers
543
return (2, num_layers // 2, num_layers - 3)