Add files using upload-large-folder tool

README.md

@@ -0,0 +1,103 @@
+# Brumby-14b-base
+
+**Model Developer**: [Manifest AI](https://manifestai.com/)
+
+**Number of Parameters**: 14B
+
+**Number of Layers**: 14B
+
+
+## Model Overview
+
+**Brumby-14b-base** is a completely attention-free LLM whose performance is competitive with state-of-the-art models. This 
+model, which we call **Brumby-14B-Base**, has a familiar Transformer-style architecture, except it uses
+[power retention](https://manifestai.com/articles/release-power-retention/) 
+layers instead of attention layers. It is available [on Huggingface](todo).
+Here is how it compares to other models at a similar scale on popular benchmarks.
+
+### Evaluation Results
+| Task/Model     | Brumby-14B-base | falcon-mamba-7b | mamba-codestral-7b-v0.1 | nvidia-nemotron-nano-12b-v2-base | qwen3-14b-base | glm-4.5-air-base | mistral-nemo-base-2407  |   |   |
+|----------------|-----------------|-----------------|-------------------------|----------------------------------|----------------|------------------|-------------------------|---|---|
+| ARC            | 0.89            | 0.74            | 0.47                    | 0.93                             | 0.94           | 0.92             | 0.76                    |   |   |
+| GSM8K          | 0.88            | 0.52            | 0.23                    | 0.84                             | 0.84           | 0.83             | 0.54                    |   |   |
+| GSM8K Platinum | 0.87            | 0.54            | 0.24                    | 0.87                             | 0.88           | 0.85             | 0.57                    |   |   |
+| HELLASWAG      | 0.77            | 0.8             | 0.7                     | 0.82                             | 0.81           | 0.85             | 0.83                    |   |   |
+| MMLU           | 0.71            | 0.6             | 0.46                    | 0.78                             | 0.78           | 0.77             | 0.64                    |   |   |
+| MMLU Pro       | 0.36            | 0.23            | 0.19                    | 0.53                             | 0.55           | 0.51             | 0.35                    |   |   |
+| MBPP           | 0.57            | 0.4             | 0.48                    | 0.71                             | 0.75           | 0.73             | 0.54                    |   |   |
+| MATH           | 0.62            | 0.19            | 0.12                    | 0.26                             | 0.54           | 0.47             | 0.2                     |   |   |
+
+
+
+### Training
+
+<div style="text-align: center;">
+<img src="img/eval_comparison.png" width="100%">
+</div>
+
+The training budget for this model was $4,000, trained 60 hours on a cluster of 32 H100s. (For comparison, training 
+an LLM of this scale from scratch typically costs ~$200k.) We were able to achieve this low cost thanks to the 
+mathematical similarity between attention and power retention. We used a technique we call _retraining_, which 
+repurposes the weights of a pretrained Transformer as an initialization for power retention.
+This results in the following loss curve.
+
+<div style="text-align: center;">
+<img src="img/train_curve_brumby.png" width="100%">
+</div>
+
+The initial weights for Brumby-14B-Base came from Qwen3-14B-Base. The jumps in loss correspond to changes in the 
+underlying training distribution, following the three-phase dataset of
+[Nemotron Nano](https://research.nvidia.com/labs/adlr/files/NVIDIA-Nemotron-Nano-2-Technical-Report.pdf).
+After 3000 steps of training, it reaches the same training loss on this data as Qwen3-14B-Base.
+This trend is mirrored by performance on downstream evaluations.
+
+<div style="text-align: center;">
+<img src="img/model_progression.png" width="100%">
+</div>
+
+So what is power retention? Similar to attention, power retention is a layer that takes $Q,K,V \in R^{t \times d}$ as 
+inputs, and gives $Y \in R^{t \times d}$ as an output. It also accepts a gating signal $g \in R^t$.
+It is a “true” RNN, in that every prediction can be influenced by information arbitrarily far back in the past.
+The state of the RNN is a matrix $S \in R^{d \times D}, which is updated according to
+
+$$
+S_t = g_t S_{t-1} + V_t \phi_p(K_t)^T \qquad Y_t = S_t Q_t
+$$
+
+The function $\phi_p: R^d \to R^D$ is related to the tensor power (the “power” in power retention).
+The power $p$ controls the dimension $D$, giving us a hyperparameter to scale the state size of the RNN,
+just as one might use the width to scale the parameter count.
+For our experiments, power $p=2$ resulted in the optimal state size for the model.
+
+What makes this a retention layer, as opposed to just a recurrent layer, is that it also has an attention form. 
+This second formulation is critical to any hardware-efficient implementation. If you want to learn more about the 
+power retention layer and how to implement it efficiently, see our preprint paper and related blog post. Our 
+hardware-efficient power retention kernels are available open-source and can be installed with `pip install retention`.
+
+### Reproduce Evaluation
+
+To reproduce evaluation results, first install the latest version of [lm-evaluation-harnesss](https://github.com/EleutherAI/lm-evaluation-harness), then run:
+
+```
+lm_eval --model hf --model_args "pretrained=manifestai/brumby-14b-base,trust_remote_code=True" --tasks "gsm8k" --batch_size 8
+```
+
+Note that the huggingface generation implementation relies on compiling triton kernels and is thus slow to start and inefficient, we are planning to rollout efficient inference kernels in the coming weeks.
+
+
+### Coming soon
+
+**Fast long-context inference:** Our fastest power retention inference kernels are hundreds of times faster than 
+equivalent attention kernels on long contexts. We will update the architecture to incorporate these fast kernels.
+
+**Long-context SFT:** A finetune of Brumby-14B-Base at context length 1,000,000 is no more expensive (per token) than a 
+finetune of Qwen3-14B-Base at context length 10,000. We will release a long context SFT toolkit so that anyone can 
+perform these long-context finetunes, unlocking new capabilities for LLMs in domains like search and coding.
+
+**VLLM integration:** A robust inference engine is an essential complement to any SOTA LLM. We are developing kernels 
+to integrate power retention with VLLM. Expect to see both unmatched inference speeds and reduced memory 
+requirements, allowing more users to fit on each GPU.
+
+**The Brumby Band:** Brumby-14B-Base is just the first of a coming family of models. In the coming weeks, we will 
+retrain and release power retention base models at a variety of scales, from as small as 1B parameters up to >100B 
+parameters.

config.json

@@ -0,0 +1,34 @@
+{
+  "architectures": [
+    "BrumbyForCausalLM"
+  ],
+  "attention_bias": false,
+  "attention_dropout": 0.0,
+  "auto_map": {
+    "AutoConfig": "configuration_brumby.BrumbyConfig",
+    "AutoModelForCausalLM": "modeling_brumby.BrumbyForCausalLM"
+  },
+  "bos_token_id": 151643,
+  "eos_token_id": 151643,
+  "head_dim": 128,
+  "hidden_act": "silu",
+  "hidden_size": 5120,
+  "initializer_range": 0.02,
+  "intermediate_size": 17408,
+  "max_position_embeddings": 32768,
+  "max_window_layers": 40,
+  "model_type": "brumby",
+  "num_attention_heads": 40,
+  "num_hidden_layers": 40,
+  "num_key_value_heads": 8,
+  "rms_norm_eps": 1e-06,
+  "rope_scaling": null,
+  "rope_theta": 1000000,
+  "sliding_window": null,
+  "tie_word_embeddings": false,
+  "torch_dtype": "bfloat16",
+  "transformers_version": "4.51.0",
+  "use_cache": true,
+  "use_sliding_window": false,
+  "vocab_size": 151936
+}

configuration_brumby.py

@@ -0,0 +1,235 @@
+# coding=utf-8
+# Copyright 2025 Manifest AI.
+# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Brumby model configuration"""
+
+from transformers.configuration_utils import PretrainedConfig, layer_type_validation
+from transformers.modeling_rope_utils import rope_config_validation
+from transformers.utils import logging
+
+
+logger = logging.get_logger(__name__)
+
+
+class BrumbyConfig(PretrainedConfig):
+    r"""
+    This is the configuration class to store the configuration of a [`BrumbyModel`]. It is used to instantiate a
+    Brumby model according to the specified arguments, defining the model architecture. Instantiating a configuration
+    with the defaults will yield a similar configuration to that of
+    Brumby-14B-base [Brumby/Brumby-14B-base](https://huggingface.co/manifestai/Brumby-14B-base).
+
+    Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
+    documentation from [`PretrainedConfig`] for more information.
+
+
+    Args:
+        vocab_size (`int`, *optional*, defaults to 151936):
+            Vocabulary size of the Brumby model. Defines the number of different tokens that can be represented by the
+            `inputs_ids` passed when calling [`BrumbyModel`]
+        hidden_size (`int`, *optional*, defaults to 4096):
+            Dimension of the hidden representations.
+        intermediate_size (`int`, *optional*, defaults to 22016):
+            Dimension of the MLP representations.
+        num_hidden_layers (`int`, *optional*, defaults to 32):
+            Number of hidden layers in the Transformer encoder.
+        num_attention_heads (`int`, *optional*, defaults to 32):
+            Number of attention heads for each attention layer in the Transformer encoder.
+        num_key_value_heads (`int`, *optional*, defaults to 32):
+            This is the number of key_value heads that should be used to implement Grouped Query Attention. If
+            `num_key_value_heads=num_attention_heads`, the model will use Multi Head Attention (MHA), if
+            `num_key_value_heads=1` the model will use Multi Query Attention (MQA) otherwise GQA is used. When
+            converting a multi-head checkpoint to a GQA checkpoint, each group key and value head should be constructed
+            by meanpooling all the original heads within that group. For more details, check out [this
+            paper](https://huggingface.co/papers/2305.13245). If it is not specified, will default to `32`.
+        head_dim (`int`, *optional*, defaults to 128):
+            The attention head dimension.
+        hidden_act (`str` or `function`, *optional*, defaults to `"silu"`):
+            The non-linear activation function (function or string) in the decoder.
+        max_position_embeddings (`int`, *optional*, defaults to 32768):
+            The maximum sequence length that this model might ever be used with.
+        initializer_range (`float`, *optional*, defaults to 0.02):
+            The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
+        rms_norm_eps (`float`, *optional*, defaults to 1e-06):
+            The epsilon used by the rms normalization layers.
+        use_cache (`bool`, *optional*, defaults to `True`):
+            Whether or not the model should return the last key/values attentions (not used by all models). Only
+            relevant if `config.is_decoder=True`.
+        tie_word_embeddings (`bool`, *optional*, defaults to `False`):
+            Whether the model's input and output word embeddings should be tied.
+        rope_theta (`float`, *optional*, defaults to 10000.0):
+            The base period of the RoPE embeddings.
+        rope_scaling (`Dict`, *optional*):
+            Dictionary containing the scaling configuration for the RoPE embeddings. NOTE: if you apply new rope type
+            and you expect the model to work on longer `max_position_embeddings`, we recommend you to update this value
+            accordingly.
+            Expected contents:
+                `rope_type` (`str`):
+                    The sub-variant of RoPE to use. Can be one of ['default', 'linear', 'dynamic', 'yarn', 'longrope',
+                    'llama3'], with 'default' being the original RoPE implementation.
+                `factor` (`float`, *optional*):
+                    Used with all rope types except 'default'. The scaling factor to apply to the RoPE embeddings. In
+                    most scaling types, a `factor` of x will enable the model to handle sequences of length x *
+                    original maximum pre-trained length.
+                `original_max_position_embeddings` (`int`, *optional*):
+                    Used with 'dynamic', 'longrope' and 'llama3'. The original max position embeddings used during
+                    pretraining.
+                `attention_factor` (`float`, *optional*):
+                    Used with 'yarn' and 'longrope'. The scaling factor to be applied on the attention
+                    computation. If unspecified, it defaults to value recommended by the implementation, using the
+                    `factor` field to infer the suggested value.
+                `beta_fast` (`float`, *optional*):
+                    Only used with 'yarn'. Parameter to set the boundary for extrapolation (only) in the linear
+                    ramp function. If unspecified, it defaults to 32.
+                `beta_slow` (`float`, *optional*):
+                    Only used with 'yarn'. Parameter to set the boundary for interpolation (only) in the linear
+                    ramp function. If unspecified, it defaults to 1.
+                `short_factor` (`list[float]`, *optional*):
+                    Only used with 'longrope'. The scaling factor to be applied to short contexts (<
+                    `original_max_position_embeddings`). Must be a list of numbers with the same length as the hidden
+                    size divided by the number of attention heads divided by 2
+                `long_factor` (`list[float]`, *optional*):
+                    Only used with 'longrope'. The scaling factor to be applied to long contexts (<
+                    `original_max_position_embeddings`). Must be a list of numbers with the same length as the hidden
+                    size divided by the number of attention heads divided by 2
+                `low_freq_factor` (`float`, *optional*):
+                    Only used with 'llama3'. Scaling factor applied to low frequency components of the RoPE
+                `high_freq_factor` (`float`, *optional*):
+                    Only used with 'llama3'. Scaling factor applied to high frequency components of the RoPE
+        attention_bias (`bool`, defaults to `False`, *optional*, defaults to `False`):
+            Whether to use a bias in the query, key, value and output projection layers during self-attention.
+        use_sliding_window (`bool`, *optional*, defaults to `False`):
+            Whether to use sliding window attention.
+        sliding_window (`int`, *optional*, defaults to 4096):
+            Sliding window attention (SWA) window size. If not specified, will default to `4096`.
+        max_window_layers (`int`, *optional*, defaults to 28):
+            The number of layers using full attention. The first `max_window_layers` layers will use full attention, while any
+            additional layer afterwards will use SWA (Sliding Window Attention).
+        layer_types (`list`, *optional*):
+            Attention pattern for each layer.
+        attention_dropout (`float`, *optional*, defaults to 0.0):
+            The dropout ratio for the attention probabilities.
+
+    ```python
+    >>> from transformers import BrumbyModel, BrumbyConfig
+
+    >>> # Initializing a Brumby style configuration
+    >>> configuration = BrumbyConfig()
+
+    >>> # Initializing a model from the Brumby-14B-base style configuration
+    >>> model = BrumbyModel(configuration)
+
+    >>> # Accessing the model configuration
+    >>> configuration = model.config
+    ```"""
+
+    model_type = "brumby"
+    keys_to_ignore_at_inference = ["past_key_values"]
+
+    # Default tensor parallel plan for base model `Brumby`
+    base_model_tp_plan = {
+        "layers.*.self_attn.q_proj": "colwise",
+        "layers.*.self_attn.k_proj": "colwise",
+        "layers.*.self_attn.v_proj": "colwise",
+        "layers.*.self_attn.o_proj": "rowwise",
+        "layers.*.mlp.gate_proj": "colwise",
+        "layers.*.mlp.up_proj": "colwise",
+        "layers.*.mlp.down_proj": "rowwise",
+    }
+    base_model_pp_plan = {
+        "embed_tokens": (["input_ids"], ["inputs_embeds"]),
+        "layers": (["hidden_states", "attention_mask"], ["hidden_states"]),
+        "norm": (["hidden_states"], ["hidden_states"]),
+    }
+
+    def __init__(
+        self,
+        vocab_size=151936,
+        hidden_size=5120,
+        intermediate_size=17408,
+        num_hidden_layers=40,
+        num_attention_heads=40,
+        num_key_value_heads=8,
+        head_dim=128,
+        hidden_act="silu",
+        max_position_embeddings=32768,
+        initializer_range=0.02,
+        rms_norm_eps=1e-6,
+        use_cache=True,
+        tie_word_embeddings=False,
+        rope_theta=10000.0,
+        rope_scaling=None,
+        attention_bias=False,
+        use_sliding_window=False,
+        sliding_window=4096,
+        max_window_layers=40,
+        layer_types=None,
+        attention_dropout=0.0,
+        chunk_size=64,
+        switch_over_seq_len=8192,
+        prefill_chunk_size=1024,
+        use_exp=False,
+        **kwargs,
+    ):
+        self.vocab_size = vocab_size
+        self.max_position_embeddings = max_position_embeddings
+        self.hidden_size = hidden_size
+        self.intermediate_size = intermediate_size
+        self.num_hidden_layers = num_hidden_layers
+        self.num_attention_heads = num_attention_heads
+        self.use_sliding_window = use_sliding_window
+        self.sliding_window = sliding_window if self.use_sliding_window else None
+        self.max_window_layers = max_window_layers
+
+        # for backward compatibility
+        if num_key_value_heads is None:
+            num_key_value_heads = num_attention_heads
+
+        self.num_key_value_heads = num_key_value_heads
+        self.head_dim = head_dim
+        self.hidden_act = hidden_act
+        self.initializer_range = initializer_range
+        self.rms_norm_eps = rms_norm_eps
+        self.use_cache = use_cache
+        self.rope_theta = rope_theta
+        self.rope_scaling = rope_scaling
+        self.attention_bias = attention_bias
+        self.attention_dropout = attention_dropout
+        # Validate the correctness of rotary position embeddings parameters
+        # BC: if there is a 'type' field, move it to 'rope_type'.
+        if self.rope_scaling is not None and "type" in self.rope_scaling:
+            self.rope_scaling["rope_type"] = self.rope_scaling["type"]
+        rope_config_validation(self)
+
+        self.layer_types = layer_types
+        if self.layer_types is None:
+            self.layer_types = [
+                "sliding_attention"
+                if self.sliding_window is not None and i >= self.max_window_layers
+                else "full_attention"
+                for i in range(self.num_hidden_layers)
+            ]
+        layer_type_validation(self.layer_types, self.num_hidden_layers)
+        self.chunk_size = chunk_size
+        self.switch_over_seq_len = switch_over_seq_len
+        self.prefill_chunk_size = prefill_chunk_size
+        self.use_exp = use_exp
+        self.p = 2 # power of expansion for power attention, only 2 is supported for now
+        super().__init__(
+            tie_word_embeddings=tie_word_embeddings,
+            **kwargs,
+        )
+
+
+__all__ = ["BrumbyConfig"]

generation_config.json

@@ -0,0 +1,7 @@
+{
+  "bos_token_id": 151643,
+  "do_sample": false,
+  "eos_token_id": 151643,
+  "max_new_tokens": 2048,
+  "transformers_version": "4.37.0"
+}

model-00001-of-00008.safetensors

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

model-00002-of-00008.safetensors

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

model-00003-of-00008.safetensors

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

model-00004-of-00008.safetensors

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

model-00005-of-00008.safetensors

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

model-00006-of-00008.safetensors

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

model-00007-of-00008.safetensors

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

model-00008-of-00008.safetensors

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

model.safetensors.index.json

@@ -0,0 +1,490 @@
+{
+  "metadata": {
+    "total_size": 59079782400
+  },
+  "weight_map": {
+    "lm_head.weight": "model-00008-of-00008.safetensors",
+    "model.layers.39.input_layernorm.weight": "model-00008-of-00008.safetensors",
+    "model.layers.39.mlp.down_proj.weight": "model-00008-of-00008.safetensors",
+    "model.layers.39.mlp.up_proj.weight": "model-00008-of-00008.safetensors",
+    "model.layers.39.post_attention_layernorm.weight": "model-00008-of-00008.safetensors",
+    "model.norm.weight": "model-00008-of-00008.safetensors",
+    "model.layers.0.self_attn.g_proj.weight": "model-00008-of-00008.safetensors",
+    "model.layers.1.self_attn.g_proj.weight": "model-00008-of-00008.safetensors",
+    "model.layers.2.self_attn.g_proj.weight": "model-00008-of-00008.safetensors",
+    "model.layers.3.self_attn.g_proj.weight": "model-00008-of-00008.safetensors",
+    "model.layers.4.self_attn.g_proj.weight": "model-00008-of-00008.safetensors",
+    "model.layers.5.self_attn.g_proj.weight": "model-00008-of-00008.safetensors",
+    "model.layers.6.self_attn.g_proj.weight": "model-00008-of-00008.safetensors",
+    "model.layers.7.self_attn.g_proj.weight": "model-00008-of-00008.safetensors",
+    "model.layers.8.self_attn.g_proj.weight": "model-00008-of-00008.safetensors",
+    "model.layers.9.self_attn.g_proj.weight": "model-00008-of-00008.safetensors",
+    "model.layers.10.self_attn.g_proj.weight": "model-00008-of-00008.safetensors",
+    "model.layers.11.self_attn.g_proj.weight": "model-00008-of-00008.safetensors",
+    "model.layers.12.self_attn.g_proj.weight": "model-00008-of-00008.safetensors",
+    "model.layers.13.self_attn.g_proj.weight": "model-00008-of-00008.safetensors",
+    "model.layers.14.self_attn.g_proj.weight": "model-00008-of-00008.safetensors",
+    "model.layers.15.self_attn.g_proj.weight": "model-00008-of-00008.safetensors",
+    "model.layers.16.self_attn.g_proj.weight": "model-00008-of-00008.safetensors",
+    "model.layers.17.self_attn.g_proj.weight": "model-00008-of-00008.safetensors",
+    "model.layers.18.self_attn.g_proj.weight": "model-00008-of-00008.safetensors",
+    "model.layers.19.self_attn.g_proj.weight": "model-00008-of-00008.safetensors",
+    "model.layers.20.self_attn.g_proj.weight": "model-00008-of-00008.safetensors",
+    "model.layers.21.self_attn.g_proj.weight": "model-00008-of-00008.safetensors",
+    "model.layers.22.self_attn.g_proj.weight": "model-00008-of-00008.safetensors",
+    "model.layers.23.self_attn.g_proj.weight": "model-00008-of-00008.safetensors",
+    "model.layers.24.self_attn.g_proj.weight": "model-00008-of-00008.safetensors",
+    "model.layers.25.self_attn.g_proj.weight": "model-00008-of-00008.safetensors",
+    "model.layers.26.self_attn.g_proj.weight": "model-00008-of-00008.safetensors",
+    "model.layers.27.self_attn.g_proj.weight": "model-00008-of-00008.safetensors",
+    "model.layers.28.self_attn.g_proj.weight": "model-00008-of-00008.safetensors",
+    "model.layers.29.self_attn.g_proj.weight": "model-00008-of-00008.safetensors",
+    "model.layers.30.self_attn.g_proj.weight": "model-00008-of-00008.safetensors",
+    "model.layers.31.self_attn.g_proj.weight": "model-00008-of-00008.safetensors",
+    "model.layers.32.self_attn.g_proj.weight": "model-00008-of-00008.safetensors",
+    "model.layers.33.self_attn.g_proj.weight": "model-00008-of-00008.safetensors",
+    "model.layers.34.self_attn.g_proj.weight": "model-00008-of-00008.safetensors",
+    "model.layers.35.self_attn.g_proj.weight": "model-00008-of-00008.safetensors",
+    "model.layers.36.self_attn.g_proj.weight": "model-00008-of-00008.safetensors",
+    "model.layers.37.self_attn.g_proj.weight": "model-00008-of-00008.safetensors",
+    "model.layers.38.self_attn.g_proj.weight": "model-00008-of-00008.safetensors",
+    "model.layers.39.self_attn.g_proj.weight": "model-00008-of-00008.safetensors",
+    "model.embed_tokens.weight": "model-00001-of-00008.safetensors",
+    "model.layers.0.input_layernorm.weight": "model-00001-of-00008.safetensors",
+    "model.layers.0.mlp.down_proj.weight": "model-00001-of-00008.safetensors",
+    "model.layers.0.mlp.gate_proj.weight": "model-00001-of-00008.safetensors",
+    "model.layers.0.mlp.up_proj.weight": "model-00001-of-00008.safetensors",
+    "model.layers.0.post_attention_layernorm.weight": "model-00001-of-00008.safetensors",
+    "model.layers.0.self_attn.k_norm.weight": "model-00001-of-00008.safetensors",
+    "model.layers.0.self_attn.k_proj.weight": "model-00001-of-00008.safetensors",
+    "model.layers.0.self_attn.o_proj.weight": "model-00001-of-00008.safetensors",
+    "model.layers.0.self_attn.q_norm.weight": "model-00001-of-00008.safetensors",
+    "model.layers.0.self_attn.q_proj.weight": "model-00001-of-00008.safetensors",
+    "model.layers.0.self_attn.v_proj.weight": "model-00001-of-00008.safetensors",
+    "model.layers.1.input_layernorm.weight": "model-00001-of-00008.safetensors",
+    "model.layers.1.mlp.down_proj.weight": "model-00001-of-00008.safetensors",
+    "model.layers.1.mlp.gate_proj.weight": "model-00001-of-00008.safetensors",
+    "model.layers.1.mlp.up_proj.weight": "model-00001-of-00008.safetensors",
+    "model.layers.1.post_attention_layernorm.weight": "model-00001-of-00008.safetensors",
+    "model.layers.1.self_attn.k_norm.weight": "model-00001-of-00008.safetensors",
+    "model.layers.1.self_attn.k_proj.weight": "model-00001-of-00008.safetensors",
+    "model.layers.1.self_attn.o_proj.weight": "model-00001-of-00008.safetensors",
+    "model.layers.1.self_attn.q_norm.weight": "model-00001-of-00008.safetensors",
+    "model.layers.1.self_attn.q_proj.weight": "model-00001-of-00008.safetensors",
+    "model.layers.1.self_attn.v_proj.weight": "model-00001-of-00008.safetensors",
+    "model.layers.2.input_layernorm.weight": "model-00001-of-00008.safetensors",
+    "model.layers.2.mlp.down_proj.weight": "model-00001-of-00008.safetensors",
+    "model.layers.2.mlp.gate_proj.weight": "model-00001-of-00008.safetensors",
+    "model.layers.2.mlp.up_proj.weight": "model-00001-of-00008.safetensors",
+    "model.layers.2.post_attention_layernorm.weight": "model-00001-of-00008.safetensors",
+    "model.layers.2.self_attn.k_norm.weight": "model-00001-of-00008.safetensors",
+    "model.layers.2.self_attn.k_proj.weight": "model-00001-of-00008.safetensors",
+    "model.layers.2.self_attn.o_proj.weight": "model-00001-of-00008.safetensors",
+    "model.layers.2.self_attn.q_norm.weight": "model-00001-of-00008.safetensors",
+    "model.layers.2.self_attn.q_proj.weight": "model-00001-of-00008.safetensors",
+    "model.layers.2.self_attn.v_proj.weight": "model-00001-of-00008.safetensors",
+    "model.layers.3.mlp.gate_proj.weight": "model-00001-of-00008.safetensors",
+    "model.layers.3.self_attn.k_norm.weight": "model-00001-of-00008.safetensors",
+    "model.layers.3.self_attn.k_proj.weight": "model-00001-of-00008.safetensors",
+    "model.layers.3.self_attn.o_proj.weight": "model-00001-of-00008.safetensors",
+    "model.layers.3.self_attn.q_norm.weight": "model-00001-of-00008.safetensors",
+    "model.layers.3.self_attn.q_proj.weight": "model-00001-of-00008.safetensors",
+    "model.layers.3.self_attn.v_proj.weight": "model-00001-of-00008.safetensors",
+    "model.layers.10.input_layernorm.weight": "model-00003-of-00008.safetensors",
+    "model.layers.10.mlp.down_proj.weight": "model-00003-of-00008.safetensors",
+    "model.layers.10.mlp.gate_proj.weight": "model-00003-of-00008.safetensors",
+    "model.layers.10.mlp.up_proj.weight": "model-00003-of-00008.safetensors",
+    "model.layers.10.post_attention_layernorm.weight": "model-00003-of-00008.safetensors",
+    "model.layers.10.self_attn.k_norm.weight": "model-00003-of-00008.safetensors",
+    "model.layers.10.self_attn.k_proj.weight": "model-00003-of-00008.safetensors",
+    "model.layers.10.self_attn.o_proj.weight": "model-00003-of-00008.safetensors",
+    "model.layers.10.self_attn.q_norm.weight": "model-00003-of-00008.safetensors",
+    "model.layers.10.self_attn.q_proj.weight": "model-00003-of-00008.safetensors",
+    "model.layers.10.self_attn.v_proj.weight": "model-00003-of-00008.safetensors",
+    "model.layers.11.input_layernorm.weight": "model-00003-of-00008.safetensors",
+    "model.layers.11.mlp.down_proj.weight": "model-00003-of-00008.safetensors",
+    "model.layers.11.mlp.gate_proj.weight": "model-00003-of-00008.safetensors",
+    "model.layers.11.mlp.up_proj.weight": "model-00003-of-00008.safetensors",
+    "model.layers.11.post_attention_layernorm.weight": "model-00003-of-00008.safetensors",
+    "model.layers.11.self_attn.k_norm.weight": "model-00003-of-00008.safetensors",
+    "model.layers.11.self_attn.k_proj.weight": "model-00003-of-00008.safetensors",
+    "model.layers.11.self_attn.o_proj.weight": "model-00003-of-00008.safetensors",
+    "model.layers.11.self_attn.q_norm.weight": "model-00003-of-00008.safetensors",
+    "model.layers.11.self_attn.q_proj.weight": "model-00003-of-00008.safetensors",
+    "model.layers.11.self_attn.v_proj.weight": "model-00003-of-00008.safetensors",
+    "model.layers.12.input_layernorm.weight": "model-00003-of-00008.safetensors",
+    "model.layers.12.mlp.down_proj.weight": "model-00003-of-00008.safetensors",
+    "model.layers.12.mlp.gate_proj.weight": "model-00003-of-00008.safetensors",
+    "model.layers.12.mlp.up_proj.weight": "model-00003-of-00008.safetensors",
+    "model.layers.12.post_attention_layernorm.weight": "model-00003-of-00008.safetensors",
+    "model.layers.12.self_attn.k_norm.weight": "model-00003-of-00008.safetensors",
+    "model.layers.12.self_attn.k_proj.weight": "model-00003-of-00008.safetensors",
+    "model.layers.12.self_attn.o_proj.weight": "model-00003-of-00008.safetensors",
+    "model.layers.12.self_attn.q_norm.weight": "model-00003-of-00008.safetensors",
+    "model.layers.12.self_attn.q_proj.weight": "model-00003-of-00008.safetensors",
+    "model.layers.12.self_attn.v_proj.weight": "model-00003-of-00008.safetensors",
+    "model.layers.13.input_layernorm.weight": "model-00003-of-00008.safetensors",
+    "model.layers.13.mlp.down_proj.weight": "model-00003-of-00008.safetensors",
+    "model.layers.13.mlp.gate_proj.weight": "model-00003-of-00008.safetensors",
+    "model.layers.13.mlp.up_proj.weight": "model-00003-of-00008.safetensors",
+    "model.layers.13.post_attention_layernorm.weight": "model-00003-of-00008.safetensors",
+    "model.layers.13.self_attn.k_norm.weight": "model-00003-of-00008.safetensors",
+    "model.layers.13.self_attn.k_proj.weight": "model-00003-of-00008.safetensors",
+    "model.layers.13.self_attn.o_proj.weight": "model-00003-of-00008.safetensors",
+    "model.layers.13.self_attn.q_norm.weight": "model-00003-of-00008.safetensors",
+    "model.layers.13.self_attn.q_proj.weight": "model-00003-of-00008.safetensors",
+    "model.layers.13.self_attn.v_proj.weight": "model-00003-of-00008.safetensors",
+    "model.layers.14.input_layernorm.weight": "model-00003-of-00008.safetensors",
+    "model.layers.14.mlp.down_proj.weight": "model-00003-of-00008.safetensors",
+    "model.layers.14.mlp.gate_proj.weight": "model-00003-of-00008.safetensors",
+    "model.layers.14.mlp.up_proj.weight": "model-00003-of-00008.safetensors",
+    "model.layers.14.post_attention_layernorm.weight": "model-00003-of-00008.safetensors",
+    "model.layers.14.self_attn.k_norm.weight": "model-00003-of-00008.safetensors",
+    "model.layers.14.self_attn.k_proj.weight": "model-00003-of-00008.safetensors",
+    "model.layers.14.self_attn.o_proj.weight": "model-00003-of-00008.safetensors",
+    "model.layers.14.self_attn.q_norm.weight": "model-00003-of-00008.safetensors",
+    "model.layers.14.self_attn.q_proj.weight": "model-00003-of-00008.safetensors",
+    "model.layers.14.self_attn.v_proj.weight": "model-00003-of-00008.safetensors",
+    "model.layers.15.mlp.gate_proj.weight": "model-00003-of-00008.safetensors",
+    "model.layers.15.self_attn.k_norm.weight": "model-00003-of-00008.safetensors",
+    "model.layers.15.self_attn.k_proj.weight": "model-00003-of-00008.safetensors",
+    "model.layers.15.self_attn.o_proj.weight": "model-00003-of-00008.safetensors",
+    "model.layers.15.self_attn.q_norm.weight": "model-00003-of-00008.safetensors",
+    "model.layers.15.self_attn.q_proj.weight": "model-00003-of-00008.safetensors",
+    "model.layers.15.self_attn.v_proj.weight": "model-00003-of-00008.safetensors",
+    "model.layers.9.input_layernorm.weight": "model-00003-of-00008.safetensors",
+    "model.layers.9.mlp.down_proj.weight": "model-00003-of-00008.safetensors",
+    "model.layers.9.mlp.up_proj.weight": "model-00003-of-00008.safetensors",
+    "model.layers.9.post_attention_layernorm.weight": "model-00003-of-00008.safetensors",
+    "model.layers.15.input_layernorm.weight": "model-00004-of-00008.safetensors",
+    "model.layers.15.mlp.down_proj.weight": "model-00004-of-00008.safetensors",
+    "model.layers.15.mlp.up_proj.weight": "model-00004-of-00008.safetensors",
+    "model.layers.15.post_attention_layernorm.weight": "model-00004-of-00008.safetensors",
+    "model.layers.16.input_layernorm.weight": "model-00004-of-00008.safetensors",
+    "model.layers.16.mlp.down_proj.weight": "model-00004-of-00008.safetensors",
+    "model.layers.16.mlp.gate_proj.weight": "model-00004-of-00008.safetensors",
+    "model.layers.16.mlp.up_proj.weight": "model-00004-of-00008.safetensors",
+    "model.layers.16.post_attention_layernorm.weight": "model-00004-of-00008.safetensors",
+    "model.layers.16.self_attn.k_norm.weight": "model-00004-of-00008.safetensors",
+    "model.layers.16.self_attn.k_proj.weight": "model-00004-of-00008.safetensors",
+    "model.layers.16.self_attn.o_proj.weight": "model-00004-of-00008.safetensors",
+    "model.layers.16.self_attn.q_norm.weight": "model-00004-of-00008.safetensors",
+    "model.layers.16.self_attn.q_proj.weight": "model-00004-of-00008.safetensors",
+    "model.layers.16.self_attn.v_proj.weight": "model-00004-of-00008.safetensors",
+    "model.layers.17.input_layernorm.weight": "model-00004-of-00008.safetensors",
+    "model.layers.17.mlp.down_proj.weight": "model-00004-of-00008.safetensors",
+    "model.layers.17.mlp.gate_proj.weight": "model-00004-of-00008.safetensors",
+    "model.layers.17.mlp.up_proj.weight": "model-00004-of-00008.safetensors",
+    "model.layers.17.post_attention_layernorm.weight": "model-00004-of-00008.safetensors",
+    "model.layers.17.self_attn.k_norm.weight": "model-00004-of-00008.safetensors",
+    "model.layers.17.self_attn.k_proj.weight": "model-00004-of-00008.safetensors",
+    "model.layers.17.self_attn.o_proj.weight": "model-00004-of-00008.safetensors",
+    "model.layers.17.self_attn.q_norm.weight": "model-00004-of-00008.safetensors",
+    "model.layers.17.self_attn.q_proj.weight": "model-00004-of-00008.safetensors",
+    "model.layers.17.self_attn.v_proj.weight": "model-00004-of-00008.safetensors",
+    "model.layers.18.input_layernorm.weight": "model-00004-of-00008.safetensors",
+    "model.layers.18.mlp.down_proj.weight": "model-00004-of-00008.safetensors",
+    "model.layers.18.mlp.gate_proj.weight": "model-00004-of-00008.safetensors",
+    "model.layers.18.mlp.up_proj.weight": "model-00004-of-00008.safetensors",
+    "model.layers.18.post_attention_layernorm.weight": "model-00004-of-00008.safetensors",
+    "model.layers.18.self_attn.k_norm.weight": "model-00004-of-00008.safetensors",
+    "model.layers.18.self_attn.k_proj.weight": "model-00004-of-00008.safetensors",
+    "model.layers.18.self_attn.o_proj.weight": "model-00004-of-00008.safetensors",
+    "model.layers.18.self_attn.q_norm.weight": "model-00004-of-00008.safetensors",
+    "model.layers.18.self_attn.q_proj.weight": "model-00004-of-00008.safetensors",
+    "model.layers.18.self_attn.v_proj.weight": "model-00004-of-00008.safetensors",
+    "model.layers.19.input_layernorm.weight": "model-00004-of-00008.safetensors",
+    "model.layers.19.mlp.down_proj.weight": "model-00004-of-00008.safetensors",
+    "model.layers.19.mlp.gate_proj.weight": "model-00004-of-00008.safetensors",
+    "model.layers.19.mlp.up_proj.weight": "model-00004-of-00008.safetensors",
+    "model.layers.19.post_attention_layernorm.weight": "model-00004-of-00008.safetensors",
+    "model.layers.19.self_attn.k_norm.weight": "model-00004-of-00008.safetensors",
+    "model.layers.19.self_attn.k_proj.weight": "model-00004-of-00008.safetensors",
+    "model.layers.19.self_attn.o_proj.weight": "model-00004-of-00008.safetensors",
+    "model.layers.19.self_attn.q_norm.weight": "model-00004-of-00008.safetensors",
+    "model.layers.19.self_attn.q_proj.weight": "model-00004-of-00008.safetensors",
+    "model.layers.19.self_attn.v_proj.weight": "model-00004-of-00008.safetensors",
+    "model.layers.20.input_layernorm.weight": "model-00004-of-00008.safetensors",
+    "model.layers.20.mlp.down_proj.weight": "model-00004-of-00008.safetensors",
+    "model.layers.20.mlp.gate_proj.weight": "model-00004-of-00008.safetensors",
+    "model.layers.20.mlp.up_proj.weight": "model-00004-of-00008.safetensors",
+    "model.layers.20.post_attention_layernorm.weight": "model-00004-of-00008.safetensors",
+    "model.layers.20.self_attn.k_norm.weight": "model-00004-of-00008.safetensors",
+    "model.layers.20.self_attn.k_proj.weight": "model-00004-of-00008.safetensors",
+    "model.layers.20.self_attn.o_proj.weight": "model-00004-of-00008.safetensors",
+    "model.layers.20.self_attn.q_norm.weight": "model-00004-of-00008.safetensors",
+    "model.layers.20.self_attn.q_proj.weight": "model-00004-of-00008.safetensors",
+    "model.layers.20.self_attn.v_proj.weight": "model-00004-of-00008.safetensors",
+    "model.layers.21.mlp.gate_proj.weight": "model-00004-of-00008.safetensors",
+    "model.layers.21.self_attn.k_norm.weight": "model-00004-of-00008.safetensors",
+    "model.layers.21.self_attn.k_proj.weight": "model-00004-of-00008.safetensors",
+    "model.layers.21.self_attn.o_proj.weight": "model-00004-of-00008.safetensors",
+    "model.layers.21.self_attn.q_norm.weight": "model-00004-of-00008.safetensors",
+    "model.layers.21.self_attn.q_proj.weight": "model-00004-of-00008.safetensors",
+    "model.layers.21.self_attn.v_proj.weight": "model-00004-of-00008.safetensors",
+    "model.layers.21.input_layernorm.weight": "model-00005-of-00008.safetensors",
+    "model.layers.21.mlp.down_proj.weight": "model-00005-of-00008.safetensors",
+    "model.layers.21.mlp.up_proj.weight": "model-00005-of-00008.safetensors",
+    "model.layers.21.post_attention_layernorm.weight": "model-00005-of-00008.safetensors",
+    "model.layers.22.input_layernorm.weight": "model-00005-of-00008.safetensors",
+    "model.layers.22.mlp.down_proj.weight": "model-00005-of-00008.safetensors",
+    "model.layers.22.mlp.gate_proj.weight": "model-00005-of-00008.safetensors",
+    "model.layers.22.mlp.up_proj.weight": "model-00005-of-00008.safetensors",
+    "model.layers.22.post_attention_layernorm.weight": "model-00005-of-00008.safetensors",
+    "model.layers.22.self_attn.k_norm.weight": "model-00005-of-00008.safetensors",
+    "model.layers.22.self_attn.k_proj.weight": "model-00005-of-00008.safetensors",
+    "model.layers.22.self_attn.o_proj.weight": "model-00005-of-00008.safetensors",
+    "model.layers.22.self_attn.q_norm.weight": "model-00005-of-00008.safetensors",
+    "model.layers.22.self_attn.q_proj.weight": "model-00005-of-00008.safetensors",
+    "model.layers.22.self_attn.v_proj.weight": "model-00005-of-00008.safetensors",
+    "model.layers.23.input_layernorm.weight": "model-00005-of-00008.safetensors",
+    "model.layers.23.mlp.down_proj.weight": "model-00005-of-00008.safetensors",
+    "model.layers.23.mlp.gate_proj.weight": "model-00005-of-00008.safetensors",
+    "model.layers.23.mlp.up_proj.weight": "model-00005-of-00008.safetensors",
+    "model.layers.23.post_attention_layernorm.weight": "model-00005-of-00008.safetensors",
+    "model.layers.23.self_attn.k_norm.weight": "model-00005-of-00008.safetensors",
+    "model.layers.23.self_attn.k_proj.weight": "model-00005-of-00008.safetensors",
+    "model.layers.23.self_attn.o_proj.weight": "model-00005-of-00008.safetensors",
+    "model.layers.23.self_attn.q_norm.weight": "model-00005-of-00008.safetensors",
+    "model.layers.23.self_attn.q_proj.weight": "model-00005-of-00008.safetensors",
+    "model.layers.23.self_attn.v_proj.weight": "model-00005-of-00008.safetensors",
+    "model.layers.24.input_layernorm.weight": "model-00005-of-00008.safetensors",
+    "model.layers.24.mlp.down_proj.weight": "model-00005-of-00008.safetensors",
+    "model.layers.24.mlp.gate_proj.weight": "model-00005-of-00008.safetensors",
+    "model.layers.24.mlp.up_proj.weight": "model-00005-of-00008.safetensors",
+    "model.layers.24.post_attention_layernorm.weight": "model-00005-of-00008.safetensors",
+    "model.layers.24.self_attn.k_norm.weight": "model-00005-of-00008.safetensors",
+    "model.layers.24.self_attn.k_proj.weight": "model-00005-of-00008.safetensors",
+    "model.layers.24.self_attn.o_proj.weight": "model-00005-of-00008.safetensors",
+    "model.layers.24.self_attn.q_norm.weight": "model-00005-of-00008.safetensors",
+    "model.layers.24.self_attn.q_proj.weight": "model-00005-of-00008.safetensors",
+    "model.layers.24.self_attn.v_proj.weight": "model-00005-of-00008.safetensors",
+    "model.layers.25.input_layernorm.weight": "model-00005-of-00008.safetensors",
+    "model.layers.25.mlp.down_proj.weight": "model-00005-of-00008.safetensors",
+    "model.layers.25.mlp.gate_proj.weight": "model-00005-of-00008.safetensors",
+    "model.layers.25.mlp.up_proj.weight": "model-00005-of-00008.safetensors",
+    "model.layers.25.post_attention_layernorm.weight": "model-00005-of-00008.safetensors",
+    "model.layers.25.self_attn.k_norm.weight": "model-00005-of-00008.safetensors",
+    "model.layers.25.self_attn.k_proj.weight": "model-00005-of-00008.safetensors",
+    "model.layers.25.self_attn.o_proj.weight": "model-00005-of-00008.safetensors",
+    "model.layers.25.self_attn.q_norm.weight": "model-00005-of-00008.safetensors",
+    "model.layers.25.self_attn.q_proj.weight": "model-00005-of-00008.safetensors",
+    "model.layers.25.self_attn.v_proj.weight": "model-00005-of-00008.safetensors",
+    "model.layers.26.input_layernorm.weight": "model-00005-of-00008.safetensors",
+    "model.layers.26.mlp.down_proj.weight": "model-00005-of-00008.safetensors",
+    "model.layers.26.mlp.gate_proj.weight": "model-00005-of-00008.safetensors",
+    "model.layers.26.mlp.up_proj.weight": "model-00005-of-00008.safetensors",
+    "model.layers.26.post_attention_layernorm.weight": "model-00005-of-00008.safetensors",
+    "model.layers.26.self_attn.k_norm.weight": "model-00005-of-00008.safetensors",
+    "model.layers.26.self_attn.k_proj.weight": "model-00005-of-00008.safetensors",
+    "model.layers.26.self_attn.o_proj.weight": "model-00005-of-00008.safetensors",
+    "model.layers.26.self_attn.q_norm.weight": "model-00005-of-00008.safetensors",
+    "model.layers.26.self_attn.q_proj.weight": "model-00005-of-00008.safetensors",
+    "model.layers.26.self_attn.v_proj.weight": "model-00005-of-00008.safetensors",
+    "model.layers.27.mlp.gate_proj.weight": "model-00005-of-00008.safetensors",
+    "model.layers.27.self_attn.k_norm.weight": "model-00005-of-00008.safetensors",
+    "model.layers.27.self_attn.k_proj.weight": "model-00005-of-00008.safetensors",
+    "model.layers.27.self_attn.o_proj.weight": "model-00005-of-00008.safetensors",
+    "model.layers.27.self_attn.q_norm.weight": "model-00005-of-00008.safetensors",
+    "model.layers.27.self_attn.q_proj.weight": "model-00005-of-00008.safetensors",
+    "model.layers.27.self_attn.v_proj.weight": "model-00005-of-00008.safetensors",
+    "model.layers.27.input_layernorm.weight": "model-00006-of-00008.safetensors",
+    "model.layers.27.mlp.down_proj.weight": "model-00006-of-00008.safetensors",
+    "model.layers.27.mlp.up_proj.weight": "model-00006-of-00008.safetensors",
+    "model.layers.27.post_attention_layernorm.weight": "model-00006-of-00008.safetensors",
+    "model.layers.28.input_layernorm.weight": "model-00006-of-00008.safetensors",
+    "model.layers.28.mlp.down_proj.weight": "model-00006-of-00008.safetensors",
+    "model.layers.28.mlp.gate_proj.weight": "model-00006-of-00008.safetensors",
+    "model.layers.28.mlp.up_proj.weight": "model-00006-of-00008.safetensors",
+    "model.layers.28.post_attention_layernorm.weight": "model-00006-of-00008.safetensors",
+    "model.layers.28.self_attn.k_norm.weight": "model-00006-of-00008.safetensors",
+    "model.layers.28.self_attn.k_proj.weight": "model-00006-of-00008.safetensors",
+    "model.layers.28.self_attn.o_proj.weight": "model-00006-of-00008.safetensors",
+    "model.layers.28.self_attn.q_norm.weight": "model-00006-of-00008.safetensors",
+    "model.layers.28.self_attn.q_proj.weight": "model-00006-of-00008.safetensors",
+    "model.layers.28.self_attn.v_proj.weight": "model-00006-of-00008.safetensors",
+    "model.layers.29.input_layernorm.weight": "model-00006-of-00008.safetensors",
+    "model.layers.29.mlp.down_proj.weight": "model-00006-of-00008.safetensors",
+    "model.layers.29.mlp.gate_proj.weight": "model-00006-of-00008.safetensors",
+    "model.layers.29.mlp.up_proj.weight": "model-00006-of-00008.safetensors",
+    "model.layers.29.post_attention_layernorm.weight": "model-00006-of-00008.safetensors",
+    "model.layers.29.self_attn.k_norm.weight": "model-00006-of-00008.safetensors",
+    "model.layers.29.self_attn.k_proj.weight": "model-00006-of-00008.safetensors",
+    "model.layers.29.self_attn.o_proj.weight": "model-00006-of-00008.safetensors",
+    "model.layers.29.self_attn.q_norm.weight": "model-00006-of-00008.safetensors",
+    "model.layers.29.self_attn.q_proj.weight": "model-00006-of-00008.safetensors",
+    "model.layers.29.self_attn.v_proj.weight": "model-00006-of-00008.safetensors",
+    "model.layers.30.input_layernorm.weight": "model-00006-of-00008.safetensors",
+    "model.layers.30.mlp.down_proj.weight": "model-00006-of-00008.safetensors",
+    "model.layers.30.mlp.gate_proj.weight": "model-00006-of-00008.safetensors",
+    "model.layers.30.mlp.up_proj.weight": "model-00006-of-00008.safetensors",
+    "model.layers.30.post_attention_layernorm.weight": "model-00006-of-00008.safetensors",
+    "model.layers.30.self_attn.k_norm.weight": "model-00006-of-00008.safetensors",
+    "model.layers.30.self_attn.k_proj.weight": "model-00006-of-00008.safetensors",
+    "model.layers.30.self_attn.o_proj.weight": "model-00006-of-00008.safetensors",
+    "model.layers.30.self_attn.q_norm.weight": "model-00006-of-00008.safetensors",
+    "model.layers.30.self_attn.q_proj.weight": "model-00006-of-00008.safetensors",
+    "model.layers.30.self_attn.v_proj.weight": "model-00006-of-00008.safetensors",
+    "model.layers.31.input_layernorm.weight": "model-00006-of-00008.safetensors",
+    "model.layers.31.mlp.down_proj.weight": "model-00006-of-00008.safetensors",
+    "model.layers.31.mlp.gate_proj.weight": "model-00006-of-00008.safetensors",
+    "model.layers.31.mlp.up_proj.weight": "model-00006-of-00008.safetensors",
+    "model.layers.31.post_attention_layernorm.weight": "model-00006-of-00008.safetensors",
+    "model.layers.31.self_attn.k_norm.weight": "model-00006-of-00008.safetensors",
+    "model.layers.31.self_attn.k_proj.weight": "model-00006-of-00008.safetensors",
+    "model.layers.31.self_attn.o_proj.weight": "model-00006-of-00008.safetensors",
+    "model.layers.31.self_attn.q_norm.weight": "model-00006-of-00008.safetensors",
+    "model.layers.31.self_attn.q_proj.weight": "model-00006-of-00008.safetensors",
+    "model.layers.31.self_attn.v_proj.weight": "model-00006-of-00008.safetensors",
+    "model.layers.32.input_layernorm.weight": "model-00006-of-00008.safetensors",
+    "model.layers.32.mlp.down_proj.weight": "model-00006-of-00008.safetensors",
+    "model.layers.32.mlp.gate_proj.weight": "model-00006-of-00008.safetensors",
+    "model.layers.32.mlp.up_proj.weight": "model-00006-of-00008.safetensors",
+    "model.layers.32.post_attention_layernorm.weight": "model-00006-of-00008.safetensors",
+    "model.layers.32.self_attn.k_norm.weight": "model-00006-of-00008.safetensors",
+    "model.layers.32.self_attn.k_proj.weight": "model-00006-of-00008.safetensors",
+    "model.layers.32.self_attn.o_proj.weight": "model-00006-of-00008.safetensors",
+    "model.layers.32.self_attn.q_norm.weight": "model-00006-of-00008.safetensors",
+    "model.layers.32.self_attn.q_proj.weight": "model-00006-of-00008.safetensors",
+    "model.layers.32.self_attn.v_proj.weight": "model-00006-of-00008.safetensors",
+    "model.layers.33.mlp.gate_proj.weight": "model-00006-of-00008.safetensors",
+    "model.layers.33.self_attn.k_norm.weight": "model-00006-of-00008.safetensors",
+    "model.layers.33.self_attn.k_proj.weight": "model-00006-of-00008.safetensors",
+    "model.layers.33.self_attn.o_proj.weight": "model-00006-of-00008.safetensors",
+    "model.layers.33.self_attn.q_norm.weight": "model-00006-of-00008.safetensors",
+    "model.layers.33.self_attn.q_proj.weight": "model-00006-of-00008.safetensors",
+    "model.layers.33.self_attn.v_proj.weight": "model-00006-of-00008.safetensors",
+    "model.layers.3.input_layernorm.weight": "model-00002-of-00008.safetensors",
+    "model.layers.3.mlp.down_proj.weight": "model-00002-of-00008.safetensors",
+    "model.layers.3.mlp.up_proj.weight": "model-00002-of-00008.safetensors",
+    "model.layers.3.post_attention_layernorm.weight": "model-00002-of-00008.safetensors",
+    "model.layers.4.input_layernorm.weight": "model-00002-of-00008.safetensors",
+    "model.layers.4.mlp.down_proj.weight": "model-00002-of-00008.safetensors",
+    "model.layers.4.mlp.gate_proj.weight": "model-00002-of-00008.safetensors",
+    "model.layers.4.mlp.up_proj.weight": "model-00002-of-00008.safetensors",
+    "model.layers.4.post_attention_layernorm.weight": "model-00002-of-00008.safetensors",
+    "model.layers.4.self_attn.k_norm.weight": "model-00002-of-00008.safetensors",
+    "model.layers.4.self_attn.k_proj.weight": "model-00002-of-00008.safetensors",
+    "model.layers.4.self_attn.o_proj.weight": "model-00002-of-00008.safetensors",
+    "model.layers.4.self_attn.q_norm.weight": "model-00002-of-00008.safetensors",
+    "model.layers.4.self_attn.q_proj.weight": "model-00002-of-00008.safetensors",
+    "model.layers.4.self_attn.v_proj.weight": "model-00002-of-00008.safetensors",
+    "model.layers.5.input_layernorm.weight": "model-00002-of-00008.safetensors",
+    "model.layers.5.mlp.down_proj.weight": "model-00002-of-00008.safetensors",
+    "model.layers.5.mlp.gate_proj.weight": "model-00002-of-00008.safetensors",
+    "model.layers.5.mlp.up_proj.weight": "model-00002-of-00008.safetensors",
+    "model.layers.5.post_attention_layernorm.weight": "model-00002-of-00008.safetensors",
+    "model.layers.5.self_attn.k_norm.weight": "model-00002-of-00008.safetensors",
+    "model.layers.5.self_attn.k_proj.weight": "model-00002-of-00008.safetensors",
+    "model.layers.5.self_attn.o_proj.weight": "model-00002-of-00008.safetensors",
+    "model.layers.5.self_attn.q_norm.weight": "model-00002-of-00008.safetensors",
+    "model.layers.5.self_attn.q_proj.weight": "model-00002-of-00008.safetensors",
+    "model.layers.5.self_attn.v_proj.weight": "model-00002-of-00008.safetensors",
+    "model.layers.6.input_layernorm.weight": "model-00002-of-00008.safetensors",
+    "model.layers.6.mlp.down_proj.weight": "model-00002-of-00008.safetensors",
+    "model.layers.6.mlp.gate_proj.weight": "model-00002-of-00008.safetensors",
+    "model.layers.6.mlp.up_proj.weight": "model-00002-of-00008.safetensors",
+    "model.layers.6.post_attention_layernorm.weight": "model-00002-of-00008.safetensors",
+    "model.layers.6.self_attn.k_norm.weight": "model-00002-of-00008.safetensors",
+    "model.layers.6.self_attn.k_proj.weight": "model-00002-of-00008.safetensors",
+    "model.layers.6.self_attn.o_proj.weight": "model-00002-of-00008.safetensors",
+    "model.layers.6.self_attn.q_norm.weight": "model-00002-of-00008.safetensors",
+    "model.layers.6.self_attn.q_proj.weight": "model-00002-of-00008.safetensors",
+    "model.layers.6.self_attn.v_proj.weight": "model-00002-of-00008.safetensors",
+    "model.layers.7.input_layernorm.weight": "model-00002-of-00008.safetensors",
+    "model.layers.7.mlp.down_proj.weight": "model-00002-of-00008.safetensors",
+    "model.layers.7.mlp.gate_proj.weight": "model-00002-of-00008.safetensors",
+    "model.layers.7.mlp.up_proj.weight": "model-00002-of-00008.safetensors",
+    "model.layers.7.post_attention_layernorm.weight": "model-00002-of-00008.safetensors",
+    "model.layers.7.self_attn.k_norm.weight": "model-00002-of-00008.safetensors",
+    "model.layers.7.self_attn.k_proj.weight": "model-00002-of-00008.safetensors",
+    "model.layers.7.self_attn.o_proj.weight": "model-00002-of-00008.safetensors",
+    "model.layers.7.self_attn.q_norm.weight": "model-00002-of-00008.safetensors",
+    "model.layers.7.self_attn.q_proj.weight": "model-00002-of-00008.safetensors",
+    "model.layers.7.self_attn.v_proj.weight": "model-00002-of-00008.safetensors",
+    "model.layers.8.input_layernorm.weight": "model-00002-of-00008.safetensors",
+    "model.layers.8.mlp.down_proj.weight": "model-00002-of-00008.safetensors",
+    "model.layers.8.mlp.gate_proj.weight": "model-00002-of-00008.safetensors",
+    "model.layers.8.mlp.up_proj.weight": "model-00002-of-00008.safetensors",
+    "model.layers.8.post_attention_layernorm.weight": "model-00002-of-00008.safetensors",
+    "model.layers.8.self_attn.k_norm.weight": "model-00002-of-00008.safetensors",
+    "model.layers.8.self_attn.k_proj.weight": "model-00002-of-00008.safetensors",
+    "model.layers.8.self_attn.o_proj.weight": "model-00002-of-00008.safetensors",
+    "model.layers.8.self_attn.q_norm.weight": "model-00002-of-00008.safetensors",
+    "model.layers.8.self_attn.q_proj.weight": "model-00002-of-00008.safetensors",
+    "model.layers.8.self_attn.v_proj.weight": "model-00002-of-00008.safetensors",
+    "model.layers.9.mlp.gate_proj.weight": "model-00002-of-00008.safetensors",
+    "model.layers.9.self_attn.k_norm.weight": "model-00002-of-00008.safetensors",
+    "model.layers.9.self_attn.k_proj.weight": "model-00002-of-00008.safetensors",
+    "model.layers.9.self_attn.o_proj.weight": "model-00002-of-00008.safetensors",
+    "model.layers.9.self_attn.q_norm.weight": "model-00002-of-00008.safetensors",
+    "model.layers.9.self_attn.q_proj.weight": "model-00002-of-00008.safetensors",
+    "model.layers.9.self_attn.v_proj.weight": "model-00002-of-00008.safetensors",
+    "model.layers.33.input_layernorm.weight": "model-00007-of-00008.safetensors",
+    "model.layers.33.mlp.down_proj.weight": "model-00007-of-00008.safetensors",
+    "model.layers.33.mlp.up_proj.weight": "model-00007-of-00008.safetensors",
+    "model.layers.33.post_attention_layernorm.weight": "model-00007-of-00008.safetensors",
+    "model.layers.34.input_layernorm.weight": "model-00007-of-00008.safetensors",
+    "model.layers.34.mlp.down_proj.weight": "model-00007-of-00008.safetensors",
+    "model.layers.34.mlp.gate_proj.weight": "model-00007-of-00008.safetensors",
+    "model.layers.34.mlp.up_proj.weight": "model-00007-of-00008.safetensors",
+    "model.layers.34.post_attention_layernorm.weight": "model-00007-of-00008.safetensors",
+    "model.layers.34.self_attn.k_norm.weight": "model-00007-of-00008.safetensors",
+    "model.layers.34.self_attn.k_proj.weight": "model-00007-of-00008.safetensors",
+    "model.layers.34.self_attn.o_proj.weight": "model-00007-of-00008.safetensors",
+    "model.layers.34.self_attn.q_norm.weight": "model-00007-of-00008.safetensors",
+    "model.layers.34.self_attn.q_proj.weight": "model-00007-of-00008.safetensors",
+    "model.layers.34.self_attn.v_proj.weight": "model-00007-of-00008.safetensors",
+    "model.layers.35.input_layernorm.weight": "model-00007-of-00008.safetensors",
+    "model.layers.35.mlp.down_proj.weight": "model-00007-of-00008.safetensors",
+    "model.layers.35.mlp.gate_proj.weight": "model-00007-of-00008.safetensors",
+    "model.layers.35.mlp.up_proj.weight": "model-00007-of-00008.safetensors",
+    "model.layers.35.post_attention_layernorm.weight": "model-00007-of-00008.safetensors",
+    "model.layers.35.self_attn.k_norm.weight": "model-00007-of-00008.safetensors",
+    "model.layers.35.self_attn.k_proj.weight": "model-00007-of-00008.safetensors",
+    "model.layers.35.self_attn.o_proj.weight": "model-00007-of-00008.safetensors",
+    "model.layers.35.self_attn.q_norm.weight": "model-00007-of-00008.safetensors",
+    "model.layers.35.self_attn.q_proj.weight": "model-00007-of-00008.safetensors",
+    "model.layers.35.self_attn.v_proj.weight": "model-00007-of-00008.safetensors",
+    "model.layers.36.input_layernorm.weight": "model-00007-of-00008.safetensors",
+    "model.layers.36.mlp.down_proj.weight": "model-00007-of-00008.safetensors",
+    "model.layers.36.mlp.gate_proj.weight": "model-00007-of-00008.safetensors",
+    "model.layers.36.mlp.up_proj.weight": "model-00007-of-00008.safetensors",
+    "model.layers.36.post_attention_layernorm.weight": "model-00007-of-00008.safetensors",
+    "model.layers.36.self_attn.k_norm.weight": "model-00007-of-00008.safetensors",
+    "model.layers.36.self_attn.k_proj.weight": "model-00007-of-00008.safetensors",
+    "model.layers.36.self_attn.o_proj.weight": "model-00007-of-00008.safetensors",
+    "model.layers.36.self_attn.q_norm.weight": "model-00007-of-00008.safetensors",
+    "model.layers.36.self_attn.q_proj.weight": "model-00007-of-00008.safetensors",
+    "model.layers.36.self_attn.v_proj.weight": "model-00007-of-00008.safetensors",
+    "model.layers.37.input_layernorm.weight": "model-00007-of-00008.safetensors",
+    "model.layers.37.mlp.down_proj.weight": "model-00007-of-00008.safetensors",
+    "model.layers.37.mlp.gate_proj.weight": "model-00007-of-00008.safetensors",
+    "model.layers.37.mlp.up_proj.weight": "model-00007-of-00008.safetensors",
+    "model.layers.37.post_attention_layernorm.weight": "model-00007-of-00008.safetensors",
+    "model.layers.37.self_attn.k_norm.weight": "model-00007-of-00008.safetensors",
+    "model.layers.37.self_attn.k_proj.weight": "model-00007-of-00008.safetensors",
+    "model.layers.37.self_attn.o_proj.weight": "model-00007-of-00008.safetensors",
+    "model.layers.37.self_attn.q_norm.weight": "model-00007-of-00008.safetensors",
+    "model.layers.37.self_attn.q_proj.weight": "model-00007-of-00008.safetensors",
+    "model.layers.37.self_attn.v_proj.weight": "model-00007-of-00008.safetensors",
+    "model.layers.38.input_layernorm.weight": "model-00007-of-00008.safetensors",
+    "model.layers.38.mlp.down_proj.weight": "model-00007-of-00008.safetensors",
+    "model.layers.38.mlp.gate_proj.weight": "model-00007-of-00008.safetensors",
+    "model.layers.38.mlp.up_proj.weight": "model-00007-of-00008.safetensors",
+    "model.layers.38.post_attention_layernorm.weight": "model-00007-of-00008.safetensors",
+    "model.layers.38.self_attn.k_norm.weight": "model-00007-of-00008.safetensors",
+    "model.layers.38.self_attn.k_proj.weight": "model-00007-of-00008.safetensors",
+    "model.layers.38.self_attn.o_proj.weight": "model-00007-of-00008.safetensors",
+    "model.layers.38.self_attn.q_norm.weight": "model-00007-of-00008.safetensors",
+    "model.layers.38.self_attn.q_proj.weight": "model-00007-of-00008.safetensors",
+    "model.layers.38.self_attn.v_proj.weight": "model-00007-of-00008.safetensors",
+    "model.layers.39.mlp.gate_proj.weight": "model-00007-of-00008.safetensors",
+    "model.layers.39.self_attn.k_norm.weight": "model-00007-of-00008.safetensors",
+    "model.layers.39.self_attn.k_proj.weight": "model-00007-of-00008.safetensors",
+    "model.layers.39.self_attn.o_proj.weight": "model-00007-of-00008.safetensors",
+    "model.layers.39.self_attn.q_norm.weight": "model-00007-of-00008.safetensors",
+    "model.layers.39.self_attn.q_proj.weight": "model-00007-of-00008.safetensors",
+    "model.layers.39.self_attn.v_proj.weight": "model-00007-of-00008.safetensors"
+  }
+}

modeling_brumby.py

@@ -0,0 +1,738 @@
+# coding=utf-8
+# Copyright 2024 HuggingFace Inc. team.
+# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""PyTorch Brumby model."""
+
+from typing import Callable, Optional, Union, Any, Dict, Tuple
+
+import torch
+from torch import nn
+
+from transformers.activations import ACT2FN
+from transformers.cache_utils import Cache, DynamicCache
+from transformers.generation import GenerationMixin
+from transformers.integrations import use_kernel_forward_from_hub
+from transformers.masking_utils import create_causal_mask, create_sliding_window_causal_mask
+from transformers.modeling_flash_attention_utils import FlashAttentionKwargs
+from transformers.modeling_layers import (
+    GenericForQuestionAnswering,
+    GenericForSequenceClassification,
+    GenericForTokenClassification,
+    GradientCheckpointingLayer,
+)
+from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast
+from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update
+from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel
+from transformers.processing_utils import Unpack
+from transformers.utils import TransformersKwargs, auto_docstring, can_return_tuple
+from transformers.utils.deprecation import deprecate_kwarg
+from transformers.utils.generic import check_model_inputs
+from .configuration_brumby import BrumbyConfig
+
+
+try:
+    from retention.triton import power_retention, power_retention_inference
+    from retention._utils import compute_expanded_dim
+except ImportError:
+    raise ImportError("Retention is required by the Brumby model. Please install it with `pip install retention`.")
+
+
+class PowerAttentionDynamicCache(DynamicCache):
+    """
+    A dynamic cache that encompasses 2 sets of caches:
+    1. attention cache with a short seq_len (determined by the chunk_size parameter):
+       - key_cache: [batch_size, num_heads, chunk_size, head_dim]
+       - value_cache: [batch_size, num_heads, chunk_size, head_dim]
+       - gating_cache: [batch_size, num_heads, chunk_size]
+    2. fixed-size state-based cache:
+       - state: [batch_size, num_heads, state_dim, head_dim]
+       - sum_of_keys: [batch_size, num_heads, state_dim]
+
+    where state_dim is determined by the power of expansion for power attention.
+    """
+    def __init__(self, config: BrumbyConfig, batch_size: int, dtype=torch.bfloat16, device=None):
+        super().__init__()
+        self.config = config
+        self.batch_size = batch_size
+        self.chunk_size = config.chunk_size
+        self.head_dim = config.hidden_size // config.num_attention_heads
+        self.p = config.p
+        self.state_dim = compute_expanded_dim(self.head_dim, deg=self.p)
+        
+        self.states = [None for _ in range(config.num_hidden_layers)]
+        self.sum_of_keys = [None for _ in range(config.num_hidden_layers)]
+        self.key_cache = [torch.tensor([[]] * batch_size, device=device, dtype=dtype) for _ in range(config.num_hidden_layers)]
+        self.value_cache = [torch.tensor([[]] * batch_size, device=device, dtype=dtype) for _ in range(config.num_hidden_layers)]
+        self.gate_cache = [torch.tensor([[]] * batch_size, device=device, dtype=torch.float32) for _ in range(config.num_hidden_layers)]
+
+    def clean_cache(self, layer_idx: int) -> None:
+        self.key_cache[layer_idx] = torch.tensor([[]] * self.batch_size, device=self.device, dtype=self.dtype)
+        self.value_cache[layer_idx] = torch.tensor([[]] * self.batch_size, device=self.device, dtype=self.dtype)
+        self.gate_cache[layer_idx] = torch.tensor([[]] * self.batch_size, device=self.device, dtype=torch.float32)
+
+    def update_cache(
+        self,
+        key_states: torch.Tensor,
+        value_states: torch.Tensor,
+        gate_states: torch.Tensor,
+        layer_idx: int,
+        cache_kwargs: Optional[Dict[str, Any]] = None,
+    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        if self.key_cache[layer_idx].shape[-1] == 0:
+            self.key_cache[layer_idx] = key_states
+            self.value_cache[layer_idx] = value_states
+            self.gate_cache[layer_idx] = gate_states
+        else:
+            self.key_cache[layer_idx] = torch.cat([self.key_cache[layer_idx], key_states], dim=2)
+            self.value_cache[layer_idx] = torch.cat([self.value_cache[layer_idx], value_states], dim=2)
+            self.gate_cache[layer_idx] = torch.cat([self.gate_cache[layer_idx], gate_states], dim=2)
+        return self.key_cache[layer_idx], self.value_cache[layer_idx], self.gate_cache[layer_idx], self.states[layer_idx], self.sum_of_keys[layer_idx]
+
+    def reorder_cache(self, beam_idx: torch.LongTensor):
+        """Reorders the cache for beam search, given the selected beam indices."""
+        for layer_idx in range(len(self.key_cache)):
+            device = self.key_cache[layer_idx].device
+            self.key_cache[layer_idx] = self.key_cache[layer_idx].index_select(0, beam_idx.to(device))
+            device = self.value_cache[layer_idx].device
+            self.value_cache[layer_idx] = self.value_cache[layer_idx].index_select(0, beam_idx.to(device))
+            device = self.gate_cache[layer_idx].device
+            self.gate_cache[layer_idx] = self.gate_cache[layer_idx].index_select(0, beam_idx.to(device))
+            device = self.states[layer_idx].device
+            self.states[layer_idx] = self.states[layer_idx].index_select(0, beam_idx.to(device))
+            device = self.sum_of_keys[layer_idx].device
+            self.sum_of_keys[layer_idx] = self.sum_of_keys[layer_idx].index_select(0, beam_idx.to(device))
+
+    def get_seq_length(self, layer_idx: Optional[int] = 0) -> int:
+        """Returns the sequence length of the cached states. A layer index can be optionally passed."""
+        # take any layer that contains cache and not empty tensor
+        if layer_idx is None:
+            layer_idx = 0
+        if layer_idx >= len(self.key_cache):
+            return 0
+        # Check if the cache for this layer is empty
+        if self.key_cache[layer_idx].numel() == 0:
+            return 0
+        return self.key_cache[layer_idx].shape[-2]
+
+    def to_legacy_cache(self) -> tuple[tuple[torch.Tensor, torch.Tensor]]:
+        raise NotImplementedError("PowerAttentionDynamicCache does not have a legacy cache equivalent.")
+
+    @classmethod
+    def from_legacy_cache(cls, past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None) -> "PowerAttentionDynamicCache":
+        raise NotImplementedError("PowerAttentionDynamicCache does not have a legacy cache equivalent.")
+
+    def update_state(self, layer_idx: int, new_state: torch.Tensor, new_sum_of_keys: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+        self.states[layer_idx] = new_state
+        self.sum_of_keys[layer_idx] = new_sum_of_keys
+        return self.states[layer_idx], self.sum_of_keys[layer_idx]
+
+    def reset(self):
+        self.states.zero_()
+        self.sum_of_keys.zero_()
+        self.key_cache.zero_()
+        self.value_cache.zero_()
+        self.gate_cache.zero_()
+
+
+
+@use_kernel_forward_from_hub("RMSNorm")
+class BrumbyRMSNorm(nn.Module):
+    def __init__(self, hidden_size, eps: float = 1e-6) -> None:
+        """
+        BrumbyRMSNorm is equivalent to T5LayerNorm
+        """
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(hidden_size))
+        self.variance_epsilon = eps
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        input_dtype = hidden_states.dtype
+        hidden_states = hidden_states.to(torch.float32)
+        variance = hidden_states.pow(2).mean(-1, keepdim=True)
+        hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
+        return self.weight * hidden_states.to(input_dtype)
+
+    def extra_repr(self):
+        return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}"
+
+
+class BrumbyMLP(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.config = config
+        self.hidden_size = config.hidden_size
+        self.intermediate_size = config.intermediate_size
+        self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
+        self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
+        self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)
+        self.act_fn = ACT2FN[config.hidden_act]
+
+    def forward(self, x):
+        down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
+        return down_proj
+
+
+def rotate_half(x):
+    """Rotates half the hidden dims of the input."""
+    x1 = x[..., : x.shape[-1] // 2]
+    x2 = x[..., x.shape[-1] // 2 :]
+    return torch.cat((-x2, x1), dim=-1)
+
+
+def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1):
+    """Applies Rotary Position Embedding to the query and key tensors.
+
+    Args:
+        q (`torch.Tensor`): The query tensor.
+        k (`torch.Tensor`): The key tensor.
+        cos (`torch.Tensor`): The cosine part of the rotary embedding.
+        sin (`torch.Tensor`): The sine part of the rotary embedding.
+        position_ids (`torch.Tensor`, *optional*):
+            Deprecated and unused.
+        unsqueeze_dim (`int`, *optional*, defaults to 1):
+            The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and
+            sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note
+            that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and
+            k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes
+            cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have
+            the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2.
+    Returns:
+        `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding.
+    """
+    cos = cos.unsqueeze(unsqueeze_dim)
+    sin = sin.unsqueeze(unsqueeze_dim)
+    q_embed = (q * cos) + (rotate_half(q) * sin)
+    k_embed = (k * cos) + (rotate_half(k) * sin)
+    return q_embed, k_embed
+
+
+def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
+    """
+    This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch,
+    num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim)
+    """
+    batch, num_key_value_heads, slen, head_dim = hidden_states.shape
+    if n_rep == 1:
+        return hidden_states
+    hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim)
+    return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)
+
+
+def eager_attention_forward(
+    module: nn.Module,
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    attention_mask: Optional[torch.Tensor],
+    scaling: float,
+    dropout: float = 0.0,
+    **kwargs: Unpack[TransformersKwargs],
+):
+    key_states = repeat_kv(key, module.num_key_value_groups)
+    value_states = repeat_kv(value, module.num_key_value_groups)
+
+    attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling
+    if attention_mask is not None:
+        causal_mask = attention_mask[:, :, :, : key_states.shape[-2]]
+        attn_weights = attn_weights + causal_mask
+
+    attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype)
+    attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training)
+    attn_output = torch.matmul(attn_weights, value_states)
+    attn_output = attn_output.transpose(1, 2).contiguous()
+
+    return attn_output, attn_weights
+
+
+class BrumbyAttention(nn.Module):
+    """Multi-headed attention from 'Attention Is All You Need' paper"""
+
+    def __init__(self, config: BrumbyConfig, layer_idx: int):
+        super().__init__()
+        self.config = config
+        self.layer_idx = layer_idx
+        self.head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads)
+        self.num_key_value_groups = config.num_attention_heads // config.num_key_value_heads
+        self.scaling = self.head_dim**-0.5
+        self.attention_dropout = config.attention_dropout
+        self.is_causal = True
+        self.use_exp = config.use_exp
+        self.prefill_chunk_size = config.prefill_chunk_size
+        self.chunk_size = config.chunk_size
+        self.switch_over_seq_len = config.switch_over_seq_len
+
+        self.q_proj = nn.Linear(
+            config.hidden_size, config.num_attention_heads * self.head_dim, bias=config.attention_bias
+        )
+        self.k_proj = nn.Linear(
+            config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias
+        )
+        self.v_proj = nn.Linear(
+            config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias
+        )
+        self.g_proj = nn.Linear(
+            config.hidden_size, config.num_key_value_heads, bias=config.attention_bias
+        )
+        self.o_proj = nn.Linear(
+            config.num_attention_heads * self.head_dim, config.hidden_size, bias=config.attention_bias
+        )
+        self.q_norm = BrumbyRMSNorm(self.head_dim, eps=config.rms_norm_eps)  # unlike olmo, only on the head dim!
+        self.k_norm = BrumbyRMSNorm(self.head_dim, eps=config.rms_norm_eps)  # thus post q_norm does not need reshape
+        self.sliding_window = config.sliding_window if config.layer_types[layer_idx] == "sliding_attention" else None
+
+    @deprecate_kwarg("past_key_value", new_name="past_key_values", version="4.58")
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        position_embeddings: tuple[torch.Tensor, torch.Tensor],
+        attention_mask: Optional[torch.Tensor],
+        past_key_values: Optional[PowerAttentionDynamicCache] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+        **kwargs: Unpack[FlashAttentionKwargs],
+    ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
+        input_shape = hidden_states.shape[:-1]
+        hidden_shape = (*input_shape, -1, self.head_dim)
+
+        query_states = self.q_norm(self.q_proj(hidden_states).view(hidden_shape)).transpose(1, 2)
+        key_states = self.k_norm(self.k_proj(hidden_states).view(hidden_shape)).transpose(1, 2)
+        value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+        gate_states = self.g_proj(hidden_states).view(hidden_shape[:-1]).transpose(1, 2)
+        gate_states = nn.functional.logsigmoid(gate_states.to(torch.float32))
+
+        cos, sin = position_embeddings
+        query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)
+
+        if past_key_values is not None:
+            # sin and cos are specific to RoPE models; cache_position needed for the static cache
+            cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position}
+            key_states, value_states, gate_states, state, sum_of_keys = past_key_values.update_cache(key_states, value_states, gate_states, self.layer_idx, cache_kwargs)
+
+        if self.use_exp:
+            attention_interface: Callable = eager_attention_forward
+            if self.config._attn_implementation != "eager":
+                attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation]
+
+            attn_output, attn_weights = attention_interface(
+                self,
+                query_states,
+                key_states,
+                value_states,
+                attention_mask,
+                dropout=0.0 if not self.training else self.attention_dropout,
+                scaling=self.scaling,
+                sliding_window=self.sliding_window,  # diff with Llama
+                **kwargs,
+            )
+        
+        elif query_states.shape[2] == 1:
+            key_len = key_states.shape[2]
+            attn_output, state, sum_of_keys = power_retention_inference(
+                query_states.transpose(1, 2),
+                key_states.transpose(1, 2),
+                value_states.transpose(1, 2),
+                gate_states.transpose(1, 2),
+                initial_state=state,
+                sum_of_keys=sum_of_keys,
+                deg=2,
+                scale=self.scaling,
+                switch_over_seq_len=self.chunk_size,
+            )
+            if self.chunk_size is not None and key_len >= self.chunk_size:
+                past_key_values.clean_cache(self.layer_idx)
+                past_key_values.update_state(self.layer_idx, state, sum_of_keys)
+
+            attn_weights = None
+
+        else:
+            key_len = key_states.shape[2]
+            attn_output = power_retention(
+                query_states.transpose(1, 2),
+                key_states.transpose(1, 2),
+                value_states.transpose(1, 2),
+                gate_states.transpose(1, 2),
+                deg=2,
+                scale=self.scaling,
+                chunk_size=self.prefill_chunk_size, # enable chunked prefilling by default
+                switch_over_seq_len=self.switch_over_seq_len,
+            )
+            attn_weights = None
+
+
+        attn_output = attn_output.reshape(*input_shape, -1).contiguous()
+        attn_output = self.o_proj(attn_output)
+        return attn_output, attn_weights
+
+
+class BrumbyDecoderLayer(GradientCheckpointingLayer):
+    def __init__(self, config: BrumbyConfig, layer_idx: int):
+        super().__init__()
+        self.hidden_size = config.hidden_size
+
+        self.self_attn = BrumbyAttention(config=config, layer_idx=layer_idx)
+
+        self.mlp = BrumbyMLP(config)
+        self.input_layernorm = BrumbyRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.post_attention_layernorm = BrumbyRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.attention_type = config.layer_types[layer_idx]
+
+    @deprecate_kwarg("past_key_value", new_name="past_key_values", version="4.58")
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[Cache] = None,
+        use_cache: Optional[bool] = False,
+        cache_position: Optional[torch.LongTensor] = None,
+        position_embeddings: Optional[tuple[torch.Tensor, torch.Tensor]] = None,  # necessary, but kept here for BC
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> torch.Tensor:
+        residual = hidden_states
+        hidden_states = self.input_layernorm(hidden_states)
+        # Self Attention
+        hidden_states, _ = self.self_attn(
+            hidden_states=hidden_states,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            use_cache=use_cache,
+            cache_position=cache_position,
+            position_embeddings=position_embeddings,
+            **kwargs,
+        )
+        hidden_states = residual + hidden_states
+
+        # Fully Connected
+        residual = hidden_states
+        hidden_states = self.post_attention_layernorm(hidden_states)
+        hidden_states = self.mlp(hidden_states)
+        hidden_states = residual + hidden_states
+        return hidden_states
+
+
+@auto_docstring
+class BrumbyPreTrainedModel(PreTrainedModel):
+    config: BrumbyConfig
+    base_model_prefix = "model"
+    supports_gradient_checkpointing = True
+    _no_split_modules = ["BrumbyDecoderLayer"]
+    _skip_keys_device_placement = ["past_key_values"]
+    _supports_flash_attn = True
+    _supports_sdpa = True
+    _supports_flex_attn = True
+
+    _can_compile_fullgraph = True
+    _supports_attention_backend = True
+    _can_record_outputs = {
+        "hidden_states": BrumbyDecoderLayer,
+        "attentions": BrumbyAttention,
+    }
+
+
+class BrumbyRotaryEmbedding(nn.Module):
+    inv_freq: torch.Tensor  # fix linting for `register_buffer`
+
+    def __init__(self, config: BrumbyConfig, device=None):
+        super().__init__()
+        # BC: "rope_type" was originally "type"
+        if hasattr(config, "rope_scaling") and isinstance(config.rope_scaling, dict):
+            self.rope_type = config.rope_scaling.get("rope_type", config.rope_scaling.get("type"))
+        else:
+            self.rope_type = "default"
+        self.max_seq_len_cached = config.max_position_embeddings
+        self.original_max_seq_len = config.max_position_embeddings
+
+        self.config = config
+        self.rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type]
+
+        inv_freq, self.attention_scaling = self.rope_init_fn(self.config, device)
+        self.register_buffer("inv_freq", inv_freq, persistent=False)
+        self.original_inv_freq = self.inv_freq
+
+    @torch.no_grad()
+    @dynamic_rope_update  # power user: used with advanced RoPE types (e.g. dynamic rope)
+    def forward(self, x, position_ids):
+        inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1).to(x.device)
+        position_ids_expanded = position_ids[:, None, :].float()
+
+        device_type = x.device.type if isinstance(x.device.type, str) and x.device.type != "mps" else "cpu"
+        with torch.autocast(device_type=device_type, enabled=False):  # Force float32
+            freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2)
+            emb = torch.cat((freqs, freqs), dim=-1)
+            cos = emb.cos() * self.attention_scaling
+            sin = emb.sin() * self.attention_scaling
+
+        return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)
+
+
+@auto_docstring
+class BrumbyModel(BrumbyPreTrainedModel):
+    def __init__(self, config: BrumbyConfig):
+        super().__init__(config)
+        self.padding_idx = config.pad_token_id
+        self.vocab_size = config.vocab_size
+
+        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)
+        self.layers = nn.ModuleList(
+            [BrumbyDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]
+        )
+        self.norm = BrumbyRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.rotary_emb = BrumbyRotaryEmbedding(config=config)
+        self.gradient_checkpointing = False
+        self.has_sliding_layers = "sliding_attention" in self.config.layer_types
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    @check_model_inputs
+    @auto_docstring
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[PowerAttentionDynamicCache] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        use_cache: Optional[bool] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> BaseModelOutputWithPast:
+        if (input_ids is None) ^ (inputs_embeds is not None):
+            raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
+
+        if inputs_embeds is None:
+            inputs_embeds = self.embed_tokens(input_ids)
+
+        if use_cache and past_key_values is None:
+            raise ValueError("Brumby requires an initialized `PowerAttentionDynamicCache` to return a cache. None was provided, so no cache will be returned.")
+
+        if cache_position is None:
+            past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
+            cache_position = torch.arange(
+                past_seen_tokens, past_seen_tokens + inputs_embeds.shape[1], device=inputs_embeds.device
+            )
+
+        if position_ids is None:
+            position_ids = cache_position.unsqueeze(0)
+
+        # It may already have been prepared by e.g. `generate`
+        if not isinstance(causal_mask_mapping := attention_mask, dict):
+            # Prepare mask arguments
+            mask_kwargs = {
+                "config": self.config,
+                "input_embeds": inputs_embeds,
+                "attention_mask": attention_mask,
+                "cache_position": cache_position,
+                "past_key_values": past_key_values,
+                "position_ids": position_ids,
+            }
+            # Create the masks
+            causal_mask_mapping = {
+                "full_attention": create_causal_mask(**mask_kwargs),
+            }
+            # The sliding window alternating layers are not always activated depending on the config
+            if self.has_sliding_layers:
+                causal_mask_mapping["sliding_attention"] = create_sliding_window_causal_mask(**mask_kwargs)
+
+        hidden_states = inputs_embeds
+
+        # create position embeddings to be shared across the decoder layers
+        position_embeddings = self.rotary_emb(hidden_states, position_ids)
+
+        for decoder_layer in self.layers[: self.config.num_hidden_layers]:
+            hidden_states = decoder_layer(
+                hidden_states,
+                attention_mask=causal_mask_mapping[decoder_layer.attention_type],
+                position_ids=position_ids,
+                past_key_values=past_key_values,
+                use_cache=use_cache,
+                cache_position=cache_position,
+                position_embeddings=position_embeddings,
+                **kwargs,
+            )
+
+        hidden_states = self.norm(hidden_states)
+        return BaseModelOutputWithPast(
+            last_hidden_state=hidden_states,
+            past_key_values=past_key_values if use_cache else None,
+        )
+
+
+@auto_docstring
+class BrumbyForCausalLM(BrumbyPreTrainedModel, GenerationMixin):
+    _tied_weights_keys = ["lm_head.weight"]
+    _tp_plan = {"lm_head": "colwise_rep"}
+    _pp_plan = {"lm_head": (["hidden_states"], ["logits"])}
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.model = BrumbyModel(config)
+        self.vocab_size = config.vocab_size
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    def prepare_inputs_for_generation(
+        self,
+        input_ids,
+        past_key_values=None,
+        attention_mask=None,
+        inputs_embeds=None,
+        cache_position=None,
+        position_ids=None,
+        use_cache=True,
+        **kwargs,
+    ):
+        # Copy from https://github.com/huggingface/transformers/blob/main/src/transformers/models/jamba/modeling_jamba.py
+        # Overwitten -- uses `past_key_values` as opposed to `past_key_values`
+        empty_past_kv = past_key_values is None
+        all_dynamic_cache = isinstance(past_key_values, DynamicCache)
+        if all_dynamic_cache:
+            past_key_values = PowerAttentionDynamicCache(self.config, input_ids.shape[0], self.dtype, device=self.device)
+
+        # If we have cache: let's slice `input_ids` through `cache_position`, to keep only the unprocessed tokens
+        # Exception 1: when passing input_embeds, input_ids may be missing entries
+        # Exception 2: some generation methods do special slicing of input_ids, so we don't need to do it here
+        # Exception 3: with synced GPUs cache_position may go out of bounds, but we only want dummy token in that case.
+        #              (we can't check exception 3 while compiling)
+        if not empty_past_kv:
+            if (
+                inputs_embeds is not None  # Exception 1
+                or cache_position[-1] >= input_ids.shape[1]  # Exception 3
+            ):
+                input_ids = input_ids[:, -cache_position.shape[0] :]
+            elif input_ids.shape[1] != cache_position.shape[0]:  # Default case (the "else", a no op, is Exception 2)
+                input_ids = input_ids[:, cache_position]
+        else:
+            past_key_values = PowerAttentionDynamicCache(
+                self.config, input_ids.shape[0], self.dtype, device=self.device
+            )
+
+        if attention_mask is not None and position_ids is None:
+            # create position_ids on the fly for batch generation
+            position_ids = attention_mask.long().cumsum(-1) - 1
+            position_ids.masked_fill_(attention_mask == 0, 1)
+            if not empty_past_kv:
+                position_ids = position_ids[:, -input_ids.shape[1] :]
+
+        # if `inputs_embeds` are passed, we only want to use them in the 1st generation step
+        if inputs_embeds is not None and empty_past_kv:
+            # TODO(pjin): workaround fix for properly extending inputs_embeds;
+            # longer term, may be better handled elsewhere in .generate().
+            if input_ids is not None and inputs_embeds.shape[1] < input_ids.shape[1]:
+                new_token_embeds = self.get_input_embeddings()(input_ids[:,inputs_embeds.shape[1]:])
+                inputs_embeds = torch.cat([inputs_embeds, new_token_embeds], dim=1)
+            model_inputs = {"inputs_embeds": inputs_embeds}
+        else:
+            model_inputs = {"input_ids": input_ids.contiguous()}  # `contiguous()` needed for compilation use cases
+
+        model_inputs.update(
+            {
+                "position_ids": position_ids,
+                "past_key_values": past_key_values,
+                "use_cache": use_cache,
+                "attention_mask": attention_mask,
+                "cache_position": cache_position,
+            }
+        )
+        return model_inputs
+
+    @can_return_tuple
+    @auto_docstring
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[PowerAttentionDynamicCache] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        use_cache: Optional[bool] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+        logits_to_keep: Union[int, torch.Tensor] = 0,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> CausalLMOutputWithPast:
+        r"""
+        labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
+            Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,
+            config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored
+            (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.
+
+        Example:
+
+        ```python
+        >>> from transformers import AutoTokenizer, BrumbyForCausalLM
+
+        >>> model = BrumbyForCausalLM.from_pretrained("Qwen/Brumby-8B")
+        >>> tokenizer = AutoTokenizer.from_pretrained("Qwen/Brumby-8B")
+
+        >>> prompt = "Hey, are you conscious? Can you talk to me?"
+        >>> inputs = tokenizer(prompt, return_tensors="pt")
+
+        >>> # Generate
+        >>> generate_ids = model.generate(inputs.input_ids, max_length=30)
+        >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
+        "Hey, are you conscious? Can you talk to me?\nI'm not conscious, but I can talk to you."
+        ```"""
+        outputs: BaseModelOutputWithPast = self.model(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            inputs_embeds=inputs_embeds,
+            use_cache=use_cache,
+            cache_position=cache_position,
+            **kwargs,
+        )
+
+        hidden_states = outputs.last_hidden_state
+        # Only compute necessary logits, and do not upcast them to float if we are not computing the loss
+        slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
+        logits = self.lm_head(hidden_states[:, slice_indices, :])
+
+        loss = None
+        if labels is not None:
+            loss = self.loss_function(logits=logits, labels=labels, vocab_size=self.config.vocab_size, **kwargs)
+
+        return CausalLMOutputWithPast(
+            loss=loss,
+            logits=logits,
+            past_key_values=outputs.past_key_values,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )
+
+
+class BrumbyForSequenceClassification(GenericForSequenceClassification, BrumbyPreTrainedModel):
+    pass
+
+
+class BrumbyForTokenClassification(GenericForTokenClassification, BrumbyPreTrainedModel):
+    pass
+
+
+class BrumbyForQuestionAnswering(GenericForQuestionAnswering, BrumbyPreTrainedModel):
+    base_model_prefix = "transformer"  # For BC, where `transformer` was used instead of `model`
+
+
+__all__ = [
+    "BrumbyForCausalLM",
+    "BrumbyForQuestionAnswering",
+    "BrumbyPreTrainedModel",
+    "BrumbyModel",
+    "BrumbyForSequenceClassification",
+    "BrumbyForTokenClassification",
+]

