Initial upload of RishAI-Base-v2: Sparse MoE multilingual model

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+tokenizer.json filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,155 @@
+# Rish AI
+
+## Model Description
+
+Rish AI is a cutting-edge Mixture of Experts (MoE) transformer model designed for efficient and scalable language understanding and generation. It features sparse routing with 7 experts per token, advanced rotary position embeddings, and optimized attention mechanisms.
+
+## Key Features
+
+- **Sparse Mixture of Experts**: 7 experts with 5 experts activated per token for optimal efficiency
+- **Rotary Position Embeddings**: Dynamic RoPE scaling for better long-context handling
+- **Grouped Query Attention**: Efficient attention with reduced key/value heads
+- **RMSNorm**: Improved normalization for stable training
+- **Load Balancing**: Automatic expert load balancing during training
+
+## Usage
+
+### Installation
+
+```bash
+pip install transformers
+```
+
+### Basic Usage
+
+```python
+from transformers import AutoTokenizer, AutoModelForCausalLM
+
+# Load model and tokenizer
+model_name = "your-org/RishAI-1B-7B"
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+model = AutoModelForCausalLM.from_pretrained(model_name)
+
+# Prepare input
+text = "Hello, how are you?"
+inputs = tokenizer(text, return_tensors="pt")
+
+# Generate response
+outputs = model.generate(**inputs, max_length=50, do_sample=True, temperature=0.7)
+response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+print(response)
+```
+
+### Advanced Usage
+
+```python
+import torch
+from transformers import AutoTokenizer, AutoModelForCausalLM
+
+# Load model with specific configuration
+model = AutoModelForCausalLM.from_pretrained(
+    "your-org/RishAI-1B-7B",
+    torch_dtype=torch.bfloat16,  # For memory efficiency
+    device_map="auto"  # Automatic device placement
+)
+
+tokenizer = AutoTokenizer.from_pretrained("your-org/RishAI-1B-7B")
+
+# Multi-turn conversation
+conversation = [
+    {"role": "user", "content": "What is machine learning?"},
+    {"role": "assistant", "content": "Machine learning is a subset of AI..."},
+    {"role": "user", "content": "Can you give a practical example?"}
+]
+
+# Format conversation
+formatted_input = tokenizer.apply_chat_template(conversation, tokenize=False)
+inputs = tokenizer(formatted_input, return_tensors="pt")
+
+# Generate with controlled parameters
+outputs = model.generate(
+    **inputs,
+    max_length=200,
+    temperature=0.8,
+    top_p=0.9,
+    do_sample=True,
+    pad_token_id=tokenizer.eos_token_id
+)
+
+response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+print(response)
+```
+
+### Model Configuration
+
+```python
+from transformers import RishAIConfig
+
+# Create custom configuration
+config = RishAIConfig(
+    vocab_size=100352,
+    hidden_size=4096,
+    num_hidden_layers=32,
+    num_attention_heads=32,
+    num_experts=7,           # Number of experts
+    num_experts_per_tok=5,   # Experts activated per token
+    max_position_embeddings=4096,
+    rope_scaling={"rope_type": "dynamic", "factor": 1.0}
+)
+
+# Initialize model with config
+from transformers import RishAIModel
+model = RishAIModel(config)
+```
+
+## Model Architecture
+
+### Sparse Mixture of Experts (MoE)
+- **Experts**: 7 specialized sub-networks
+- **Routing**: Top-5 expert selection per token
+- **Load Balancing**: Automatic expert utilization optimization
+
+### Attention Mechanism
+- **Grouped Query Attention**: Efficient key/value head reduction
+- **Rotary Embeddings**: Position-aware attention with dynamic scaling
+- **RMSNorm**: Stable layer normalization
+
+### Training Features
+- **Gradient Checkpointing**: Memory-efficient training
+- **Flash Attention**: Optimized attention computation
+- **Expert Parallelism**: Distributed expert training
+
+## Performance
+
+### Speed
+- **Inference**: Optimized for fast generation
+- **Training**: Efficient MoE routing and load balancing
+- **Memory**: Sparse activation reduces memory footprint
+
+### Quality
+- **Perplexity**: Competitive with state-of-the-art models
+- **Long Context**: Effective handling of 4K+ token sequences
+- **Multitask**: Strong performance across diverse tasks
+
+## Limitations
+
+- Requires significant computational resources for training
+- Memory usage scales with number of active experts
+- Best performance on modern GPUs with ample VRAM
+
+## Citation
+
+```bibtex
+@misc{rishailabs_2026,
+    author       = { RishAILabs },
+    title        = { RLLM-Base (Revision 552ee30) },
+    year         = 2026,
+    url          = { https://huggingface.co/RishAILabs/RLLM-Base },
+    doi          = { 10.57967/hf/7560 },
+    publisher    = { Hugging Face }
+}
+```
+
+## License
+
+This model is released under the Apache 2.0 license.

__init__.py

@@ -0,0 +1,75 @@
+# Copyright 2024 The HuggingFace Team. 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.
+
+from typing import TYPE_CHECKING
+
+from transformers.utils import (
+    OptionalDependencyNotAvailable,
+    _LazyModule,
+    is_tokenizers_available,
+    is_torch_available,
+)
+
+
+_import_structure = {
+    "configuration_rish_ai": ["RishAIConfig"],
+}
+
+try:
+    if not is_torch_available():
+        raise OptionalDependencyNotAvailable()
+except OptionalDependencyNotAvailable:
+    pass
+else:
+    _import_structure["modeling_rish_ai"] = [
+        "RishAICausalLM",
+        "RishAIModel",
+        "RishAIPreTrainedModel",
+    ]
+
+try:
+    if not is_tokenizers_available():
+        raise OptionalDependencyNotAvailable()
+except OptionalDependencyNotAvailable:
+    pass
+else:
+    _import_structure["tokenization_rish_ai"] = ["RishAITokenizer"]
+
+if TYPE_CHECKING:
+    from .configuration_rish_ai import RishAIConfig
+
+    try:
+        if not is_torch_available():
+            raise OptionalDependencyNotAvailable()
+    except OptionalDependencyNotAvailable:
+        pass
+    else:
+        from .modeling_rish_ai import (
+            RishAICausalLM,
+            RishAIModel,
+            RishAIPreTrainedModel,
+        )
+
+    try:
+        if not is_tokenizers_available():
+            raise OptionalDependencyNotAvailable()
+    except OptionalDependencyNotAvailable:
+        pass
+    else:
+        from .tokenization_rish_ai import RishAITokenizer
+
+else:
+    import sys
+
+    sys.modules[__name__] = _LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)

chat_template.jinja

@@ -0,0 +1 @@
+{% if messages[0]['role'] == 'system' %}{% set loop_messages = messages[1:] %}{% set system_message = messages[0]['content'] %}{% else %}{% set loop_messages = messages %}{% set system_message = 'You are Rish AI, a powerful large language model built by Rish AI Labs. You are helpful, creative, accurate, and always aim to provide the best possible assistance.' %}{% endif %}{{ '<|im_start|>system\n' + system_message + '<|im_end|>\n' }}{% for message in loop_messages %}{% if (message['role'] == 'user') != (loop.index0 % 2 == 0) %}{{ raise_exception('Conversation roles must alternate user/assistant/user/assistant/...') }}{% endif %}{{ '<|im_start|>' + message['role'] + '\n' + message['content'] + '<|im_end|>\n' }}{% endfor %}{% if add_generation_prompt %}{{ '<|im_start|>assistant\n' }}{% endif %}

config.json

@@ -0,0 +1,31 @@
+{
+  "attention_bias": false,
+  "attention_dropout": 0.0,
+  "eos_token_id": 151645,
+  "hidden_act": "silu",
+  "hidden_size": 2048,
+  "initializer_range": 0.02,
+  "intermediate_size": 4864,
+  "max_position_embeddings": 4096,
+  "model_type": "rish_ai",
+  "norm_topk_prob": false,
+  "num_attention_heads": 16,
+  "num_experts": 7,
+  "num_experts_per_tok": 5,
+  "num_hidden_layers": 24,
+  "num_key_value_heads": 16,
+  "output_router_logits": false,
+  "pad_token_id": 151643,
+  "rms_norm_eps": 1e-06,
+  "rope_parameters": {
+    "factor": 1.0,
+    "rope_theta": 500000.0,
+    "rope_type": "dynamic"
+  },
+  "rope_theta": 500000.0,
+  "router_aux_loss_coef": 0.01,
+  "tie_word_embeddings": false,
+  "transformers_version": "5.0.0rc1",
+  "use_cache": true,
+  "vocab_size": 151680
+}

configuration_rish_ai.py

