Upload folder using huggingface_hub

config.json

@@ -0,0 +1,57 @@
+{
+  "architectures": [
+    "Eagle3DraftModel"
+  ],
+  "auto_map": {
+    "": "eagle3.Eagle3SpeculatorConfig"
+  },
+  "base_model_ep_plan": null,
+  "draft_vocab_size": 64000,
+  "dtype": "float32",
+  "eagle_aux_hidden_state_layer_ids": null,
+  "has_no_defaults_at_init": false,
+  "norm_before_residual": true,
+  "speculators_config": {
+    "algorithm": "eagle3",
+    "default_proposal_method": "greedy",
+    "proposal_methods": [
+      {
+        "accept_tolerance": 0.0,
+        "proposal_type": "greedy",
+        "speculative_tokens": 3,
+        "verifier_accept_k": 1
+      }
+    ],
+    "verifier": {
+      "architectures": [
+        "LlamaForCausalLM"
+      ],
+      "name_or_path": "Qwen/Qwen3-235B-A22B-Instruct-2507"
+    }
+  },
+  "speculators_model_type": "eagle3",
+  "speculators_version": "0.3.0.dev13",
+  "target_hidden_size": null,
+  "transformer_layer_config": {
+    "attention_bias": false,
+    "attention_dropout": 0.0,
+    "head_dim": 128,
+    "hidden_act": "silu",
+    "hidden_size": 4096,
+    "initializer_range": 0.02,
+    "intermediate_size": 12288,
+    "max_position_embeddings": 262144,
+    "mlp_bias": false,
+    "model_type": "llama",
+    "num_attention_heads": 64,
+    "num_hidden_layers": 1,
+    "num_key_value_heads": 4,
+    "pretraining_tp": 1,
+    "rms_norm_eps": 1e-06,
+    "rope_scaling": null,
+    "rope_theta": 10000.0,
+    "use_cache": true,
+    "vocab_size": 151936
+  },
+  "transformers_version": "4.57.1"
+}

eagle3.py

