Upload folder using huggingface_hub

.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,85 @@
+---
+license: mit
+language:
+- zh
+- en
+- yue
+pipeline_tag: automatic-speech-recognition
+tags:
+- safetensors
+- fp8
+- quantization
+- speech-recognition
+base_model: XiaomiMiMo/MiMo-V2.5-ASR
+---
+
+# MiMo-V2.5-ASR — FP8 (e4m3fn)
+
+FP8-quantized build of [XiaomiMiMo/MiMo-V2.5-ASR](https://huggingface.co/XiaomiMiMo/MiMo-V2.5-ASR),
+the Xiaomi MiMo end-to-end ASR model with native Mandarin/English code-switching,
+Chinese dialects, song lyrics, noisy/multi-speaker robustness, and native punctuation.
+
+## What this is
+
+- **Weights:** `float8_e4m3fn`, per-output-channel absmax scaling (one fp32 scale per row), baked at save time.
+- **Activations:** dynamic per-tensor fp8 quantization each forward pass.
+- **Matmul:** `torch._scaled_mm` (FP8 tensor cores on Ada / Hopper / Blackwell).
+- **Skipped (kept bf16):** embeddings, RMSNorm/LayerNorm, biases. 417 Linear layers converted.
+- The audio encoder/tokenizer ([MiMo-Audio-Tokenizer](https://huggingface.co/XiaomiMiMo/MiMo-Audio-Tokenizer)) is **not** quantized; download it separately for inference.
+
+This roughly halves the LLM weight footprint (~32 GB bf16 → ~16-17 GB on disk).
+
+## Important: this is NOT a drop-in `from_pretrained` checkpoint
+
+`model.safetensors` stores custom `FP8Linear` buffers (`*.weight_fp8`, `*.weight_scale`),
+not standard HF Linear weights. It must be loaded through the matching `FP8Linear`
+modules. Use the loader below.
+
+## Usage
+
+```bash
+git clone https://github.com/XiaomiMiMo/MiMo-V2.5-ASR.git
+cd MiMo-V2.5-ASR
+pip install -r requirements.txt
+pip install flash-attn==2.7.4.post1   # required by the audio tokenizer
+hf download XiaomiMiMo/MiMo-Audio-Tokenizer --local-dir ./models/MiMo-Audio-Tokenizer
+hf download Infatoshi/MiMo-V2.5-ASR-FP8 --local-dir ./MiMo-V2.5-ASR-FP8
+```
+
+Then load with the `FP8Linear` loader (`quantize_fp8.py`, included here as `quantize_fp8.py`):
+
+```python
+from quantize_fp8 import load_fp8_model
+mimo = load_fp8_model(
+    fp8_dir="./MiMo-V2.5-ASR-FP8",
+    tokenizer_path="./models/MiMo-Audio-Tokenizer",
+    repo_root=".",                     # the cloned MiMo-V2.5-ASR repo
+)
+print(mimo.asr_sft("audio.wav", audio_tag="<english>"))
+```
+
+## Quantization fidelity
+
+Per-output-channel absmax dequant error vs the original fp32 weights, sampled across
+depth (layers 0/17/35), all attn+mlp projections, lm_head, and the audio local transformer:
+
+- relative Frobenius error: **~0.026, uniform** across every sampled layer (max 0.027 on lm_head)
+- no corrupted or outlier layers
+
+This is the expected magnitude for fp8 e4m3 with per-channel scaling (3 mantissa bits).
+
+## Requirements
+
+- CUDA GPU with FP8 tensor cores (Ada / Hopper / Blackwell), CUDA >= 12.0
+- torch >= 2.6, safetensors
+- **Blackwell (sm_120, e.g. RTX PRO 6000 / RTX 50xx):** use a torch build with CUDA 12.8+
+  (torch >= 2.7, `cu128`). torch 2.6 `cu124` ships no sm_120 kernels and will fail with
+  "no kernel image is available for execution on the device".
+
+## Notes / caveats
+
+- FP8 e4m3fn weight-only-style quantization is lossy; expect small WER deltas vs bf16.
+- Per-tensor dynamic activation scaling is simple and fast but less accurate than
+  per-token scaling on activations with large outliers.
+
+Derivative of an MIT-licensed model; original credit to the Xiaomi MiMo team.

added_tokens.json

@@ -0,0 +1,32 @@
+{
+  "</tool_call>": 151658,
+  "<tool_call>": 151657,
+  "<|Human|>": 151668,
+  "<|SpeechLM|>": 151669,
+  "<|box_end|>": 151649,
+  "<|box_start|>": 151648,
+  "<|empty|>": 151667,
+  "<|endoftext|>": 151643,
+  "<|eosp|>": 151666,
+  "<|eostm|>": 151671,
+  "<|eot|>": 151672,
+  "<|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,
+  "<|sosp|>": 151665,
+  "<|sostm|>": 151670,
+  "<|video_pad|>": 151656,
+  "<|vision_end|>": 151653,
+  "<|vision_pad|>": 151654,
+  "<|vision_start|>": 151652
+}

chat_template.jinja

@@ -0,0 +1,170 @@
+{%- if not add_generation_prompt is defined -%}
+    {%- set add_generation_prompt = false -%}
+{%- endif -%}
+{%- if not enable_thinking is defined -%}
+    {%- set enable_thinking = false -%}
+{%- endif -%}
+{%- if not keep_all_reasoning is defined -%}
+    {%- set keep_all_reasoning = false -%}
+{%- endif -%}
+{%- macro render_extra_keys(json_dict, handled_keys) -%}
+    {%- if json_dict is mapping %}
+        {%- for json_key in json_dict if json_key not in handled_keys %}
+            {%- if json_dict[json_key] is mapping or (json_dict[json_key] is sequence and json_dict[json_key] is not string) %}
+                {{- '\n<' ~ json_key ~ '>' ~ (json_dict[json_key] | tojson | safe) ~ '</' ~ json_key ~ '>' }}
+            {%- else %}
+                {{-'\n<' ~ json_key ~ '>' ~ (json_dict[json_key] | string) ~ '</' ~ json_key ~ '>' }}
+            {%- endif %}
+        {%- endfor %}
+    {%- endif %}
+{%- endmacro -%}
+{%- macro render_content(message_content) -%}
+    {%- if message_content is string -%}
+        {{- message_content -}}
+    {%- else -%}
+        {%- for content in message_content -%}
+            {%- if content['type'] == 'image' or 'image' in content or 'image_url' in content -%}
+                {{- '<|vision_start|><|image_pad|><|vision_end|>' -}}
+            {%- elif content['type'] == 'audio' or 'audio' in content or 'audio_url' in content -%}
+                {{- '<|sosp|><|empty|><|eosp|>' -}}
+            {%- elif content['type'] == 'video' or 'video' in content or 'video_url' in content -%}
+                {{- '<|vision_start|><|video_pad|><|vision_end|>' -}}
+            {%- elif 'text' in content -%}
+                {{- content['text'] -}}
+            {%- endif -%}
+        {%- endfor -%}
+    {%- endif -%}
+{%- endmacro -%}
+{%- if messages[0]["role"] == "system" %}
+    {%- set system_message = messages[0]["content"] %}
+    {%- set loop_messages = messages[1:] %}
+{%- else %}
+    {%- set loop_messages = messages %}
+{%- endif %}
+{%- set ns = namespace(last_user_index=-1, assistant_is_last=false) %}
+{%- for m in loop_messages %}
+    {%- if m.role == 'user' %}
+        {%- set ns.last_user_index = loop.index0 -%}
+    {%- endif %}
+{%- endfor %}
+{%- if not tools is defined %}
+    {%- set tools = [] %}
+{%- endif %}
+{%- set has_system = false %}
+{%- if system_message is defined %}
+    {{- "<|im_start|>system\n" + system_message }}
+    {%- set has_system = true %}
+{%- endif %}
+{%- if tools is iterable and tools | length > 0 %}
+    {%- if not has_system %}
+        {{- "<|im_start|>system\n" }}
+        {%- set has_system = true %}
+    {%- endif %}
+    {{- "\n\n# Tools\n\nYou may call one or more functions to assist with the user query.\n\nYou have access to the following functions:\n\n" }}
+    {{- "<tools>" }}
+    {%- for tool in tools %}
+        {%- if tool.function is defined %}
+            {%- set tool = tool.function %}
+        {%- endif %}
+        {{- "\n<function>\n<name>" ~ tool.name ~ "</name>" }}
+        {%- if tool.description is defined %}
+            {{- '\n<description>' ~ (tool.description | trim) ~ '</description>' }}
+        {%- endif %}
+        {{- '\n<parameters>' }}
+        {%- if tool.parameters is defined and tool.parameters is mapping and tool.parameters.properties is defined and tool.parameters.properties is mapping %}
+            {%- for param_name, param_fields in tool.parameters.properties|items %}
+                {{- '\n<parameter>' }}
+                {{- '\n<name>' ~ param_name ~ '</name>' }}
+                {%- if param_fields.type is defined %}
+                    {{- '\n<type>' ~ (param_fields.type | string) ~ '</type>' }}
+                {%- endif %}
+                {%- if param_fields.description is defined %}
+                    {{- '\n<description>' ~ (param_fields.description | trim) ~ '</description>' }}
+                {%- endif %}
+                {%- set handled_keys = ['name', 'type', 'description'] %}
+                {{- render_extra_keys(param_fields, handled_keys) }}
+                {{- '\n</parameter>' }}
+            {%- endfor %}
+        {%- endif %}
+        {%- set handled_keys = ['type', 'properties'] %}
+        {{- render_extra_keys(tool.parameters, handled_keys) }}
+        {{- '\n</parameters>' }}
+        {%- set handled_keys = ['type', 'name', 'description', 'parameters'] %}
+        {{- render_extra_keys(tool, handled_keys) }}
+        {{- '\n</function>' }}
+    {%- endfor %}
+    {{- "\n</tools>" }}
+    {{- '\n\nFor each function call, output the function name and arguments in the following format:\n<tool_call>\n<function=example_function_name>\n<parameter=example_parameter_1>value_1</parameter>\n<parameter=example_parameter_2>This is the value for the second parameter\nthat can span\nmultiple lines</parameter>\n</function>\n</tool_call>\n\n<IMPORTANT>\n- Function calls MUST follow the specified format: an inner <function=...></function> block must be nested within <tool_call></tool_call> XML tags\n- DO NOT use function calls inside <think></think> tags.\n- The value enclosed between parameter tags is preserved exactly as-is, including newlines and spaces.\n</IMPORTANT>' }}
+{%- endif %}
+{%- if has_system %}
+    {{- '<|im_end|>\n' }}
+{%- endif %}
+{%- for message in loop_messages %}
+    {%- if message.content is string %}
+        {%- set content = message.content %}
+    {%- else %}
+        {%- set content = render_content(message.content) %}
+    {%- endif %}
+    {%- if message.role == "assistant" %}
+        {%- if message.reasoning_content is string %}
+            {%- set reasoning_content = message.reasoning_content %}
+        {%- else %}
+            {%- set reasoning_content = '' %}
+            {%- if '</think>' in content %}
+                {%- set reasoning_content = content.split('</think>')[0].split('<think>')[-1] %}
+                {%- set content = content.split('</think>')[-1] %}
+            {%- endif %}
+        {%- endif %}
+        {%- if (keep_all_reasoning or loop.index0 > ns.last_user_index) and reasoning_content -%}
+            {{- '<|im_start|>' + message.role + '\n<think>' + reasoning_content + '</think>' + content }}
+        {%- else %}
+            {{- '<|im_start|>' + message.role + '\n<think></think>' + content }}
+        {%- endif %}
+        {%- if message.tool_calls is defined and message.tool_calls is iterable and message.tool_calls | length > 0 %}
+            {%- for tool_call in message.tool_calls %}
+                {%- if tool_call.function is defined %}
+                    {%- set tool_call = tool_call.function %}
+                {%- endif %}
+                {{- '<tool_call>\n<function=' + tool_call.name + '>\n' }}
+                {%- if tool_call.arguments is defined %}
+                    {%- for args_name, args_value in tool_call.arguments|items %}
+                        {{- '<parameter=' + args_name + '>' }}
+                        {%- set args_value = args_value | tojson | safe if args_value is mapping or (args_value is sequence and args_value is not string) else args_value | string %}
+                        {{- args_value }}
+                        {{- '</parameter>\n' }}
+                    {%- endfor %}
+                {%- endif %}
+                {{- '</function>\n</tool_call>' }}
+            {%- endfor %}
+        {%- endif %}
+        {%- if loop.last %}
+            {%- set ns.assistant_is_last = true %}
+        {%- else %}
+            {{- '<|im_end|>\n' }}
+        {%- endif %}
+    {%- elif message.role == "user" %}
+        {{- '<|im_start|>' + message.role + '\n' + render_content(message.content) + '<|im_end|>\n' }}
+    {%- elif message.role == "system" %}
+        {{- '<|im_start|>' + message.role + '\n' + message.content + '<|im_end|>\n' }}
+    {%- elif message.role == "tool" %}
+        {%- if loop.previtem and loop.previtem.role != "tool" %}
+            {{- '<|im_start|>tool\n' }}
+        {%- endif %}
+        {{- '<tool_response>\n' }}
+        {{- render_content(message.content) }}
+        {{- '\n</tool_response>\n' }}
+        {%- if not loop.last and loop.nextitem.role != "tool" %}
+            {{- '<|im_end|>\n' }}
+        {%- elif loop.last %}
+            {{- '<|im_end|>\n' }}
+        {%- endif %}
+    {%- else %}
+        {{- '<|im_start|>' + message.role + '\n' + message.content + '<|im_end|>\n' }}
+    {%- endif %}
+{%- endfor %}
+{%- if add_generation_prompt and not ns.assistant_is_last %}
+    {{- '<|im_start|>assistant\n' }}
+    {%- if not enable_thinking -%}
+        {{- '<think>\n\n</think>\n' -}}
+    {%- endif -%}
+{%- endif %}

config.json

@@ -0,0 +1,108 @@
+{
+  "add_input_local_transformer": true,
+  "add_speech_sosp_eosp": false,
+  "architectures": [
+    "MiMoV2ASRForCausalLM"
+  ],
+  "attention_bias": true,
+  "attention_dropout": 0.0,
+  "audio_channels": 8,
+  "delay_pattern": "0-1-2-3-4-5-6-7",
+  "dtype": "bfloat16",
+  "empty_loss_weight": 0.01,
+  "group_size": 4,
+  "head_dim": 128,
+  "hidden_act": "silu",
+  "hidden_size": 4096,
+  "initializer_range": 0.02,
+  "input_full_attention": true,
+  "input_local_dim": 1024,
+  "input_local_layers": 6,
+  "intermediate_size": 11008,
+  "layer_types": [
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention"
+  ],
+  "local_attn_dropout": 0.1,
+  "local_attn_heads": 64,
+  "local_dim": 1024,
+  "local_ffn_dim": 4096,
+  "local_hidden_dropout": 0.1,
+  "local_layers": 16,
+  "local_rotary_base": 640000,
+  "max_position_embeddings": 8192,
+  "max_window_layers": 28,
+  "mlp_layers": 1,
+  "model_type": "qwen2",
+  "n_rvq": 20,
+  "no_speech_loss": false,
+  "no_text_loss": false,
+  "num_attention_heads": 32,
+  "num_hidden_layers": 36,
+  "num_key_value_heads": 8,
+  "rms_norm_eps": 1e-06,
+  "rope_scaling": null,
+  "rope_theta": 640000,
+  "sliding_window": null,
+  "speech_vocab_size": "1025-1025-129-129-129-129-129-129",
+  "speech_zeroemb_idx": "1024-1024-128-128-128-128-128-128",
+  "tie_word_embeddings": false,
+  "transformers_version": "4.57.1",
+  "use_cache": true,
+  "use_sliding_window": false,
+  "vocab_size": 151680,
+  "audio_config": {
+    "tokenizer_version": "v1",
+    "speech_vocab_size": "1025-1025-129-129-129-129-129-129",
+    "speech_zeroemb_idx": "1024-1024-128-128-128-128-128-128",
+    "group_size": 4,
+    "audio_channels": 8,
+    "input_local_layers": 6,
+    "input_local_dim": 1024,
+    "input_full_attention": true,
+    "input_local_attn_heads": 64,
+    "input_local_head_dim": 16,
+    "input_local_intermediate_size": 4096,
+    "input_local_hidden_dropout": 0.1,
+    "out_hidden_size": 4096,
+    "rope_theta": 640000,
+    "partial_rotary_factor": 1.0,
+    "projection_layers": 1,
+    "add_post_norm": true,
+    "audio_segment_size": 6000
+  }
+}

fp8_meta.json

@@ -0,0 +1,12 @@
+{
+  "dtype": "float8_e4m3fn",
+  "weight_scaling": "per_channel_absmax",
+  "activation_scaling": "dynamic_per_tensor",
+  "matmul_op": "torch._scaled_mm",
+  "output_dtype": "bfloat16",
+  "converted_layers": 417,
+  "weight_gb_before": 32.074,
+  "weight_gb_after": 8.65,
+  "compression_ratio": 3.708,
+  "quantizer": "streaming_cpu"
+}

generation_config.json

@@ -0,0 +1,6 @@
+{
+  "do_sample": true,
+  "temperature": 0.6,
+  "top_k": -1,
+  "top_p": 0.95
+}

model.safetensors

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

quantize_fp8.py

@@ -0,0 +1,294 @@
+"""
+MiMo-V2.5-ASR -> FP8 e4m3fn  (per-channel weight quant, dynamic activation quant)
+
+Quantize entrypoint loads MiMoAudioForCausalLM directly (no audio tokenizer / no
+flash-attn needed -- the LLM is pure Qwen2). Verify/load paths still go through the
+full MimoAudio stack and DO require flash-attn + the audio tokenizer.
+"""
+
+import os
+import sys
+import json
+import shutil
+import argparse
+from pathlib import Path
+
+import torch
+import torch.nn as nn
+from safetensors.torch import save_file, safe_open
+
+REPO_ROOT = Path(__file__).resolve().parent
+if str(REPO_ROOT) not in sys.path:
+    sys.path.insert(0, str(REPO_ROOT))
+
+# ----- constants -----
+FP8_DTYPE   = torch.float8_e4m3fn
+FP8_MAX     = torch.finfo(FP8_DTYPE).max   # 448.0
+SCALE_DTYPE = torch.float32
+SKIP_TYPES  = (nn.Embedding, nn.LayerNorm, nn.GroupNorm,
+               nn.BatchNorm1d, nn.BatchNorm2d, nn.RMSNorm)
+
+CONFIG_FILES = [
+    "config.json", "tokenizer_config.json", "tokenizer.json",
+    "special_tokens_map.json", "generation_config.json",
+    "added_tokens.json", "merges.txt", "vocab.json", "chat_template.jinja",
+]
+
+SPECIAL_TOKENS = ["<|sosp|>", "<|eosp|>", "<|empty|>", "<|Human|>",
+                  "<|SpeechLM|>", "<|sostm|>", "<|eostm|>", "<|eot|>"]
+
+
+# ----- weight quantization -----
+def quantize_weight_per_channel(weight: torch.Tensor):
+    """Per output-channel absmax scaling. weight: [out, in]"""
+    w     = weight.float()
+    amax  = w.abs().amax(dim=1, keepdim=True).clamp(min=1e-12)
+    scale = (amax / FP8_MAX).to(SCALE_DTYPE)
+    w_fp8 = (w / scale).clamp(-FP8_MAX, FP8_MAX).to(FP8_DTYPE)
+    return w_fp8, scale
+
+
+class FP8Linear(nn.Module):
+    """FP8 e4m3fn weights + per-channel scales; dynamic per-tensor activation quant."""
+
+    def __init__(self, linear: nn.Linear):
+        super().__init__()
+        with torch.no_grad():
+            w_fp8, w_scale = quantize_weight_per_channel(linear.weight)
+        self.register_buffer("weight_fp8",   w_fp8.contiguous())   # [out, in]
+        self.register_buffer("weight_scale", w_scale.squeeze(1))   # [out]
+        if linear.bias is not None:
+            self.register_buffer("bias", linear.bias.detach().clone())
+        else:
+            self.bias = None
+        self.in_features  = linear.in_features
+        self.out_features = linear.out_features
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        leading = x.shape[:-1]
+        x2d     = x.reshape(-1, self.in_features)
+        x_scale = (x2d.float().abs().max().clamp(min=1e-12) / FP8_MAX).to(SCALE_DTYPE)
+        x_fp8   = (x2d.float() / x_scale).clamp(-FP8_MAX, FP8_MAX).to(FP8_DTYPE)
+        w_scale_scalar = self.weight_scale.max().to(SCALE_DTYPE)
+        out = torch._scaled_mm(
+            x_fp8, self.weight_fp8.t(),
+            scale_a=x_scale, scale_b=w_scale_scalar,
+            out_dtype=torch.bfloat16, use_fast_accum=True,
+        )
+        correction = (self.weight_scale / w_scale_scalar).to(torch.bfloat16)
+        out = out * correction.unsqueeze(0)
+        if self.bias is not None:
+            out = out + self.bias.to(out.dtype)
+        return out.reshape(*leading, self.out_features)
+
+    def extra_repr(self):
+        return f"in={self.in_features}, out={self.out_features}, fp8=e4m3fn"
+
+
+# ----- model walk -----
+def quantize_model(model: nn.Module, verbose: bool = True):
+    stats = {"converted": 0, "skipped": 0, "bytes_before": 0, "bytes_after": 0}
+
+    def _walk(parent, prefix=""):
+        for name, module in list(parent.named_children()):
+            full = f"{prefix}.{name}" if prefix else name
+            if isinstance(module, nn.Linear) and not isinstance(module, SKIP_TYPES):
+                b_before = module.weight.numel() * module.weight.element_size()
+                if module.bias is not None:
+                    b_before += module.bias.numel() * module.bias.element_size()
+                fp8mod = FP8Linear(module)
+                b_after = fp8mod.weight_fp8.numel() + fp8mod.weight_scale.numel() * 4
+                if fp8mod.bias is not None:
+                    b_after += fp8mod.bias.numel() * fp8mod.bias.element_size()
+                setattr(parent, name, fp8mod)
+                stats["converted"]    += 1
+                stats["bytes_before"] += b_before
+                stats["bytes_after"]  += b_after
+                if verbose:
+                    print(f"  [FP8] {full:<70} {b_before/max(b_after,1):.1f}x")
+            elif isinstance(module, SKIP_TYPES):
+                stats["skipped"] += 1
+            else:
+                _walk(module, full)
+
+    _walk(model)
+    return model, stats
+
+
+# ----- save -----
+def save_fp8(model, out_dir: Path, stats: dict, model_path: Path):
+    out_dir.mkdir(parents=True, exist_ok=True)
+    state = {k: v.contiguous().cpu() for k, v in model.state_dict().items()}
+    st_path = out_dir / "model.safetensors"
+    save_file(state, str(st_path), metadata={"format": "pt"})
+
+    copied = []
+    for cfg in CONFIG_FILES:
+        src = model_path / cfg
+        if src.exists():
+            shutil.copy2(src, out_dir / cfg)
+            copied.append(cfg)
+    if copied:
+        print(f"  Copied config: {', '.join(copied)}")
+
+    gb_before = stats["bytes_before"] / 1e9
+    gb_after  = stats["bytes_after"]  / 1e9
+    ratio     = round(stats["bytes_before"] / max(stats["bytes_after"], 1), 3)
+    meta = {
+        "dtype": "float8_e4m3fn", "weight_scaling": "per_channel_absmax",
+        "activation_scaling": "dynamic_per_tensor", "matmul_op": "torch._scaled_mm",
+        "output_dtype": "bfloat16", "converted_layers": stats["converted"],
+        "skipped_layers": stats["skipped"], "weight_gb_before": round(gb_before, 3),
+        "weight_gb_after": round(gb_after, 3), "compression_ratio": ratio,
+    }
+    with open(out_dir / "fp8_meta.json", "w") as f:
+        json.dump(meta, f, indent=2)
+
+    actual_gb = st_path.stat().st_size / 1e9
+    print(f"\nOK {st_path}  ({actual_gb:.2f} GB on disk)")
+    print(f"   weight bytes: {gb_before:.2f} GB -> {gb_after:.2f} GB ({ratio}x)")
+    print(f"   {stats['converted']} layers converted, {stats['skipped']} skipped")
+
+
+def _build_args_and_tokenizer(model_path: str):
+    from transformers import AutoTokenizer
+    from src.mimo_audio.modeling_mimo_audio import MiMoAudioArguments
+    tok = AutoTokenizer.from_pretrained(model_path)
+    for t in SPECIAL_TOKENS:
+        if t not in tok.get_vocab():
+            tok.add_tokens([t], special_tokens=True)
+    gid = lambda t: tok.convert_tokens_to_ids(t)
+    args = MiMoAudioArguments(
+        model_name_or_path=model_path,
+        sosp_idx=gid("<|sosp|>"), eosp_idx=gid("<|eosp|>"),
+        empty_idx=gid("<|empty|>"), sostm_idx=gid("<|sostm|>"),
+        eostm_idx=gid("<|eostm|>"), eot_idx=gid("<|eot|>"),
+    )
+    return args, tok
+
+
+# ----- quantize entrypoint (direct LLM load, no audio tokenizer / flash-attn) -----
+def run_quantize(args):
+    from src.mimo_audio.modeling_mimo_audio import MiMoAudioForCausalLM
+
+    print(f"Loading MiMoAudioForCausalLM from {args.model_path} on {args.device} ...")
+    model_args, _ = _build_args_and_tokenizer(args.model_path)
+    model = MiMoAudioForCausalLM.from_pretrained(
+        args.model_path,
+        args=model_args,
+        torch_dtype=torch.bfloat16,
+        device_map={"": args.device},
+        attn_implementation="sdpa",
+    )
+    model.eval()
+    print("OK loaded\n")
+
+    print("Quantizing to FP8 e4m3fn ...")
+    with torch.no_grad():
+        model, stats = quantize_model(model, verbose=not args.quiet)
+    save_fp8(model, Path(args.out_dir), stats, Path(args.model_path))
+    print("\nDone.")
+
+
+# ----- load FP8 model (for inference / verify) -----
+def load_fp8_model(fp8_dir: str, tokenizer_path: str, repo_root: str, device: str = "cuda"):
+    """
+    Load the FP8 checkpoint for inference, returning a MimoAudio wrapper exposing .asr_sft().
+
+    Strategy: instantiate the real architecture via from_pretrained on the ORIGINAL repo
+    weights is NOT required -- instead we build the bf16 architecture from config (correct
+    rotary init), replace Linears with FP8Linear shells, then load the FP8 state dict.
+
+    repo_root must be the cloned MiMo-V2.5-ASR repo (contains src/).
+    fp8_dir must contain model.safetensors + config/tokenizer files.
+    """
+    rr = Path(repo_root).resolve()
+    if str(rr) not in sys.path:
+        sys.path.insert(0, str(rr))
+
+    from src.mimo_audio.mimo_audio import MimoAudio
+    from src.mimo_audio.modeling_mimo_audio import MiMoAudioForCausalLM
+    from src.mimo_audio_tokenizer import MiMoAudioTokenizer
+    from transformers import AutoTokenizer, AutoConfig, GenerationConfig
+
+    fp8_dir = Path(fp8_dir)
+    model_args, tokenizer = _build_args_and_tokenizer(str(fp8_dir))
+
+    # Build architecture with real init (correct rotary inv_freq), no pretrained shards.
+    print("Building architecture (config init) ...")
+    cfg = AutoConfig.from_pretrained(str(fp8_dir))
+    model = MiMoAudioForCausalLM(cfg, model_args).to(torch.bfloat16)
+
+    print("Installing FP8 modules ...")
+    with torch.no_grad():
+        quantize_model(model, verbose=False)
+    model = model.to(device)
+
+    print("Loading FP8 weights ...")
+    state = {}
+    with safe_open(str(fp8_dir / "model.safetensors"), framework="pt", device=device) as f:
+        for key in f.keys():
+            state[key] = f.get_tensor(key)
+    model.load_state_dict(state, strict=True)
+    model.eval()
+
+    # Wrap in MimoAudio without re-running its __init__ (which would reload weights).
+    mimo = object.__new__(MimoAudio)
+    mimo.device = device
+    mimo.path = str(fp8_dir)
+    mimo.mimo_audio_tokenizer_path = tokenizer_path
+    mimo.tokenizer = tokenizer
+    mimo.padding_idx = int(tokenizer.pad_token_id)
+    mimo.sosp_idx = model_args.sosp_idx
+    mimo.eosp_idx = model_args.eosp_idx
+    mimo.empty_token = model_args.empty_idx
+    mimo.sostm_idx = model_args.sostm_idx
+    mimo.eostm_idx = model_args.eostm_idx
+    mimo.eot_idx = model_args.eot_idx
+    mimo.im_start_idx = tokenizer.convert_tokens_to_ids("<|im_start|>")
+    mimo.im_end_idx = tokenizer.convert_tokens_to_ids("<|im_end|>")
+    mimo.model = model
+    mimo.group_size = model.config.group_size
+    mimo.audio_channels = model.config.audio_channels
+    mimo.delay_pattern = model.config.delay_pattern
+    mimo.vocab_size = model.config.vocab_size
+    mimo.speech_zeroemb_idx = model.speech_empty_ids
+
+    from src.mimo_audio.modeling_mimo_audio import MiMoSampler
+    mimo.default_global_sampler = MiMoSampler(do_sample=True, temperature=0.6, top_k=50, top_p=0.95)
+    mimo.default_local_sampler = MiMoSampler(do_sample=True, temperature=0.9, top_k=50, top_p=0.95)
+    mimo.task_sampler_configs = {
+        "asr": {"global": MiMoSampler(do_sample=False, temperature=1.0, top_p=1.0),
+                "local": MiMoSampler(do_sample=True, temperature=0.9, top_p=0.95)},
+    }
+    mimo.generate_kwargs = {
+        "max_length": 8192,
+        "eos_token_id": tokenizer.eos_token_id,
+        "pad_token_id": tokenizer.pad_token_id,
+    }
+
+    mimo.mimo_audio_tokenizer = MiMoAudioTokenizer.from_pretrained(tokenizer_path)
+    mimo.mimo_audio_tokenizer.eval().bfloat16().to(device)
+    from torchaudio.transforms import MelSpectrogram
+    tcfg = mimo.mimo_audio_tokenizer.config
+    mimo.mel_transform = MelSpectrogram(
+        sample_rate=tcfg.sampling_rate, n_fft=tcfg.nfft, hop_length=tcfg.hop_length,
+        win_length=tcfg.window_size, f_min=tcfg.fmin, f_max=tcfg.fmax,
+        n_mels=tcfg.n_mels, power=1.0, center=True,
+    ).to(device)
+    print("FP8 model ready\n")
+    return mimo
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--model-path",     required=True)
+    ap.add_argument("--out-dir",        default="./MiMo-V2.5-ASR-FP8")
+    ap.add_argument("--device",         default="cuda")
+    ap.add_argument("--quiet",          action="store_true")
+    args = ap.parse_args()
+    run_quantize(args)
+
+
+if __name__ == "__main__":
+    main()

quantize_fp8_stream.py

@@ -0,0 +1,121 @@
+"""
+Streaming CPU FP8 quantizer for MiMo-V2.5-ASR.
+
+Avoids loading the full model into RAM and never uses CUDA (works around the lack of
+sm_120 / Blackwell kernels in torch 2.6+cu124). Reads the original safetensors shards,
+quantizes every nn.Linear weight to float8_e4m3fn (per-out-channel absmax), leaves
+embeddings / norms / biases in bf16, and writes a single model.safetensors whose keys
+match the FP8Linear loader (`*.weight_fp8`, `*.weight_scale`).
+"""
+import sys, json, shutil, argparse
+from pathlib import Path
+import torch
+import torch.nn as nn
+from safetensors.torch import save_file, safe_open
+
+REPO = Path(__file__).resolve().parent
+sys.path.insert(0, str(REPO))
+from quantize_fp8 import (
+    quantize_weight_per_channel, CONFIG_FILES, _build_args_and_tokenizer, FP8_DTYPE,
+)
+
+
+def linear_weight_names(model_path: str) -> set:
+    """Build the model on meta and collect '<path>.weight' for every nn.Linear."""
+    from src.mimo_audio.modeling_mimo_audio import MiMoAudioForCausalLM
+    from transformers import AutoConfig
+    args, _ = _build_args_and_tokenizer(model_path)
+    cfg = AutoConfig.from_pretrained(model_path)
+    with torch.device("meta"):
+        model = MiMoAudioForCausalLM(cfg, args)
+    names = set()
+    for name, mod in model.named_modules():
+        if isinstance(mod, nn.Linear):
+            names.add(name + ".weight")
+    return names
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--model-path", required=True)
+    ap.add_argument("--out-dir", required=True)
+    ap.add_argument("--dry-run", action="store_true")
+    a = ap.parse_args()
+    mp = Path(a.model_path)
+    out = Path(a.out_dir)
+
+    lin = linear_weight_names(str(mp))
+    print(f"quantizable Linear weights: {len(lin)}")
+
+    idx = json.loads((mp / "model.safetensors.index.json").read_text())
+    weight_map = idx["weight_map"]
+    shards = {}
+    for k, sh in weight_map.items():
+        shards.setdefault(sh, []).append(k)
+    all_keys = set(weight_map)
+    matched = lin & all_keys
+    missing = lin - all_keys
+    print(f"checkpoint keys: {len(all_keys)} | linear matched: {len(matched)} | linear missing in ckpt: {len(missing)}")
+    if missing:
+        print("  MISSING (first 10):", sorted(missing)[:10])
+
+    if a.dry_run:
+        # show a few non-linear keys for sanity
+        nonlin = sorted(all_keys - lin)
+        print("sample NON-linear keys kept as-is:", nonlin[:8])
+        print("total params:", len(all_keys), "-> fp8:", len(matched), "kept:", len(all_keys) - len(matched))
+        return
+
+    out.mkdir(parents=True, exist_ok=True)
+    new_state = {}
+    bytes_before = bytes_after = 0
+    conv = 0
+    for shard in sorted(shards):
+        with safe_open(str(mp / shard), framework="pt", device="cpu") as f:
+            for k in shards[shard]:
+                t = f.get_tensor(k)
+                if k in lin:
+                    w_fp8, scale = quantize_weight_per_channel(t)   # scale [out,1]
+                    base = k[:-len(".weight")]
+                    new_state[base + ".weight_fp8"] = w_fp8.contiguous()
+                    new_state[base + ".weight_scale"] = scale.squeeze(1).contiguous()
+                    bytes_before += t.numel() * t.element_size()
+                    bytes_after += w_fp8.numel() + scale.numel() * 4
+                    conv += 1
+                else:
+                    bytes_before += t.numel() * t.element_size()
+                    # kept tensors (embeddings/norms/biases): store bf16 to match the
+                    # bf16 model the FP8Linear loader builds, and to shrink the file.
+                    if t.is_floating_point():
+                        t = t.to(torch.bfloat16)
+                    new_state[k] = t.contiguous()
+                    bytes_after += t.numel() * t.element_size()
+        print(f"  done {shard}  (running fp8 layers: {conv})")
+
+    st = out / "model.safetensors"
+    save_file(new_state, str(st), metadata={"format": "pt"})
+
+    copied = []
+    for cfg in CONFIG_FILES:
+        src = mp / cfg
+        if src.exists():
+            shutil.copy2(src, out / cfg); copied.append(cfg)
+
+    ratio = round(bytes_before / max(bytes_after, 1), 3)
+    meta = {
+        "dtype": "float8_e4m3fn", "weight_scaling": "per_channel_absmax",
+        "activation_scaling": "dynamic_per_tensor", "matmul_op": "torch._scaled_mm",
+        "output_dtype": "bfloat16", "converted_layers": conv,
+        "weight_gb_before": round(bytes_before / 1e9, 3),
+        "weight_gb_after": round(bytes_after / 1e9, 3), "compression_ratio": ratio,
+        "quantizer": "streaming_cpu",
+    }
+    (out / "fp8_meta.json").write_text(json.dumps(meta, indent=2))
+    gb = st.stat().st_size / 1e9
+    print(f"\nOK {st} ({gb:.2f} GB on disk)")
+    print(f"   {conv} layers -> fp8 | {round(bytes_before/1e9,2)}GB -> {round(bytes_after/1e9,2)}GB ({ratio}x)")
+    print(f"   copied config: {', '.join(copied)}")
+
+
+if __name__ == "__main__":
+    main()

special_tokens_map.json

@@ -0,0 +1,31 @@
+{
+  "additional_special_tokens": [
+    "<|im_start|>",
+    "<|im_end|>",
+    "<|object_ref_start|>",
+    "<|object_ref_end|>",
+    "<|box_start|>",
+    "<|box_end|>",
+    "<|quad_start|>",
+    "<|quad_end|>",
+    "<|vision_start|>",
+    "<|vision_end|>",
+    "<|vision_pad|>",
+    "<|image_pad|>",
+    "<|video_pad|>"
+  ],
+  "eos_token": {
+    "content": "<|endoftext|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "pad_token": {
+    "content": "<|endoftext|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  }
+}

tokenizer.json

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

tokenizer_config.json

@@ -0,0 +1,271 @@
+{
+  "add_bos_token": false,
+  "add_prefix_space": false,
+  "added_tokens_decoder": {
+    "151643": {
+      "content": "<|endoftext|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151644": {
+      "content": "<|im_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151645": {
+      "content": "<|im_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151646": {
+      "content": "<|object_ref_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151647": {
+      "content": "<|object_ref_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151648": {
+      "content": "<|box_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151649": {
+      "content": "<|box_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151650": {
+      "content": "<|quad_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151651": {
+      "content": "<|quad_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151652": {
+      "content": "<|vision_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151653": {
+      "content": "<|vision_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151654": {
+      "content": "<|vision_pad|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151655": {
+      "content": "<|image_pad|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151656": {
+      "content": "<|video_pad|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151657": {
+      "content": "<tool_call>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151658": {
+      "content": "</tool_call>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151659": {
+      "content": "<|fim_prefix|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151660": {
+      "content": "<|fim_middle|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151661": {
+      "content": "<|fim_suffix|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151662": {
+      "content": "<|fim_pad|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151663": {
+      "content": "<|repo_name|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151664": {
+      "content": "<|file_sep|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151665": {
+      "content": "<|sosp|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151666": {
+      "content": "<|eosp|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151667": {
+      "content": "<|empty|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151668": {
+      "content": "<|Human|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151669": {
+      "content": "<|SpeechLM|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151670": {
+      "content": "<|sostm|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151671": {
+      "content": "<|eostm|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151672": {
+      "content": "<|eot|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    }
+  },
+  "additional_special_tokens": [
+    "<|im_start|>",
+    "<|im_end|>",
+    "<|object_ref_start|>",
+    "<|object_ref_end|>",
+    "<|box_start|>",
+    "<|box_end|>",
+    "<|quad_start|>",
+    "<|quad_end|>",
+    "<|vision_start|>",
+    "<|vision_end|>",
+    "<|vision_pad|>",
+    "<|image_pad|>",
+    "<|video_pad|>"
+  ],
+  "bos_token": null,
+  "clean_up_tokenization_spaces": false,
+  "eos_token": "<|im_end|>",
+  "errors": "replace",
+  "extra_special_tokens": {},
+  "model_max_length": 131072,
+  "pad_token": "<|endoftext|>",
+  "split_special_tokens": false,
+  "tokenizer_class": "Qwen2Tokenizer",
+  "unk_token": null
+}