tokenizer_config.json

@@ -0,0 +1,239 @@
+{
+  "add_bos_token": false,
+  "add_prefix_space": false,
+  "added_tokens_decoder": {
+    "151643": {
+      "content": "<|endoftext|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151644": {
+      "content": "<|im_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151645": {
+      "content": "<|im_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151646": {
+      "content": "<|object_ref_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151647": {
+      "content": "<|object_ref_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151648": {
+      "content": "<|box_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151649": {
+      "content": "<|box_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151650": {
+      "content": "<|quad_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151651": {
+      "content": "<|quad_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151652": {
+      "content": "<|vision_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151653": {
+      "content": "<|vision_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151654": {
+      "content": "<|vision_pad|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151655": {
+      "content": "<|image_pad|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151656": {
+      "content": "<|video_pad|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151657": {
+      "content": "<tool_call>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151658": {
+      "content": "</tool_call>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151659": {
+      "content": "<|fim_prefix|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151660": {
+      "content": "<|fim_middle|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151661": {
+      "content": "<|fim_suffix|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151662": {
+      "content": "<|fim_pad|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151663": {
+      "content": "<|repo_name|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151664": {
+      "content": "<|file_sep|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151665": {
+      "content": "<tool_response>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151666": {
+      "content": "</tool_response>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151667": {
+      "content": "<think>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151668": {
+      "content": "</think>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    }
+  },
+  "additional_special_tokens": [
+    "<|im_start|>",
+    "<|im_end|>",
+    "<|object_ref_start|>",
+    "<|object_ref_end|>",
+    "<|box_start|>",
+    "<|box_end|>",
+    "<|quad_start|>",
+    "<|quad_end|>",
+    "<|vision_start|>",
+    "<|vision_end|>",
+    "<|vision_pad|>",
+    "<|image_pad|>",
+    "<|video_pad|>"
+  ],
+  "bos_token": null,
+  "chat_template": "{%- if tools %}\n    {{- '<|im_start|>system\\n' }}\n    {%- if messages[0].role == 'system' %}\n        {{- messages[0].content + '\\n\\n' }}\n    {%- endif %}\n    {{- \"# Tools\\n\\nYou may call one or more functions to assist with the user query.\\n\\nYou are provided with function signatures within <tools></tools> XML tags:\\n<tools>\" }}\n    {%- for tool in tools %}\n        {{- \"\\n\" }}\n        {{- tool | tojson }}\n    {%- endfor %}\n    {{- \"\\n</tools>\\n\\nFor each function call, return a json object with function name and arguments within <tool_call></tool_call> XML tags:\\n<tool_call>\\n{\\\"name\\\": <function-name>, \\\"arguments\\\": <args-json-object>}\\n</tool_call><|im_end|>\\n\" }}\n{%- else %}\n    {%- if messages[0].role == 'system' %}\n        {{- '<|im_start|>system\\n' + messages[0].content + '<|im_end|>\\n' }}\n    {%- endif %}\n{%- endif %}\n{%- set ns = namespace(multi_step_tool=true, last_query_index=messages|length - 1) %}\n{%- for message in messages[::-1] %}\n    {%- set index = (messages|length - 1) - loop.index0 %}\n    {%- if ns.multi_step_tool and message.role == \"user\" and not(message.content.startswith('<tool_response>') and message.content.endswith('</tool_response>')) %}\n        {%- set ns.multi_step_tool = false %}\n        {%- set ns.last_query_index = index %}\n    {%- endif %}\n{%- endfor %}\n{%- for message in messages %}\n    {%- if (message.role == \"user\") or (message.role == \"system\" and not loop.first) %}\n        {{- '<|im_start|>' + message.role + '\\n' + message.content + '<|im_end|>' + '\\n' }}\n    {%- elif message.role == \"assistant\" %}\n        {%- set content = message.content %}\n        {%- set reasoning_content = '' %}\n        {%- if message.reasoning_content is defined and message.reasoning_content is not none %}\n            {%- set reasoning_content = message.reasoning_content %}\n        {%- else %}\n            {%- if '</think>' in message.content %}\n                {%- set content = message.content.split('</think>')[-1].lstrip('\\n') %}\n                {%- set reasoning_content = message.content.split('</think>')[0].rstrip('\\n').split('<think>')[-1].lstrip('\\n') %}\n            {%- endif %}\n        {%- endif %}\n        {%- if loop.index0 > ns.last_query_index %}\n            {%- if loop.last or (not loop.last and reasoning_content) %}\n                {{- '<|im_start|>' + message.role + '\\n<think>\\n' + reasoning_content.strip('\\n') + '\\n</think>\\n\\n' + content.lstrip('\\n') }}\n            {%- else %}\n                {{- '<|im_start|>' + message.role + '\\n' + content }}\n            {%- endif %}\n        {%- else %}\n            {{- '<|im_start|>' + message.role + '\\n' + content }}\n        {%- endif %}\n        {%- if message.tool_calls %}\n            {%- for tool_call in message.tool_calls %}\n                {%- if (loop.first and content) or (not loop.first) %}\n                    {{- '\\n' }}\n                {%- endif %}\n                {%- if tool_call.function %}\n                    {%- set tool_call = tool_call.function %}\n                {%- endif %}\n                {{- '<tool_call>\\n{\"name\": \"' }}\n                {{- tool_call.name }}\n                {{- '\", \"arguments\": ' }}\n                {%- if tool_call.arguments is string %}\n                    {{- tool_call.arguments }}\n                {%- else %}\n                    {{- tool_call.arguments | tojson }}\n                {%- endif %}\n                {{- '}\\n</tool_call>' }}\n            {%- endfor %}\n        {%- endif %}\n        {{- '<|im_end|>\\n' }}\n    {%- elif message.role == \"tool\" %}\n        {%- if loop.first or (messages[loop.index0 - 1].role != \"tool\") %}\n            {{- '<|im_start|>user' }}\n        {%- endif %}\n        {{- '\\n<tool_response>\\n' }}\n        {{- message.content }}\n        {{- '\\n</tool_response>' }}\n        {%- if loop.last or (messages[loop.index0 + 1].role != \"tool\") %}\n            {{- '<|im_end|>\\n' }}\n        {%- endif %}\n    {%- endif %}\n{%- endfor %}\n{%- if add_generation_prompt %}\n    {{- '<|im_start|>assistant\\n' }}\n    {%- if enable_thinking is defined and enable_thinking is false %}\n        {{- '<think>\\n\\n</think>\\n\\n' }}\n    {%- endif %}\n{%- endif %}",
+  "clean_up_tokenization_spaces": false,
+  "eos_token": "<|endoftext|>",
+  "errors": "replace",
+  "model_max_length": 131072,
+  "pad_token": "<|endoftext|>",
+  "split_special_tokens": false,
+  "tokenizer_class": "Qwen2Tokenizer",
+  "unk_token": null
+}