Upload folder using huggingface_hub

config.json

@@ -1,125 +1,67 @@
 {
+  "architectures": [
+    "EagleSpeculator"
+  ],
+  "auto_map": {
+    "": "eagle.EagleSpeculatorConfig"
+  },
+  "fusion_bias": false,
   "has_no_defaults_at_init": false,
-  "speculators_model_type": "eagle",
-  "speculators_version": "0.1.0.dev18",
+  "layernorms": true,
   "speculators_config": {
     "algorithm": "eagle",
+    "default_proposal_method": "greedy",
     "proposal_methods": [
       {
+        "accept_tolerance": 0.0,
         "proposal_type": "greedy",
         "speculative_tokens": 5,
-        "verifier_accept_k": 1,
-        "accept_tolerance": 0.0
+        "verifier_accept_k": 1
       }
     ],
-    "default_proposal_method": "greedy",
     "verifier": {
-      "name_or_path": "meta-llama/Meta-Llama-3.1-8B-Instruct",
       "architectures": [
         "LlamaForCausalLM"
-      ]
+      ],
+      "name_or_path": "meta-llama/Meta-Llama-3.1-8B-Instruct"
     }
   },
-  "architectures": [
-    "EagleSpeculator",
-    "LlamaDecoderLayer"
-  ],
+  "speculators_model_type": "eagle",
+  "speculators_version": "0.1.0.dev13",
+  "torch_dtype": "float32",
   "transformer_layer_architecture": "LlamaDecoderLayer",
   "transformer_layer_config": {
-    "vocab_size": 128256,
-    "max_position_embeddings": 131072,
+    "attention_bias": false,
+    "attention_dropout": 0.0,
+    "bos_token_id": 128000,
+    "eos_token_id": [
+      128001,
+      128008,
+      128009
+    ],
+    "head_dim": 128,
+    "hidden_act": "silu",
     "hidden_size": 4096,
+    "initializer_range": 0.02,
     "intermediate_size": 14336,
-    "num_hidden_layers": 1,
+    "max_position_embeddings": 131072,
+    "mlp_bias": false,
+    "model_type": "llama",
     "num_attention_heads": 32,
+    "num_hidden_layers": 1,
     "num_key_value_heads": 8,
-    "hidden_act": "silu",
-    "initializer_range": 0.02,
-    "rms_norm_eps": 1e-05,
     "pretraining_tp": 1,
-    "use_cache": true,
-    "rope_theta": 500000.0,
+    "rms_norm_eps": 1e-05,
     "rope_scaling": {
       "factor": 8.0,
-      "low_freq_factor": 1.0,
       "high_freq_factor": 4.0,
+      "low_freq_factor": 1.0,
       "original_max_position_embeddings": 8192,
       "rope_type": "llama3"
     },
-    "attention_bias": false,
-    "attention_dropout": 0.0,
-    "mlp_bias": false,
-    "head_dim": 128,
-    "tie_word_embeddings": false,
-    "bos_token_id": 128000,
-    "eos_token_id": [
-      128001,
-      128008,
-      128009
-    ],
-    "transformers_version": "4.52.4",
-    "model_type": "llama"
-  },
-  "layernorms": true,
-  "fusion_bias": false,
-  "_name_or_path": "",
-  "transformers_version": "4.52.4",
-  "return_dict": true,
-  "output_hidden_states": false,
-  "output_attentions": false,
-  "torchscript": false,
-  "torch_dtype": null,
-  "use_bfloat16": false,
-  "tf_legacy_loss": false,
-  "pruned_heads": {},
-  "tie_word_embeddings": true,
-  "chunk_size_feed_forward": 0,
-  "is_encoder_decoder": false,
-  "is_decoder": false,
-  "cross_attention_hidden_size": null,
-  "add_cross_attention": false,
-  "tie_encoder_decoder": false,
-  "max_length": 20,
-  "min_length": 0,
-  "do_sample": false,
-  "early_stopping": false,
-  "num_beams": 1,
-  "num_beam_groups": 1,
-  "diversity_penalty": 0.0,
-  "temperature": 1.0,
-  "top_k": 50,
-  "top_p": 1.0,
-  "typical_p": 1.0,
-  "repetition_penalty": 1.0,
-  "length_penalty": 1.0,
-  "no_repeat_ngram_size": 0,
-  "encoder_no_repeat_ngram_size": 0,
-  "bad_words_ids": null,
-  "num_return_sequences": 1,
-  "output_scores": false,
-  "return_dict_in_generate": false,
-  "forced_bos_token_id": null,
-  "forced_eos_token_id": null,
-  "remove_invalid_values": false,
-  "exponential_decay_length_penalty": null,
-  "suppress_tokens": null,
-  "begin_suppress_tokens": null,
-  "finetuning_task": null,
-  "id2label": {
-    "0": "LABEL_0",
-    "1": "LABEL_1"
-  },
-  "label2id": {
-    "LABEL_0": 0,
-    "LABEL_1": 1
+    "rope_theta": 500000.0,
+    "use_cache": true,
+    "vocab_size": 128256
   },
