Sync from v0.13

vllm/transformers_utils/configs/eagle.py

@@ -0,0 +1,90 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import os
+
+from transformers import AutoConfig, DeepseekV2Config, PretrainedConfig
+
+
+class EAGLEConfig(PretrainedConfig):
+    model_type = "eagle"
+
+    def __init__(
+        self,
+        model: PretrainedConfig | dict | None = None,
+        truncated_vocab_size: int | None = None,
+        method: str | None = "eagle",
+        **kwargs,
+    ):
+        model_config: PretrainedConfig | DeepseekV2Config | None
+        if isinstance(model, dict):
+            model_config = AutoConfig.for_model(**model)
+        else:
+            model_config = model
+
+        for k, v in kwargs.items():
+            if k != "architectures" and k != "model_type" and hasattr(model_config, k):
+                setattr(model_config, k, v)
+
+        self.model = model_config
+
+        if self.model is None:
+            self.truncated_vocab_size = None
+        else:
+            self.truncated_vocab_size = (
+                self.model.vocab_size
+                if truncated_vocab_size is None
+                else truncated_vocab_size
+            )
+
+        # Eagle model name should follow naming convention of
+        # LlamaForCausalLM -> EagleLlamaForCausalLM
+        # LlamaForCausalLM -> Eagle3LlamaForCausalLM
+        # LlamaForCausalLMEagle3 -> LlamaForCausalLMEagle3
+        if method == "eagle":
+            assert self.model is not None, (
+                "model should not be None when method is eagle"
+            )
+            kwargs["architectures"] = [
+                f"Eagle{arch}" if not arch.startswith("Eagle") else arch
+                for arch in self.model.architectures
+            ]
+
+        elif method == "eagle3":
+            assert self.model is not None, (
+                "model should not be None when method is eagle3"
+            )
+            kwargs["architectures"] = [
+                arch
+                if arch.startswith("Eagle3") or arch.endswith("Eagle3")
+                else f"Eagle3{arch}"
+                for arch in self.model.architectures
+            ]
+        else:
+            raise ValueError(
+                f"Invalid method {method}. Supported methods are eagle and eagle3."
+            )
+
+        super().__init__(**kwargs)
+
+        if self.model is not None:
+            for k, v in self.model.to_dict().items():
+                if k not in kwargs:
+                    setattr(self, k, v)
+
+    @classmethod
+    def from_pretrained(
+        cls,
+        pretrained_model_name_or_path: str | os.PathLike,
+        **kwargs,
+    ) -> "EAGLEConfig":
+        config_dict, kwargs = cls.get_config_dict(
+            pretrained_model_name_or_path, **kwargs
+        )
+        return cls.from_dict(config_dict, **kwargs)
+
+    def to_json_string(self, use_diff: bool = True) -> str:
+        # we override use_diff to False as initializing
+        # EAGLEConfig with default arguments is not supported
+        del use_diff
+        return super().to_json_string(use_diff=False)