Upload MiMo-V2-Flash-JANG_2S

.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,51 @@
+---
+base_model: XiaomiMiMo/MiMo-V2-Flash
+library_name: mlx
+pipeline_tag: text-generation
+license: apache-2.0
+tags:
+- mlx
+- jang
+- jang-quantized
+- JANG_2S
+- mixed-precision
+- apple-silicon
+---
+
+# MiMo-V2-Flash-JANG_2S
+
+JANG adaptive mixed-precision MLX quantization produced via [vmlx / jang-tools](https://github.com/jjang-ai/jangq).
+
+- **Quantization:** 3.06b avg, profile JANG_2S, method mse, calibration weights
+- **Profile:** JANG_2S
+- **Format:** JANG v2 MLX safetensors
+- **Compatible with:** vmlx, MLX Studio, oMLX (with JANG patch)
+
+## Usage
+
+### vmlx (recommended)
+
+```bash
+pip install 'vmlx[jang]'
+vmlx serve bearzi/MiMo-V2-Flash-JANG_2S
+```
+
+### Python
+
+```python
+from jang_tools.loader import load_jang_model
+from mlx_lm import generate
+
+model, tokenizer = load_jang_model("bearzi/MiMo-V2-Flash-JANG_2S")
+messages = [{"role": "user", "content": "Hello"}]
+prompt = tokenizer.apply_chat_template(messages, add_generation_prompt=True, tokenize=False)
+print(generate(model, tokenizer, prompt=prompt, max_tokens=512, verbose=True))
+```
+
+## About JANG
+
+JANG (Jang Adaptive N-bit Grading) assigns different bit widths to different layer types — attention layers get more bits, MLP/expert layers compress harder. This preserves model coherence at aggressive compression levels where uniform quantization breaks down.
+
+See [JANG documentation](https://github.com/jjang-ai/jangq) and scores at [jangq.ai](https://jangq.ai).
+
+Comparative benchmarks and feedback welcome — please open a discussion.

added_tokens.json

@@ -0,0 +1,28 @@
+{
+  "</think>": 151668,
+  "</tool_call>": 151658,
+  "</tool_response>": 151666,
+  "<think>": 151667,
+  "<tool_call>": 151657,
+  "<tool_response>": 151665,
+  "<|box_end|>": 151649,
+  "<|box_start|>": 151648,
+  "<|endoftext|>": 151643,
+  "<|file_sep|>": 151664,
+  "<|fim_middle|>": 151660,
+  "<|fim_pad|>": 151662,
+  "<|fim_prefix|>": 151659,
+  "<|fim_suffix|>": 151661,
+  "<|im_end|>": 151645,
+  "<|im_start|>": 151644,
+  "<|image_pad|>": 151655,
+  "<|object_ref_end|>": 151647,
+  "<|object_ref_start|>": 151646,
+  "<|quad_end|>": 151651,
+  "<|quad_start|>": 151650,
+  "<|repo_name|>": 151663,
+  "<|video_pad|>": 151656,
+  "<|vision_end|>": 151653,
+  "<|vision_pad|>": 151654,
+  "<|vision_start|>": 151652
+}

config.json

@@ -0,0 +1,156 @@
+{
+  "architectures": [
+    "MiMoV2FlashForCausalLM"
+  ],
+  "auto_map": {
+    "AutoConfig": "configuration_mimo_v2_flash.MiMoV2FlashConfig",
+    "AutoModel": "modeling_mimo_v2_flash.MiMoV2FlashModel",
+    "AutoModelForCausalLM": "modeling_mimo_v2_flash.MiMoV2FlashForCausalLM"
+  },
+  "attention_dropout": 0.0,
+  "attention_value_scale": 0.707,
+  "hidden_act": "silu",
+  "hidden_size": 4096,
+  "initializer_range": 0.02,
+  "intermediate_size": 16384,
+  "max_position_embeddings": 262144,
+  "model_type": "mimo_v2_flash",
+  "num_attention_heads": 64,
+  "head_dim": 192,
+  "num_hidden_layers": 48,
+  "num_key_value_heads": 4,
+  "layernorm_epsilon": 1e-05,
+  "rope_theta": 5000000,
+  "tie_word_embeddings": false,
+  "torch_dtype": "bfloat16",
+  "transformers_version": "4.40.1",
+  "use_cache": true,
+  "vocab_size": 152576,
+  "partial_rotary_factor": 0.334,
+  "sliding_window": 128,
+  "swa_rope_theta": 10000,
+  "attention_bias": false,
+  "v_head_dim": 128,
+  "hybrid_layer_pattern": [
+    0,
+    1,
+    1,
+    1,
+    1,
+    0,
+    1,
+    1,
+    1,
+    1,
+    1,
+    0,
+    1,
+    1,
+    1,
+    1,
+    1,
+    0,
+    1,
+    1,
+    1,
+    1,
+    1,
+    0,
+    1,
+    1,
+    1,
+    1,
+    1,
+    0,
+    1,
+    1,
+    1,
+    1,
+    1,
+    0,
+    1,
+    1,
+    1,
+    1,
+    1,
+    0,
+    1,
+    1,
+    1,
+    1,
+    1,
+    0
+  ],
+  "add_swa_attention_sink_bias": true,
+  "add_full_attention_sink_bias": false,
+  "sliding_window_size": 128,
+  "attention_chunk_size": 128,
+  "moe_layer_freq": [
+    0,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1
+  ],
+  "moe_intermediate_size": 2048,
+  "n_routed_experts": 256,
+  "n_shared_experts": null,
+  "num_experts_per_tok": 8,
+  "norm_topk_prob": true,
+  "scoring_func": "sigmoid",
+  "n_group": 1,
+  "topk_group": 1,
+  "topk_method": "noaux_tc",
+  "routed_scaling_factor": null,
+  "swa_num_attention_heads": 64,
+  "swa_num_key_value_heads": 8,
+  "swa_head_dim": 192,
+  "swa_v_head_dim": 128,
+  "quantization": {
+    "group_size": 128,
+    "bits": 2
+  }
+}

configuration_mimo_v2_flash.py

@@ -0,0 +1,109 @@
+# coding=utf-8
+#
+# Copyright 2025 Xiaomi Corporation.
+# Copyright 2025 The HuggingFace Inc. team.
+#
+# 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
+from transformers.modeling_rope_utils import rope_config_validation
+from transformers.utils import logging
+
+
+logger = logging.get_logger(__name__)
+
+
+class MiMoV2FlashConfig(PretrainedConfig):
+
+    model_type = ""
+    keys_to_ignore_at_inference = ["past_key_values"]
+
+    # Default tensor parallel plan for base model `Hybrid`
+    base_model_tp_plan = {
+        "layers.*.self_attn.q_proj": "colwise",
+        "layers.*.self_attn.k_proj": "colwise",
+        "layers.*.self_attn.v_proj": "colwise",
+        "layers.*.self_attn.o_proj": "rowwise",
+        "layers.*.mlp.gate_proj": "colwise",
+        "layers.*.mlp.up_proj": "colwise",
+        "layers.*.mlp.down_proj": "rowwise",
+    }
+    base_model_pp_plan = {
+        "embed_tokens": (["input_ids"], ["inputs_embeds"]),
+        "layers": (["hidden_states", "attention_mask"], ["hidden_states"]),
+        "norm": (["hidden_states"], ["hidden_states"]),
+    }
+
+    attribute_map = {
+        "num_local_experts": "n_routed_experts",
+    }
+
+    def __init__(
+        self,
+        vocab_size=151936,
+        hidden_size=4096,
+        intermediate_size=22016,
+        num_hidden_layers=32,
+        num_attention_heads=32,
+        num_key_value_heads=32,
+        hidden_act="silu",
+        max_position_embeddings=32768,
+        initializer_range=0.02,
+        layernorm_epsilon=1e-6,
+        use_cache=True,
+        tie_word_embeddings=False,
+        rope_theta=10000.0,
+        rope_scaling=None,
+        attention_dropout=0.0,
+        hybrid_block_size=None,
+        hybrid_layer_pattern=None,
+        partial_rotary_factor=1.0,
+        **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.layernorm_epsilon = layernorm_epsilon
+        self.use_cache = use_cache
+        self.rope_theta = rope_theta
+        self.rope_scaling = rope_scaling
+        self.attention_dropout = attention_dropout
+
+        if hybrid_block_size is not None and hybrid_layer_pattern is None:
+            hybrid_layer_pattern = [0 if ((i + 1) % hybrid_block_size == 0) else 1 for i in range(num_hidden_layers)]
+        self.hybrid_block_size = hybrid_block_size
+        self.hybrid_layer_pattern = hybrid_layer_pattern
+
+        self.partial_rotary_factor = partial_rotary_factor
+
+        # Validate the correctness of rotary position embeddings parameters
+        # BC: if there is a 'type' field, move it to 'rope_type'.
+        if self.rope_scaling is not None and "type" in self.rope_scaling:
+            self.rope_scaling["rope_type"] = self.rope_scaling["type"]
+        rope_config_validation(self)
+
+        super().__init__(
+            tie_word_embeddings=tie_word_embeddings,
+            **kwargs,
+        )

jang_config.json

@@ -0,0 +1,49 @@
+{
+  "quantization": {
+    "method": "jang-importance",
+    "profile": "JANG_2S",
+    "target_bits": 2,
+    "actual_bits": 3.06,
+    "block_size": 128,
+    "calibration_method": "weights",
+    "quantization_method": "mse",
+    "scoring_method": "weight-magnitude",
+    "bit_widths_used": [
+      2,
+      3,
+      4,
+      6
+    ],
+    "quantization_scheme": "asymmetric",
+    "quantization_backend": "mx.quantize",
+    "hadamard_rotation": false
+  },
+  "source_model": {
+    "name": "MiMo-V2-Flash",
+    "dtype": "bfloat16",
+    "parameters": "13.5B"
+  },
+  "architecture": {
+    "type": "moe",
+    "attention": "gqa",
+    "has_vision": false,
+    "has_ssm": false,
+    "has_moe": true
+  },
+  "runtime": {
+    "total_weight_bytes": 171442176,
+    "total_weight_gb": 0.16
+  },
+  "capabilities": {
+    "reasoning_parser": "deepseek_r1",
+    "tool_parser": "qwen",
+    "think_in_template": true,
+    "supports_tools": true,
+    "supports_thinking": true,
+    "family": "mimo_v2_flash",
+    "modality": "text",
+    "cache_type": "kv"
+  },
+  "format": "jang",
+  "format_version": "2.0"
+}

model-00001-of-00144.safetensors

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

model-00002-of-00144.safetensors

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

model-00003-of-00144.safetensors

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

model-00004-of-00144.safetensors

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

model-00005-of-00144.safetensors

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

model-00006-of-00144.safetensors

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

model-00007-of-00144.safetensors

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

model-00008-of-00144.safetensors

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

model-00009-of-00144.safetensors

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

model-00010-of-00144.safetensors

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

model-00011-of-00144.safetensors

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

model-00012-of-00144.safetensors

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

model-00013-of-00144.safetensors

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

model-00014-of-00144.safetensors

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

model-00015-of-00144.safetensors

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

model-00016-of-00144.safetensors

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

model-00017-of-00144.safetensors

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

model-00018-of-00144.safetensors

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

model-00019-of-00144.safetensors

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

model-00020-of-00144.safetensors

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

model-00021-of-00144.safetensors

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

model-00022-of-00144.safetensors

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

model-00023-of-00144.safetensors

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

model-00024-of-00144.safetensors

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

model-00025-of-00144.safetensors

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

model-00026-of-00144.safetensors

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

model-00027-of-00144.safetensors

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

model-00028-of-00144.safetensors

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

model-00029-of-00144.safetensors

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

model-00030-of-00144.safetensors

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

model-00031-of-00144.safetensors

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

model-00032-of-00144.safetensors

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

model-00033-of-00144.safetensors

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

model-00034-of-00144.safetensors

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

model-00035-of-00144.safetensors

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

model-00036-of-00144.safetensors

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

model-00037-of-00144.safetensors

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

model-00038-of-00144.safetensors

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

model-00039-of-00144.safetensors

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

model-00040-of-00144.safetensors

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

model-00041-of-00144.safetensors

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

model-00042-of-00144.safetensors

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

model-00043-of-00144.safetensors

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