@@ -0,0 +1,169 @@
+# Copyright 2024 The HuggingFace Team. 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.
+
+from transformers.configuration_utils import PretrainedConfig
+
+
+class RishAIConfig(PretrainedConfig):
+    r"""
+    Configuration class for RishAI models.
+
+    Args:
+        vocab_size (`int`, *optional*, defaults to 100352):
+            Vocabulary size of the RishAI model. Defines the number of different tokens that can be represented by the
+            `inputs_ids` passed when calling [`RishAIModel`]
+        hidden_size (`int`, *optional*, defaults to 4096):
+            Dimension of the hidden representations.
+        intermediate_size (`int`, *optional*, defaults to 11008):
+            Dimension of the MLP representations.
+        num_hidden_layers (`int`, *optional*, defaults to 32):
+            Number of hidden layers in the Transformer decoder.
+        num_attention_heads (`int`, *optional*, defaults to 32):
+            Number of attention heads for each attention layer in the Transformer decoder.
+        num_key_value_heads (`int`, *optional*):
+            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
+            `num_attention_heads`.
+        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 4096):
+            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`.
+        pad_token_id (`int`, *optional*, defaults to 100277):
+            Padding token id.
+        bos_token_id (`int`, *optional*):
+            Beginning of stream token id.
+        eos_token_id (`int`, *optional*, defaults to 100257):
+            End of stream token id.
+        tie_word_embeddings (`bool`, *optional*, defaults to `False`):
+            Whether to tie weight embeddings
+        rope_theta (`float`, *optional*, defaults to 500000.0):
+            The base period of the RoPE embeddings.
+        rope_scaling (`Dict`, *optional*):
+            Dictionary containing the scaling configuration for the RoPE embeddings. Currently supports two scaling
+            strategies: linear and dynamic. Their scaling factor must be a float greater than 1. The expected format is
+            `{"type": strategy name, "factor": scaling factor}`. When using this flag, don't update
+            `max_position_embeddings` to the expected new maximum. See the following thread for more information on how
+            these scaling strategies behave:
+            https://www.reddit.com/r/LocalLLaMA/comments/14mrgpr/dynamically_scaled_rope_further_increases/. This is an
+            experimental feature, subject to breaking API changes in future versions.
+        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.
+        attention_dropout (`float`, *optional*, defaults to 0.0):
+            The dropout ratio for the attention probabilities.
+        num_experts_per_tok (`int`, *optional*, defaults to 5):
+            Number of selected experts.
+        num_experts (`int`, *optional*, defaults to 7):
+            Number of routed experts.
+        output_router_logits (`bool`, *optional*, defaults to `False`):
+            Whether or not the router logits should be returned by the model. Enabling this will also
+            allow the model to output the auxiliary loss, including load balancing loss and router z-loss.
+        router_aux_loss_coef (`float`, *optional*, defaults to 0.01):
+            The aux loss factor for the total loss.
+        norm_topk_prob (`bool`, *optional*, defaults to `False`):
+            Whether to normalize the topk probabilities.
+
+    Example:
+    ```python
+    >>> from transformers import RishAIConfig, RishAIModel
+
+    >>> # Initializing a RishAI rish_ai style configuration
+    >>> configuration = RishAIConfig()
+
+    >>> # Initializing a model from the RishAI style configuration
+    >>> model = RishAIModel(configuration)
+
+    >>> # Accessing the model configuration
+    >>> configuration = model.config
+    ```
+    """
+
+    model_type = "rish_ai"
+    keys_to_ignore_at_inference = ["past_key_values"]
+
+    def __init__(
+        self,
+        vocab_size=100352,
+        hidden_size=4096,
+        intermediate_size=11008,
+        num_hidden_layers=32,
+        num_attention_heads=32,
+        num_key_value_heads=None,
+        hidden_act="silu",
+        max_position_embeddings=4096,
+        initializer_range=0.02,
+        rms_norm_eps=1e-06,
+        use_cache=True,
+        pad_token_id=100277,
+        bos_token_id=None,
+        eos_token_id=100257,
+        tie_word_embeddings=False,
+        rope_theta=500000.0,
+        rope_scaling=None,
+        attention_bias=False,
+        attention_dropout=0.0,
+        num_experts_per_tok=5,
+        num_experts=7,
+        output_router_logits=False,
+        router_aux_loss_coef=0.01,
+        norm_topk_prob=False,
+        **kwargs,
+    ):
+        super().__init__(
+            pad_token_id=pad_token_id,
+            bos_token_id=bos_token_id,
+            eos_token_id=eos_token_id,
+            tie_word_embeddings=tie_word_embeddings,
+            **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
+
+        # 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.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
+        self.num_experts_per_tok = num_experts_per_tok
+        self.num_experts = num_experts
+        self.output_router_logits = output_router_logits
+        self.router_aux_loss_coef = router_aux_loss_coef
+        self.norm_topk_prob = norm_topk_prob
+
+        # 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"]

model-00001.safetensors

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

model-00002.safetensors

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

model-00003.safetensors

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

model-00004.safetensors

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

model-00005.safetensors

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

model-00006.safetensors

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

model-00007.safetensors

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

model-00008.safetensors

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

model-00009.safetensors

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

model-00010.safetensors

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

model-00011.safetensors

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

model-00012.safetensors

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

model-00013.safetensors

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

model-00014.safetensors

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

model-00015.safetensors

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

model-00016.safetensors

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

model-00017.safetensors

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

model-00018.safetensors

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

model-00019.safetensors

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

model-00020.safetensors

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

model-00021.safetensors

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

model-00022.safetensors

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

model.safetensors.index.json

@@ -0,0 +1,730 @@
+{
+  "metadata": {
+    "total_size": 12090118144
+  },
+  "weight_map": {
+    "model.embed_tokens.weight": "model-00001.safetensors",
+    "model.layers.0.self_attn.q_proj.weight": "model-00002.safetensors",
+    "model.layers.0.self_attn.k_proj.weight": "model-00002.safetensors",
+    "model.layers.0.self_attn.v_proj.weight": "model-00002.safetensors",
+    "model.layers.0.self_attn.o_proj.weight": "model-00002.safetensors",
+    "model.layers.0.self_attn.q_norm.weight": "model-00002.safetensors",
+    "model.layers.0.self_attn.k_norm.weight": "model-00002.safetensors",
+    "model.layers.0.mlp.gate.weight": "model-00002.safetensors",
+    "model.layers.0.mlp.experts.0.gate_proj.weight": "model-00002.safetensors",
+    "model.layers.0.mlp.experts.0.up_proj.weight": "model-00002.safetensors",
+    "model.layers.0.mlp.experts.0.down_proj.weight": "model-00002.safetensors",
+    "model.layers.0.mlp.experts.1.gate_proj.weight": "model-00002.safetensors",
+    "model.layers.0.mlp.experts.1.up_proj.weight": "model-00002.safetensors",
+    "model.layers.0.mlp.experts.1.down_proj.weight": "model-00002.safetensors",
+    "model.layers.0.mlp.experts.2.gate_proj.weight": "model-00002.safetensors",
+    "model.layers.0.mlp.experts.2.up_proj.weight": "model-00002.safetensors",
+    "model.layers.0.mlp.experts.2.down_proj.weight": "model-00002.safetensors",
+    "model.layers.0.mlp.experts.3.gate_proj.weight": "model-00002.safetensors",
+    "model.layers.0.mlp.experts.3.up_proj.weight": "model-00002.safetensors",
+    "model.layers.0.mlp.experts.3.down_proj.weight": "model-00002.safetensors",
+    "model.layers.0.mlp.experts.4.gate_proj.weight": "model-00002.safetensors",
+    "model.layers.0.mlp.experts.4.up_proj.weight": "model-00002.safetensors",
+    "model.layers.0.mlp.experts.4.down_proj.weight": "model-00002.safetensors",
+    "model.layers.0.mlp.experts.5.gate_proj.weight": "model-00002.safetensors",
+    "model.layers.0.mlp.experts.5.up_proj.weight": "model-00002.safetensors",
+    "model.layers.0.mlp.experts.5.down_proj.weight": "model-00002.safetensors",
+    "model.layers.0.mlp.experts.6.gate_proj.weight": "model-00002.safetensors",
+    "model.layers.0.mlp.experts.6.up_proj.weight": "model-00002.safetensors",
+    "model.layers.0.mlp.experts.6.down_proj.weight": "model-00002.safetensors",
+    "model.layers.0.post_attention_layernorm.weight": "model-00002.safetensors",
+    "model.layers.0.post_feedforward_layernorm.weight": "model-00002.safetensors",
+    "model.layers.1.self_attn.q_proj.weight": "model-00002.safetensors",
+    "model.layers.1.self_attn.k_proj.weight": "model-00002.safetensors",
+    "model.layers.1.self_attn.v_proj.weight": "model-00002.safetensors",
+    "model.layers.1.self_attn.o_proj.weight": "model-00002.safetensors",
+    "model.layers.1.self_attn.q_norm.weight": "model-00002.safetensors",
+    "model.layers.1.self_attn.k_norm.weight": "model-00002.safetensors",
+    "model.layers.1.mlp.gate.weight": "model-00002.safetensors",
+    "model.layers.1.mlp.experts.0.gate_proj.weight": "model-00002.safetensors",
+    "model.layers.1.mlp.experts.0.up_proj.weight": "model-00002.safetensors",
+    "model.layers.1.mlp.experts.0.down_proj.weight": "model-00003.safetensors",
+    "model.layers.1.mlp.experts.1.gate_proj.weight": "model-00003.safetensors",
+    "model.layers.1.mlp.experts.1.up_proj.weight": "model-00003.safetensors",
+    "model.layers.1.mlp.experts.1.down_proj.weight": "model-00003.safetensors",
+    "model.layers.1.mlp.experts.2.gate_proj.weight": "model-00003.safetensors",
+    "model.layers.1.mlp.experts.2.up_proj.weight": "model-00003.safetensors",
+    "model.layers.1.mlp.experts.2.down_proj.weight": "model-00003.safetensors",
+    "model.layers.1.mlp.experts.3.gate_proj.weight": "model-00003.safetensors",
+    "model.layers.1.mlp.experts.3.up_proj.weight": "model-00003.safetensors",
+    "model.layers.1.mlp.experts.3.down_proj.weight": "model-00003.safetensors",
+    "model.layers.1.mlp.experts.4.gate_proj.weight": "model-00003.safetensors",
+    "model.layers.1.mlp.experts.4.up_proj.weight": "model-00003.safetensors",
+    "model.layers.1.mlp.experts.4.down_proj.weight": "model-00003.safetensors",
+    "model.layers.1.mlp.experts.5.gate_proj.weight": "model-00003.safetensors",
+    "model.layers.1.mlp.experts.5.up_proj.weight": "model-00003.safetensors",
+    "model.layers.1.mlp.experts.5.down_proj.weight": "model-00003.safetensors",
+    "model.layers.1.mlp.experts.6.gate_proj.weight": "model-00003.safetensors",
+    "model.layers.1.mlp.experts.6.up_proj.weight": "model-00003.safetensors",
+    "model.layers.1.mlp.experts.6.down_proj.weight": "model-00003.safetensors",
+    "model.layers.1.post_attention_layernorm.weight": "model-00003.safetensors",
+    "model.layers.1.post_feedforward_layernorm.weight": "model-00003.safetensors",
+    "model.layers.2.self_attn.q_proj.weight": "model-00003.safetensors",
+    "model.layers.2.self_attn.k_proj.weight": "model-00003.safetensors",
+    "model.layers.2.self_attn.v_proj.weight": "model-00003.safetensors",
+    "model.layers.2.self_attn.o_proj.weight": "model-00003.safetensors",
+    "model.layers.2.self_attn.q_norm.weight": "model-00003.safetensors",
+    "model.layers.2.self_attn.k_norm.weight": "model-00003.safetensors",
+    "model.layers.2.mlp.gate.weight": "model-00003.safetensors",
+    "model.layers.2.mlp.experts.0.gate_proj.weight": "model-00003.safetensors",
+    "model.layers.2.mlp.experts.0.up_proj.weight": "model-00003.safetensors",
+    "model.layers.2.mlp.experts.0.down_proj.weight": "model-00003.safetensors",
+    "model.layers.2.mlp.experts.1.gate_proj.weight": "model-00003.safetensors",
+    "model.layers.2.mlp.experts.1.up_proj.weight": "model-00003.safetensors",
+    "model.layers.2.mlp.experts.1.down_proj.weight": "model-00003.safetensors",
+    "model.layers.2.mlp.experts.2.gate_proj.weight": "model-00004.safetensors",
+    "model.layers.2.mlp.experts.2.up_proj.weight": "model-00004.safetensors",
+    "model.layers.2.mlp.experts.2.down_proj.weight": "model-00004.safetensors",
+    "model.layers.2.mlp.experts.3.gate_proj.weight": "model-00004.safetensors",
+    "model.layers.2.mlp.experts.3.up_proj.weight": "model-00004.safetensors",
+    "model.layers.2.mlp.experts.3.down_proj.weight": "model-00004.safetensors",
+    "model.layers.2.mlp.experts.4.gate_proj.weight": "model-00004.safetensors",
+    "model.layers.2.mlp.experts.4.up_proj.weight": "model-00004.safetensors",
+    "model.layers.2.mlp.experts.4.down_proj.weight": "model-00004.safetensors",
+    "model.layers.2.mlp.experts.5.gate_proj.weight": "model-00004.safetensors",
+    "model.layers.2.mlp.experts.5.up_proj.weight": "model-00004.safetensors",
+    "model.layers.2.mlp.experts.5.down_proj.weight": "model-00004.safetensors",
+    "model.layers.2.mlp.experts.6.gate_proj.weight": "model-00004.safetensors",
+    "model.layers.2.mlp.experts.6.up_proj.weight": "model-00004.safetensors",
+    "model.layers.2.mlp.experts.6.down_proj.weight": "model-00004.safetensors",
+    "model.layers.2.post_attention_layernorm.weight": "model-00004.safetensors",
+    "model.layers.2.post_feedforward_layernorm.weight": "model-00004.safetensors",
+    "model.layers.3.self_attn.q_proj.weight": "model-00004.safetensors",
+    "model.layers.3.self_attn.k_proj.weight": "model-00004.safetensors",
+    "model.layers.3.self_attn.v_proj.weight": "model-00004.safetensors",
+    "model.layers.3.self_attn.o_proj.weight": "model-00004.safetensors",
+    "model.layers.3.self_attn.q_norm.weight": "model-00004.safetensors",
+    "model.layers.3.self_attn.k_norm.weight": "model-00004.safetensors",
+    "model.layers.3.mlp.gate.weight": "model-00004.safetensors",
+    "model.layers.3.mlp.experts.0.gate_proj.weight": "model-00004.safetensors",
+    "model.layers.3.mlp.experts.0.up_proj.weight": "model-00004.safetensors",
+    "model.layers.3.mlp.experts.0.down_proj.weight": "model-00004.safetensors",
+    "model.layers.3.mlp.experts.1.gate_proj.weight": "model-00004.safetensors",
+    "model.layers.3.mlp.experts.1.up_proj.weight": "model-00004.safetensors",
+    "model.layers.3.mlp.experts.1.down_proj.weight": "model-00004.safetensors",
+    "model.layers.3.mlp.experts.2.gate_proj.weight": "model-00004.safetensors",
+    "model.layers.3.mlp.experts.2.up_proj.weight": "model-00004.safetensors",
+    "model.layers.3.mlp.experts.2.down_proj.weight": "model-00004.safetensors",
+    "model.layers.3.mlp.experts.3.gate_proj.weight": "model-00004.safetensors",
+    "model.layers.3.mlp.experts.3.up_proj.weight": "model-00005.safetensors",
+    "model.layers.3.mlp.experts.3.down_proj.weight": "model-00005.safetensors",
+    "model.layers.3.mlp.experts.4.gate_proj.weight": "model-00005.safetensors",
+    "model.layers.3.mlp.experts.4.up_proj.weight": "model-00005.safetensors",
+    "model.layers.3.mlp.experts.4.down_proj.weight": "model-00005.safetensors",
+    "model.layers.3.mlp.experts.5.gate_proj.weight": "model-00005.safetensors",
+    "model.layers.3.mlp.experts.5.up_proj.weight": "model-00005.safetensors",
+    "model.layers.3.mlp.experts.5.down_proj.weight": "model-00005.safetensors",
+    "model.layers.3.mlp.experts.6.gate_proj.weight": "model-00005.safetensors",
+    "model.layers.3.mlp.experts.6.up_proj.weight": "model-00005.safetensors",
+    "model.layers.3.mlp.experts.6.down_proj.weight": "model-00005.safetensors",
+    "model.layers.3.post_attention_layernorm.weight": "model-00005.safetensors",
+    "model.layers.3.post_feedforward_layernorm.weight": "model-00005.safetensors",
+    "model.layers.4.self_attn.q_proj.weight": "model-00005.safetensors",
+    "model.layers.4.self_attn.k_proj.weight": "model-00005.safetensors",
+    "model.layers.4.self_attn.v_proj.weight": "model-00005.safetensors",
+    "model.layers.4.self_attn.o_proj.weight": "model-00005.safetensors",
+    "model.layers.4.self_attn.q_norm.weight": "model-00005.safetensors",
+    "model.layers.4.self_attn.k_norm.weight": "model-00005.safetensors",
+    "model.layers.4.mlp.gate.weight": "model-00005.safetensors",
+    "model.layers.4.mlp.experts.0.gate_proj.weight": "model-00005.safetensors",
+    "model.layers.4.mlp.experts.0.up_proj.weight": "model-00005.safetensors",
+    "model.layers.4.mlp.experts.0.down_proj.weight": "model-00005.safetensors",
+    "model.layers.4.mlp.experts.1.gate_proj.weight": "model-00005.safetensors",
+    "model.layers.4.mlp.experts.1.up_proj.weight": "model-00005.safetensors",
+    "model.layers.4.mlp.experts.1.down_proj.weight": "model-00005.safetensors",
+    "model.layers.4.mlp.experts.2.gate_proj.weight": "model-00005.safetensors",
+    "model.layers.4.mlp.experts.2.up_proj.weight": "model-00005.safetensors",
+    "model.layers.4.mlp.experts.2.down_proj.weight": "model-00005.safetensors",
+    "model.layers.4.mlp.experts.3.gate_proj.weight": "model-00005.safetensors",
+    "model.layers.4.mlp.experts.3.up_proj.weight": "model-00005.safetensors",
+    "model.layers.4.mlp.experts.3.down_proj.weight": "model-00005.safetensors",
+    "model.layers.4.mlp.experts.4.gate_proj.weight": "model-00005.safetensors",
+    "model.layers.4.mlp.experts.4.up_proj.weight": "model-00005.safetensors",
+    "model.layers.4.mlp.experts.4.down_proj.weight": "model-00006.safetensors",
+    "model.layers.4.mlp.experts.5.gate_proj.weight": "model-00006.safetensors",
+    "model.layers.4.mlp.experts.5.up_proj.weight": "model-00006.safetensors",
+    "model.layers.4.mlp.experts.5.down_proj.weight": "model-00006.safetensors",
+    "model.layers.4.mlp.experts.6.gate_proj.weight": "model-00006.safetensors",
+    "model.layers.4.mlp.experts.6.up_proj.weight": "model-00006.safetensors",
+    "model.layers.4.mlp.experts.6.down_proj.weight": "model-00006.safetensors",
+    "model.layers.4.post_attention_layernorm.weight": "model-00006.safetensors",
+    "model.layers.4.post_feedforward_layernorm.weight": "model-00006.safetensors",
+    "model.layers.5.self_attn.q_proj.weight": "model-00006.safetensors",
+    "model.layers.5.self_attn.k_proj.weight": "model-00006.safetensors",
+    "model.layers.5.self_attn.v_proj.weight": "model-00006.safetensors",
+    "model.layers.5.self_attn.o_proj.weight": "model-00006.safetensors",
+    "model.layers.5.self_attn.q_norm.weight": "model-00006.safetensors",
+    "model.layers.5.self_attn.k_norm.weight": "model-00006.safetensors",
+    "model.layers.5.mlp.gate.weight": "model-00006.safetensors",
+    "model.layers.5.mlp.experts.0.gate_proj.weight": "model-00006.safetensors",
+    "model.layers.5.mlp.experts.0.up_proj.weight": "model-00006.safetensors",
+    "model.layers.5.mlp.experts.0.down_proj.weight": "model-00006.safetensors",
+    "model.layers.5.mlp.experts.1.gate_proj.weight": "model-00006.safetensors",
+    "model.layers.5.mlp.experts.1.up_proj.weight": "model-00006.safetensors",
+    "model.layers.5.mlp.experts.1.down_proj.weight": "model-00006.safetensors",
+    "model.layers.5.mlp.experts.2.gate_proj.weight": "model-00006.safetensors",
+    "model.layers.5.mlp.experts.2.up_proj.weight": "model-00006.safetensors",
+    "model.layers.5.mlp.experts.2.down_proj.weight": "model-00006.safetensors",
+    "model.layers.5.mlp.experts.3.gate_proj.weight": "model-00006.safetensors",
+    "model.layers.5.mlp.experts.3.up_proj.weight": "model-00006.safetensors",
+    "model.layers.5.mlp.experts.3.down_proj.weight": "model-00006.safetensors",
+    "model.layers.5.mlp.experts.4.gate_proj.weight": "model-00006.safetensors",
+    "model.layers.5.mlp.experts.4.up_proj.weight": "model-00006.safetensors",
+    "model.layers.5.mlp.experts.4.down_proj.weight": "model-00006.safetensors",
+    "model.layers.5.mlp.experts.5.gate_proj.weight": "model-00006.safetensors",
+    "model.layers.5.mlp.experts.5.up_proj.weight": "model-00006.safetensors",
+    "model.layers.5.mlp.experts.5.down_proj.weight": "model-00006.safetensors",
+    "model.layers.5.mlp.experts.6.gate_proj.weight": "model-00007.safetensors",
+    "model.layers.5.mlp.experts.6.up_proj.weight": "model-00007.safetensors",
+    "model.layers.5.mlp.experts.6.down_proj.weight": "model-00007.safetensors",
+    "model.layers.5.post_attention_layernorm.weight": "model-00007.safetensors",
+    "model.layers.5.post_feedforward_layernorm.weight": "model-00007.safetensors",
+    "model.layers.6.self_attn.q_proj.weight": "model-00007.safetensors",
+    "model.layers.6.self_attn.k_proj.weight": "model-00007.safetensors",
+    "model.layers.6.self_attn.v_proj.weight": "model-00007.safetensors",
+    "model.layers.6.self_attn.o_proj.weight": "model-00007.safetensors",
+    "model.layers.6.self_attn.q_norm.weight": "model-00007.safetensors",
+    "model.layers.6.self_attn.k_norm.weight": "model-00007.safetensors",
+    "model.layers.6.mlp.gate.weight": "model-00007.safetensors",
+    "model.layers.6.mlp.experts.0.gate_proj.weight": "model-00007.safetensors",
+    "model.layers.6.mlp.experts.0.up_proj.weight": "model-00007.safetensors",
+    "model.layers.6.mlp.experts.0.down_proj.weight": "model-00007.safetensors",
+    "model.layers.6.mlp.experts.1.gate_proj.weight": "model-00007.safetensors",
+    "model.layers.6.mlp.experts.1.up_proj.weight": "model-00007.safetensors",
+    "model.layers.6.mlp.experts.1.down_proj.weight": "model-00007.safetensors",
+    "model.layers.6.mlp.experts.2.gate_proj.weight": "model-00007.safetensors",
+    "model.layers.6.mlp.experts.2.up_proj.weight": "model-00007.safetensors",
+    "model.layers.6.mlp.experts.2.down_proj.weight": "model-00007.safetensors",
+    "model.layers.6.mlp.experts.3.gate_proj.weight": "model-00007.safetensors",
+    "model.layers.6.mlp.experts.3.up_proj.weight": "model-00007.safetensors",
+    "model.layers.6.mlp.experts.3.down_proj.weight": "model-00007.safetensors",
+    "model.layers.6.mlp.experts.4.gate_proj.weight": "model-00007.safetensors",
+    "model.layers.6.mlp.experts.4.up_proj.weight": "model-00007.safetensors",
+    "model.layers.6.mlp.experts.4.down_proj.weight": "model-00007.safetensors",
+    "model.layers.6.mlp.experts.5.gate_proj.weight": "model-00007.safetensors",
+    "model.layers.6.mlp.experts.5.up_proj.weight": "model-00007.safetensors",
+    "model.layers.6.mlp.experts.5.down_proj.weight": "model-00007.safetensors",
+    "model.layers.6.mlp.experts.6.gate_proj.weight": "model-00007.safetensors",
+    "model.layers.6.mlp.experts.6.up_proj.weight": "model-00007.safetensors",
+    "model.layers.6.mlp.experts.6.down_proj.weight": "model-00007.safetensors",
+    "model.layers.6.post_attention_layernorm.weight": "model-00007.safetensors",
+    "model.layers.6.post_feedforward_layernorm.weight": "model-00007.safetensors",
+    "model.layers.7.self_attn.q_proj.weight": "model-00007.safetensors",
+    "model.layers.7.self_attn.k_proj.weight": "model-00007.safetensors",
+    "model.layers.7.self_attn.v_proj.weight": "model-00008.safetensors",
+    "model.layers.7.self_attn.o_proj.weight": "model-00008.safetensors",
+    "model.layers.7.self_attn.q_norm.weight": "model-00008.safetensors",
+    "model.layers.7.self_attn.k_norm.weight": "model-00008.safetensors",
+    "model.layers.7.mlp.gate.weight": "model-00008.safetensors",
+    "model.layers.7.mlp.experts.0.gate_proj.weight": "model-00008.safetensors",
+    "model.layers.7.mlp.experts.0.up_proj.weight": "model-00008.safetensors",
+    "model.layers.7.mlp.experts.0.down_proj.weight": "model-00008.safetensors",
+    "model.layers.7.mlp.experts.1.gate_proj.weight": "model-00008.safetensors",
+    "model.layers.7.mlp.experts.1.up_proj.weight": "model-00008.safetensors",
+    "model.layers.7.mlp.experts.1.down_proj.weight": "model-00008.safetensors",
+    "model.layers.7.mlp.experts.2.gate_proj.weight": "model-00008.safetensors",
+    "model.layers.7.mlp.experts.2.up_proj.weight": "model-00008.safetensors",
+    "model.layers.7.mlp.experts.2.down_proj.weight": "model-00008.safetensors",
+    "model.layers.7.mlp.experts.3.gate_proj.weight": "model-00008.safetensors",
+    "model.layers.7.mlp.experts.3.up_proj.weight": "model-00008.safetensors",
+    "model.layers.7.mlp.experts.3.down_proj.weight": "model-00008.safetensors",
+    "model.layers.7.mlp.experts.4.gate_proj.weight": "model-00008.safetensors",
+    "model.layers.7.mlp.experts.4.up_proj.weight": "model-00008.safetensors",
+    "model.layers.7.mlp.experts.4.down_proj.weight": "model-00008.safetensors",
+    "model.layers.7.mlp.experts.5.gate_proj.weight": "model-00008.safetensors",
+    "model.layers.7.mlp.experts.5.up_proj.weight": "model-00008.safetensors",
+    "model.layers.7.mlp.experts.5.down_proj.weight": "model-00008.safetensors",
+    "model.layers.7.mlp.experts.6.gate_proj.weight": "model-00008.safetensors",
+    "model.layers.7.mlp.experts.6.up_proj.weight": "model-00008.safetensors",
+    "model.layers.7.mlp.experts.6.down_proj.weight": "model-00008.safetensors",
+    "model.layers.7.post_attention_layernorm.weight": "model-00008.safetensors",
+    "model.layers.7.post_feedforward_layernorm.weight": "model-00008.safetensors",
+    "model.layers.8.self_attn.q_proj.weight": "model-00008.safetensors",
+    "model.layers.8.self_attn.k_proj.weight": "model-00008.safetensors",
+    "model.layers.8.self_attn.v_proj.weight": "model-00008.safetensors",
+    "model.layers.8.self_attn.o_proj.weight": "model-00008.safetensors",
+    "model.layers.8.self_attn.q_norm.weight": "model-00008.safetensors",
+    "model.layers.8.self_attn.k_norm.weight": "model-00008.safetensors",
+    "model.layers.8.mlp.gate.weight": "model-00008.safetensors",
+    "model.layers.8.mlp.experts.0.gate_proj.weight": "model-00008.safetensors",
+    "model.layers.8.mlp.experts.0.up_proj.weight": "model-00008.safetensors",
+    "model.layers.8.mlp.experts.0.down_proj.weight": "model-00008.safetensors",
+    "model.layers.8.mlp.experts.1.gate_proj.weight": "model-00009.safetensors",
+    "model.layers.8.mlp.experts.1.up_proj.weight": "model-00009.safetensors",
+    "model.layers.8.mlp.experts.1.down_proj.weight": "model-00009.safetensors",
+    "model.layers.8.mlp.experts.2.gate_proj.weight": "model-00009.safetensors",
+    "model.layers.8.mlp.experts.2.up_proj.weight": "model-00009.safetensors",
+    "model.layers.8.mlp.experts.2.down_proj.weight": "model-00009.safetensors",
+    "model.layers.8.mlp.experts.3.gate_proj.weight": "model-00009.safetensors",
+    "model.layers.8.mlp.experts.3.up_proj.weight": "model-00009.safetensors",
+    "model.layers.8.mlp.experts.3.down_proj.weight": "model-00009.safetensors",
+    "model.layers.8.mlp.experts.4.gate_proj.weight": "model-00009.safetensors",
+    "model.layers.8.mlp.experts.4.up_proj.weight": "model-00009.safetensors",
+    "model.layers.8.mlp.experts.4.down_proj.weight": "model-00009.safetensors",
+    "model.layers.8.mlp.experts.5.gate_proj.weight": "model-00009.safetensors",
+    "model.layers.8.mlp.experts.5.up_proj.weight": "model-00009.safetensors",
+    "model.layers.8.mlp.experts.5.down_proj.weight": "model-00009.safetensors",
+    "model.layers.8.mlp.experts.6.gate_proj.weight": "model-00009.safetensors",
+    "model.layers.8.mlp.experts.6.up_proj.weight": "model-00009.safetensors",
+    "model.layers.8.mlp.experts.6.down_proj.weight": "model-00009.safetensors",
+    "model.layers.8.post_attention_layernorm.weight": "model-00009.safetensors",
+    "model.layers.8.post_feedforward_layernorm.weight": "model-00009.safetensors",
+    "model.layers.9.self_attn.q_proj.weight": "model-00009.safetensors",
+    "model.layers.9.self_attn.k_proj.weight": "model-00009.safetensors",
+    "model.layers.9.self_attn.v_proj.weight": "model-00009.safetensors",
+    "model.layers.9.self_attn.o_proj.weight": "model-00009.safetensors",
+    "model.layers.9.self_attn.q_norm.weight": "model-00009.safetensors",
+    "model.layers.9.self_attn.k_norm.weight": "model-00009.safetensors",
+    "model.layers.9.mlp.gate.weight": "model-00009.safetensors",
+    "model.layers.9.mlp.experts.0.gate_proj.weight": "model-00009.safetensors",
+    "model.layers.9.mlp.experts.0.up_proj.weight": "model-00009.safetensors",
+    "model.layers.9.mlp.experts.0.down_proj.weight": "model-00009.safetensors",
+    "model.layers.9.mlp.experts.1.gate_proj.weight": "model-00009.safetensors",
+    "model.layers.9.mlp.experts.1.up_proj.weight": "model-00009.safetensors",
+    "model.layers.9.mlp.experts.1.down_proj.weight": "model-00009.safetensors",
+    "model.layers.9.mlp.experts.2.gate_proj.weight": "model-00009.safetensors",
+    "model.layers.9.mlp.experts.2.up_proj.weight": "model-00010.safetensors",
+    "model.layers.9.mlp.experts.2.down_proj.weight": "model-00010.safetensors",
+    "model.layers.9.mlp.experts.3.gate_proj.weight": "model-00010.safetensors",
+    "model.layers.9.mlp.experts.3.up_proj.weight": "model-00010.safetensors",
+    "model.layers.9.mlp.experts.3.down_proj.weight": "model-00010.safetensors",
+    "model.layers.9.mlp.experts.4.gate_proj.weight": "model-00010.safetensors",
+    "model.layers.9.mlp.experts.4.up_proj.weight": "model-00010.safetensors",
+    "model.layers.9.mlp.experts.4.down_proj.weight": "model-00010.safetensors",
+    "model.layers.9.mlp.experts.5.gate_proj.weight": "model-00010.safetensors",
+    "model.layers.9.mlp.experts.5.up_proj.weight": "model-00010.safetensors",
+    "model.layers.9.mlp.experts.5.down_proj.weight": "model-00010.safetensors",
+    "model.layers.9.mlp.experts.6.gate_proj.weight": "model-00010.safetensors",
+    "model.layers.9.mlp.experts.6.up_proj.weight": "model-00010.safetensors",
+    "model.layers.9.mlp.experts.6.down_proj.weight": "model-00010.safetensors",
+    "model.layers.9.post_attention_layernorm.weight": "model-00010.safetensors",
+    "model.layers.9.post_feedforward_layernorm.weight": "model-00010.safetensors",
+    "model.layers.10.self_attn.q_proj.weight": "model-00010.safetensors",
+    "model.layers.10.self_attn.k_proj.weight": "model-00010.safetensors",
+    "model.layers.10.self_attn.v_proj.weight": "model-00010.safetensors",
+    "model.layers.10.self_attn.o_proj.weight": "model-00010.safetensors",
+    "model.layers.10.self_attn.q_norm.weight": "model-00010.safetensors",
+    "model.layers.10.self_attn.k_norm.weight": "model-00010.safetensors",
+    "model.layers.10.mlp.gate.weight": "model-00010.safetensors",
+    "model.layers.10.mlp.experts.0.gate_proj.weight": "model-00010.safetensors",
+    "model.layers.10.mlp.experts.0.up_proj.weight": "model-00010.safetensors",
+    "model.layers.10.mlp.experts.0.down_proj.weight": "model-00010.safetensors",
+    "model.layers.10.mlp.experts.1.gate_proj.weight": "model-00010.safetensors",
+    "model.layers.10.mlp.experts.1.up_proj.weight": "model-00010.safetensors",
+    "model.layers.10.mlp.experts.1.down_proj.weight": "model-00010.safetensors",
+    "model.layers.10.mlp.experts.2.gate_proj.weight": "model-00010.safetensors",
+    "model.layers.10.mlp.experts.2.up_proj.weight": "model-00010.safetensors",
+    "model.layers.10.mlp.experts.2.down_proj.weight": "model-00010.safetensors",
+    "model.layers.10.mlp.experts.3.gate_proj.weight": "model-00010.safetensors",
+    "model.layers.10.mlp.experts.3.up_proj.weight": "model-00010.safetensors",
+    "model.layers.10.mlp.experts.3.down_proj.weight": "model-00011.safetensors",
+    "model.layers.10.mlp.experts.4.gate_proj.weight": "model-00011.safetensors",
+    "model.layers.10.mlp.experts.4.up_proj.weight": "model-00011.safetensors",
+    "model.layers.10.mlp.experts.4.down_proj.weight": "model-00011.safetensors",
+    "model.layers.10.mlp.experts.5.gate_proj.weight": "model-00011.safetensors",
+    "model.layers.10.mlp.experts.5.up_proj.weight": "model-00011.safetensors",
+    "model.layers.10.mlp.experts.5.down_proj.weight": "model-00011.safetensors",
+    "model.layers.10.mlp.experts.6.gate_proj.weight": "model-00011.safetensors",
+    "model.layers.10.mlp.experts.6.up_proj.weight": "model-00011.safetensors",
+    "model.layers.10.mlp.experts.6.down_proj.weight": "model-00011.safetensors",
+    "model.layers.10.post_attention_layernorm.weight": "model-00011.safetensors",
+    "model.layers.10.post_feedforward_layernorm.weight": "model-00011.safetensors",
+    "model.layers.11.self_attn.q_proj.weight": "model-00011.safetensors",
+    "model.layers.11.self_attn.k_proj.weight": "model-00011.safetensors",
+    "model.layers.11.self_attn.v_proj.weight": "model-00011.safetensors",
+    "model.layers.11.self_attn.o_proj.weight": "model-00011.safetensors",
+    "model.layers.11.self_attn.q_norm.weight": "model-00011.safetensors",
+    "model.layers.11.self_attn.k_norm.weight": "model-00011.safetensors",
+    "model.layers.11.mlp.gate.weight": "model-00011.safetensors",
+    "model.layers.11.mlp.experts.0.gate_proj.weight": "model-00011.safetensors",
+    "model.layers.11.mlp.experts.0.up_proj.weight": "model-00011.safetensors",
+    "model.layers.11.mlp.experts.0.down_proj.weight": "model-00011.safetensors",
+    "model.layers.11.mlp.experts.1.gate_proj.weight": "model-00011.safetensors",
+    "model.layers.11.mlp.experts.1.up_proj.weight": "model-00011.safetensors",
+    "model.layers.11.mlp.experts.1.down_proj.weight": "model-00011.safetensors",
+    "model.layers.11.mlp.experts.2.gate_proj.weight": "model-00011.safetensors",
+    "model.layers.11.mlp.experts.2.up_proj.weight": "model-00011.safetensors",
+    "model.layers.11.mlp.experts.2.down_proj.weight": "model-00011.safetensors",
+    "model.layers.11.mlp.experts.3.gate_proj.weight": "model-00011.safetensors",
+    "model.layers.11.mlp.experts.3.up_proj.weight": "model-00011.safetensors",
+    "model.layers.11.mlp.experts.3.down_proj.weight": "model-00011.safetensors",
+    "model.layers.11.mlp.experts.4.gate_proj.weight": "model-00011.safetensors",
+    "model.layers.11.mlp.experts.4.up_proj.weight": "model-00011.safetensors",
+    "model.layers.11.mlp.experts.4.down_proj.weight": "model-00011.safetensors",
+    "model.layers.11.mlp.experts.5.gate_proj.weight": "model-00012.safetensors",
+    "model.layers.11.mlp.experts.5.up_proj.weight": "model-00012.safetensors",
+    "model.layers.11.mlp.experts.5.down_proj.weight": "model-00012.safetensors",
+    "model.layers.11.mlp.experts.6.gate_proj.weight": "model-00012.safetensors",
+    "model.layers.11.mlp.experts.6.up_proj.weight": "model-00012.safetensors",
+    "model.layers.11.mlp.experts.6.down_proj.weight": "model-00012.safetensors",
+    "model.layers.11.post_attention_layernorm.weight": "model-00012.safetensors",
+    "model.layers.11.post_feedforward_layernorm.weight": "model-00012.safetensors",
+    "model.layers.12.self_attn.q_proj.weight": "model-00012.safetensors",
+    "model.layers.12.self_attn.k_proj.weight": "model-00012.safetensors",
+    "model.layers.12.self_attn.v_proj.weight": "model-00012.safetensors",
+    "model.layers.12.self_attn.o_proj.weight": "model-00012.safetensors",
+    "model.layers.12.self_attn.q_norm.weight": "model-00012.safetensors",
+    "model.layers.12.self_attn.k_norm.weight": "model-00012.safetensors",
+    "model.layers.12.mlp.gate.weight": "model-00012.safetensors",
+    "model.layers.12.mlp.experts.0.gate_proj.weight": "model-00012.safetensors",
+    "model.layers.12.mlp.experts.0.up_proj.weight": "model-00012.safetensors",
+    "model.layers.12.mlp.experts.0.down_proj.weight": "model-00012.safetensors",
+    "model.layers.12.mlp.experts.1.gate_proj.weight": "model-00012.safetensors",
+    "model.layers.12.mlp.experts.1.up_proj.weight": "model-00012.safetensors",
+    "model.layers.12.mlp.experts.1.down_proj.weight": "model-00012.safetensors",
+    "model.layers.12.mlp.experts.2.gate_proj.weight": "model-00012.safetensors",
+    "model.layers.12.mlp.experts.2.up_proj.weight": "model-00012.safetensors",
+    "model.layers.12.mlp.experts.2.down_proj.weight": "model-00012.safetensors",
+    "model.layers.12.mlp.experts.3.gate_proj.weight": "model-00012.safetensors",
+    "model.layers.12.mlp.experts.3.up_proj.weight": "model-00012.safetensors",
+    "model.layers.12.mlp.experts.3.down_proj.weight": "model-00012.safetensors",
+    "model.layers.12.mlp.experts.4.gate_proj.weight": "model-00012.safetensors",
+    "model.layers.12.mlp.experts.4.up_proj.weight": "model-00012.safetensors",
+    "model.layers.12.mlp.experts.4.down_proj.weight": "model-00012.safetensors",
+    "model.layers.12.mlp.experts.5.gate_proj.weight": "model-00012.safetensors",
+    "model.layers.12.mlp.experts.5.up_proj.weight": "model-00012.safetensors",
+    "model.layers.12.mlp.experts.5.down_proj.weight": "model-00012.safetensors",
+    "model.layers.12.mlp.experts.6.gate_proj.weight": "model-00012.safetensors",
+    "model.layers.12.mlp.experts.6.up_proj.weight": "model-00013.safetensors",
+    "model.layers.12.mlp.experts.6.down_proj.weight": "model-00013.safetensors",
+    "model.layers.12.post_attention_layernorm.weight": "model-00013.safetensors",
+    "model.layers.12.post_feedforward_layernorm.weight": "model-00013.safetensors",
+    "model.layers.13.self_attn.q_proj.weight": "model-00013.safetensors",
+    "model.layers.13.self_attn.k_proj.weight": "model-00013.safetensors",
+    "model.layers.13.self_attn.v_proj.weight": "model-00013.safetensors",
+    "model.layers.13.self_attn.o_proj.weight": "model-00013.safetensors",
+    "model.layers.13.self_attn.q_norm.weight": "model-00013.safetensors",
+    "model.layers.13.self_attn.k_norm.weight": "model-00013.safetensors",
+    "model.layers.13.mlp.gate.weight": "model-00013.safetensors",
+    "model.layers.13.mlp.experts.0.gate_proj.weight": "model-00013.safetensors",
+    "model.layers.13.mlp.experts.0.up_proj.weight": "model-00013.safetensors",
+    "model.layers.13.mlp.experts.0.down_proj.weight": "model-00013.safetensors",
+    "model.layers.13.mlp.experts.1.gate_proj.weight": "model-00013.safetensors",
+    "model.layers.13.mlp.experts.1.up_proj.weight": "model-00013.safetensors",
+    "model.layers.13.mlp.experts.1.down_proj.weight": "model-00013.safetensors",
+    "model.layers.13.mlp.experts.2.gate_proj.weight": "model-00013.safetensors",
+    "model.layers.13.mlp.experts.2.up_proj.weight": "model-00013.safetensors",
+    "model.layers.13.mlp.experts.2.down_proj.weight": "model-00013.safetensors",
+    "model.layers.13.mlp.experts.3.gate_proj.weight": "model-00013.safetensors",
+    "model.layers.13.mlp.experts.3.up_proj.weight": "model-00013.safetensors",
+    "model.layers.13.mlp.experts.3.down_proj.weight": "model-00013.safetensors",
+    "model.layers.13.mlp.experts.4.gate_proj.weight": "model-00013.safetensors",
+    "model.layers.13.mlp.experts.4.up_proj.weight": "model-00013.safetensors",
+    "model.layers.13.mlp.experts.4.down_proj.weight": "model-00013.safetensors",
+    "model.layers.13.mlp.experts.5.gate_proj.weight": "model-00013.safetensors",
+    "model.layers.13.mlp.experts.5.up_proj.weight": "model-00013.safetensors",
+    "model.layers.13.mlp.experts.5.down_proj.weight": "model-00013.safetensors",
+    "model.layers.13.mlp.experts.6.gate_proj.weight": "model-00013.safetensors",
+    "model.layers.13.mlp.experts.6.up_proj.weight": "model-00013.safetensors",
+    "model.layers.13.mlp.experts.6.down_proj.weight": "model-00013.safetensors",
+    "model.layers.13.post_attention_layernorm.weight": "model-00013.safetensors",
+    "model.layers.13.post_feedforward_layernorm.weight": "model-00013.safetensors",
+    "model.layers.14.self_attn.q_proj.weight": "model-00013.safetensors",
+    "model.layers.14.self_attn.k_proj.weight": "model-00013.safetensors",
+    "model.layers.14.self_attn.v_proj.weight": "model-00013.safetensors",
+    "model.layers.14.self_attn.o_proj.weight": "model-00013.safetensors",
+    "model.layers.14.self_attn.q_norm.weight": "model-00014.safetensors",
+    "model.layers.14.self_attn.k_norm.weight": "model-00014.safetensors",
+    "model.layers.14.mlp.gate.weight": "model-00014.safetensors",
+    "model.layers.14.mlp.experts.0.gate_proj.weight": "model-00014.safetensors",
+    "model.layers.14.mlp.experts.0.up_proj.weight": "model-00014.safetensors",
+    "model.layers.14.mlp.experts.0.down_proj.weight": "model-00014.safetensors",
+    "model.layers.14.mlp.experts.1.gate_proj.weight": "model-00014.safetensors",
+    "model.layers.14.mlp.experts.1.up_proj.weight": "model-00014.safetensors",
+    "model.layers.14.mlp.experts.1.down_proj.weight": "model-00014.safetensors",
+    "model.layers.14.mlp.experts.2.gate_proj.weight": "model-00014.safetensors",
+    "model.layers.14.mlp.experts.2.up_proj.weight": "model-00014.safetensors",
+    "model.layers.14.mlp.experts.2.down_proj.weight": "model-00014.safetensors",
+    "model.layers.14.mlp.experts.3.gate_proj.weight": "model-00014.safetensors",
+    "model.layers.14.mlp.experts.3.up_proj.weight": "model-00014.safetensors",
+    "model.layers.14.mlp.experts.3.down_proj.weight": "model-00014.safetensors",
+    "model.layers.14.mlp.experts.4.gate_proj.weight": "model-00014.safetensors",
+    "model.layers.14.mlp.experts.4.up_proj.weight": "model-00014.safetensors",
+    "model.layers.14.mlp.experts.4.down_proj.weight": "model-00014.safetensors",
+    "model.layers.14.mlp.experts.5.gate_proj.weight": "model-00014.safetensors",
+    "model.layers.14.mlp.experts.5.up_proj.weight": "model-00014.safetensors",
+    "model.layers.14.mlp.experts.5.down_proj.weight": "model-00014.safetensors",
+    "model.layers.14.mlp.experts.6.gate_proj.weight": "model-00014.safetensors",
+    "model.layers.14.mlp.experts.6.up_proj.weight": "model-00014.safetensors",
+    "model.layers.14.mlp.experts.6.down_proj.weight": "model-00014.safetensors",
+    "model.layers.14.post_attention_layernorm.weight": "model-00014.safetensors",
+    "model.layers.14.post_feedforward_layernorm.weight": "model-00014.safetensors",
+    "model.layers.15.self_attn.q_proj.weight": "model-00014.safetensors",
+    "model.layers.15.self_attn.k_proj.weight": "model-00014.safetensors",
+    "model.layers.15.self_attn.v_proj.weight": "model-00014.safetensors",
+    "model.layers.15.self_attn.o_proj.weight": "model-00014.safetensors",
+    "model.layers.15.self_attn.q_norm.weight": "model-00014.safetensors",
+    "model.layers.15.self_attn.k_norm.weight": "model-00014.safetensors",
+    "model.layers.15.mlp.gate.weight": "model-00014.safetensors",
+    "model.layers.15.mlp.experts.0.gate_proj.weight": "model-00014.safetensors",
+    "model.layers.15.mlp.experts.0.up_proj.weight": "model-00014.safetensors",
+    "model.layers.15.mlp.experts.0.down_proj.weight": "model-00014.safetensors",
+    "model.layers.15.mlp.experts.1.gate_proj.weight": "model-00014.safetensors",
+    "model.layers.15.mlp.experts.1.up_proj.weight": "model-00015.safetensors",
+    "model.layers.15.mlp.experts.1.down_proj.weight": "model-00015.safetensors",
+    "model.layers.15.mlp.experts.2.gate_proj.weight": "model-00015.safetensors",
+    "model.layers.15.mlp.experts.2.up_proj.weight": "model-00015.safetensors",
+    "model.layers.15.mlp.experts.2.down_proj.weight": "model-00015.safetensors",
+    "model.layers.15.mlp.experts.3.gate_proj.weight": "model-00015.safetensors",
+    "model.layers.15.mlp.experts.3.up_proj.weight": "model-00015.safetensors",
+    "model.layers.15.mlp.experts.3.down_proj.weight": "model-00015.safetensors",
+    "model.layers.15.mlp.experts.4.gate_proj.weight": "model-00015.safetensors",
+    "model.layers.15.mlp.experts.4.up_proj.weight": "model-00015.safetensors",
+    "model.layers.15.mlp.experts.4.down_proj.weight": "model-00015.safetensors",
+    "model.layers.15.mlp.experts.5.gate_proj.weight": "model-00015.safetensors",
+    "model.layers.15.mlp.experts.5.up_proj.weight": "model-00015.safetensors",
+    "model.layers.15.mlp.experts.5.down_proj.weight": "model-00015.safetensors",
+    "model.layers.15.mlp.experts.6.gate_proj.weight": "model-00015.safetensors",
+    "model.layers.15.mlp.experts.6.up_proj.weight": "model-00015.safetensors",
+    "model.layers.15.mlp.experts.6.down_proj.weight": "model-00015.safetensors",
+    "model.layers.15.post_attention_layernorm.weight": "model-00015.safetensors",
+    "model.layers.15.post_feedforward_layernorm.weight": "model-00015.safetensors",
+    "model.layers.16.self_attn.q_proj.weight": "model-00015.safetensors",
+    "model.layers.16.self_attn.k_proj.weight": "model-00015.safetensors",
+    "model.layers.16.self_attn.v_proj.weight": "model-00015.safetensors",
+    "model.layers.16.self_attn.o_proj.weight": "model-00015.safetensors",
+    "model.layers.16.self_attn.q_norm.weight": "model-00015.safetensors",
+    "model.layers.16.self_attn.k_norm.weight": "model-00015.safetensors",
+    "model.layers.16.mlp.gate.weight": "model-00015.safetensors",
+    "model.layers.16.mlp.experts.0.gate_proj.weight": "model-00015.safetensors",
+    "model.layers.16.mlp.experts.0.up_proj.weight": "model-00015.safetensors",
+    "model.layers.16.mlp.experts.0.down_proj.weight": "model-00015.safetensors",
+    "model.layers.16.mlp.experts.1.gate_proj.weight": "model-00015.safetensors",
+    "model.layers.16.mlp.experts.1.up_proj.weight": "model-00015.safetensors",
+    "model.layers.16.mlp.experts.1.down_proj.weight": "model-00015.safetensors",
+    "model.layers.16.mlp.experts.2.gate_proj.weight": "model-00015.safetensors",
+    "model.layers.16.mlp.experts.2.up_proj.weight": "model-00015.safetensors",
+    "model.layers.16.mlp.experts.2.down_proj.weight": "model-00016.safetensors",
+    "model.layers.16.mlp.experts.3.gate_proj.weight": "model-00016.safetensors",
+    "model.layers.16.mlp.experts.3.up_proj.weight": "model-00016.safetensors",
+    "model.layers.16.mlp.experts.3.down_proj.weight": "model-00016.safetensors",
+    "model.layers.16.mlp.experts.4.gate_proj.weight": "model-00016.safetensors",
+    "model.layers.16.mlp.experts.4.up_proj.weight": "model-00016.safetensors",
+    "model.layers.16.mlp.experts.4.down_proj.weight": "model-00016.safetensors",
+    "model.layers.16.mlp.experts.5.gate_proj.weight": "model-00016.safetensors",
+    "model.layers.16.mlp.experts.5.up_proj.weight": "model-00016.safetensors",
+    "model.layers.16.mlp.experts.5.down_proj.weight": "model-00016.safetensors",
+    "model.layers.16.mlp.experts.6.gate_proj.weight": "model-00016.safetensors",
+    "model.layers.16.mlp.experts.6.up_proj.weight": "model-00016.safetensors",
+    "model.layers.16.mlp.experts.6.down_proj.weight": "model-00016.safetensors",
+    "model.layers.16.post_attention_layernorm.weight": "model-00016.safetensors",
+    "model.layers.16.post_feedforward_layernorm.weight": "model-00016.safetensors",
+    "model.layers.17.self_attn.q_proj.weight": "model-00016.safetensors",
+    "model.layers.17.self_attn.k_proj.weight": "model-00016.safetensors",
+    "model.layers.17.self_attn.v_proj.weight": "model-00016.safetensors",
+    "model.layers.17.self_attn.o_proj.weight": "model-00016.safetensors",
+    "model.layers.17.self_attn.q_norm.weight": "model-00016.safetensors",
+    "model.layers.17.self_attn.k_norm.weight": "model-00016.safetensors",
+    "model.layers.17.mlp.gate.weight": "model-00016.safetensors",
+    "model.layers.17.mlp.experts.0.gate_proj.weight": "model-00016.safetensors",
+    "model.layers.17.mlp.experts.0.up_proj.weight": "model-00016.safetensors",
+    "model.layers.17.mlp.experts.0.down_proj.weight": "model-00016.safetensors",
+    "model.layers.17.mlp.experts.1.gate_proj.weight": "model-00016.safetensors",
+    "model.layers.17.mlp.experts.1.up_proj.weight": "model-00016.safetensors",
+    "model.layers.17.mlp.experts.1.down_proj.weight": "model-00016.safetensors",
+    "model.layers.17.mlp.experts.2.gate_proj.weight": "model-00016.safetensors",
+    "model.layers.17.mlp.experts.2.up_proj.weight": "model-00016.safetensors",
+    "model.layers.17.mlp.experts.2.down_proj.weight": "model-00016.safetensors",
+    "model.layers.17.mlp.experts.3.gate_proj.weight": "model-00016.safetensors",
+    "model.layers.17.mlp.experts.3.up_proj.weight": "model-00016.safetensors",
+    "model.layers.17.mlp.experts.3.down_proj.weight": "model-00016.safetensors",
+    "model.layers.17.mlp.experts.4.gate_proj.weight": "model-00017.safetensors",
+    "model.layers.17.mlp.experts.4.up_proj.weight": "model-00017.safetensors",
+    "model.layers.17.mlp.experts.4.down_proj.weight": "model-00017.safetensors",
+    "model.layers.17.mlp.experts.5.gate_proj.weight": "model-00017.safetensors",
+    "model.layers.17.mlp.experts.5.up_proj.weight": "model-00017.safetensors",
+    "model.layers.17.mlp.experts.5.down_proj.weight": "model-00017.safetensors",
+    "model.layers.17.mlp.experts.6.gate_proj.weight": "model-00017.safetensors",
+    "model.layers.17.mlp.experts.6.up_proj.weight": "model-00017.safetensors",
+    "model.layers.17.mlp.experts.6.down_proj.weight": "model-00017.safetensors",
+    "model.layers.17.post_attention_layernorm.weight": "model-00017.safetensors",
+    "model.layers.17.post_feedforward_layernorm.weight": "model-00017.safetensors",
+    "model.layers.18.self_attn.q_proj.weight": "model-00017.safetensors",
+    "model.layers.18.self_attn.k_proj.weight": "model-00017.safetensors",
+    "model.layers.18.self_attn.v_proj.weight": "model-00017.safetensors",
+    "model.layers.18.self_attn.o_proj.weight": "model-00017.safetensors",
+    "model.layers.18.self_attn.q_norm.weight": "model-00017.safetensors",
+    "model.layers.18.self_attn.k_norm.weight": "model-00017.safetensors",
+    "model.layers.18.mlp.gate.weight": "model-00017.safetensors",
+    "model.layers.18.mlp.experts.0.gate_proj.weight": "model-00017.safetensors",
+    "model.layers.18.mlp.experts.0.up_proj.weight": "model-00017.safetensors",
+    "model.layers.18.mlp.experts.0.down_proj.weight": "model-00017.safetensors",
+    "model.layers.18.mlp.experts.1.gate_proj.weight": "model-00017.safetensors",
+    "model.layers.18.mlp.experts.1.up_proj.weight": "model-00017.safetensors",
+    "model.layers.18.mlp.experts.1.down_proj.weight": "model-00017.safetensors",
+    "model.layers.18.mlp.experts.2.gate_proj.weight": "model-00017.safetensors",
+    "model.layers.18.mlp.experts.2.up_proj.weight": "model-00017.safetensors",
+    "model.layers.18.mlp.experts.2.down_proj.weight": "model-00017.safetensors",
+    "model.layers.18.mlp.experts.3.gate_proj.weight": "model-00017.safetensors",
+    "model.layers.18.mlp.experts.3.up_proj.weight": "model-00017.safetensors",
+    "model.layers.18.mlp.experts.3.down_proj.weight": "model-00017.safetensors",
+    "model.layers.18.mlp.experts.4.gate_proj.weight": "model-00017.safetensors",
+    "model.layers.18.mlp.experts.4.up_proj.weight": "model-00017.safetensors",
+    "model.layers.18.mlp.experts.4.down_proj.weight": "model-00017.safetensors",
+    "model.layers.18.mlp.experts.5.gate_proj.weight": "model-00017.safetensors",
+    "model.layers.18.mlp.experts.5.up_proj.weight": "model-00018.safetensors",
+    "model.layers.18.mlp.experts.5.down_proj.weight": "model-00018.safetensors",
+    "model.layers.18.mlp.experts.6.gate_proj.weight": "model-00018.safetensors",
+    "model.layers.18.mlp.experts.6.up_proj.weight": "model-00018.safetensors",
+    "model.layers.18.mlp.experts.6.down_proj.weight": "model-00018.safetensors",
+    "model.layers.18.post_attention_layernorm.weight": "model-00018.safetensors",
+    "model.layers.18.post_feedforward_layernorm.weight": "model-00018.safetensors",
+    "model.layers.19.self_attn.q_proj.weight": "model-00018.safetensors",
+    "model.layers.19.self_attn.k_proj.weight": "model-00018.safetensors",
+    "model.layers.19.self_attn.v_proj.weight": "model-00018.safetensors",
+    "model.layers.19.self_attn.o_proj.weight": "model-00018.safetensors",
+    "model.layers.19.self_attn.q_norm.weight": "model-00018.safetensors",
+    "model.layers.19.self_attn.k_norm.weight": "model-00018.safetensors",
+    "model.layers.19.mlp.gate.weight": "model-00018.safetensors",
+    "model.layers.19.mlp.experts.0.gate_proj.weight": "model-00018.safetensors",
+    "model.layers.19.mlp.experts.0.up_proj.weight": "model-00018.safetensors",
+    "model.layers.19.mlp.experts.0.down_proj.weight": "model-00018.safetensors",
+    "model.layers.19.mlp.experts.1.gate_proj.weight": "model-00018.safetensors",
+    "model.layers.19.mlp.experts.1.up_proj.weight": "model-00018.safetensors",
+    "model.layers.19.mlp.experts.1.down_proj.weight": "model-00018.safetensors",
+    "model.layers.19.mlp.experts.2.gate_proj.weight": "model-00018.safetensors",
+    "model.layers.19.mlp.experts.2.up_proj.weight": "model-00018.safetensors",
+    "model.layers.19.mlp.experts.2.down_proj.weight": "model-00018.safetensors",
+    "model.layers.19.mlp.experts.3.gate_proj.weight": "model-00018.safetensors",
+    "model.layers.19.mlp.experts.3.up_proj.weight": "model-00018.safetensors",
+    "model.layers.19.mlp.experts.3.down_proj.weight": "model-00018.safetensors",
+    "model.layers.19.mlp.experts.4.gate_proj.weight": "model-00018.safetensors",
+    "model.layers.19.mlp.experts.4.up_proj.weight": "model-00018.safetensors",
+    "model.layers.19.mlp.experts.4.down_proj.weight": "model-00018.safetensors",
+    "model.layers.19.mlp.experts.5.gate_proj.weight": "model-00018.safetensors",
+    "model.layers.19.mlp.experts.5.up_proj.weight": "model-00018.safetensors",
+    "model.layers.19.mlp.experts.5.down_proj.weight": "model-00018.safetensors",
+    "model.layers.19.mlp.experts.6.gate_proj.weight": "model-00018.safetensors",
+    "model.layers.19.mlp.experts.6.up_proj.weight": "model-00018.safetensors",
+    "model.layers.19.mlp.experts.6.down_proj.weight": "model-00019.safetensors",
+    "model.layers.19.post_attention_layernorm.weight": "model-00019.safetensors",
+    "model.layers.19.post_feedforward_layernorm.weight": "model-00019.safetensors",
+    "model.layers.20.self_attn.q_proj.weight": "model-00019.safetensors",
+    "model.layers.20.self_attn.k_proj.weight": "model-00019.safetensors",
+    "model.layers.20.self_attn.v_proj.weight": "model-00019.safetensors",
+    "model.layers.20.self_attn.o_proj.weight": "model-00019.safetensors",
+    "model.layers.20.self_attn.q_norm.weight": "model-00019.safetensors",
+    "model.layers.20.self_attn.k_norm.weight": "model-00019.safetensors",
+    "model.layers.20.mlp.gate.weight": "model-00019.safetensors",
+    "model.layers.20.mlp.experts.0.gate_proj.weight": "model-00019.safetensors",
+    "model.layers.20.mlp.experts.0.up_proj.weight": "model-00019.safetensors",
+    "model.layers.20.mlp.experts.0.down_proj.weight": "model-00019.safetensors",
+    "model.layers.20.mlp.experts.1.gate_proj.weight": "model-00019.safetensors",
+    "model.layers.20.mlp.experts.1.up_proj.weight": "model-00019.safetensors",
+    "model.layers.20.mlp.experts.1.down_proj.weight": "model-00019.safetensors",
+    "model.layers.20.mlp.experts.2.gate_proj.weight": "model-00019.safetensors",
+    "model.layers.20.mlp.experts.2.up_proj.weight": "model-00019.safetensors",
+    "model.layers.20.mlp.experts.2.down_proj.weight": "model-00019.safetensors",
+    "model.layers.20.mlp.experts.3.gate_proj.weight": "model-00019.safetensors",
+    "model.layers.20.mlp.experts.3.up_proj.weight": "model-00019.safetensors",
+    "model.layers.20.mlp.experts.3.down_proj.weight": "model-00019.safetensors",
+    "model.layers.20.mlp.experts.4.gate_proj.weight": "model-00019.safetensors",
+    "model.layers.20.mlp.experts.4.up_proj.weight": "model-00019.safetensors",
+    "model.layers.20.mlp.experts.4.down_proj.weight": "model-00019.safetensors",
+    "model.layers.20.mlp.experts.5.gate_proj.weight": "model-00019.safetensors",
+    "model.layers.20.mlp.experts.5.up_proj.weight": "model-00019.safetensors",
+    "model.layers.20.mlp.experts.5.down_proj.weight": "model-00019.safetensors",
+    "model.layers.20.mlp.experts.6.gate_proj.weight": "model-00019.safetensors",
+    "model.layers.20.mlp.experts.6.up_proj.weight": "model-00019.safetensors",
+    "model.layers.20.mlp.experts.6.down_proj.weight": "model-00019.safetensors",
+    "model.layers.20.post_attention_layernorm.weight": "model-00019.safetensors",
+    "model.layers.20.post_feedforward_layernorm.weight": "model-00019.safetensors",
+    "model.layers.21.self_attn.q_proj.weight": "model-00019.safetensors",
+    "model.layers.21.self_attn.k_proj.weight": "model-00019.safetensors",
+    "model.layers.21.self_attn.v_proj.weight": "model-00019.safetensors",
+    "model.layers.21.self_attn.o_proj.weight": "model-00019.safetensors",
+    "model.layers.21.self_attn.q_norm.weight": "model-00019.safetensors",
+    "model.layers.21.self_attn.k_norm.weight": "model-00019.safetensors",
+    "model.layers.21.mlp.gate.weight": "model-00019.safetensors",
+    "model.layers.21.mlp.experts.0.gate_proj.weight": "model-00019.safetensors",
+    "model.layers.21.mlp.experts.0.up_proj.weight": "model-00020.safetensors",
+    "model.layers.21.mlp.experts.0.down_proj.weight": "model-00020.safetensors",
+    "model.layers.21.mlp.experts.1.gate_proj.weight": "model-00020.safetensors",
+    "model.layers.21.mlp.experts.1.up_proj.weight": "model-00020.safetensors",
+    "model.layers.21.mlp.experts.1.down_proj.weight": "model-00020.safetensors",
+    "model.layers.21.mlp.experts.2.gate_proj.weight": "model-00020.safetensors",
+    "model.layers.21.mlp.experts.2.up_proj.weight": "model-00020.safetensors",
+    "model.layers.21.mlp.experts.2.down_proj.weight": "model-00020.safetensors",
+    "model.layers.21.mlp.experts.3.gate_proj.weight": "model-00020.safetensors",
+    "model.layers.21.mlp.experts.3.up_proj.weight": "model-00020.safetensors",
+    "model.layers.21.mlp.experts.3.down_proj.weight": "model-00020.safetensors",
+    "model.layers.21.mlp.experts.4.gate_proj.weight": "model-00020.safetensors",
+    "model.layers.21.mlp.experts.4.up_proj.weight": "model-00020.safetensors",
+    "model.layers.21.mlp.experts.4.down_proj.weight": "model-00020.safetensors",
+    "model.layers.21.mlp.experts.5.gate_proj.weight": "model-00020.safetensors",
+    "model.layers.21.mlp.experts.5.up_proj.weight": "model-00020.safetensors",
+    "model.layers.21.mlp.experts.5.down_proj.weight": "model-00020.safetensors",
+    "model.layers.21.mlp.experts.6.gate_proj.weight": "model-00020.safetensors",
+    "model.layers.21.mlp.experts.6.up_proj.weight": "model-00020.safetensors",
+    "model.layers.21.mlp.experts.6.down_proj.weight": "model-00020.safetensors",
+    "model.layers.21.post_attention_layernorm.weight": "model-00020.safetensors",
+    "model.layers.21.post_feedforward_layernorm.weight": "model-00020.safetensors",
+    "model.layers.22.self_attn.q_proj.weight": "model-00020.safetensors",
+    "model.layers.22.self_attn.k_proj.weight": "model-00020.safetensors",
+    "model.layers.22.self_attn.v_proj.weight": "model-00020.safetensors",
+    "model.layers.22.self_attn.o_proj.weight": "model-00020.safetensors",
+    "model.layers.22.self_attn.q_norm.weight": "model-00020.safetensors",
+    "model.layers.22.self_attn.k_norm.weight": "model-00020.safetensors",
+    "model.layers.22.mlp.gate.weight": "model-00020.safetensors",
+    "model.layers.22.mlp.experts.0.gate_proj.weight": "model-00020.safetensors",
+    "model.layers.22.mlp.experts.0.up_proj.weight": "model-00020.safetensors",
+    "model.layers.22.mlp.experts.0.down_proj.weight": "model-00020.safetensors",
+    "model.layers.22.mlp.experts.1.gate_proj.weight": "model-00020.safetensors",
+    "model.layers.22.mlp.experts.1.up_proj.weight": "model-00020.safetensors",
+    "model.layers.22.mlp.experts.1.down_proj.weight": "model-00021.safetensors",
+    "model.layers.22.mlp.experts.2.gate_proj.weight": "model-00021.safetensors",
+    "model.layers.22.mlp.experts.2.up_proj.weight": "model-00021.safetensors",
+    "model.layers.22.mlp.experts.2.down_proj.weight": "model-00021.safetensors",
+    "model.layers.22.mlp.experts.3.gate_proj.weight": "model-00021.safetensors",
+    "model.layers.22.mlp.experts.3.up_proj.weight": "model-00021.safetensors",
+    "model.layers.22.mlp.experts.3.down_proj.weight": "model-00021.safetensors",
+    "model.layers.22.mlp.experts.4.gate_proj.weight": "model-00021.safetensors",
+    "model.layers.22.mlp.experts.4.up_proj.weight": "model-00021.safetensors",
+    "model.layers.22.mlp.experts.4.down_proj.weight": "model-00021.safetensors",
+    "model.layers.22.mlp.experts.5.gate_proj.weight": "model-00021.safetensors",
+    "model.layers.22.mlp.experts.5.up_proj.weight": "model-00021.safetensors",
+    "model.layers.22.mlp.experts.5.down_proj.weight": "model-00021.safetensors",
+    "model.layers.22.mlp.experts.6.gate_proj.weight": "model-00021.safetensors",
+    "model.layers.22.mlp.experts.6.up_proj.weight": "model-00021.safetensors",
+    "model.layers.22.mlp.experts.6.down_proj.weight": "model-00021.safetensors",
+    "model.layers.22.post_attention_layernorm.weight": "model-00021.safetensors",
+    "model.layers.22.post_feedforward_layernorm.weight": "model-00021.safetensors",
+    "model.layers.23.self_attn.q_proj.weight": "model-00021.safetensors",
+    "model.layers.23.self_attn.k_proj.weight": "model-00021.safetensors",
+    "model.layers.23.self_attn.v_proj.weight": "model-00021.safetensors",
+    "model.layers.23.self_attn.o_proj.weight": "model-00021.safetensors",
+    "model.layers.23.self_attn.q_norm.weight": "model-00021.safetensors",
+    "model.layers.23.self_attn.k_norm.weight": "model-00021.safetensors",
+    "model.layers.23.mlp.gate.weight": "model-00021.safetensors",
+    "model.layers.23.mlp.experts.0.gate_proj.weight": "model-00021.safetensors",
+    "model.layers.23.mlp.experts.0.up_proj.weight": "model-00021.safetensors",
+    "model.layers.23.mlp.experts.0.down_proj.weight": "model-00021.safetensors",
+    "model.layers.23.mlp.experts.1.gate_proj.weight": "model-00021.safetensors",
+    "model.layers.23.mlp.experts.1.up_proj.weight": "model-00021.safetensors",
+    "model.layers.23.mlp.experts.1.down_proj.weight": "model-00021.safetensors",
+    "model.layers.23.mlp.experts.2.gate_proj.weight": "model-00021.safetensors",
+    "model.layers.23.mlp.experts.2.up_proj.weight": "model-00021.safetensors",
+    "model.layers.23.mlp.experts.2.down_proj.weight": "model-00021.safetensors",
+    "model.layers.23.mlp.experts.3.gate_proj.weight": "model-00022.safetensors",
+    "model.layers.23.mlp.experts.3.up_proj.weight": "model-00022.safetensors",
+    "model.layers.23.mlp.experts.3.down_proj.weight": "model-00022.safetensors",
+    "model.layers.23.mlp.experts.4.gate_proj.weight": "model-00022.safetensors",
+    "model.layers.23.mlp.experts.4.up_proj.weight": "model-00022.safetensors",
+    "model.layers.23.mlp.experts.4.down_proj.weight": "model-00022.safetensors",
+    "model.layers.23.mlp.experts.5.gate_proj.weight": "model-00022.safetensors",
+    "model.layers.23.mlp.experts.5.up_proj.weight": "model-00022.safetensors",
+    "model.layers.23.mlp.experts.5.down_proj.weight": "model-00022.safetensors",
+    "model.layers.23.mlp.experts.6.gate_proj.weight": "model-00022.safetensors",
+    "model.layers.23.mlp.experts.6.up_proj.weight": "model-00022.safetensors",
+    "model.layers.23.mlp.experts.6.down_proj.weight": "model-00022.safetensors",
+    "model.layers.23.post_attention_layernorm.weight": "model-00022.safetensors",
+    "model.layers.23.post_feedforward_layernorm.weight": "model-00022.safetensors",
+    "model.norm.weight": "model-00022.safetensors",
+    "lm_head.weight": "model-00022.safetensors"
+  }
+}

modeling_rish_ai.py

@@ -0,0 +1,617 @@
+# Copyright 2024 The HuggingFace Team. 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.
+
+from collections.abc import Callable
+
+import torch
+import torch.nn.functional as F
+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
+from transformers.modeling_layers import GradientCheckpointingLayer
+from transformers.modeling_outputs import MoeCausalLMOutputWithPast, MoeModelOutputWithPast
+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, logging
+from transformers.utils.deprecation import deprecate_kwarg
+from transformers.utils.generic import OutputRecorder, check_model_inputs
+
+from .configuration_rish_ai import RishAIConfig
+
+
+logger = logging.get_logger(__name__)
+
+
+@use_kernel_forward_from_hub("RMSNorm")
+class RishAIRMSNorm(nn.Module):
+    def __init__(self, hidden_size, eps=1e-6):
+        """
+        RishAIRMSNorm is equivalent to T5LayerNorm
+        """
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(hidden_size))
+        self.variance_epsilon = eps
+
+    def forward(self, hidden_states):
+        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 RishAIRotaryEmbedding(nn.Module):
+    inv_freq: torch.Tensor  # fix linting for `register_buffer`
+
+    def __init__(self, config: RishAIConfig, 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:
+            # Use the rope_type from config if available, otherwise default to 'dynamic'
+            self.rope_type = getattr(config, "rope_type", "dynamic")
+
+        # Ensure we have a valid rope_type
+        if self.rope_type not in ROPE_INIT_FUNCTIONS:
+            self.rope_type = "dynamic"  # fallback to dynamic if not found
+
+        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, sin
+
+
+class RishAIMLP(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 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: torch.Tensor | None,
+    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
+
+
+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."""
+    q_type, k_type = q.dtype, k.dtype
+    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.to(q_type), k_embed.to(k_type)
+
+
+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)
+
+
+class RishAIAttention(nn.Module):
+    """Multi-headed attention from 'Attention Is All You Need' paper"""
+
+    def __init__(self, config: RishAIConfig, layer_idx: int | None = None):
+        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.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.o_proj = nn.Linear(
+            config.num_attention_heads * self.head_dim, config.hidden_size, bias=config.attention_bias
+        )
+        self.q_norm = RishAIRMSNorm(config.num_attention_heads * self.head_dim, config.rms_norm_eps)
+        self.k_norm = RishAIRMSNorm(config.num_key_value_heads * self.head_dim, config.rms_norm_eps)
+
+    @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: torch.Tensor | None,
+        past_key_values: Cache | None = None,
+        cache_position: torch.LongTensor | None = None,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> tuple[torch.Tensor, torch.Tensor | None]:
+        input_shape = hidden_states.shape[:-1]
+        hidden_shape = (*input_shape, -1, self.head_dim)
+
+        query_states = self.q_norm(self.q_proj(hidden_states))
+        key_states = self.k_norm(self.k_proj(hidden_states))
+        value_states = self.v_proj(hidden_states)
+
+        query_states = query_states.view(hidden_shape).transpose(1, 2)
+        key_states = key_states.view(hidden_shape).transpose(1, 2)
+        value_states = value_states.view(hidden_shape).transpose(1, 2)
+
+        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 = past_key_values.update(key_states, value_states, self.layer_idx, cache_kwargs)
+
+        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,
+            **kwargs,
+        )
+
+        attn_output = attn_output.reshape(*input_shape, -1).contiguous()
+        attn_output = self.o_proj(attn_output)
+        return attn_output, attn_weights
+
+
+class RishAISparseMoeBlock(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.num_experts = config.num_experts
+        self.top_k = config.num_experts_per_tok
+        self.norm_topk_prob = config.norm_topk_prob
+        self.gate = nn.Linear(config.hidden_size, self.num_experts, bias=False)
+        self.experts = nn.ModuleList([RishAIMLP(config) for _ in range(self.num_experts)])
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        batch_size, sequence_length, hidden_dim = hidden_states.shape
+        hidden_states = hidden_states.view(-1, hidden_dim)
+        # router_logits: (batch * sequence_length, n_experts)
+        router_logits = self.gate(hidden_states)
+
+        routing_weights = F.softmax(router_logits, dim=1, dtype=torch.float)
+        routing_weights, selected_experts = torch.topk(routing_weights, self.top_k, dim=-1)
+        if self.norm_topk_prob:
+            routing_weights /= routing_weights.sum(dim=-1, keepdim=True)
+        # we cast back to the input dtype
+        routing_weights = routing_weights.to(hidden_states.dtype)
+
+        final_hidden_states = torch.zeros(
+            (batch_size * sequence_length, hidden_dim), dtype=hidden_states.dtype, device=hidden_states.device
+        )
+
+        # One hot encode the selected experts to create an expert mask
+        # this will be used to easily index which expert is going to be selected
+        expert_mask = torch.nn.functional.one_hot(selected_experts, num_classes=self.num_experts).permute(2, 1, 0)
+
+        # Loop over all available experts in the model and perform the computation on each expert
+        for expert_idx in range(self.num_experts):
+            expert_layer = self.experts[expert_idx]
+            idx, top_x = torch.where(expert_mask[expert_idx])
+
+            # Index the correct hidden states and compute the expert hidden state for
+            # the current expert. We need to make sure to multiply the output hidden
+            # states by `routing_weights` on the corresponding tokens (top-1 and top-2)
+            current_state = hidden_states[None, top_x].reshape(-1, hidden_dim)
+            current_hidden_states = expert_layer(current_state) * routing_weights[top_x, idx, None]
+
+            # However `index_add_` only support torch tensors for indexing so we'll use
+            # the `top_x` tensor here.
+            final_hidden_states.index_add_(0, top_x, current_hidden_states.to(hidden_states.dtype))
+        final_hidden_states = final_hidden_states.reshape(batch_size, sequence_length, hidden_dim)
+        return final_hidden_states, router_logits
+
+
+class RishAIDecoderLayer(GradientCheckpointingLayer):
+    def __init__(self, config: RishAIConfig, layer_idx: int):
+        super().__init__()
+        self.hidden_size = config.hidden_size
+        self.self_attn = RishAIAttention(config=config, layer_idx=layer_idx)
+
+        self.mlp = RishAISparseMoeBlock(config)
+        self.post_attention_layernorm = RishAIRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.post_feedforward_layernorm = RishAIRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+
+    @deprecate_kwarg("past_key_value", new_name="past_key_values", version="4.58")
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: torch.Tensor | None = None,
+        position_ids: torch.LongTensor | None = None,
+        past_key_values: Cache | None = None,
+        cache_position: torch.LongTensor | None = None,
+        position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None,
+        **kwargs,
+    ) -> torch.FloatTensor:
+        residual = 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,
+            cache_position=cache_position,
+            position_embeddings=position_embeddings,
+            **kwargs,
+        )
+        hidden_states = self.post_attention_layernorm(hidden_states)
+        hidden_states = residual + hidden_states
+
+        # Fully Connected
+        residual = hidden_states
+        hidden_states, _ = self.mlp(hidden_states)
+        hidden_states = self.post_feedforward_layernorm(hidden_states)
+        hidden_states = residual + hidden_states
+        return hidden_states
+
+
+@auto_docstring
+class RishAIPreTrainedModel(PreTrainedModel):
+    config: RishAIConfig
+    base_model_prefix = "model"
+    supports_gradient_checkpointing = True
+    _no_split_modules = ["RishAIDecoderLayer"]
+    _skip_keys_device_placement = ["past_key_values"]
+    _supports_flash_attn = True
+    _supports_sdpa = True
+    _supports_flex_attn = True
+    _can_compile_fullgraph = False  # MoE models don't work with torch.compile (`torch.where(condition)` not supported)
+    _supports_attention_backend = True
+    _can_record_outputs = {
+        "router_logits": OutputRecorder(RishAISparseMoeBlock, index=1),
+        "hidden_states": RishAIDecoderLayer,
+        "attentions": RishAIAttention,
+    }
+
+
+@auto_docstring
+class RishAIModel(RishAIPreTrainedModel):
+    def __init__(self, config: RishAIConfig):
+        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(
+            [RishAIDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]
+        )
+        self.norm = RishAIRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.rotary_emb = RishAIRotaryEmbedding(config=config)
+        self.gradient_checkpointing = False
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    @check_model_inputs()
+    @auto_docstring
+    def forward(
+        self,
+        input_ids: torch.LongTensor | None = None,
+        attention_mask: torch.Tensor | None = None,
+        position_ids: torch.LongTensor | None = None,
+        past_key_values: Cache | None = None,
+        inputs_embeds: torch.FloatTensor | None = None,
+        use_cache: bool | None = None,
+        cache_position: torch.LongTensor | None = None,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> MoeModelOutputWithPast:
+        if (input_ids is None) ^ (inputs_embeds is not None):
+            raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
+
+        if use_cache and past_key_values is None:
+            past_key_values = DynamicCache(config=self.config)
+
+        if inputs_embeds is None:
+            inputs_embeds = self.embed_tokens(input_ids)
+
+        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)
+
+        causal_mask = create_causal_mask(
+            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,
+        )
+
+        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,
+                position_embeddings=position_embeddings,
+                attention_mask=causal_mask,
+                position_ids=position_ids,
+                past_key_values=past_key_values,
+                use_cache=use_cache,
+                cache_position=cache_position,
+                **kwargs,
+            )
+
+        hidden_states = self.norm(hidden_states)
+
+        return MoeModelOutputWithPast(
+            last_hidden_state=hidden_states,
+            past_key_values=past_key_values,
+        )
+
+
+def load_balancing_loss_func(
+    gate_logits: torch.Tensor | tuple[torch.Tensor] | None,
+    num_experts: int | None = None,
+    top_k=2,
+    attention_mask: torch.Tensor | None = None,
+) -> torch.Tensor | int:
+    r"""
+    Computes the load balancing loss for the MoE router.
+
+    Args:
+        gate_logits:
+            Logits from the `gate`, should be a tuple of model.config.num_hidden_layers tensors of
+            shape [batch_size X sequence_length, num_experts].
+        num_experts:
+            Number of experts
+        top_k:
+            The number of experts to route per-token, can be also interpreted as the `top-k` routing
+            parameter.
+        attention_mask (`torch.Tensor`, *optional*):
+            The attention_mask used in forward function
+            shape [batch_size X sequence_length] if not None.
+
+    Returns:
+        The auxiliary loss.
+    """
+    if gate_logits is None or not isinstance(gate_logits, tuple):
+        return 0
+
+    if isinstance(gate_logits, tuple):
+        compute_device = gate_logits[0].device
+        concatenated_gate_logits = torch.cat([layer_gate.to(compute_device) for layer_gate in gate_logits], dim=0)
+
+    routing_weights = torch.nn.functional.softmax(concatenated_gate_logits, dim=-1)
+
+    _, selected_experts = torch.topk(routing_weights, top_k, dim=-1)
+
+    expert_mask = torch.nn.functional.one_hot(selected_experts, num_experts)
+
+    if attention_mask is None:
+        # Compute the percentage of tokens routed to each experts
+        tokens_per_expert = torch.mean(expert_mask.float(), dim=0)
+
+        # Compute the average probability of routing to these experts
+        router_prob_per_expert = torch.mean(routing_weights, dim=0)
+    else:
+        batch_size, sequence_length = attention_mask.shape
+        num_hidden_layers = concatenated_gate_logits.shape[0] // (batch_size * sequence_length)
+
+        # Compute the mask that masks all padding tokens as 0 with the same shape of expert_mask
+        expert_attention_mask = (
+            attention_mask[None, :, :, None, None]
+            .expand((num_hidden_layers, batch_size, sequence_length, top_k, num_experts))
+            .reshape(-1, top_k, num_experts)
+            .to(compute_device)
+        )
+
+        # Compute the percentage of tokens routed to each experts
+        tokens_per_expert = torch.sum(expert_mask.float() * expert_attention_mask, dim=0) / torch.sum(
+            expert_attention_mask, dim=0
+        )
+
+        # Compute the mask that masks all padding tokens as 0 with the same shape of tokens_per_expert
+        router_per_expert_attention_mask = (
+            attention_mask[None, :, :, None]
+            .expand((num_hidden_layers, batch_size, sequence_length, routing_weights.shape[1]))
+            .reshape(-1, routing_weights.shape[1])
+            .to(compute_device)
+        )
+
+        # Compute the average probability of routing to these experts
+        router_prob_per_expert = torch.sum(routing_weights * router_per_expert_attention_mask, dim=0) / torch.sum(
+            router_per_expert_attention_mask, dim=0
+        )
+
+    device_index = routing_weights.device.index if routing_weights.device.index is not None else 0
+    rank = routing_weights.shape[1] * int(device_index)
+    overall_loss = torch.sum(
+        tokens_per_expert[:, rank : rank + routing_weights.shape[1]] * router_prob_per_expert.unsqueeze(0)
+    )
+    return overall_loss * num_experts
+
+
+class RishAICausalLM(RishAIPreTrainedModel, GenerationMixin):
+    _tied_weights_keys = ["lm_head.weight"]
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.model = RishAIModel(config)
+        self.vocab_size = config.vocab_size
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+
+        self.router_aux_loss_coef = config.router_aux_loss_coef
+        self.num_experts = config.num_experts
+        self.num_experts_per_tok = config.num_experts_per_tok
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    @auto_docstring
+    def forward(
+        self,
+        input_ids: torch.LongTensor | None = None,
+        attention_mask: torch.Tensor | None = None,
+        position_ids: torch.LongTensor | None = None,
+        past_key_values: Cache | None = None,
+        inputs_embeds: torch.FloatTensor | None = None,
+        labels: torch.LongTensor | None = None,
+        use_cache: bool | None = None,
+        output_attentions: bool | None = None,
+        output_hidden_states: bool | None = None,
+        output_router_logits: bool | None = None,
+        return_dict: bool | None = None,
+        cache_position: torch.LongTensor | None = None,
+        logits_to_keep: int | torch.Tensor = 0,
+        **kwargs,
+    ) -> tuple | MoeCausalLMOutputWithPast:
+        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+        output_router_logits = (
+            output_router_logits if output_router_logits is not None else self.config.output_router_logits
+        )
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        )
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
+        outputs = 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,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            output_router_logits=output_router_logits,
+            return_dict=return_dict,
+            cache_position=cache_position,
+        )
+
+        hidden_states = outputs[0]
+        # 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, labels, self.vocab_size, **kwargs)
+
+        aux_loss = None
+        if output_router_logits:
+            aux_loss = load_balancing_loss_func(
+                outputs.router_logits if return_dict else outputs[-1],
+                self.num_experts,
+                self.num_experts_per_tok,
+                attention_mask,
+            )
+            if labels is not None:
+                loss += self.router_aux_loss_coef * aux_loss.to(loss.device)  # make sure to reside in the same device
+
+        if not return_dict:
+            output = (logits,) + outputs[1:]
+            if output_router_logits:
+                output = (aux_loss,) + output
+            return (loss,) + output if loss is not None else output
+
+        return MoeCausalLMOutputWithPast(
+            loss=loss,
+            aux_loss=aux_loss,
+            logits=logits,
+            past_key_values=outputs.past_key_values,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+            router_logits=outputs.router_logits,
+        )

tokenization_rish_ai.py

@@ -0,0 +1,178 @@
+# Copyright 2024 The HuggingFace Team. 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.
+
+"""Tokenization class for RishAI."""
+
+import json
+
+from transformers.tokenization_utils_base import PreTrainedTokenizerBase
+from transformers.utils import add_end_docstrings, logging
+
+
+logger = logging.get_logger(__name__)
+
+
+@add_end_docstrings
+class RishAITokenizer(PreTrainedTokenizerBase):
+    """
+    Construct a RishAI tokenizer. Based on byte-level Byte-Pair-Encoding.
+
+    This tokenizer inherits from [`PreTrainedTokenizerBase`] which contains most of the main methods.
+    Users should refer to this superclass for more information regarding those methods.
+
+    Args:
+        vocab_file (`str`):
+            Path to the vocabulary file.
+        merges_file (`str`):
+            Path to the merges file.
+        errors (`str`, *optional*, defaults to `"replace"`):
+            Paradigm to follow when decoding bytes to UTF-8. See
+            [bytes.decode](https://docs.python.org/3/library/stdtypes.html#bytes.decode) for more information.
+        unk_token (`str`, *optional*, defaults to `"<|endoftext|>"`):
+            The unknown token. A token that is not in the vocabulary cannot be converted to an ID and is set to be this
+            token instead.
+        bos_token (`str`, *optional*, defaults to `"<|endoftext|>"`):
+            The beginning of sequence token.
+        eos_token (`str`, *optional*, defaults to `"<|endoftext|>"`):
+            The end of sequence token.
+        pad_token (`str`, *optional*):
+            The token used for padding, for example when batching sequences of different lengths.
+        clean_up_tokenization_spaces (`bool`, *optional*, defaults to `False`):
+            Whether or not to cleanup spaces after decoding, cleanup consists in removing potential artifacts like
+            extra spaces.
+        split_special_tokens (`bool`, *optional*, defaults to `False`):
+            Whether or not the special tokens should be split during the encoding.
+    """
+
+    vocab_files_names = {
+        "vocab_file": "vocab.json",
+        "merges_file": "merges.txt",
+    }
+    pretrained_vocab_files_map = {
+        "vocab_file": {},
+        "merges_file": {},
+    }
+    max_model_input_sizes = {"default": 4096}
+    model_input_names = ["input_ids", "attention_mask"]
+
+    def __init__(
+        self,
+        vocab_file=None,
+        merges_file=None,
+        errors="replace",
+        unk_token="<|endoftext|>",
+        bos_token="<|endoftext|>",
+        eos_token="<|endoftext|>",
+        pad_token=None,
+        clean_up_tokenization_spaces=False,
+        split_special_tokens=False,
+        **kwargs,
+    ):
+        # Set default special tokens if not provided
+        if pad_token is None:
+            pad_token = "<|endoftext|>"
+
+        super().__init__(
+            errors=errors,
+            unk_token=unk_token,
+            bos_token=bos_token,
+            eos_token=eos_token,
+            pad_token=pad_token,
+            clean_up_tokenization_spaces=clean_up_tokenization_spaces,
+            split_special_tokens=split_special_tokens,
+            **kwargs,
+        )
+
+        self.vocab_file = vocab_file
+        self.merges_file = merges_file
+
+        # Initialize vocabulary
+        self._vocab = {}
+        self._merges = []
+        self._bpe_ranks = {}
+
+        if vocab_file is not None and merges_file is not None:
+            self._load_vocab_and_merges(vocab_file, merges_file)
+
+    def _load_vocab_and_merges(self, vocab_file, merges_file):
+        """Load vocabulary and merges from files."""
+        # Load vocabulary
+        with open(vocab_file, "r", encoding="utf-8") as f:
+            self._vocab = json.load(f)
+
+        # Load merges
+        with open(merges_file, "r", encoding="utf-8") as f:
+            self._merges = f.read().split("\n")
+            self._merges = [merge for merge in self._merges if merge.strip()]
+
+        # Build BPE ranks
+        self._bpe_ranks = {merge: i for i, merge in enumerate(self._merges)}
+
+    @property
+    def vocab_size(self) -> int:
+        """Returns vocab size."""
+        return len(self._vocab)
+
+    def get_vocab(self) -> dict[str, int]:
+        """Returns vocab as a dict."""
+        return dict(self._vocab)
+
+    def _tokenize(self, text: str, **kwargs) -> list[str]:
+        """Tokenize a string."""
+        # Simple whitespace tokenization for now
+        # In a real implementation, this would use BPE
+        return text.split()
+
+    def _convert_token_to_id(self, token: str) -> int:
+        """Converts a token (str) to an id using the vocab."""
+        return self._vocab.get(token, self._vocab.get(self.unk_token, 0))
+
+    def _convert_id_to_token(self, index: int) -> str:
+        """Converts an index (integer) to a token (str) using the vocab."""
+        for token, idx in self._vocab.items():
+            if idx == index:
+                return token
+        return self.unk_token
+
+    def convert_tokens_to_string(self, tokens: list[str]) -> str:
+        """Converts a sequence of tokens (string) in a single string."""
+        # Simple detokenization - join with spaces
+        return " ".join(tokens)
+
+    def save_vocabulary(self, save_directory: str, filename_prefix: str | None = None) -> tuple[str, str]:
+        """Save the vocabulary and merges files to a directory."""
+        if not self.can_save_slow_tokenizer:
+            raise ValueError(
+                "Your tokenizer does not have the necessary information to save the vocabulary. "
+                "Please use a tokenizer that has been trained with the correct parameters."
+            )
+
+        vocab_file = (filename_prefix + "-" if filename_prefix else "") + "vocab.json"
+        merges_file = (filename_prefix + "-" if filename_prefix else "") + "merges.txt"
+
+        vocab_file_path = f"{save_directory}/{vocab_file}"
+        merges_file_path = f"{save_directory}/{merges_file}"
+
+        with open(vocab_file_path, "w", encoding="utf-8") as f:
+            json.dump(self._vocab, f, ensure_ascii=False, indent=2)
+
+        with open(merges_file_path, "w", encoding="utf-8") as f:
+            f.write("\n".join(self._merges))
+
+        return vocab_file_path, merges_file_path
+
+    @property
+    def can_save_slow_tokenizer(self) -> bool:
+        """Check if the tokenizer can be saved."""
+        return self._vocab is not None and self._merges is not None

tokenizer.json

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

tokenizer_config.json

@@ -0,0 +1,30 @@
+{
+  "add_prefix_space": false,
+  "additional_special_tokens": null,
+  "backend": "tokenizers",
+  "bos_token": null,
+  "clean_up_tokenization_spaces": false,
+  "eos_token": "<|im_end|>",
+  "errors": "replace",
+  "extra_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|>"
+  ],
+  "is_local": false,
+  "model_max_length": 32768,
+  "pad_token": "<|endoftext|>",
+  "split_special_tokens": false,
+  "tokenizer_class": "Qwen2Tokenizer",
+  "unk_token": null
+}