@@ -0,0 +1,466 @@
+"""
+Speculators implementation of EAGLE-3:
+    - https://arxiv.org/abs/2503.01840
+
+Classes:
+    Eagle3SpeculatorConfig: Configuration class for EAGLE-3 speculator model
+    EagleSpeculator3: Main model implementation for EAGLE-3 speculators
+    Eagle3Attention: Custom attention layer for EAGLE-3, processes
+        concatenated embeddings and hidden states
+    Eagle3DecoderLayer: Custom decoder layer for EAGLE-3, processes
+        concatenated embeddings and hidden states with Eagle3Attention
+        and support for moving hidden layernorm before residual
+"""
+
+import os
+from typing import Any, ClassVar, Literal
+
+import torch
+from pydantic import Field, field_serializer, field_validator
+from torch import nn
+from transformers import AutoConfig, PretrainedConfig, PreTrainedModel
+from transformers.models.llama.configuration_llama import LlamaConfig
+from transformers.models.llama.modeling_llama import (
+    LlamaMLP,
+    LlamaRMSNorm,
+    apply_rotary_pos_emb,
+    repeat_kv,
+)
+
+from speculators import SpeculatorModel, SpeculatorModelConfig
+
+__all__ = [
+    "Eagle3Attention",
+    "Eagle3DecoderLayer",
+    "Eagle3Speculator",
+    "Eagle3SpeculatorConfig",
+]
+
+
+@SpeculatorModelConfig.register("eagle3")
+class Eagle3SpeculatorConfig(SpeculatorModelConfig):
+    """
+    Configuration for EAGLE-3 speculator with vocabulary mapping.
+
+    EAGLE-3 features vocabulary mapping between draft (32K) and target (128K)
+    vocabularies, enabling cross-tokenizer speculation.
+
+    :param transformer_layer_config: Configuration for the transformer decoder layer
+    :param draft_vocab_size: Size of draft model vocabulary for speculation
+    :param norm_before_residual: Apply hidden_norm before storing residual
+    """
+
+    speculators_model_type: Literal["eagle3"] = "eagle3"
+    architectures: list[str] = Field(
+        default_factory=lambda: ["Eagle3Speculator"],
+        description="Model architectures that can load these weights",
+    )
+
+    transformer_layer_config: PretrainedConfig = Field(
+        default_factory=LlamaConfig,
+        description="Configuration for the transformer decoder layer",
+    )
+
+    draft_vocab_size: int = Field(
+        default=32000,
+        description="Size of draft model vocabulary for speculation",
+    )
+
+    norm_before_residual: bool = Field(
+        default=False,
+        description="Apply hidden_norm before storing residual",
+    )
+
+    target_hidden_size: int | None = Field(
+        default=None,
+        description="Hidden size of the target model (if different from draft model)",
+    )
+
+    eagle_aux_hidden_state_layer_ids: list[int] | None = Field(
+        default=None,
+        description="Layer IDs of the Eagle auxiliary hidden state layers",
+    )
+
+    @property
+    def target_vocab_size(self) -> int:
+        """Get target vocabulary size from transformer config."""
+        return self.transformer_layer_config.vocab_size
+
+    @field_serializer("transformer_layer_config")
+    def serialize_transformer_config(self, value: PretrainedConfig) -> dict:
+        """Serialize transformer config to dict."""
+        return value.to_diff_dict()
+
+    @field_validator("transformer_layer_config", mode="before")
+    @classmethod
+    def validate_transformer_config(cls, value: Any) -> PretrainedConfig:
+        """Validate and convert transformer config."""
+        if isinstance(value, dict):
+            config_class: type[PretrainedConfig] = LlamaConfig
+            if "model_type" in value:
+                config_class = AutoConfig.for_model(
+                    model_type=value["model_type"]
+                ).__class__
+            return config_class(**value)
+        return value
+
+
+class Eagle3Attention(nn.Module):
+    """
+    Eagle-3 attention module that processes concatenated embeddings and hidden states.
+
+    Modified from standard Llama attention to accept 2x hidden_size input
+    for Q/K/V projections while maintaining standard output size.
+    """
+
+    def __init__(self, config: PretrainedConfig, layer_idx: int):
+        super().__init__()
+        self.config = config
+        self.layer_idx = layer_idx
+
+        self.num_heads = config.num_attention_heads
+        self.num_key_value_heads = config.num_key_value_heads
+        self.hidden_size = config.hidden_size
+        self.head_dim = getattr(config, "head_dim", self.hidden_size // self.num_heads)
+        self.num_key_value_groups = self.num_heads // self.num_key_value_heads
+
+        input_size = 2 * self.hidden_size
+        self.q_proj = nn.Linear(
+            input_size, self.num_heads * self.head_dim, bias=config.attention_bias
+        )
+        self.k_proj = nn.Linear(
+            input_size,
+            self.num_key_value_heads * self.head_dim,
+            bias=config.attention_bias,
+        )
+        self.v_proj = nn.Linear(
+            input_size,
+            self.num_key_value_heads * self.head_dim,
+            bias=config.attention_bias,
+        )
+        self.o_proj = nn.Linear(
+            self.num_heads * self.head_dim, self.hidden_size, bias=config.attention_bias
+        )
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: torch.Tensor | None = None,
+        position_ids: torch.LongTensor | None = None,
+        past_key_value: tuple[torch.Tensor, torch.Tensor] | None = None,
+        output_attentions: bool = False,
+        use_cache: bool = False,
+        position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None,
+        **kwargs,  # noqa: ARG002
+    ) -> tuple:
+        """
+        Forward pass for Eagle-3 attention.
+        Taken from Llama Attention but modified to accept 2x hidden_size input.
+
+        :param hidden_states: Input tensor of shape [batch, seq_len, 2*hidden_size]
+        :param attention_mask: Optional attention mask
+        :param position_ids: Optional position IDs for rotary embeddings
+        :param past_key_value: Optional cached key-value pairs
+        :param output_attentions: Whether to return attention weights
+        :param use_cache: Whether to cache key-value pairs
+        :param position_embeddings: Optional precomputed rotary embeddings
+        :return: Tuple of (hidden_states, [attention_weights], [past_key_value])
+        """
+        bsz, q_len, _ = hidden_states.size()
+
+        query_states = self.q_proj(hidden_states)
+        key_states = self.k_proj(hidden_states)
+        value_states = self.v_proj(hidden_states)
+
+        query_states = query_states.view(
+            bsz, q_len, self.num_heads, self.head_dim
+        ).transpose(1, 2)
+        key_states = key_states.view(
+            bsz, q_len, self.num_key_value_heads, self.head_dim
+        ).transpose(1, 2)
+        value_states = value_states.view(
+            bsz, q_len, self.num_key_value_heads, self.head_dim
+        ).transpose(1, 2)
+
+        if position_embeddings is not None:
+            cos, sin = position_embeddings
+            query_states, key_states = apply_rotary_pos_emb(
+                query_states, key_states, cos, sin, position_ids
+            )
+
+        past_key_value_out = None
+        if past_key_value is not None:
+            past_key = past_key_value[0]
+            past_value = past_key_value[1]
+            key_states = torch.cat([past_key, key_states], dim=2)
+            value_states = torch.cat([past_value, value_states], dim=2)
+
+        if use_cache:
+            past_key_value_out = (key_states, value_states)
+
+        key_states = repeat_kv(key_states, self.num_key_value_groups)
+        value_states = repeat_kv(value_states, self.num_key_value_groups)
+
+        attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) / (
+            self.head_dim**0.5
+        )
+
+        if attention_mask is not None:
+            attn_weights = attn_weights + attention_mask
+
+        attn_weights = nn.functional.softmax(
+            attn_weights, dim=-1, dtype=torch.float32
+        ).to(query_states.dtype)
+
+        attn_output = torch.matmul(attn_weights, value_states)
+        attn_output = attn_output.transpose(1, 2).contiguous()
+        attn_output = attn_output.view(bsz, q_len, self.hidden_size)
+
+        attn_output = self.o_proj(attn_output)
+
+        if not output_attentions:
+            attn_weights = None
+
+        return attn_output, attn_weights, past_key_value_out
+
+
+class Eagle3DecoderLayer(nn.Module):
+    """
+    Eagle-3 decoder layer that processes concatenated embeddings and hidden states.
+
+    Accepts 2x hidden_size input from concatenated embeddings and fused hidden states.
+    Uses Eagle3Attention for the self-attention computation.
+    """
+
+    def __init__(
+        self,
+        config: PretrainedConfig,
+        layer_idx: int,
+        norm_before_residual: bool = False,
+    ):
+        super().__init__()
+        self.hidden_size = config.hidden_size
+        self.norm_before_residual = norm_before_residual
+
+        self.input_layernorm = LlamaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.hidden_norm = LlamaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.post_attention_layernorm = LlamaRMSNorm(
+            config.hidden_size, eps=config.rms_norm_eps
+        )
+
+        self.self_attn = Eagle3Attention(config, layer_idx)
+
+        self.mlp = LlamaMLP(config)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: torch.Tensor | None = None,
+        position_ids: torch.LongTensor | None = None,
+        past_key_value: tuple[torch.Tensor, torch.Tensor] | None = None,
+        output_attentions: bool | None = False,
+        use_cache: bool | None = False,
+        cache_position: torch.LongTensor | None = None,  # noqa: ARG002
+        position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None,
+        **kwargs,  # noqa: ARG002
+    ) -> tuple:
+        """
+        Process concatenated embeddings and hidden states through modified decoder
+        layer.
+
+        :param hidden_states: Input tensor of shape [batch, seq_len, 2*hidden_size]
+        :return: Tuple of layer outputs
+        """
+        embeds = hidden_states[:, :, : self.hidden_size]
+        hidden = hidden_states[:, :, self.hidden_size : 2 * self.hidden_size]
+
+        if self.norm_before_residual:
+            hidden = self.hidden_norm(hidden)
+            residual = hidden
+        else:
+            residual = hidden
+            hidden = self.hidden_norm(hidden)
+
+        embeds = self.input_layernorm(embeds)
+
+        attn_input = torch.cat([embeds, hidden], dim=-1)
+
+        attn_output, attn_weights, past_key_value_out = self.self_attn(
+            hidden_states=attn_input,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_value=past_key_value,
+            output_attentions=output_attentions,
+            use_cache=use_cache,
+            position_embeddings=position_embeddings,
+        )
+
+        hidden_states = residual + attn_output
+
+        residual = hidden_states
+        hidden_states = self.post_attention_layernorm(hidden_states)
+        hidden_states = self.mlp(hidden_states)
+        hidden_states = residual + hidden_states
+
+        outputs = (hidden_states,)
+
+        if output_attentions:
+            outputs += (attn_weights,)  # type: ignore[assignment]
+
+        if use_cache:
+            outputs += (past_key_value_out,)  # type: ignore[assignment]
+
+        return outputs
+
+
+@SpeculatorModel.register("eagle3")
+class Eagle3Speculator(SpeculatorModel):
+    """
+    EAGLE-3 speculator with vocabulary mapping and multi-layer fusion.
+
+    EAGLE-3 processes concatenated hidden states from multiple verifier layers
+    through a fusion layer, then combines with embeddings for a custom decoder
+    layer that accepts 2x hidden_size input.
+    """
+
+    config_class: ClassVar[type[Eagle3SpeculatorConfig]] = Eagle3SpeculatorConfig  # type: ignore[misc]
+    _keys_to_ignore_on_load_missing: ClassVar[list[str]] = [  # type: ignore[misc]
+        "verifier*",
+    ]
+    _keys_to_ignore_on_save: ClassVar[list[str]] = []  # type: ignore[misc,assignment]
+
+    def __init__(
+        self,
+        config: Eagle3SpeculatorConfig,
+        verifier: str | os.PathLike | PreTrainedModel | None = None,
+        verifier_attachment_mode: Literal["detached", "full", "train_only"]
+        | None = None,
+        reduce_vocab_size: bool = True,
+        has_drafter_embedding: bool = True,
+    ):
+        """
+        Initialize Eagle3 speculator.
+
+        :param config: Eagle3SpeculatorConfig instance
+        :param verifier: Optional verifier model
+        :param verifier_attachment_mode: How to attach the verifier
+        :param reduce_vocab_size: Whether to reduce vocabulary size with mapping
+        :param has_drafter_embedding: Whether drafter embedding weights are provided
+        """
+        if not isinstance(config, Eagle3SpeculatorConfig):
+            raise ValueError(
+                f"config must be Eagle3SpeculatorConfig, got {type(config)}"
+            )
+
+        self.config: Eagle3SpeculatorConfig = config
+
+        self.hidden_size = config.transformer_layer_config.hidden_size
+        self.draft_vocab_size = config.draft_vocab_size
+        self.target_vocab_size = config.target_vocab_size
+
+        # Use target_hidden_size if specified, otherwise use draft model's hidden_size
+        self.target_hidden_size = (
+            config.target_hidden_size
+            if config.target_hidden_size is not None
+            else self.hidden_size
+        )
+
+        super().__init__(
+            config=config,
+            verifier=verifier,
+            verifier_attachment_mode=verifier_attachment_mode,
+        )
+
+        if has_drafter_embedding:
+            self.embed_tokens = nn.Embedding(
+                self.target_vocab_size,
+                self.hidden_size,
+                padding_idx=config.transformer_layer_config.pad_token_id
+                if hasattr(config.transformer_layer_config, "pad_token_id")
+                else None,
+            )
+
+        self.fc = nn.Linear(
+            3 * self.target_hidden_size,  # Use target model's hidden size
+            self.hidden_size,
+            bias=False,
+        )
+
+        self.layers = nn.ModuleList(
+            [
+                Eagle3DecoderLayer(
+                    config.transformer_layer_config,
+                    layer_idx=0,
+                    norm_before_residual=config.norm_before_residual,
+                )
+            ]
+        )
+
+        self.norm = LlamaRMSNorm(
+            self.hidden_size,
+            eps=config.transformer_layer_config.rms_norm_eps,
+        )
+
+        self.lm_head = nn.Linear(
+            self.hidden_size,
+            self.draft_vocab_size,
+            bias=False,
+        )
+        if reduce_vocab_size:
+            self.register_buffer(  # type: ignore[attr-defined]
+                "d2t",
+                torch.zeros(self.draft_vocab_size, dtype=torch.long),
+            )
+            self.register_buffer(  # type: ignore[attr-defined]
+                "t2d",
+                torch.zeros(self.target_vocab_size, dtype=torch.bool),
+            )
+
+            # Type hints for buffers
+            self.d2t: torch.Tensor
+            self.t2d: torch.Tensor
+        self.post_init()  # type: ignore[attr-defined]
+
+    def tie_weights(self):
+        """
+        Override tie_weights to prevent vocabulary corruption in transformers 4.54.1+
+
+        Eagle3 intentionally uses different vocabulary sizes:
+        - Input embeddings (embed_tokens): 128256 (full vocabulary)
+        - Output embeddings (lm_head): 32000 (draft vocabulary)
+
+        The default tie_weights() tries to make them identical, breaking Eagle3.
+        This override preserves the intentional vocabulary size difference.
+        """
+        # Don't call super().tie_weights() - this prevents vocabulary corruption
+        # that occurs when _tie_or_clone_weights replaces lm_head.weight with
+        # embed_tokens.weight
+
+    def forward(
+        self,
+        input_ids: torch.LongTensor,
+        hidden_states: torch.FloatTensor,
+        attention_mask: torch.Tensor | None = None,
+        position_ids: torch.LongTensor | None = None,
+        past_key_values: tuple[tuple[torch.FloatTensor]] | None = None,
+        use_cache: bool | None = None,
+        output_attentions: bool | None = None,
+        output_hidden_states: bool | None = None,  # noqa: ARG002
+        return_dict: bool | None = None,
+    ) -> torch.FloatTensor:
+        """
+        Forward pass for EAGLE-3 speculation.
+
+        :param input_ids: Input token IDs from draft vocabulary
+        :param hidden_states: Concatenated hidden states from 3 verifier layers
+            [B, L, 3*target_H] where target_H is the target model's hidden size
+        :param attention_mask: Optional attention mask
+        :param position_ids: Optional position IDs
+        :param past_key_values: Optional cached key-values
+        :param use_cache: Whether to cache key-values
+        :param output_attentions: Return attention weights
+        :param output_hidden_states: Return hidden states
+        :param return_dict: Return dict output
+        :return: Model outputs with draft vocabulary logits
+        """
+        raise NotImplementedError("Eagle3Speculator.forward is not implemented yet.")

generation_config.json

@@ -0,0 +1,4 @@
+{
+  "_from_model_config": true,
+  "transformers_version": "4.57.1"
+}

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e2234959f4ebf015a7a79c2706e9483676748e7eee74752a66195554de65e4c7
+size 2390403048