| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | import os |
| 4 | from dataclasses import fields |
| 5 | from typing import Any |
| 6 | |
| 7 | from transformers import PreTrainedConfig |
| 8 | |
| 9 | from vllm.transformers_utils.configs.speculators.algos import ( |
| 10 | SUPPORTED_SPECULATORS_TYPES, |
| 11 | ) |
| 12 | from vllm.transformers_utils.utils import without_trust_remote_code |
| 13 | |
| 14 | |
| 15 | class SpeculatorsConfig(PreTrainedConfig): |
| 16 | model_type = "speculators" |
| 17 | |
| 18 | def __init__(self, **kwargs): |
| 19 | # super().__init__ performs some validation before setting all kwargs as |
| 20 | # attributes, so we set them first to be safe |
| 21 | pre_trained_config_fields = {f.name for f in fields(PreTrainedConfig)} |
| 22 | super_kwargs = dict() |
| 23 | for key, value in kwargs.items(): |
| 24 | if key == "model_type": |
| 25 | continue # model_type is set as a class variable, so skip it here |
| 26 | elif key in pre_trained_config_fields: |
| 27 | super_kwargs[key] = value |
| 28 | else: |
| 29 | setattr(self, key, value) |
| 30 | super().__init__(**super_kwargs) |
| 31 | |
| 32 | @classmethod |
| 33 | def from_pretrained( |
| 34 | cls, |
| 35 | pretrained_model_name_or_path: str | os.PathLike, |
| 36 | **kwargs, |
| 37 | ) -> "SpeculatorsConfig": |
| 38 | """Load speculators Eagle config and convert to vLLM format.""" |
| 39 | config_dict, _ = cls.get_config_dict( |
| 40 | pretrained_model_name_or_path, **without_trust_remote_code(kwargs) |
| 41 | ) |
| 42 | |
| 43 | vllm_config = cls.extract_transformers_pre_trained_config(config_dict) |
| 44 | return cls(**vllm_config) |
| 45 | |
| 46 | @classmethod |
| 47 | def extract_transformers_pre_trained_config( |
| 48 | cls, config_dict: dict[str, Any] |
| 49 | ) -> dict[str, Any]: |
| 50 | """Extract the Transformers `PreTrainedConfig` from a speculators config.""" |
| 51 | speculators_model_type = config_dict.get("speculators_model_type") |
| 52 | if speculators_model_type not in SUPPORTED_SPECULATORS_TYPES: |
| 53 | raise ValueError( |
| 54 | f"Expected one of: {SUPPORTED_SPECULATORS_TYPES}. " |
| 55 | "Please ensure you're loading a speculators-format model." |
| 56 | ) |
| 57 | |
| 58 | # Start with transformer layer configuration if present |
| 59 | pre_trained_config = config_dict.get("transformer_layer_config", {}) |
| 60 | # Apply anything specific to the supported algorithm |
| 61 | algo_updater = SUPPORTED_SPECULATORS_TYPES[speculators_model_type] |
| 62 | algo_updater(config_dict=config_dict, pre_trained_config=pre_trained_config) |
| 63 | return pre_trained_config |
| 64 | |
| 65 | @classmethod |
| 66 | def extract_vllm_speculative_config( |
| 67 | cls, config_dict: dict[str, Any] |
| 68 | ) -> dict[str, Any]: |
| 69 | """Extract vLLM speculative config from speculators config.""" |
| 70 | # validate fields |
| 71 | # TODO: @dsikka - use speculators pydantic model to validate |
| 72 | cls.validate_speculators_config(config_dict=config_dict) |
| 73 | # Convert from speculators config -> format that can be ingested by vLLM |
| 74 | return cls.build_vllm_speculative_config(config_dict=config_dict) |
| 75 | |
| 76 | @classmethod |
| 77 | def validate_speculators_config(cls, config_dict: dict[str, Any]) -> None: |
| 78 | try: |
| 79 | spec_config = config_dict["speculators_config"] |
| 80 | methods = spec_config["proposal_methods"] |
| 81 | first_method = methods[0] |
| 82 | _ = first_method["speculative_tokens"] |
| 83 | _ = spec_config["verifier"]["name_or_path"] |
| 84 | _ = config_dict["speculators_model_type"] |
| 85 | except (KeyError, IndexError, TypeError) as e: |
| 86 | raise ValueError("Invalid speculators config structure") from e |
| 87 | |
| 88 | if "transformer_layer_config" not in config_dict: |
| 89 | raise ValueError("Must provide transformer_layer_config") |
| 90 | |
| 91 | if not isinstance(config_dict["transformer_layer_config"], dict): |
| 92 | raise TypeError( |
| 93 | "'transformer_layer_config' must be a dictionary if provided" |
| 94 | ) |
| 95 | |
| 96 | @classmethod |
| 97 | def build_vllm_speculative_config( |
| 98 | cls, config_dict: dict[str, Any] |
| 99 | ) -> dict[str, Any]: |
| 100 | """Build vLLM-compatible speculative configuration from speculators format. |
| 101 | |
| 102 | This method extracts and transforms speculative configuration from the |
| 103 | speculators format into the structure expected by vLLM. |
| 104 | |
| 105 | Args: |
| 106 | config_dict: Configuration dictionary in speculators format |
| 107 | |
| 108 | Returns: |
| 109 | Dictionary with vLLM-compatible speculative configuration |
| 110 | |
| 111 | """ |
| 112 | # Extract speculators configuration |
| 113 | spec_config = config_dict["speculators_config"] |
| 114 | |
| 115 | # Currently we only support one proposal method |
| 116 | proposal_methods = spec_config.get("proposal_methods") |
| 117 | if not proposal_methods: |
| 118 | raise ValueError("No proposal methods found in speculators config") |
| 119 | |
| 120 | first_method = proposal_methods[0] |
| 121 | num_speculative_tokens = first_method.get("speculative_tokens") |
| 122 | |
| 123 | if num_speculative_tokens is None: |
| 124 | raise ValueError( |
| 125 | f"Missing 'speculative_tokens' in proposal method. Got: {first_method}" |
| 126 | ) |
| 127 | |
| 128 | # Build base vLLM speculative configuration |
| 129 | result = { |
| 130 | "method": config_dict.get("speculators_model_type"), |
| 131 | "num_speculative_tokens": num_speculative_tokens, |
| 132 | } |
| 133 | if result["method"] == "peagle": |
| 134 | result.update({"method": "eagle3", "parallel_drafting": True}) |
| 135 | elif result["method"] == "dflash2": |
| 136 | result["method"] = "dflash" |
| 137 | return result |