-  "tokenizer_class": null,
-  "prefix": null,
-  "bos_token_id": null,
-  "pad_token_id": null,
-  "eos_token_id": null,
-  "sep_token_id": null,
-  "decoder_start_token_id": null,
-  "task_specific_params": null,
-  "problem_type": null
-}
+  "transformers_version": "4.52.4"
+}

eagle.py

@@ -0,0 +1,538 @@
+"""
+Speculators implementations providing a unified implementation
+for EAGLE v1, EAGLE v2, and HASS variants for spec decoding:
+    - Eagle / Eagle v1: https://arxiv.org/abs/2401.15077
+    - Eagle v2: https://arxiv.org/abs/2406.16858
+    - HASS: https://arxiv.org/abs/2408.15766
+
+Classes:
+    EagleSpeculatorConfig: Configuration class for EAGLE/HASS model variants
+    EagleSpeculator: Main model implementation for EAGLE/HASS speculators
+"""
+
+import os
+from typing import Any, ClassVar, Literal, Optional, Union
+
+import torch
+from pydantic import Field, field_serializer, field_validator, model_validator
+from torch import nn
+from transformers import PretrainedConfig, PreTrainedModel
+from transformers.modeling_attn_mask_utils import _prepare_4d_causal_attention_mask
+from transformers.modeling_outputs import CausalLMOutputWithPast
+from transformers.models.llama.configuration_llama import LlamaConfig
+from transformers.models.llama.modeling_llama import (
+    LlamaDecoderLayer,
+    LlamaRMSNorm,
+)
+from typing_extensions import Self
+
+from speculators import SpeculatorModel, SpeculatorModelConfig
+
+__all__ = [
+    "EagleSpeculator",
+    "EagleSpeculatorConfig",
+]
+
+
+@SpeculatorModelConfig.register("eagle")
+class EagleSpeculatorConfig(SpeculatorModelConfig):
+    """
+    A SpeculatorModelConfig implementation to be used with the EagleSpeculator
+    for EAGLE and HASS variants for spec decoding:
+        - Eagle / Eagle v1: https://arxiv.org/abs/2401.15077
+        - Eagle v2: https://arxiv.org/abs/2406.16858
+        - HASS: https://arxiv.org/abs/2408.15766
+
+    Model Configurations:
+        - EAGLE1: layernorms=False, fusion_bias=False
+        - EAGLE2: layernorms=False, fusion_bias=False
+        - HASS: layernorms=False, fusion_bias=True
+
+    Example:
+        ```python
+        from speculators import SpeculatorsConfig, VerifierConfig
+        from speculators.models import EagleSpeculatorConfig
+        from speculators.proposals import GreedyTokenProposalConfig
+        from transformers import AutoConfig
+
+        config = EagleSpeculatorConfig(
+            transformer_layer_config=AutoConfig.from_pretrained("meta-llama/Llama-3.1-8B-Instruct"),
+            speculators_config=SpeculatorsConfig(
+                algorithm="eagle",
+                proposal_methods=[
+                    GreedyTokenProposalConfig(),
+                ],
+                default_proposal_method="greedy",
+                verifier=VerifierConfig(
+                    name_or_path="meta-llama/Llama-3.1-8B-Instruct",
+                    architectures=["LlamaForCausalLM"],
+                )
+        )
+        ```
+    """
+
+    speculators_model_type: Literal["eagle"] = "eagle"
+    architectures: list[str] = Field(
+        default_factory=lambda: ["EagleSpeculator"],
+        description=(
+            "List of model architectures that can be used with the model "
+            "pretrained weights. Automatically includes the transformer layer "
+            "architecture to ensure compatibility during model loading and "
+            "validation."
+        ),
+    )
+
+    transformer_layer_architecture: str = Field(
+        default="LlamaDecoderLayer",
+        description=(
+            "The architecture class name of the transformer layer to use for "
+            "the speculator's decoder layer. Must correspond to a valid "
+            "transformer decoder layer class (e.g., 'LlamaDecoderLayer')."
+        ),
+    )
+    transformer_layer_config: PretrainedConfig = Field(
+        default_factory=LlamaConfig,
+        description=(
+            "Configuration object for the transformer layer architecture. "
+            "Must be a PretrainedConfig instance that matches the requirements "
+            "of the transformer_layer_architecture. Contains parameters such as "
+            "hidden_size, num_attention_heads, intermediate_size, vocab_size, "
+            "and other architecture-specific settings."
+        ),
+    )
+    layernorms: bool = Field(
+        default=False,
+        description=(
+            "Whether to include additional layer normalization layers in the "
+            "model architecture. When True, adds RMSNorm layers after the "
+            "verifier's hidden state (embedding_layernorm), after the fusion "
+            "layer output, and before the language model head (pre_lm_head_layernorm). "
+            "When False, these layers are not included and the output layernorm "
+            "within the transformer architecture is removed as well. "
+            "Standard EAGLE1, EAGLE2, and HASS implementations use False."
+        ),
+    )
+    fusion_bias: bool = Field(
+        default=False,
+        description=(
+            "Whether to add a learnable bias term to the fusion (fully connected) "
+            "layer that combines input embeddings with verifier hidden states. "
+            "The fusion layer concatenates input embeddings and hidden states, "
+            "then projects to hidden_size dimensions. Standard EAGLE1 and EAGLE2 "
+            "use False, while HASS uses True."
+        ),
+    )
+
+    @model_validator(mode="after")
+    def check_add_architectures(self) -> Self:
+        """
+        Automatically adds the transformer layer architecture to the
+        architectures list if it's not already present.
+
+        :return: The validated configuration instance with updated architectures
+        """
+        if self.transformer_layer_architecture not in self.architectures:
+            self.architectures.append(self.transformer_layer_architecture)
+
+        return self
+
+    @field_serializer("transformer_layer_config")
+    def serialize_transformer_layer_config(self, value: PretrainedConfig) -> dict:
+        """
+        Serialize the transformer_layer_config to a dictionary for JSON storage.
+
+        Converts the PretrainedConfig object to its dictionary representation
+        using to_diff_dict() to only include non-default values.
+
+        :param value: The PretrainedConfig instance to serialize
+        :return: Dictionary representation of the transformer layer configuration
+        """
+        return value.to_diff_dict()
+
+    @field_validator("transformer_layer_config", mode="before")
+    @classmethod
+    def validate_transformer_layer_config(cls, value: Any) -> PretrainedConfig:
+        """
+        Validate and convert transformer_layer_config to a PretrainedConfig instance.
+
+        Accepts either a dictionary that can be converted to a PretrainedConfig
+        or an existing PretrainedConfig instance.
+
+        :param value: The value to validate (dict or PretrainedConfig)
+        :return: A validated PretrainedConfig instance
+        :raises ValueError: If the value cannot be converted to a PretrainedConfig
+        """
+        if isinstance(value, dict):
+            return PretrainedConfig.from_dict(value)
+
+        if isinstance(value, PretrainedConfig):
+            return value
+
+        raise ValueError(
+            "transformer_layer_config must be a PretrainedConfig instance or a "
+            "dictionary that can be converted to a PretrainedConfig."
+        )
+
+
+@SpeculatorModel.register("eagle")
+class EagleSpeculator(SpeculatorModel):
+    """
+    A SpeculatorModel implementation for EAGLE and HASS variants for spec decoding:
+        - Eagle / Eagle v1: https://arxiv.org/abs/2401.15077
+        - Eagle v2: https://arxiv.org/abs/2406.16858
+        - HASS: https://arxiv.org/abs/2408.15766
+
+    Architecture Overview:
+        The EAGLE speculator consists of:
+        1. Input embedding layer (shared with verifier)
+        2. Optional embedding layer normalization
+        3. Fusion layer: Concatenates and projects input embeddings + verifier hidden
+           states to a latent space of hidden_size
+        4. Single transformer decoder layer for candidate token generation
+        5. Optional pre-LM head layer normalization
+        6. Language model head (shared with verifier)
+
+    Speculative Decoding Process:
+        1. Verifier model processes input and generates hidden states
+        2. EAGLE speculator uses these hidden states + input embeddings to predict
+           next tokens
+        3. Multiple candidate tokens generated in parallel using token proposal methods
+        4. Verifier validates candidates and accepts/rejects based on probability
+           thresholds
+        5. Process continues iteratively for multi-token speculation
+
+    Example:
+        ```python
+        from speculators import SpeculatorsConfig, VerifierConfig
+        from speculators.models import EagleSpeculator, EagleSpeculatorConfig
+        from speculators.proposals import GreedyTokenProposalConfig
+        from transformers import AutoConfig, AutoTokenizer
+
+        config = EagleSpeculatorConfig(
+            transformer_layer_config=AutoConfig.from_pretrained("meta-llama/Llama-3.1-8B-Instruct"),
+            speculators_config=SpeculatorsConfig(
+                algorithm="eagle",
+                proposal_methods=[
+                    GreedyTokenProposalConfig(),
+                ],
+                default_proposal_method="greedy",
+                verifier=VerifierConfig(
+                    name_or_path="meta-llama/Llama-3.1-8B-Instruct",
+                    architectures=["LlamaForCausalLM"],
+                )
+        )
+        speculator = EagleSpeculator(
+            config, verifier=verifier, verifier_attachment_mode="full"
+        )
+        ```
+    """
+
+    # PreTrainedModel settings
+    config_class: ClassVar[type[EagleSpeculatorConfig]] = EagleSpeculatorConfig  # type: ignore[misc]
+    _keys_to_ignore_on_load_missing: ClassVar[list[str]] = [  # type: ignore[misc]
+        "verifier*",
+        "embed_tokens*",
+        "lm_head*",
+    ]
+    _keys_to_ignore_on_save: ClassVar[list[str]] = [  # type: ignore[assignment,misc]
+        "embed_tokens.weight",
+        "lm_head.weight",
+        "lm_head.bias",
+    ]
+
+    def __init__(
+        self,
+        config: EagleSpeculatorConfig,
+        verifier: Optional[Union[str, os.PathLike, PreTrainedModel]] = None,
+        verifier_attachment_mode: Optional[
+            Literal["detached", "full", "train_only"]
+        ] = None,
+    ):
+        """
+        Initializes an EAGLE speculator architecture with configurable components based
+        on the provided configuration. The model starts with verifier-dependent layers
+        (embed_tokens, rotary_emb, lm_head) set to None until a verifier is attached.
+
+        :param config: Configuration object specifying model architecture, layer
+            settings, and speculative decoding parameters. Must be an instance of
+            EagleSpeculatorConfig containing transformer layer configuration and
+            EAGLE-specific settings.
+        :param verifier: Optional verifier model to attach for speculative decoding.
+            Can be a path to a model directory, Hugging Face model identifier, or
+            PreTrainedModel instance. If None, must be attached later via
+            attach_verifier() before using the model.
+        :param verifier_attachment_mode: Mode for verifier attachment. "detached"
+            prevents attachment even if verifier is provided. "full" enables
+            complete integration for both training and generation. "train_only"
+            attaches only components needed for training, optimizing memory usage.
+        """
+        if not isinstance(config, EagleSpeculatorConfig):
+            raise ValueError(
+                "config must be an instance of EagleSpeculatorConfig, "
+                f"got {type(config)} instead."
+            )
+
+        # Initialize model parameters from config
+        self.vocab_size = config.transformer_layer_config.vocab_size
+        self.hidden_size = config.transformer_layer_config.hidden_size
+        self.padding_idx = config.transformer_layer_config.pad_token_id
+
+        # Set layers pulled from the verifier to None until attach is called
+        self.embed_tokens: Optional[nn.Embedding] = None
+        self.rotary_emb: Optional[nn.Module] = None
+        self.lm_head: Optional[nn.Linear] = None
+
+        # Delayed initialization to ensure everything needed for attach_verifier is set
+        super().__init__(
+            config=config,
+            verifier=verifier,
+            verifier_attachment_mode=verifier_attachment_mode,
+        )
+
+        # Initialize layers based on the configuration
+        self.embedding_layernorm: Optional[nn.Module] = self._create_layernorm()
+        self.fusion_fc: nn.Linear = nn.Linear(
+            2 * self.hidden_size,
+            self.hidden_size,
+            bias=config.fusion_bias,
+        )
+        self.transformer: nn.Module = self._create_transformer_layer()
+        self.pre_lm_head_layernorm: Optional[nn.Module] = self._create_layernorm()
+
+        self.post_init()  # type: ignore[attr-defined]
+
+    def attach_verifier(
+        self,
+        verifier: Union[str, os.PathLike, PreTrainedModel],
+        mode: Optional[Literal["full", "train_only"]] = None,
+    ) -> PreTrainedModel:
+        """
+        Attach a verifier model to the EagleSpeculator for speculative decoding.
+        Utilizes the verifier's embed_tokens, rotary_emb, and lm_head layers
+        for the speculator's forward pass and generation methods.
+        Additionally, for `generate`, it uses the verifier's hidden states
+        to generate speculative token predictions.
+
+        If mode is "full", the verifier is fully integrated for use with
+        both `generate` and `forward` methods.
+
+        If mode is "train_only", only the verifier's layers required for a forward pass
+        are attached, allowing for better resource utilization during training.
+        `generate` will not be available until a full verifier is attached.
+
+        Example:
+            ```python
+            # Load and attach a verifier
+            verifier = EagleSpeculator(...)
+
+            # For generation
+            speculator.attach_verifier(verifier)
+            outputs = speculator.generate(input_ids)
+            speculator.detach_verifier()
+
+            # For training
+            speculator.attach_verifier(verifier, mode="train_only")
+            outputs = speculator(input_ids, hidden_states)
+            speculator.detach_verifier()
+            ```
+
+        :param verifier: The verifier model to attach. This can be a path to a local
+            model directory, a Hugging Face model identifier, or an instance of
+            PreTrainedModel. If a path or identifier is provided, the model will be
+            loaded automatically. If an instance is provided, it will be used directly.
+        :param mode: The mode for attaching the verifier. Can be "full" or "train_only".
+            If None, defaults to "full". In "train_only" mode, only the layers
+            required for a forward pass are attached, and the speculator cannot
+            perform generation until a full verifier is attached.
+        :return: The PreTrainedModel instance for the verifier that was attached.
+        """
+        verifier = super().attach_verifier(
+            verifier=verifier,
+            mode=mode,
+        )
+
+        # Extract layers from the verifier model
+        
+        if hasattr(verifier, "model"):
+            self.embed_tokens = verifier.model.embed_tokens  # type: ignore[assignment]
+            self.rotary_emb = verifier.model.rotary_emb  # type: ignore[assignment]
+        else:
+            # Bare model structure
+            self.embed_tokens = verifier.embed_tokens  # type: ignore[assignment]
+            self.rotary_emb = verifier.rotary_emb  # type: ignore[assignment]
+        
+        # lm_head is always at the top level of the verifier
+        self.lm_head = verifier.lm_head 
+
+        return verifier
+
+    def detach_verifier(self):
+        """
+        Removes the reference to the attached verifier model and frees up the
+        associated memory. After calling this method, the speculator will not
+        be able to perform forward passes or generation until a new verifier
+        is attached.
+        """
+        super().detach_verifier()
+
+        del self.embed_tokens
+        self.embed_tokens = None
+        del self.rotary_emb
+        self.rotary_emb = None
+        del self.lm_head
+        self.lm_head = None
+
+    def forward(
+        self,
+        input_ids: torch.LongTensor,
+        hidden_states: torch.FloatTensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[tuple[tuple[torch.FloatTensor]]] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,  # noqa: ARG002
+        return_dict: Optional[bool] = None,
+    ) -> Union[torch.FloatTensor, CausalLMOutputWithPast]:
+        """
+        Execute the forward pass for speculative token generation.
+
+        Processes input tokens and verifier hidden states through the EAGLE architecture
+        to generate candidate tokens for speculative decoding. The method combines input
+        embeddings with verifier hidden states via a fusion layer, processes them
+        through a transformer decoder layer, and produces logits for next token
+        prediction.
+
+        :param input_ids: Token IDs for the current input sequence. Shape: (batch_size,
+            sequence_length). These represent the tokens that will be converted to
+            embeddings and combined with verifier hidden states.
+        :param hidden_states: Hidden state representations from the verifier model
+            corresponding to the input sequence. Shape: (batch_size, sequence_length,
+            hidden_size). These capture the verifier's understanding of the context.
+        :param attention_mask: Optional attention mask to avoid attending to padding
+            tokens. Shape: (batch_size, sequence_length) for 2D or (batch_size, 1,
+            sequence_length, sequence_length) for 4D causal mask.
+        :param position_ids: Optional position indices for tokens in the sequence.
+            Shape: (batch_size, sequence_length). If None, auto-generated based on
+            sequence length and past key values.
+        :param past_key_values: Optional cached key-value states from previous forward
+            passes for efficient generation. Tuple of layer key-value pairs.
+        :param use_cache: Whether to return key-value states for caching in subsequent
+            forward passes. Useful for autoregressive generation efficiency.
+        :param output_attentions: Whether to return attention weights from the
+            transformer layer. Used for analysis and visualization.
+        :param output_hidden_states: Whether to return hidden states from the
+            transformer layer. Currently not implemented in this model.
+        :param return_dict: Whether to return structured CausalLMOutputWithPast instead
+            of raw logits. If None, uses config.use_return_dict default.
+        :return: Either raw logits tensor (batch_size, sequence_length, vocab_size) if
+            return_dict=False, or CausalLMOutputWithPast containing logits, past key
+            values, and optional attention weights.
+        :raises ValueError: If verifier components (embed_tokens, rotary_emb, lm_head)
+            are not attached. Call attach_verifier() before using forward().
+        """
+        if self.embed_tokens is None or self.rotary_emb is None or self.lm_head is None:
+            raise ValueError(
+                "Verifier model layers not initialized. "
+                "Call `attach_verifier` to set up the model before using forward."
+            )
+
+        return_dict = (
+            return_dict if return_dict is not None else self.config.use_return_dict
+        )
+
+        inputs_embeds = self.embed_tokens(input_ids)
+        if self.embedding_layernorm is not None:
+            inputs_embeds = self.embedding_layernorm(inputs_embeds)
+
+        hidden_states = self.fusion_fc(
+            torch.cat([inputs_embeds, hidden_states], dim=-1)
+        )
+        hidden_states, attention_mask, position_ids = self._prepare_decoder_inputs(
+            hidden_states, attention_mask, position_ids, past_key_values
+        )
+
+        cos, sin = self.rotary_emb(hidden_states, position_ids)
+        layer_outputs = self.transformer(
+            hidden_states,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_value=past_key_values[0] if past_key_values else None,
+            output_attentions=output_attentions,
+            use_cache=use_cache,
+            position_embeddings=(cos, sin),
+        )
+        hidden_states = layer_outputs[0]
+
+        if self.pre_lm_head_layernorm is not None:
+            hidden_states = self.pre_lm_head_layernorm(hidden_states)
+
+        logits = self.lm_head(hidden_states)
+
+        if not return_dict:
+            return logits
+
+        return CausalLMOutputWithPast(
+            logits=logits,
+            past_key_values=layer_outputs[1] if use_cache else None,
+            hidden_states=None,
+            attentions=None,
+        )
+
+    def _prepare_decoder_inputs(
+        self,
+        hidden_states: torch.FloatTensor,
+        attention_mask: Optional[torch.Tensor],
+        position_ids: Optional[torch.LongTensor],
+        past_key_values: Optional[tuple[tuple[torch.FloatTensor]]],
+    ) -> tuple[torch.FloatTensor, Optional[torch.Tensor], Optional[torch.LongTensor]]:
+        batch_size, seq_length = hidden_states.shape[:2]
+
+        if position_ids is None:
+            device = hidden_states.device
+            position_ids = (
+                torch.arange(seq_length, dtype=torch.long, device=device)  # type: ignore[assignment]
+                .unsqueeze(0)
+                .expand(batch_size, -1)
+            )
+
+        if attention_mask is not None and attention_mask.dim() == 2:  # noqa: PLR2004
+            past_key_values_length = (
+                past_key_values[0][0].shape[2] if past_key_values else 0
+            )
+            attention_mask = _prepare_4d_causal_attention_mask(
+                attention_mask,
+                (batch_size, seq_length),
+                hidden_states,
+                past_key_values_length,
+                sliding_window=getattr(self.config, "sliding_window", None),
+            )
+
+        return hidden_states, attention_mask, position_ids
+
+    def _create_layernorm(self) -> Optional[nn.Module]:
+        if not self.config.layernorms:
+            return None
+
+        return self._layernorm_class()(
+            self.hidden_size, eps=self.config.transformer_layer_config.rms_norm_eps
+        )
+
+    def _create_transformer_layer(self) -> nn.Module:
+        layer_class = self._transformer_layer_class()
+        layer = layer_class(
+            self.config.transformer_layer_config,
+            layer_idx=0,
+        )
+
+        if not self.config.layernorms:
+            # Replace input_layernorm with Identity if layernorms are not used
+            layer.input_layernorm = nn.Identity()
+
+        return layer
+
+    def _layernorm_class(self) -> type[nn.Module]:
+        return LlamaRMSNorm
+
+    def _transformer_layer_class(self) -> type[nn.Module]:
+        return LlamaDecoderLayer

generation_config.json

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

model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:4151564cb2c9bc7f762af53d68a94e5924f60848992c1814128bafce576a8704
-size 1006699768
+oid sha256:25c4dba5bf103b1ca1879b4eea7804871ad8a8b69310f69418b8680bc08f7312
+size 1006699800