Hugging Face's logo

mlx-community
/
MiMo-V2.6-Pro-RL-mxfp4-q8

Text Generation
MLX
Safetensors
English
Chinese
mimo_v2
agent
long-context
conversational
custom_code
4-bit precision
Model card Files
xet
 Community
kernelpool commited on
Commit
f7edb41
·
verified ·
1 Parent(s): ba06603

Add files using upload-large-folder tool

Browse files
This view is limited to 50 files because it contains too many changes.   See raw diff
Files changed (50) hide show
  1. .gitattributes +1 -0
  2. README.md +66 -0
  3. chat_template.jinja +95 -0
  4. config.json +3166 -0
  5. configuration_mimo_v2.py +247 -0
  6. generation_config.json +9 -0
  7. model-00001-of-00116.safetensors +3 -0
  8. model-00004-of-00116.safetensors +3 -0
  9. model-00005-of-00116.safetensors +3 -0
  10. model-00013-of-00116.safetensors +3 -0
  11. model-00016-of-00116.safetensors +3 -0
  12. model-00019-of-00116.safetensors +3 -0
  13. model-00021-of-00116.safetensors +3 -0
  14. model-00024-of-00116.safetensors +3 -0
  15. model-00032-of-00116.safetensors +3 -0
  16. model-00033-of-00116.safetensors +3 -0
  17. model-00036-of-00116.safetensors +3 -0
  18. model-00037-of-00116.safetensors +3 -0
  19. model-00038-of-00116.safetensors +3 -0
  20. model-00039-of-00116.safetensors +3 -0
  21. model-00040-of-00116.safetensors +3 -0
  22. model-00044-of-00116.safetensors +3 -0
  23. model-00045-of-00116.safetensors +3 -0
  24. model-00046-of-00116.safetensors +3 -0
  25. model-00052-of-00116.safetensors +3 -0
  26. model-00057-of-00116.safetensors +3 -0
  27. model-00058-of-00116.safetensors +3 -0
  28. model-00060-of-00116.safetensors +3 -0
  29. model-00065-of-00116.safetensors +3 -0
  30. model-00072-of-00116.safetensors +3 -0
  31. model-00073-of-00116.safetensors +3 -0
  32. model-00076-of-00116.safetensors +3 -0
  33. model-00077-of-00116.safetensors +3 -0
  34. model-00078-of-00116.safetensors +3 -0
  35. model-00079-of-00116.safetensors +3 -0
  36. model-00080-of-00116.safetensors +3 -0
  37. model-00081-of-00116.safetensors +3 -0
  38. model-00084-of-00116.safetensors +3 -0
  39. model-00085-of-00116.safetensors +3 -0
  40. model-00093-of-00116.safetensors +3 -0
  41. model-00096-of-00116.safetensors +3 -0
  42. model-00099-of-00116.safetensors +3 -0
  43. model-00100-of-00116.safetensors +3 -0
  44. model-00105-of-00116.safetensors +3 -0
  45. model-00112-of-00116.safetensors +3 -0
  46. model-00113-of-00116.safetensors +3 -0
  47. model-00116-of-00116.safetensors +3 -0
  48. model.safetensors.index.json +0 -0
  49. modeling_mimo_v2.py +1878 -0
  50. tokenizer.json +3 -0
.gitattributes CHANGED
@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
33
  *.zip filter=lfs diff=lfs merge=lfs -text
34
  *.zst filter=lfs diff=lfs merge=lfs -text
35
  *tfevents* filter=lfs diff=lfs merge=lfs -text
 
 
33
  *.zip filter=lfs diff=lfs merge=lfs -text
34
  *.zst filter=lfs diff=lfs merge=lfs -text
35
  *tfevents* filter=lfs diff=lfs merge=lfs -text
36
+ tokenizer.json filter=lfs diff=lfs merge=lfs -text
README.md ADDED
@@ -0,0 +1,66 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ ---
2
+ license: mit
3
+ language:
4
+ - en
5
+ - zh
6
+ base_model: XiaomiMiMo/MiMo-V2.6-Pro-RL
7
+ library_name: mlx
8
+ pipeline_tag: text-generation
9
+ tags:
10
+ - mlx
11
+ - mimo_v2
12
+ - agent
13
+ - long-context
14
+ ---
15
+
16
+ # mlx-community/MiMo-V2.6-Pro-RL-mxfp4-q8
17
+
18
+ This model [mlx-community/MiMo-V2.6-Pro-RL-mxfp4-q8](https://huggingface.co/mlx-community/MiMo-V2.6-Pro-RL-mxfp4-q8) was
19
+ converted to MLX format from [XiaomiMiMo/MiMo-V2.6-Pro-RL](https://huggingface.co/XiaomiMiMo/MiMo-V2.6-Pro-RL)
20
+ using mlx-lm version **0.32.0** ([PR #1219](https://github.com/ml-explore/mlx-lm/pull/1219)).
21
+
22
+ ## Quantization
23
+
24
+ - MoE expert weights are the original checkpoint's native **MXFP4** (4-bit, group size 32), loaded directly without requantization.
25
+ - Attention, dense MLP, embeddings and `lm_head` are **8-bit affine**, group size 64.
26
+ - 4.339 bits per weight overall, 516 GB on disk.
27
+
28
+ This is a text-only conversion: the vision and audio encoders and the MTP/DFlash draft weights are not included.
29
+
30
+ ## Requirements
31
+
32
+ MiMo-V2 support is in [mlx-lm PR #1219](https://github.com/ml-explore/mlx-lm/pull/1219). Until it is merged, install mlx-lm from that branch:
33
+
34
+ ```bash
35
+ pip install git+https://github.com/kernelpool/mlx-lm.git@add-mimo-v2
36
+ ```
37
+
38
+ At 1.02T total parameters (42B active) this model does not fit on a single Mac. It is meant to run tensor-parallel across two 512 GB machines with mlx-lm's distributed support, for example:
39
+
40
+ ```bash
41
+ mlx.launch --backend jaccl --hostfile hosts.json -- \
42
+ python -m mlx_lm.examples.sharded_generate \
43
+ --model mlx-community/MiMo-V2.6-Pro-RL-mxfp4-q8 --prompt "hello" -m 256
44
+ ```
45
+
46
+ See the mlx-lm distributed inference documentation for the hostfile format and backend setup.
47
+
48
+ ## Use with mlx
49
+
50
+ ```python
51
+ from mlx_lm import load, generate
52
+
53
+ model, tokenizer = load("mlx-community/MiMo-V2.6-Pro-RL-mxfp4-q8")
54
+
55
+ prompt = "hello"
56
+
57
+ if tokenizer.chat_template is not None:
58
+ messages = [{"role": "user", "content": prompt}]
59
+ prompt = tokenizer.apply_chat_template(
60
+ messages, add_generation_prompt=True, return_dict=False,
61
+ )
62
+
63
+ response = generate(model, tokenizer, prompt=prompt, verbose=True)
64
+ ```
65
+
66
+ Thinking is enabled by default in the chat template; pass `enable_thinking=False` to `apply_chat_template` to disable it. Tool calls use the Qwen3-Coder format (`<tool_call><function=...>`), which the `qwen3_coder` tool parser in mlx-lm handles.
chat_template.jinja ADDED
@@ -0,0 +1,95 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {%- macro render_value(value) -%}
2
+ {%- if value is string -%}
3
+ {{- value -}}
4
+ {%- else -%}
5
+ {{- value | tojson(ensure_ascii=False) -}}
6
+ {%- endif -%}
7
+ {%- endmacro -%}
8
+
9
+ {%- macro render_content(message_content) -%}
10
+ {%- if message_content is string -%}
11
+ {{- message_content -}}
12
+ {%- elif message_content is iterable -%}
13
+ {%- for part in message_content -%}
14
+ {%- if part is not mapping -%}
15
+ {{- part -}}
16
+ {%- elif part['type'] == 'image' or 'image' in part or 'image_url' in part -%}
17
+ {{- '<|vision_start|><|image_pad|><|vision_end|>' -}}
18
+ {%- elif part['type'] == 'audio' or part['type'] == 'input_audio' or 'audio' in part or 'audio_url' in part or 'input_audio' in part -%}
19
+ {{- '<|mimo_audio_start|><|audio_pad|><|mimo_audio_end|>' -}}
20
+ {%- elif part['type'] == 'video' or 'video' in part or 'video_url' in part -%}
21
+ {{- '<|vision_start|><|video_pad|><|vision_end|>' -}}
22
+ {%- elif 'text' in part -%}
23
+ {{- part['text'] -}}
24
+ {%- endif -%}
25
+ {%- endfor -%}
26
+ {%- endif -%}
27
+ {%- endmacro -%}
28
+
29
+ {%- macro render_tools(tools) -%}
30
+ {{- 'You are provided with the following tools:\n\n<tools>' -}}
31
+ {%- for tool in tools -%}
32
+ {{- '\n' ~ (tool | tojson(ensure_ascii=False)) -}}
33
+ {%- endfor -%}
34
+ {{- '\n</tools>' -}}
35
+ {%- endmacro -%}
36
+
37
+ {%- macro render_tool_calls(tool_calls) -%}
38
+ {%- for tool_call in tool_calls -%}
39
+ {%- if tool_call.function is defined -%}
40
+ {%- set tool_call = tool_call.function -%}
41
+ {%- elif tool_call.custom is defined -%}
42
+ {%- set tool_call = tool_call.custom -%}
43
+ {%- endif -%}
44
+ {{- '<tool_call><function=' ~ tool_call.name ~ '>' -}}
45
+ {%- if tool_call.input is defined and tool_call.input is string -%}
46
+ {{- tool_call.input -}}
47
+ {%- elif tool_call.arguments -%}
48
+ {%- if tool_call.arguments is string -%}
49
+ {{- tool_call.arguments -}}
50
+ {%- else -%}
51
+ {%- for args_name, args_value in tool_call.arguments | items -%}
52
+ {{- '<parameter=' ~ args_name ~ '>' ~ render_value(args_value) ~ '</parameter>' -}}
53
+ {%- endfor -%}
54
+ {%- endif -%}
55
+ {%- endif -%}
56
+ {{- '</function></tool_call>' -}}
57
+ {%- endfor -%}
58
+ {%- endmacro -%}
59
+
60
+ {%- macro render_assistant_message(message) -%}
61
+ {%- set content = render_content(message.content) -%}
62
+ {%- set reasoning = message.reasoning_content if message.reasoning_content is string else '' -%}
63
+ {{- '<|im_start|>assistant\n<think>' ~ reasoning ~ '</think>' ~ content -}}
64
+ {%- if message.tool_calls is defined and message.tool_calls is iterable and message.tool_calls | length > 0 -%}
65
+ {{- render_tool_calls(message.tool_calls) -}}
66
+ {%- endif -%}
67
+ {{- '<|im_end|>' -}}
68
+ {%- endmacro -%}
69
+
70
+ {%- if tools is defined and tools is iterable and tools | length > 0 -%}
71
+ {{- '<|im_start|>system\n' ~ render_tools(tools) ~ '<|im_end|>' -}}
72
+ {%- endif -%}
73
+
74
+ {%- for message in messages -%}
75
+ {%- if message.role == 'assistant' -%}
76
+ {{- render_assistant_message(message) -}}
77
+ {%- else -%}
78
+ {%- set body = render_content(message.content) -%}
79
+ {{- '<|im_start|>' ~ message.role ~ '\n' ~ body -}}
80
+ {%- if message.tools is defined and message.tools is iterable and message.tools | length > 0 -%}
81
+ {%- if body -%}
82
+ {{- '\n\n' -}}
83
+ {%- endif -%}
84
+ {{- render_tools(message.tools) -}}
85
+ {%- endif -%}
86
+ {{- '<|im_end|>' -}}
87
+ {%- endif -%}
88
+ {%- endfor -%}
89
+
90
+ {%- if add_generation_prompt -%}
91
+ {{- '<|im_start|>assistant\n' -}}
92
+ {%- if enable_thinking is false -%}
93
+ {{- '<think></think>' -}}
94
+ {%- endif -%}
95
+ {%- endif -%}
config.json ADDED
@@ -0,0 +1,3166 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "add_full_attention_sink_bias": false,
3
+ "add_swa_attention_sink_bias": true,
4
+ "architectures": [
5
+ "MiMoV2ForCausalLM"
6
+ ],
7
+ "attention_bias": false,
8
+ "attention_chunk_size": 128,
9
+ "attention_dropout": 0.0,
10
+ "attention_projection_layout": "fused_qkv",
11
+ "attention_value_scale": 0.612,
12
+ "audio_config": {
13
+ "add_post_norm": true,
14
+ "audio_channels": 20,
15
+ "audio_segment_size": 6000,
16
+ "group_size": 4,
17
+ "input_full_attention": true,
18
+ "input_local_attn_heads": 16,
19
+ "input_local_dim": 1024,
20
+ "input_local_head_dim": 64,
21
+ "input_local_hidden_dropout": 0.0,
22
+ "input_local_intermediate_size": 4096,
23
+ "input_local_layers": 6,
24
+ "out_hidden_size": 6144,
25
+ "partial_rotary_factor": 1.0,
26
+ "projection_layers": 2,
27
+ "rope_theta": 640000,
28
+ "speech_vocab_size": "1280",
29
+ "speech_zeroemb_idx": "1024"
30
+ },
31
+ "audio_end_token_id": 151674,
32
+ "audio_start_token_id": 151673,
33
+ "audio_token_id": 151669,
34
+ "auto_map": {
35
+ "AutoConfig": "configuration_mimo_v2.MiMoV2Config",
36
+ "AutoModel": "modeling_mimo_v2.MiMoV2Model",
37
+ "AutoModelForCausalLM": "modeling_mimo_v2.MiMoV2ForCausalLM"
38
+ },
39
+ "bos_token_id": null,
40
+ "dtype": "bfloat16",
41
+ "eos_token_id": [
42
+ 151643,
43
+ 151645,
44
+ 151672
45
+ ],
46
+ "head_dim": 192,
47
+ "hidden_act": "silu",
48
+ "hidden_size": 6144,
49
+ "hybrid_block_size": null,
50
+ "hybrid_layer_pattern": [
51
+ 0,
52
+ 1,
53
+ 1,
54
+ 1,
55
+ 1,
56
+ 1,
57
+ 1,
58
+ 0,
59
+ 1,
60
+ 1,
61
+ 1,
62
+ 1,
63
+ 1,
64
+ 1,
65
+ 1,
66
+ 0,
67
+ 1,
68
+ 1,
69
+ 1,
70
+ 1,
71
+ 1,
72
+ 1,
73
+ 1,
74
+ 0,
75
+ 1,
76
+ 1,
77
+ 1,
78
+ 1,
79
+ 1,
80
+ 1,
81
+ 1,
82
+ 0,
83
+ 1,
84
+ 1,
85
+ 1,
86
+ 1,
87
+ 1,
88
+ 1,
89
+ 1,
90
+ 0,
91
+ 1,
92
+ 1,
93
+ 1,
94
+ 1,
95
+ 1,
96
+ 1,
97
+ 1,
98
+ 0,
99
+ 1,
100
+ 1,
101
+ 1,
102
+ 1,
103
+ 1,
104
+ 1,
105
+ 1,
106
+ 0,
107
+ 1,
108
+ 1,
109
+ 1,
110
+ 1,
111
+ 1,
112
+ 1,
113
+ 0,
114
+ 1,
115
+ 1,
116
+ 1,
117
+ 1,
118
+ 1,
119
+ 1,
120
+ 0
121
+ ],
122
+ "image_token_id": 151655,
123
+ "initializer_range": 0.02,
124
+ "intermediate_size": 16384,
125
+ "layernorm_epsilon": 1e-05,
126
+ "max_position_embeddings": 1048576,
127
+ "model_type": "mimo_v2",
128
+ "moe_intermediate_size": 2048,
129
+ "moe_layer_freq": [
130
+ 0,
131
+ 1,
132
+ 1,
133
+ 1,
134
+ 1,
135
+ 1,
136
+ 1,
137
+ 1,
138
+ 1,
139
+ 1,
140
+ 1,
141
+ 1,
142
+ 1,
143
+ 1,
144
+ 1,
145
+ 1,
146
+ 1,
147
+ 1,
148
+ 1,
149
+ 1,
150
+ 1,
151
+ 1,
152
+ 1,
153
+ 1,
154
+ 1,
155
+ 1,
156
+ 1,
157
+ 1,
158
+ 1,
159
+ 1,
160
+ 1,
161
+ 1,
162
+ 1,
163
+ 1,
164
+ 1,
165
+ 1,
166
+ 1,
167
+ 1,
168
+ 1,
169
+ 1,
170
+ 1,
171
+ 1,
172
+ 1,
173
+ 1,
174
+ 1,
175
+ 1,
176
+ 1,
177
+ 1,
178
+ 1,
179
+ 1,
180
+ 1,
181
+ 1,
182
+ 1,
183
+ 1,
184
+ 1,
185
+ 1,
186
+ 1,
187
+ 1,
188
+ 1,
189
+ 1,
190
+ 1,
191
+ 1,
192
+ 1,
193
+ 1,
194
+ 1,
195
+ 1,
196
+ 1,
197
+ 1,
198
+ 1,
199
+ 1
200
+ ],
201
+ "moe_router_dtype": "bfloat16",
202
+ "n_group": 1,
203
+ "n_routed_experts": 384,
204
+ "n_shared_experts": null,
205
+ "norm_topk_prob": true,
206
+ "num_attention_heads": 128,
207
+ "num_experts_per_tok": 8,
208
+ "num_hidden_layers": 70,
209
+ "num_key_value_heads": 8,
210
+ "pad_token_id": 151643,
211
+ "partial_rotary_factor": 0.334,
212
+ "processor_config": {
213
+ "audio_avg_pooler": 2,
214
+ "audio_channels": 20,
215
+ "audio_end_token_id": 151674,
216
+ "audio_fmax": null,
217
+ "audio_fmin": 0,
218
+ "audio_group_size": 4,
219
+ "audio_hop_length": 240,
220
+ "audio_input_id_per_second": 25.0,
221
+ "audio_kernel_size": 3,
222
+ "audio_n_mels": 128,
223
+ "audio_nfft": 960,
224
+ "audio_sampling_rate": 24000,
225
+ "audio_segment_size": 6000,
226
+ "audio_start_token_id": 151673,
227
+ "audio_stride_size": 2,
228
+ "audio_token_id": 151669,
229
+ "audio_window_size": 960,
230
+ "audio_zeroemb_idx": [
231
+ 1024,
232
+ 1024,
233
+ 1024,
234
+ 1024,
235
+ 1024,
236
+ 1024,
237
+ 1024,
238
+ 1024,
239
+ 1024,
240
+ 1024,
241
+ 1024,
242
+ 1024,
243
+ 1024,
244
+ 1024,
245
+ 1024,
246
+ 1024,
247
+ 1024,
248
+ 1024,
249
+ 1024,
250
+ 1024
251
+ ],
252
+ "fps": 1.0,
253
+ "image_max_pixels": 8388608,
254
+ "image_min_pixels": 8192,
255
+ "image_token_id": 151655,
256
+ "max_frames": 3600,
257
+ "merge_size": 2,
258
+ "min_frames": null,
259
+ "num_frames": null,
260
+ "pad_token_id": 151643,
261
+ "patch_size": 16,
262
+ "rope_type": "rope",
263
+ "temporal_compression_ratio": 1,
264
+ "temporal_patch_size": 2,
265
+ "use_per_grid_t_timestamps": false,
266
+ "use_video_timestamps": true,
267
+ "video_audio_interleave_length": 0.0,
268
+ "video_end_token_id": 151671,
269
+ "video_max_pixels": 8388608,
270
+ "video_min_pixels": 8192,
271
+ "video_process_num_threads": 16,
272
+ "video_start_token_id": 151670,
273
+ "video_token_id": 151656,
274
+ "video_tokens_per_second": 2,
275
+ "video_total_max_pixels": 268435456,
276
+ "vision_end_token_id": 151653,
277
+ "vision_start_token_id": 151652
278
+ },
279
+ "quantization": {
280
+ "group_size": 32,
281
+ "bits": 4,
282
+ "mode": "mxfp4",
283
+ "model.embed_tokens": {
284
+ "group_size": 64,
285
+ "bits": 8,
286
+ "mode": "affine"
287
+ },
288
+ "model.layers.0.self_attn.q_proj": {
289
+ "group_size": 64,
290
+ "bits": 8,
291
+ "mode": "affine"
292
+ },
293
+ "model.layers.0.self_attn.k_proj": {
294
+ "group_size": 64,
295
+ "bits": 8,
296
+ "mode": "affine"
297
+ },
298
+ "model.layers.0.self_attn.v_proj": {
299
+ "group_size": 64,
300
+ "bits": 8,
301
+ "mode": "affine"
302
+ },
303
+ "model.layers.0.self_attn.o_proj": {
304
+ "group_size": 64,
305
+ "bits": 8,
306
+ "mode": "affine"
307
+ },
308
+ "model.layers.0.mlp.gate_proj": {
309
+ "group_size": 64,
310
+ "bits": 8,
311
+ "mode": "affine"
312
+ },
313
+ "model.layers.0.mlp.up_proj": {
314
+ "group_size": 64,
315
+ "bits": 8,
316
+ "mode": "affine"
317
+ },
318
+ "model.layers.0.mlp.down_proj": {
319
+ "group_size": 64,
320
+ "bits": 8,
321
+ "mode": "affine"
322
+ },
323
+ "model.layers.1.self_attn.q_proj": {
324
+ "group_size": 64,
325
+ "bits": 8,
326
+ "mode": "affine"
327
+ },
328
+ "model.layers.1.self_attn.k_proj": {
329
+ "group_size": 64,
330
+ "bits": 8,
331
+ "mode": "affine"
332
+ },
333
+ "model.layers.1.self_attn.v_proj": {
334
+ "group_size": 64,
335
+ "bits": 8,
336
+ "mode": "affine"
337
+ },
338
+ "model.layers.1.self_attn.o_proj": {
339
+ "group_size": 64,
340
+ "bits": 8,
341
+ "mode": "affine"
342
+ },
343
+ "model.layers.2.self_attn.q_proj": {
344
+ "group_size": 64,
345
+ "bits": 8,
346
+ "mode": "affine"
347
+ },
348
+ "model.layers.2.self_attn.k_proj": {
349
+ "group_size": 64,
350
+ "bits": 8,
351
+ "mode": "affine"
352
+ },
353
+ "model.layers.2.self_attn.v_proj": {
354
+ "group_size": 64,
355
+ "bits": 8,
356
+ "mode": "affine"
357
+ },
358
+ "model.layers.2.self_attn.o_proj": {
359
+ "group_size": 64,
360
+ "bits": 8,
361
+ "mode": "affine"
362
+ },
363
+ "model.layers.3.self_attn.q_proj": {
364
+ "group_size": 64,
365
+ "bits": 8,
366
+ "mode": "affine"
367
+ },
368
+ "model.layers.3.self_attn.k_proj": {
369
+ "group_size": 64,
370
+ "bits": 8,
371
+ "mode": "affine"
372
+ },
373
+ "model.layers.3.self_attn.v_proj": {
374
+ "group_size": 64,
375
+ "bits": 8,
376
+ "mode": "affine"
377
+ },
378
+ "model.layers.3.self_attn.o_proj": {
379
+ "group_size": 64,
380
+ "bits": 8,
381
+ "mode": "affine"
382
+ },
383
+ "model.layers.4.self_attn.q_proj": {
384
+ "group_size": 64,
385
+ "bits": 8,
386
+ "mode": "affine"
387
+ },
388
+ "model.layers.4.self_attn.k_proj": {
389
+ "group_size": 64,
390
+ "bits": 8,
391
+ "mode": "affine"
392
+ },
393
+ "model.layers.4.self_attn.v_proj": {
394
+ "group_size": 64,
395
+ "bits": 8,
396
+ "mode": "affine"
397
+ },
398
+ "model.layers.4.self_attn.o_proj": {
399
+ "group_size": 64,
400
+ "bits": 8,
401
+ "mode": "affine"
402
+ },
403
+ "model.layers.5.self_attn.q_proj": {
404
+ "group_size": 64,
405
+ "bits": 8,
406
+ "mode": "affine"
407
+ },
408
+ "model.layers.5.self_attn.k_proj": {
409
+ "group_size": 64,
410
+ "bits": 8,
411
+ "mode": "affine"
412
+ },
413
+ "model.layers.5.self_attn.v_proj": {
414
+ "group_size": 64,
415
+ "bits": 8,
416
+ "mode": "affine"
417
+ },
418
+ "model.layers.5.self_attn.o_proj": {
419
+ "group_size": 64,
420
+ "bits": 8,
421
+ "mode": "affine"
422
+ },
423
+ "model.layers.6.self_attn.q_proj": {
424
+ "group_size": 64,
425
+ "bits": 8,
426
+ "mode": "affine"
427
+ },
428
+ "model.layers.6.self_attn.k_proj": {
429
+ "group_size": 64,
430
+ "bits": 8,
431
+ "mode": "affine"
432
+ },
433
+ "model.layers.6.self_attn.v_proj": {
434
+ "group_size": 64,
435
+ "bits": 8,
436
+ "mode": "affine"
437
+ },
438
+ "model.layers.6.self_attn.o_proj": {
439
+ "group_size": 64,
440
+ "bits": 8,
441
+ "mode": "affine"
442
+ },
443
+ "model.layers.7.self_attn.q_proj": {
444
+ "group_size": 64,
445
+ "bits": 8,
446
+ "mode": "affine"
447
+ },
448
+ "model.layers.7.self_attn.k_proj": {
449
+ "group_size": 64,
450
+ "bits": 8,
451
+ "mode": "affine"
452
+ },
453
+ "model.layers.7.self_attn.v_proj": {
454
+ "group_size": 64,
455
+ "bits": 8,
456
+ "mode": "affine"
457
+ },
458
+ "model.layers.7.self_attn.o_proj": {
459
+ "group_size": 64,
460
+ "bits": 8,
461
+ "mode": "affine"
462
+ },
463
+ "model.layers.8.self_attn.q_proj": {
464
+ "group_size": 64,
465
+ "bits": 8,
466
+ "mode": "affine"
467
+ },
468
+ "model.layers.8.self_attn.k_proj": {
469
+ "group_size": 64,
470
+ "bits": 8,
471
+ "mode": "affine"
472
+ },
473
+ "model.layers.8.self_attn.v_proj": {
474
+ "group_size": 64,
475
+ "bits": 8,
476
+ "mode": "affine"
477
+ },
478
+ "model.layers.8.self_attn.o_proj": {
479
+ "group_size": 64,
480
+ "bits": 8,
481
+ "mode": "affine"
482
+ },
483
+ "model.layers.9.self_attn.q_proj": {
484
+ "group_size": 64,
485
+ "bits": 8,
486
+ "mode": "affine"
487
+ },
488
+ "model.layers.9.self_attn.k_proj": {
489
+ "group_size": 64,
490
+ "bits": 8,
491
+ "mode": "affine"
492
+ },
493
+ "model.layers.9.self_attn.v_proj": {
494
+ "group_size": 64,
495
+ "bits": 8,
496
+ "mode": "affine"
497
+ },
498
+ "model.layers.9.self_attn.o_proj": {
499
+ "group_size": 64,
500
+ "bits": 8,
501
+ "mode": "affine"
502
+ },
503
+ "model.layers.10.self_attn.q_proj": {
504
+ "group_size": 64,
505
+ "bits": 8,
506
+ "mode": "affine"
507
+ },
508
+ "model.layers.10.self_attn.k_proj": {
509
+ "group_size": 64,
510
+ "bits": 8,
511
+ "mode": "affine"
512
+ },
513
+ "model.layers.10.self_attn.v_proj": {
514
+ "group_size": 64,
515
+ "bits": 8,
516
+ "mode": "affine"
517
+ },
518
+ "model.layers.10.self_attn.o_proj": {
519
+ "group_size": 64,
520
+ "bits": 8,
521
+ "mode": "affine"
522
+ },
523
+ "model.layers.11.self_attn.q_proj": {
524
+ "group_size": 64,
525
+ "bits": 8,
526
+ "mode": "affine"
527
+ },
528
+ "model.layers.11.self_attn.k_proj": {
529
+ "group_size": 64,
530
+ "bits": 8,
531
+ "mode": "affine"
532
+ },
533
+ "model.layers.11.self_attn.v_proj": {
534
+ "group_size": 64,
535
+ "bits": 8,
536
+ "mode": "affine"
537
+ },
538
+ "model.layers.11.self_attn.o_proj": {
539
+ "group_size": 64,
540
+ "bits": 8,
541
+ "mode": "affine"
542
+ },
543
+ "model.layers.12.self_attn.q_proj": {
544
+ "group_size": 64,
545
+ "bits": 8,
546
+ "mode": "affine"
547
+ },
548
+ "model.layers.12.self_attn.k_proj": {
549
+ "group_size": 64,
550
+ "bits": 8,
551
+ "mode": "affine"
552
+ },
553
+ "model.layers.12.self_attn.v_proj": {
554
+ "group_size": 64,
555
+ "bits": 8,
556
+ "mode": "affine"
557
+ },
558
+ "model.layers.12.self_attn.o_proj": {
559
+ "group_size": 64,
560
+ "bits": 8,
561
+ "mode": "affine"
562
+ },
563
+ "model.layers.13.self_attn.q_proj": {
564
+ "group_size": 64,
565
+ "bits": 8,
566
+ "mode": "affine"
567
+ },
568
+ "model.layers.13.self_attn.k_proj": {
569
+ "group_size": 64,
570
+ "bits": 8,
571
+ "mode": "affine"
572
+ },
573
+ "model.layers.13.self_attn.v_proj": {
574
+ "group_size": 64,
575
+ "bits": 8,
576
+ "mode": "affine"
577
+ },
578
+ "model.layers.13.self_attn.o_proj": {
579
+ "group_size": 64,
580
+ "bits": 8,
581
+ "mode": "affine"
582
+ },
583
+ "model.layers.14.self_attn.q_proj": {
584
+ "group_size": 64,
585
+ "bits": 8,
586
+ "mode": "affine"
587
+ },
588
+ "model.layers.14.self_attn.k_proj": {
589
+ "group_size": 64,
590
+ "bits": 8,
591
+ "mode": "affine"
592
+ },
593
+ "model.layers.14.self_attn.v_proj": {
594
+ "group_size": 64,
595
+ "bits": 8,
596
+ "mode": "affine"
597
+ },
598
+ "model.layers.14.self_attn.o_proj": {
599
+ "group_size": 64,
600
+ "bits": 8,
601
+ "mode": "affine"
602
+ },
603
+ "model.layers.15.self_attn.q_proj": {
604
+ "group_size": 64,
605
+ "bits": 8,
606
+ "mode": "affine"
607
+ },
608
+ "model.layers.15.self_attn.k_proj": {
609
+ "group_size": 64,
610
+ "bits": 8,
611
+ "mode": "affine"
612
+ },
613
+ "model.layers.15.self_attn.v_proj": {
614
+ "group_size": 64,
615
+ "bits": 8,
616
+ "mode": "affine"
617
+ },
618
+ "model.layers.15.self_attn.o_proj": {
619
+ "group_size": 64,
620
+ "bits": 8,
621
+ "mode": "affine"
622
+ },
623
+ "model.layers.16.self_attn.q_proj": {
624
+ "group_size": 64,
625
+ "bits": 8,
626
+ "mode": "affine"
627
+ },
628
+ "model.layers.16.self_attn.k_proj": {
629
+ "group_size": 64,
630
+ "bits": 8,
631
+ "mode": "affine"
632
+ },
633
+ "model.layers.16.self_attn.v_proj": {
634
+ "group_size": 64,
635
+ "bits": 8,
636
+ "mode": "affine"
637
+ },
638
+ "model.layers.16.self_attn.o_proj": {
639
+ "group_size": 64,
640
+ "bits": 8,
641
+ "mode": "affine"
642
+ },
643
+ "model.layers.17.self_attn.q_proj": {
644
+ "group_size": 64,
645
+ "bits": 8,
646
+ "mode": "affine"
647
+ },
648
+ "model.layers.17.self_attn.k_proj": {
649
+ "group_size": 64,
650
+ "bits": 8,
651
+ "mode": "affine"
652
+ },
653
+ "model.layers.17.self_attn.v_proj": {
654
+ "group_size": 64,
655
+ "bits": 8,
656
+ "mode": "affine"
657
+ },
658
+ "model.layers.17.self_attn.o_proj": {
659
+ "group_size": 64,
660
+ "bits": 8,
661
+ "mode": "affine"
662
+ },
663
+ "model.layers.18.self_attn.q_proj": {
664
+ "group_size": 64,
665
+ "bits": 8,
666
+ "mode": "affine"
667
+ },
668
+ "model.layers.18.self_attn.k_proj": {
669
+ "group_size": 64,
670
+ "bits": 8,
671
+ "mode": "affine"
672
+ },
673
+ "model.layers.18.self_attn.v_proj": {
674
+ "group_size": 64,
675
+ "bits": 8,
676
+ "mode": "affine"
677
+ },
678
+ "model.layers.18.self_attn.o_proj": {
679
+ "group_size": 64,
680
+ "bits": 8,
681
+ "mode": "affine"
682
+ },
683
+ "model.layers.19.self_attn.q_proj": {
684
+ "group_size": 64,
685
+ "bits": 8,
686
+ "mode": "affine"
687
+ },
688
+ "model.layers.19.self_attn.k_proj": {
689
+ "group_size": 64,
690
+ "bits": 8,
691
+ "mode": "affine"
692
+ },
693
+ "model.layers.19.self_attn.v_proj": {
694
+ "group_size": 64,
695
+ "bits": 8,
696
+ "mode": "affine"
697
+ },
698
+ "model.layers.19.self_attn.o_proj": {
699
+ "group_size": 64,
700
+ "bits": 8,
701
+ "mode": "affine"
702
+ },
703
+ "model.layers.20.self_attn.q_proj": {
704
+ "group_size": 64,
705
+ "bits": 8,
706
+ "mode": "affine"
707
+ },
708
+ "model.layers.20.self_attn.k_proj": {
709
+ "group_size": 64,
710
+ "bits": 8,
711
+ "mode": "affine"
712
+ },
713
+ "model.layers.20.self_attn.v_proj": {
714
+ "group_size": 64,
715
+ "bits": 8,
716
+ "mode": "affine"
717
+ },
718
+ "model.layers.20.self_attn.o_proj": {
719
+ "group_size": 64,
720
+ "bits": 8,
721
+ "mode": "affine"
722
+ },
723
+ "model.layers.21.self_attn.q_proj": {
724
+ "group_size": 64,
725
+ "bits": 8,
726
+ "mode": "affine"
727
+ },
728
+ "model.layers.21.self_attn.k_proj": {
729
+ "group_size": 64,
730
+ "bits": 8,
731
+ "mode": "affine"
732
+ },
733
+ "model.layers.21.self_attn.v_proj": {
734
+ "group_size": 64,
735
+ "bits": 8,
736
+ "mode": "affine"
737
+ },
738
+ "model.layers.21.self_attn.o_proj": {
739
+ "group_size": 64,
740
+ "bits": 8,
741
+ "mode": "affine"
742
+ },
743
+ "model.layers.22.self_attn.q_proj": {
744
+ "group_size": 64,
745
+ "bits": 8,
746
+ "mode": "affine"
747
+ },
748
+ "model.layers.22.self_attn.k_proj": {
749
+ "group_size": 64,
750
+ "bits": 8,
751
+ "mode": "affine"
752
+ },
753
+ "model.layers.22.self_attn.v_proj": {
754
+ "group_size": 64,
755
+ "bits": 8,
756
+ "mode": "affine"
757
+ },
758
+ "model.layers.22.self_attn.o_proj": {
759
+ "group_size": 64,
760
+ "bits": 8,
761
+ "mode": "affine"
762
+ },
763
+ "model.layers.23.self_attn.q_proj": {
764
+ "group_size": 64,
765
+ "bits": 8,
766
+ "mode": "affine"
767
+ },
768
+ "model.layers.23.self_attn.k_proj": {
769
+ "group_size": 64,
770
+ "bits": 8,
771
+ "mode": "affine"
772
+ },
773
+ "model.layers.23.self_attn.v_proj": {
774
+ "group_size": 64,
775
+ "bits": 8,
776
+ "mode": "affine"
777
+ },
778
+ "model.layers.23.self_attn.o_proj": {
779
+ "group_size": 64,
780
+ "bits": 8,
781
+ "mode": "affine"
782
+ },
783
+ "model.layers.24.self_attn.q_proj": {
784
+ "group_size": 64,
785
+ "bits": 8,
786
+ "mode": "affine"
787
+ },
788
+ "model.layers.24.self_attn.k_proj": {
789
+ "group_size": 64,
790
+ "bits": 8,
791
+ "mode": "affine"
792
+ },
793
+ "model.layers.24.self_attn.v_proj": {
794
+ "group_size": 64,
795
+ "bits": 8,
796
+ "mode": "affine"
797
+ },
798
+ "model.layers.24.self_attn.o_proj": {
799
+ "group_size": 64,
800
+ "bits": 8,
801
+ "mode": "affine"
802
+ },
803
+ "model.layers.25.self_attn.q_proj": {
804
+ "group_size": 64,
805
+ "bits": 8,
806
+ "mode": "affine"
807
+ },
808
+ "model.layers.25.self_attn.k_proj": {
809
+ "group_size": 64,
810
+ "bits": 8,
811
+ "mode": "affine"
812
+ },
813
+ "model.layers.25.self_attn.v_proj": {
814
+ "group_size": 64,
815
+ "bits": 8,
816
+ "mode": "affine"
817
+ },
818
+ "model.layers.25.self_attn.o_proj": {
819
+ "group_size": 64,
820
+ "bits": 8,
821
+ "mode": "affine"
822
+ },
823
+ "model.layers.26.self_attn.q_proj": {
824
+ "group_size": 64,
825
+ "bits": 8,
826
+ "mode": "affine"
827
+ },
828
+ "model.layers.26.self_attn.k_proj": {
829
+ "group_size": 64,
830
+ "bits": 8,
831
+ "mode": "affine"
832
+ },
833
+ "model.layers.26.self_attn.v_proj": {
834
+ "group_size": 64,
835
+ "bits": 8,
836
+ "mode": "affine"
837
+ },
838
+ "model.layers.26.self_attn.o_proj": {
839
+ "group_size": 64,
840
+ "bits": 8,
841
+ "mode": "affine"
842
+ },
843
+ "model.layers.27.self_attn.q_proj": {
844
+ "group_size": 64,
845
+ "bits": 8,
846
+ "mode": "affine"
847
+ },
848
+ "model.layers.27.self_attn.k_proj": {
849
+ "group_size": 64,
850
+ "bits": 8,
851
+ "mode": "affine"
852
+ },
853
+ "model.layers.27.self_attn.v_proj": {
854
+ "group_size": 64,
855
+ "bits": 8,
856
+ "mode": "affine"
857
+ },
858
+ "model.layers.27.self_attn.o_proj": {
859
+ "group_size": 64,
860
+ "bits": 8,
861
+ "mode": "affine"
862
+ },
863
+ "model.layers.28.self_attn.q_proj": {
864
+ "group_size": 64,
865
+ "bits": 8,
866
+ "mode": "affine"
867
+ },
868
+ "model.layers.28.self_attn.k_proj": {
869
+ "group_size": 64,
870
+ "bits": 8,
871
+ "mode": "affine"
872
+ },
873
+ "model.layers.28.self_attn.v_proj": {
874
+ "group_size": 64,
875
+ "bits": 8,
876
+ "mode": "affine"
877
+ },
878
+ "model.layers.28.self_attn.o_proj": {
879
+ "group_size": 64,
880
+ "bits": 8,
881
+ "mode": "affine"
882
+ },
883
+ "model.layers.29.self_attn.q_proj": {
884
+ "group_size": 64,
885
+ "bits": 8,
886
+ "mode": "affine"
887
+ },
888
+ "model.layers.29.self_attn.k_proj": {
889
+ "group_size": 64,
890
+ "bits": 8,
891
+ "mode": "affine"
892
+ },
893
+ "model.layers.29.self_attn.v_proj": {
894
+ "group_size": 64,
895
+ "bits": 8,
896
+ "mode": "affine"
897
+ },
898
+ "model.layers.29.self_attn.o_proj": {
899
+ "group_size": 64,
900
+ "bits": 8,
901
+ "mode": "affine"
902
+ },
903
+ "model.layers.30.self_attn.q_proj": {
904
+ "group_size": 64,
905
+ "bits": 8,
906
+ "mode": "affine"
907
+ },
908
+ "model.layers.30.self_attn.k_proj": {
909
+ "group_size": 64,
910
+ "bits": 8,
911
+ "mode": "affine"
912
+ },
913
+ "model.layers.30.self_attn.v_proj": {
914
+ "group_size": 64,
915
+ "bits": 8,
916
+ "mode": "affine"
917
+ },
918
+ "model.layers.30.self_attn.o_proj": {
919
+ "group_size": 64,
920
+ "bits": 8,
921
+ "mode": "affine"
922
+ },
923
+ "model.layers.31.self_attn.q_proj": {
924
+ "group_size": 64,
925
+ "bits": 8,
926
+ "mode": "affine"
927
+ },
928
+ "model.layers.31.self_attn.k_proj": {
929
+ "group_size": 64,
930
+ "bits": 8,
931
+ "mode": "affine"
932
+ },
933
+ "model.layers.31.self_attn.v_proj": {
934
+ "group_size": 64,
935
+ "bits": 8,
936
+ "mode": "affine"
937
+ },
938
+ "model.layers.31.self_attn.o_proj": {
939
+ "group_size": 64,
940
+ "bits": 8,
941
+ "mode": "affine"
942
+ },
943
+ "model.layers.32.self_attn.q_proj": {
944
+ "group_size": 64,
945
+ "bits": 8,
946
+ "mode": "affine"
947
+ },
948
+ "model.layers.32.self_attn.k_proj": {
949
+ "group_size": 64,
950
+ "bits": 8,
951
+ "mode": "affine"
952
+ },
953
+ "model.layers.32.self_attn.v_proj": {
954
+ "group_size": 64,
955
+ "bits": 8,
956
+ "mode": "affine"
957
+ },
958
+ "model.layers.32.self_attn.o_proj": {
959
+ "group_size": 64,
960
+ "bits": 8,
961
+ "mode": "affine"
962
+ },
963
+ "model.layers.33.self_attn.q_proj": {
964
+ "group_size": 64,
965
+ "bits": 8,
966
+ "mode": "affine"
967
+ },
968
+ "model.layers.33.self_attn.k_proj": {
969
+ "group_size": 64,
970
+ "bits": 8,
971
+ "mode": "affine"
972
+ },
973
+ "model.layers.33.self_attn.v_proj": {
974
+ "group_size": 64,
975
+ "bits": 8,
976
+ "mode": "affine"
977
+ },
978
+ "model.layers.33.self_attn.o_proj": {
979
+ "group_size": 64,
980
+ "bits": 8,
981
+ "mode": "affine"
982
+ },
983
+ "model.layers.34.self_attn.q_proj": {
984
+ "group_size": 64,
985
+ "bits": 8,
986
+ "mode": "affine"
987
+ },
988
+ "model.layers.34.self_attn.k_proj": {
989
+ "group_size": 64,
990
+ "bits": 8,
991
+ "mode": "affine"
992
+ },
993
+ "model.layers.34.self_attn.v_proj": {
994
+ "group_size": 64,
995
+ "bits": 8,
996
+ "mode": "affine"
997
+ },
998
+ "model.layers.34.self_attn.o_proj": {
999
+ "group_size": 64,
1000
+ "bits": 8,
1001
+ "mode": "affine"
1002
+ },
1003
+ "model.layers.35.self_attn.q_proj": {
1004
+ "group_size": 64,
1005
+ "bits": 8,
1006
+ "mode": "affine"
1007
+ },
1008
+ "model.layers.35.self_attn.k_proj": {
1009
+ "group_size": 64,
1010
+ "bits": 8,
1011
+ "mode": "affine"
1012
+ },
1013
+ "model.layers.35.self_attn.v_proj": {
1014
+ "group_size": 64,
1015
+ "bits": 8,
1016
+ "mode": "affine"
1017
+ },
1018
+ "model.layers.35.self_attn.o_proj": {
1019
+ "group_size": 64,
1020
+ "bits": 8,
1021
+ "mode": "affine"
1022
+ },
1023
+ "model.layers.36.self_attn.q_proj": {
1024
+ "group_size": 64,
1025
+ "bits": 8,
1026
+ "mode": "affine"
1027
+ },
1028
+ "model.layers.36.self_attn.k_proj": {
1029
+ "group_size": 64,
1030
+ "bits": 8,
1031
+ "mode": "affine"
1032
+ },
1033
+ "model.layers.36.self_attn.v_proj": {
1034
+ "group_size": 64,
1035
+ "bits": 8,
1036
+ "mode": "affine"
1037
+ },
1038
+ "model.layers.36.self_attn.o_proj": {
1039
+ "group_size": 64,
1040
+ "bits": 8,
1041
+ "mode": "affine"
1042
+ },
1043
+ "model.layers.37.self_attn.q_proj": {
1044
+ "group_size": 64,
1045
+ "bits": 8,
1046
+ "mode": "affine"
1047
+ },
1048
+ "model.layers.37.self_attn.k_proj": {
1049
+ "group_size": 64,
1050
+ "bits": 8,
1051
+ "mode": "affine"
1052
+ },
1053
+ "model.layers.37.self_attn.v_proj": {
1054
+ "group_size": 64,
1055
+ "bits": 8,
1056
+ "mode": "affine"
1057
+ },
1058
+ "model.layers.37.self_attn.o_proj": {
1059
+ "group_size": 64,
1060
+ "bits": 8,
1061
+ "mode": "affine"
1062
+ },
1063
+ "model.layers.38.self_attn.q_proj": {
1064
+ "group_size": 64,
1065
+ "bits": 8,
1066
+ "mode": "affine"
1067
+ },
1068
+ "model.layers.38.self_attn.k_proj": {
1069
+ "group_size": 64,
1070
+ "bits": 8,
1071
+ "mode": "affine"
1072
+ },
1073
+ "model.layers.38.self_attn.v_proj": {
1074
+ "group_size": 64,
1075
+ "bits": 8,
1076
+ "mode": "affine"
1077
+ },
1078
+ "model.layers.38.self_attn.o_proj": {
1079
+ "group_size": 64,
1080
+ "bits": 8,
1081
+ "mode": "affine"
1082
+ },
1083
+ "model.layers.39.self_attn.q_proj": {
1084
+ "group_size": 64,
1085
+ "bits": 8,
1086
+ "mode": "affine"
1087
+ },
1088
+ "model.layers.39.self_attn.k_proj": {
1089
+ "group_size": 64,
1090
+ "bits": 8,
1091
+ "mode": "affine"
1092
+ },
1093
+ "model.layers.39.self_attn.v_proj": {
1094
+ "group_size": 64,
1095
+ "bits": 8,
1096
+ "mode": "affine"
1097
+ },
1098
+ "model.layers.39.self_attn.o_proj": {
1099
+ "group_size": 64,
1100
+ "bits": 8,
1101
+ "mode": "affine"
1102
+ },
1103
+ "model.layers.40.self_attn.q_proj": {
1104
+ "group_size": 64,
1105
+ "bits": 8,
1106
+ "mode": "affine"
1107
+ },
1108
+ "model.layers.40.self_attn.k_proj": {
1109
+ "group_size": 64,
1110
+ "bits": 8,
1111
+ "mode": "affine"
1112
+ },
1113
+ "model.layers.40.self_attn.v_proj": {
1114
+ "group_size": 64,
1115
+ "bits": 8,
1116
+ "mode": "affine"
1117
+ },
1118
+ "model.layers.40.self_attn.o_proj": {
1119
+ "group_size": 64,
1120
+ "bits": 8,
1121
+ "mode": "affine"
1122
+ },
1123
+ "model.layers.41.self_attn.q_proj": {
1124
+ "group_size": 64,
1125
+ "bits": 8,
1126
+ "mode": "affine"
1127
+ },
1128
+ "model.layers.41.self_attn.k_proj": {
1129
+ "group_size": 64,
1130
+ "bits": 8,
1131
+ "mode": "affine"
1132
+ },
1133
+ "model.layers.41.self_attn.v_proj": {
1134
+ "group_size": 64,
1135
+ "bits": 8,
1136
+ "mode": "affine"
1137
+ },
1138
+ "model.layers.41.self_attn.o_proj": {
1139
+ "group_size": 64,
1140
+ "bits": 8,
1141
+ "mode": "affine"
1142
+ },
1143
+ "model.layers.42.self_attn.q_proj": {
1144
+ "group_size": 64,
1145
+ "bits": 8,
1146
+ "mode": "affine"
1147
+ },
1148
+ "model.layers.42.self_attn.k_proj": {
1149
+ "group_size": 64,
1150
+ "bits": 8,
1151
+ "mode": "affine"
1152
+ },
1153
+ "model.layers.42.self_attn.v_proj": {
1154
+ "group_size": 64,
1155
+ "bits": 8,
1156
+ "mode": "affine"
1157
+ },
1158
+ "model.layers.42.self_attn.o_proj": {
1159
+ "group_size": 64,
1160
+ "bits": 8,
1161
+ "mode": "affine"
1162
+ },
1163
+ "model.layers.43.self_attn.q_proj": {
1164
+ "group_size": 64,
1165
+ "bits": 8,
1166
+ "mode": "affine"
1167
+ },
1168
+ "model.layers.43.self_attn.k_proj": {
1169
+ "group_size": 64,
1170
+ "bits": 8,
1171
+ "mode": "affine"
1172
+ },
1173
+ "model.layers.43.self_attn.v_proj": {
1174
+ "group_size": 64,
1175
+ "bits": 8,
1176
+ "mode": "affine"
1177
+ },
1178
+ "model.layers.43.self_attn.o_proj": {
1179
+ "group_size": 64,
1180
+ "bits": 8,
1181
+ "mode": "affine"
1182
+ },
1183
+ "model.layers.44.self_attn.q_proj": {
1184
+ "group_size": 64,
1185
+ "bits": 8,
1186
+ "mode": "affine"
1187
+ },
1188
+ "model.layers.44.self_attn.k_proj": {
1189
+ "group_size": 64,
1190
+ "bits": 8,
1191
+ "mode": "affine"
1192
+ },
1193
+ "model.layers.44.self_attn.v_proj": {
1194
+ "group_size": 64,
1195
+ "bits": 8,
1196
+ "mode": "affine"
1197
+ },
1198
+ "model.layers.44.self_attn.o_proj": {
1199
+ "group_size": 64,
1200
+ "bits": 8,
1201
+ "mode": "affine"
1202
+ },
1203
+ "model.layers.45.self_attn.q_proj": {
1204
+ "group_size": 64,
1205
+ "bits": 8,
1206
+ "mode": "affine"
1207
+ },
1208
+ "model.layers.45.self_attn.k_proj": {
1209
+ "group_size": 64,
1210
+ "bits": 8,
1211
+ "mode": "affine"
1212
+ },
1213
+ "model.layers.45.self_attn.v_proj": {
1214
+ "group_size": 64,
1215
+ "bits": 8,
1216
+ "mode": "affine"
1217
+ },
1218
+ "model.layers.45.self_attn.o_proj": {
1219
+ "group_size": 64,
1220
+ "bits": 8,
1221
+ "mode": "affine"
1222
+ },
1223
+ "model.layers.46.self_attn.q_proj": {
1224
+ "group_size": 64,
1225
+ "bits": 8,
1226
+ "mode": "affine"
1227
+ },
1228
+ "model.layers.46.self_attn.k_proj": {
1229
+ "group_size": 64,
1230
+ "bits": 8,
1231
+ "mode": "affine"
1232
+ },
1233
+ "model.layers.46.self_attn.v_proj": {
1234
+ "group_size": 64,
1235
+ "bits": 8,
1236
+ "mode": "affine"
1237
+ },
1238
+ "model.layers.46.self_attn.o_proj": {
1239
+ "group_size": 64,
1240
+ "bits": 8,
1241
+ "mode": "affine"
1242
+ },
1243
+ "model.layers.47.self_attn.q_proj": {
1244
+ "group_size": 64,
1245
+ "bits": 8,
1246
+ "mode": "affine"
1247
+ },
1248
+ "model.layers.47.self_attn.k_proj": {
1249
+ "group_size": 64,
1250
+ "bits": 8,
1251
+ "mode": "affine"
1252
+ },
1253
+ "model.layers.47.self_attn.v_proj": {
1254
+ "group_size": 64,
1255
+ "bits": 8,
1256
+ "mode": "affine"
1257
+ },
1258
+ "model.layers.47.self_attn.o_proj": {
1259
+ "group_size": 64,
1260
+ "bits": 8,
1261
+ "mode": "affine"
1262
+ },
1263
+ "model.layers.48.self_attn.q_proj": {
1264
+ "group_size": 64,
1265
+ "bits": 8,
1266
+ "mode": "affine"
1267
+ },
1268
+ "model.layers.48.self_attn.k_proj": {
1269
+ "group_size": 64,
1270
+ "bits": 8,
1271
+ "mode": "affine"
1272
+ },
1273
+ "model.layers.48.self_attn.v_proj": {
1274
+ "group_size": 64,
1275
+ "bits": 8,
1276
+ "mode": "affine"
1277
+ },
1278
+ "model.layers.48.self_attn.o_proj": {
1279
+ "group_size": 64,
1280
+ "bits": 8,
1281
+ "mode": "affine"
1282
+ },
1283
+ "model.layers.49.self_attn.q_proj": {
1284
+ "group_size": 64,
1285
+ "bits": 8,
1286
+ "mode": "affine"
1287
+ },
1288
+ "model.layers.49.self_attn.k_proj": {
1289
+ "group_size": 64,
1290
+ "bits": 8,
1291
+ "mode": "affine"
1292
+ },
1293
+ "model.layers.49.self_attn.v_proj": {
1294
+ "group_size": 64,
1295
+ "bits": 8,
1296
+ "mode": "affine"
1297
+ },
1298
+ "model.layers.49.self_attn.o_proj": {
1299
+ "group_size": 64,
1300
+ "bits": 8,
1301
+ "mode": "affine"
1302
+ },
1303
+ "model.layers.50.self_attn.q_proj": {
1304
+ "group_size": 64,
1305
+ "bits": 8,
1306
+ "mode": "affine"
1307
+ },
1308
+ "model.layers.50.self_attn.k_proj": {
1309
+ "group_size": 64,
1310
+ "bits": 8,
1311
+ "mode": "affine"
1312
+ },
1313
+ "model.layers.50.self_attn.v_proj": {
1314
+ "group_size": 64,
1315
+ "bits": 8,
1316
+ "mode": "affine"
1317
+ },
1318
+ "model.layers.50.self_attn.o_proj": {
1319
+ "group_size": 64,
1320
+ "bits": 8,
1321
+ "mode": "affine"
1322
+ },
1323
+ "model.layers.51.self_attn.q_proj": {
1324
+ "group_size": 64,
1325
+ "bits": 8,
1326
+ "mode": "affine"
1327
+ },
1328
+ "model.layers.51.self_attn.k_proj": {
1329
+ "group_size": 64,
1330
+ "bits": 8,
1331
+ "mode": "affine"
1332
+ },
1333
+ "model.layers.51.self_attn.v_proj": {
1334
+ "group_size": 64,
1335
+ "bits": 8,
1336
+ "mode": "affine"
1337
+ },
1338
+ "model.layers.51.self_attn.o_proj": {
1339
+ "group_size": 64,
1340
+ "bits": 8,
1341
+ "mode": "affine"
1342
+ },
1343
+ "model.layers.52.self_attn.q_proj": {
1344
+ "group_size": 64,
1345
+ "bits": 8,
1346
+ "mode": "affine"
1347
+ },
1348
+ "model.layers.52.self_attn.k_proj": {
1349
+ "group_size": 64,
1350
+ "bits": 8,
1351
+ "mode": "affine"
1352
+ },
1353
+ "model.layers.52.self_attn.v_proj": {
1354
+ "group_size": 64,
1355
+ "bits": 8,
1356
+ "mode": "affine"
1357
+ },
1358
+ "model.layers.52.self_attn.o_proj": {
1359
+ "group_size": 64,
1360
+ "bits": 8,
1361
+ "mode": "affine"
1362
+ },
1363
+ "model.layers.53.self_attn.q_proj": {
1364
+ "group_size": 64,
1365
+ "bits": 8,
1366
+ "mode": "affine"
1367
+ },
1368
+ "model.layers.53.self_attn.k_proj": {
1369
+ "group_size": 64,
1370
+ "bits": 8,
1371
+ "mode": "affine"
1372
+ },
1373
+ "model.layers.53.self_attn.v_proj": {
1374
+ "group_size": 64,
1375
+ "bits": 8,
1376
+ "mode": "affine"
1377
+ },
1378
+ "model.layers.53.self_attn.o_proj": {
1379
+ "group_size": 64,
1380
+ "bits": 8,
1381
+ "mode": "affine"
1382
+ },
1383
+ "model.layers.54.self_attn.q_proj": {
1384
+ "group_size": 64,
1385
+ "bits": 8,
1386
+ "mode": "affine"
1387
+ },
1388
+ "model.layers.54.self_attn.k_proj": {
1389
+ "group_size": 64,
1390
+ "bits": 8,
1391
+ "mode": "affine"
1392
+ },
1393
+ "model.layers.54.self_attn.v_proj": {
1394
+ "group_size": 64,
1395
+ "bits": 8,
1396
+ "mode": "affine"
1397
+ },
1398
+ "model.layers.54.self_attn.o_proj": {
1399
+ "group_size": 64,
1400
+ "bits": 8,
1401
+ "mode": "affine"
1402
+ },
1403
+ "model.layers.55.self_attn.q_proj": {
1404
+ "group_size": 64,
1405
+ "bits": 8,
1406
+ "mode": "affine"
1407
+ },
1408
+ "model.layers.55.self_attn.k_proj": {
1409
+ "group_size": 64,
1410
+ "bits": 8,
1411
+ "mode": "affine"
1412
+ },
1413
+ "model.layers.55.self_attn.v_proj": {
1414
+ "group_size": 64,
1415
+ "bits": 8,
1416
+ "mode": "affine"
1417
+ },
1418
+ "model.layers.55.self_attn.o_proj": {
1419
+ "group_size": 64,
1420
+ "bits": 8,
1421
+ "mode": "affine"
1422
+ },
1423
+ "model.layers.56.self_attn.q_proj": {
1424
+ "group_size": 64,
1425
+ "bits": 8,
1426
+ "mode": "affine"
1427
+ },
1428
+ "model.layers.56.self_attn.k_proj": {
1429
+ "group_size": 64,
1430
+ "bits": 8,
1431
+ "mode": "affine"
1432
+ },
1433
+ "model.layers.56.self_attn.v_proj": {
1434
+ "group_size": 64,
1435
+ "bits": 8,
1436
+ "mode": "affine"
1437
+ },
1438
+ "model.layers.56.self_attn.o_proj": {
1439
+ "group_size": 64,
1440
+ "bits": 8,
1441
+ "mode": "affine"
1442
+ },
1443
+ "model.layers.57.self_attn.q_proj": {
1444
+ "group_size": 64,
1445
+ "bits": 8,
1446
+ "mode": "affine"
1447
+ },
1448
+ "model.layers.57.self_attn.k_proj": {
1449
+ "group_size": 64,
1450
+ "bits": 8,
1451
+ "mode": "affine"
1452
+ },
1453
+ "model.layers.57.self_attn.v_proj": {
1454
+ "group_size": 64,
1455
+ "bits": 8,
1456
+ "mode": "affine"
1457
+ },
1458
+ "model.layers.57.self_attn.o_proj": {
1459
+ "group_size": 64,
1460
+ "bits": 8,
1461
+ "mode": "affine"
1462
+ },
1463
+ "model.layers.58.self_attn.q_proj": {
1464
+ "group_size": 64,
1465
+ "bits": 8,
1466
+ "mode": "affine"
1467
+ },
1468
+ "model.layers.58.self_attn.k_proj": {
1469
+ "group_size": 64,
1470
+ "bits": 8,
1471
+ "mode": "affine"
1472
+ },
1473
+ "model.layers.58.self_attn.v_proj": {
1474
+ "group_size": 64,
1475
+ "bits": 8,
1476
+ "mode": "affine"
1477
+ },
1478
+ "model.layers.58.self_attn.o_proj": {
1479
+ "group_size": 64,
1480
+ "bits": 8,
1481
+ "mode": "affine"
1482
+ },
1483
+ "model.layers.59.self_attn.q_proj": {
1484
+ "group_size": 64,
1485
+ "bits": 8,
1486
+ "mode": "affine"
1487
+ },
1488
+ "model.layers.59.self_attn.k_proj": {
1489
+ "group_size": 64,
1490
+ "bits": 8,
1491
+ "mode": "affine"
1492
+ },
1493
+ "model.layers.59.self_attn.v_proj": {
1494
+ "group_size": 64,
1495
+ "bits": 8,
1496
+ "mode": "affine"
1497
+ },
1498
+ "model.layers.59.self_attn.o_proj": {
1499
+ "group_size": 64,
1500
+ "bits": 8,
1501
+ "mode": "affine"
1502
+ },
1503
+ "model.layers.60.self_attn.q_proj": {
1504
+ "group_size": 64,
1505
+ "bits": 8,
1506
+ "mode": "affine"
1507
+ },
1508
+ "model.layers.60.self_attn.k_proj": {
1509
+ "group_size": 64,
1510
+ "bits": 8,
1511
+ "mode": "affine"
1512
+ },
1513
+ "model.layers.60.self_attn.v_proj": {
1514
+ "group_size": 64,
1515
+ "bits": 8,
1516
+ "mode": "affine"
1517
+ },
1518
+ "model.layers.60.self_attn.o_proj": {
1519
+ "group_size": 64,
1520
+ "bits": 8,
1521
+ "mode": "affine"
1522
+ },
1523
+ "model.layers.61.self_attn.q_proj": {
1524
+ "group_size": 64,
1525
+ "bits": 8,
1526
+ "mode": "affine"
1527
+ },
1528
+ "model.layers.61.self_attn.k_proj": {
1529
+ "group_size": 64,
1530
+ "bits": 8,
1531
+ "mode": "affine"
1532
+ },
1533
+ "model.layers.61.self_attn.v_proj": {
1534
+ "group_size": 64,
1535
+ "bits": 8,
1536
+ "mode": "affine"
1537
+ },
1538
+ "model.layers.61.self_attn.o_proj": {
1539
+ "group_size": 64,
1540
+ "bits": 8,
1541
+ "mode": "affine"
1542
+ },
1543
+ "model.layers.62.self_attn.q_proj": {
1544
+ "group_size": 64,
1545
+ "bits": 8,
1546
+ "mode": "affine"
1547
+ },
1548
+ "model.layers.62.self_attn.k_proj": {
1549
+ "group_size": 64,
1550
+ "bits": 8,
1551
+ "mode": "affine"
1552
+ },
1553
+ "model.layers.62.self_attn.v_proj": {
1554
+ "group_size": 64,
1555
+ "bits": 8,
1556
+ "mode": "affine"
1557
+ },
1558
+ "model.layers.62.self_attn.o_proj": {
1559
+ "group_size": 64,
1560
+ "bits": 8,
1561
+ "mode": "affine"
1562
+ },
1563
+ "model.layers.63.self_attn.q_proj": {
1564
+ "group_size": 64,
1565
+ "bits": 8,
1566
+ "mode": "affine"
1567
+ },
1568
+ "model.layers.63.self_attn.k_proj": {
1569
+ "group_size": 64,
1570
+ "bits": 8,
1571
+ "mode": "affine"
1572
+ },
1573
+ "model.layers.63.self_attn.v_proj": {
1574
+ "group_size": 64,
1575
+ "bits": 8,
1576
+ "mode": "affine"
1577
+ },
1578
+ "model.layers.63.self_attn.o_proj": {
1579
+ "group_size": 64,
1580
+ "bits": 8,
1581
+ "mode": "affine"
1582
+ },
1583
+ "model.layers.64.self_attn.q_proj": {
1584
+ "group_size": 64,
1585
+ "bits": 8,
1586
+ "mode": "affine"
1587
+ },
1588
+ "model.layers.64.self_attn.k_proj": {
1589
+ "group_size": 64,
1590
+ "bits": 8,
1591
+ "mode": "affine"
1592
+ },
1593
+ "model.layers.64.self_attn.v_proj": {
1594
+ "group_size": 64,
1595
+ "bits": 8,
1596
+ "mode": "affine"
1597
+ },
1598
+ "model.layers.64.self_attn.o_proj": {
1599
+ "group_size": 64,
1600
+ "bits": 8,
1601
+ "mode": "affine"
1602
+ },
1603
+ "model.layers.65.self_attn.q_proj": {
1604
+ "group_size": 64,
1605
+ "bits": 8,
1606
+ "mode": "affine"
1607
+ },
1608
+ "model.layers.65.self_attn.k_proj": {
1609
+ "group_size": 64,
1610
+ "bits": 8,
1611
+ "mode": "affine"
1612
+ },
1613
+ "model.layers.65.self_attn.v_proj": {
1614
+ "group_size": 64,
1615
+ "bits": 8,
1616
+ "mode": "affine"
1617
+ },
1618
+ "model.layers.65.self_attn.o_proj": {
1619
+ "group_size": 64,
1620
+ "bits": 8,
1621
+ "mode": "affine"
1622
+ },
1623
+ "model.layers.66.self_attn.q_proj": {
1624
+ "group_size": 64,
1625
+ "bits": 8,
1626
+ "mode": "affine"
1627
+ },
1628
+ "model.layers.66.self_attn.k_proj": {
1629
+ "group_size": 64,
1630
+ "bits": 8,
1631
+ "mode": "affine"
1632
+ },
1633
+ "model.layers.66.self_attn.v_proj": {
1634
+ "group_size": 64,
1635
+ "bits": 8,
1636
+ "mode": "affine"
1637
+ },
1638
+ "model.layers.66.self_attn.o_proj": {
1639
+ "group_size": 64,
1640
+ "bits": 8,
1641
+ "mode": "affine"
1642
+ },
1643
+ "model.layers.67.self_attn.q_proj": {
1644
+ "group_size": 64,
1645
+ "bits": 8,
1646
+ "mode": "affine"
1647
+ },
1648
+ "model.layers.67.self_attn.k_proj": {
1649
+ "group_size": 64,
1650
+ "bits": 8,
1651
+ "mode": "affine"
1652
+ },
1653
+ "model.layers.67.self_attn.v_proj": {
1654
+ "group_size": 64,
1655
+ "bits": 8,
1656
+ "mode": "affine"
1657
+ },
1658
+ "model.layers.67.self_attn.o_proj": {
1659
+ "group_size": 64,
1660
+ "bits": 8,
1661
+ "mode": "affine"
1662
+ },
1663
+ "model.layers.68.self_attn.q_proj": {
1664
+ "group_size": 64,
1665
+ "bits": 8,
1666
+ "mode": "affine"
1667
+ },
1668
+ "model.layers.68.self_attn.k_proj": {
1669
+ "group_size": 64,
1670
+ "bits": 8,
1671
+ "mode": "affine"
1672
+ },
1673
+ "model.layers.68.self_attn.v_proj": {
1674
+ "group_size": 64,
1675
+ "bits": 8,
1676
+ "mode": "affine"
1677
+ },
1678
+ "model.layers.68.self_attn.o_proj": {
1679
+ "group_size": 64,
1680
+ "bits": 8,
1681
+ "mode": "affine"
1682
+ },
1683
+ "model.layers.69.self_attn.q_proj": {
1684
+ "group_size": 64,
1685
+ "bits": 8,
1686
+ "mode": "affine"
1687
+ },
1688
+ "model.layers.69.self_attn.k_proj": {
1689
+ "group_size": 64,
1690
+ "bits": 8,
1691
+ "mode": "affine"
1692
+ },
1693
+ "model.layers.69.self_attn.v_proj": {
1694
+ "group_size": 64,
1695
+ "bits": 8,
1696
+ "mode": "affine"
1697
+ },
1698
+ "model.layers.69.self_attn.o_proj": {
1699
+ "group_size": 64,
1700
+ "bits": 8,
1701
+ "mode": "affine"
1702
+ },
1703
+ "lm_head": {
1704
+ "group_size": 64,
1705
+ "bits": 8,
1706
+ "mode": "affine"
1707
+ }
1708
+ },
1709
+ "quantization_config": {
1710
+ "group_size": 32,
1711
+ "bits": 4,
1712
+ "mode": "mxfp4",
1713
+ "model.embed_tokens": {
1714
+ "group_size": 64,
1715
+ "bits": 8,
1716
+ "mode": "affine"
1717
+ },
1718
+ "model.layers.0.self_attn.q_proj": {
1719
+ "group_size": 64,
1720
+ "bits": 8,
1721
+ "mode": "affine"
1722
+ },
1723
+ "model.layers.0.self_attn.k_proj": {
1724
+ "group_size": 64,
1725
+ "bits": 8,
1726
+ "mode": "affine"
1727
+ },
1728
+ "model.layers.0.self_attn.v_proj": {
1729
+ "group_size": 64,
1730
+ "bits": 8,
1731
+ "mode": "affine"
1732
+ },
1733
+ "model.layers.0.self_attn.o_proj": {
1734
+ "group_size": 64,
1735
+ "bits": 8,
1736
+ "mode": "affine"
1737
+ },
1738
+ "model.layers.0.mlp.gate_proj": {
1739
+ "group_size": 64,
1740
+ "bits": 8,
1741
+ "mode": "affine"
1742
+ },
1743
+ "model.layers.0.mlp.up_proj": {
1744
+ "group_size": 64,
1745
+ "bits": 8,
1746
+ "mode": "affine"
1747
+ },
1748
+ "model.layers.0.mlp.down_proj": {
1749
+ "group_size": 64,
1750
+ "bits": 8,
1751
+ "mode": "affine"
1752
+ },
1753
+ "model.layers.1.self_attn.q_proj": {
1754
+ "group_size": 64,
1755
+ "bits": 8,
1756
+ "mode": "affine"
1757
+ },
1758
+ "model.layers.1.self_attn.k_proj": {
1759
+ "group_size": 64,
1760
+ "bits": 8,
1761
+ "mode": "affine"
1762
+ },
1763
+ "model.layers.1.self_attn.v_proj": {
1764
+ "group_size": 64,
1765
+ "bits": 8,
1766
+ "mode": "affine"
1767
+ },
1768
+ "model.layers.1.self_attn.o_proj": {
1769
+ "group_size": 64,
1770
+ "bits": 8,
1771
+ "mode": "affine"
1772
+ },
1773
+ "model.layers.2.self_attn.q_proj": {
1774
+ "group_size": 64,
1775
+ "bits": 8,
1776
+ "mode": "affine"
1777
+ },
1778
+ "model.layers.2.self_attn.k_proj": {
1779
+ "group_size": 64,
1780
+ "bits": 8,
1781
+ "mode": "affine"
1782
+ },
1783
+ "model.layers.2.self_attn.v_proj": {
1784
+ "group_size": 64,
1785
+ "bits": 8,
1786
+ "mode": "affine"
1787
+ },
1788
+ "model.layers.2.self_attn.o_proj": {
1789
+ "group_size": 64,
1790
+ "bits": 8,
1791
+ "mode": "affine"
1792
+ },
1793
+ "model.layers.3.self_attn.q_proj": {
1794
+ "group_size": 64,
1795
+ "bits": 8,
1796
+ "mode": "affine"
1797
+ },
1798
+ "model.layers.3.self_attn.k_proj": {
1799
+ "group_size": 64,
1800
+ "bits": 8,
1801
+ "mode": "affine"
1802
+ },
1803
+ "model.layers.3.self_attn.v_proj": {
1804
+ "group_size": 64,
1805
+ "bits": 8,
1806
+ "mode": "affine"
1807
+ },
1808
+ "model.layers.3.self_attn.o_proj": {
1809
+ "group_size": 64,
1810
+ "bits": 8,
1811
+ "mode": "affine"
1812
+ },
1813
+ "model.layers.4.self_attn.q_proj": {
1814
+ "group_size": 64,
1815
+ "bits": 8,
1816
+ "mode": "affine"
1817
+ },
1818
+ "model.layers.4.self_attn.k_proj": {
1819
+ "group_size": 64,
1820
+ "bits": 8,
1821
+ "mode": "affine"
1822
+ },
1823
+ "model.layers.4.self_attn.v_proj": {
1824
+ "group_size": 64,
1825
+ "bits": 8,
1826
+ "mode": "affine"
1827
+ },
1828
+ "model.layers.4.self_attn.o_proj": {
1829
+ "group_size": 64,
1830
+ "bits": 8,
1831
+ "mode": "affine"
1832
+ },
1833
+ "model.layers.5.self_attn.q_proj": {
1834
+ "group_size": 64,
1835
+ "bits": 8,
1836
+ "mode": "affine"
1837
+ },
1838
+ "model.layers.5.self_attn.k_proj": {
1839
+ "group_size": 64,
1840
+ "bits": 8,
1841
+ "mode": "affine"
1842
+ },
1843
+ "model.layers.5.self_attn.v_proj": {
1844
+ "group_size": 64,
1845
+ "bits": 8,
1846
+ "mode": "affine"
1847
+ },
1848
+ "model.layers.5.self_attn.o_proj": {
1849
+ "group_size": 64,
1850
+ "bits": 8,
1851
+ "mode": "affine"
1852
+ },
1853
+ "model.layers.6.self_attn.q_proj": {
1854
+ "group_size": 64,
1855
+ "bits": 8,
1856
+ "mode": "affine"
1857
+ },
1858
+ "model.layers.6.self_attn.k_proj": {
1859
+ "group_size": 64,
1860
+ "bits": 8,
1861
+ "mode": "affine"
1862
+ },
1863
+ "model.layers.6.self_attn.v_proj": {
1864
+ "group_size": 64,
1865
+ "bits": 8,
1866
+ "mode": "affine"
1867
+ },
1868
+ "model.layers.6.self_attn.o_proj": {
1869
+ "group_size": 64,
1870
+ "bits": 8,
1871
+ "mode": "affine"
1872
+ },
1873
+ "model.layers.7.self_attn.q_proj": {
1874
+ "group_size": 64,
1875
+ "bits": 8,
1876
+ "mode": "affine"
1877
+ },
1878
+ "model.layers.7.self_attn.k_proj": {
1879
+ "group_size": 64,
1880
+ "bits": 8,
1881
+ "mode": "affine"
1882
+ },
1883
+ "model.layers.7.self_attn.v_proj": {
1884
+ "group_size": 64,
1885
+ "bits": 8,
1886
+ "mode": "affine"
1887
+ },
1888
+ "model.layers.7.self_attn.o_proj": {
1889
+ "group_size": 64,
1890
+ "bits": 8,
1891
+ "mode": "affine"
1892
+ },
1893
+ "model.layers.8.self_attn.q_proj": {
1894
+ "group_size": 64,
1895
+ "bits": 8,
1896
+ "mode": "affine"
1897
+ },
1898
+ "model.layers.8.self_attn.k_proj": {
1899
+ "group_size": 64,
1900
+ "bits": 8,
1901
+ "mode": "affine"
1902
+ },
1903
+ "model.layers.8.self_attn.v_proj": {
1904
+ "group_size": 64,
1905
+ "bits": 8,
1906
+ "mode": "affine"
1907
+ },
1908
+ "model.layers.8.self_attn.o_proj": {
1909
+ "group_size": 64,
1910
+ "bits": 8,
1911
+ "mode": "affine"
1912
+ },
1913
+ "model.layers.9.self_attn.q_proj": {
1914
+ "group_size": 64,
1915
+ "bits": 8,
1916
+ "mode": "affine"
1917
+ },
1918
+ "model.layers.9.self_attn.k_proj": {
1919
+ "group_size": 64,
1920
+ "bits": 8,
1921
+ "mode": "affine"
1922
+ },
1923
+ "model.layers.9.self_attn.v_proj": {
1924
+ "group_size": 64,
1925
+ "bits": 8,
1926
+ "mode": "affine"
1927
+ },
1928
+ "model.layers.9.self_attn.o_proj": {
1929
+ "group_size": 64,
1930
+ "bits": 8,
1931
+ "mode": "affine"
1932
+ },
1933
+ "model.layers.10.self_attn.q_proj": {
1934
+ "group_size": 64,
1935
+ "bits": 8,
1936
+ "mode": "affine"
1937
+ },
1938
+ "model.layers.10.self_attn.k_proj": {
1939
+ "group_size": 64,
1940
+ "bits": 8,
1941
+ "mode": "affine"
1942
+ },
1943
+ "model.layers.10.self_attn.v_proj": {
1944
+ "group_size": 64,
1945
+ "bits": 8,
1946
+ "mode": "affine"
1947
+ },
1948
+ "model.layers.10.self_attn.o_proj": {
1949
+ "group_size": 64,
1950
+ "bits": 8,
1951
+ "mode": "affine"
1952
+ },
1953
+ "model.layers.11.self_attn.q_proj": {
1954
+ "group_size": 64,
1955
+ "bits": 8,
1956
+ "mode": "affine"
1957
+ },
1958
+ "model.layers.11.self_attn.k_proj": {
1959
+ "group_size": 64,
1960
+ "bits": 8,
1961
+ "mode": "affine"
1962
+ },
1963
+ "model.layers.11.self_attn.v_proj": {
1964
+ "group_size": 64,
1965
+ "bits": 8,
1966
+ "mode": "affine"
1967
+ },
1968
+ "model.layers.11.self_attn.o_proj": {
1969
+ "group_size": 64,
1970
+ "bits": 8,
1971
+ "mode": "affine"
1972
+ },
1973
+ "model.layers.12.self_attn.q_proj": {
1974
+ "group_size": 64,
1975
+ "bits": 8,
1976
+ "mode": "affine"
1977
+ },
1978
+ "model.layers.12.self_attn.k_proj": {
1979
+ "group_size": 64,
1980
+ "bits": 8,
1981
+ "mode": "affine"
1982
+ },
1983
+ "model.layers.12.self_attn.v_proj": {
1984
+ "group_size": 64,
1985
+ "bits": 8,
1986
+ "mode": "affine"
1987
+ },
1988
+ "model.layers.12.self_attn.o_proj": {
1989
+ "group_size": 64,
1990
+ "bits": 8,
1991
+ "mode": "affine"
1992
+ },
1993
+ "model.layers.13.self_attn.q_proj": {
1994
+ "group_size": 64,
1995
+ "bits": 8,
1996
+ "mode": "affine"
1997
+ },
1998
+ "model.layers.13.self_attn.k_proj": {
1999
+ "group_size": 64,
2000
+ "bits": 8,
2001
+ "mode": "affine"
2002
+ },
2003
+ "model.layers.13.self_attn.v_proj": {
2004
+ "group_size": 64,
2005
+ "bits": 8,
2006
+ "mode": "affine"
2007
+ },
2008
+ "model.layers.13.self_attn.o_proj": {
2009
+ "group_size": 64,
2010
+ "bits": 8,
2011
+ "mode": "affine"
2012
+ },
2013
+ "model.layers.14.self_attn.q_proj": {
2014
+ "group_size": 64,
2015
+ "bits": 8,
2016
+ "mode": "affine"
2017
+ },
2018
+ "model.layers.14.self_attn.k_proj": {
2019
+ "group_size": 64,
2020
+ "bits": 8,
2021
+ "mode": "affine"
2022
+ },
2023
+ "model.layers.14.self_attn.v_proj": {
2024
+ "group_size": 64,
2025
+ "bits": 8,
2026
+ "mode": "affine"
2027
+ },
2028
+ "model.layers.14.self_attn.o_proj": {
2029
+ "group_size": 64,
2030
+ "bits": 8,
2031
+ "mode": "affine"
2032
+ },
2033
+ "model.layers.15.self_attn.q_proj": {
2034
+ "group_size": 64,
2035
+ "bits": 8,
2036
+ "mode": "affine"
2037
+ },
2038
+ "model.layers.15.self_attn.k_proj": {
2039
+ "group_size": 64,
2040
+ "bits": 8,
2041
+ "mode": "affine"
2042
+ },
2043
+ "model.layers.15.self_attn.v_proj": {
2044
+ "group_size": 64,
2045
+ "bits": 8,
2046
+ "mode": "affine"
2047
+ },
2048
+ "model.layers.15.self_attn.o_proj": {
2049
+ "group_size": 64,
2050
+ "bits": 8,
2051
+ "mode": "affine"
2052
+ },
2053
+ "model.layers.16.self_attn.q_proj": {
2054
+ "group_size": 64,
2055
+ "bits": 8,
2056
+ "mode": "affine"
2057
+ },
2058
+ "model.layers.16.self_attn.k_proj": {
2059
+ "group_size": 64,
2060
+ "bits": 8,
2061
+ "mode": "affine"
2062
+ },
2063
+ "model.layers.16.self_attn.v_proj": {
2064
+ "group_size": 64,
2065
+ "bits": 8,
2066
+ "mode": "affine"
2067
+ },
2068
+ "model.layers.16.self_attn.o_proj": {
2069
+ "group_size": 64,
2070
+ "bits": 8,
2071
+ "mode": "affine"
2072
+ },
2073
+ "model.layers.17.self_attn.q_proj": {
2074
+ "group_size": 64,
2075
+ "bits": 8,
2076
+ "mode": "affine"
2077
+ },
2078
+ "model.layers.17.self_attn.k_proj": {
2079
+ "group_size": 64,
2080
+ "bits": 8,
2081
+ "mode": "affine"
2082
+ },
2083
+ "model.layers.17.self_attn.v_proj": {
2084
+ "group_size": 64,
2085
+ "bits": 8,
2086
+ "mode": "affine"
2087
+ },
2088
+ "model.layers.17.self_attn.o_proj": {
2089
+ "group_size": 64,
2090
+ "bits": 8,
2091
+ "mode": "affine"
2092
+ },
2093
+ "model.layers.18.self_attn.q_proj": {
2094
+ "group_size": 64,
2095
+ "bits": 8,
2096
+ "mode": "affine"
2097
+ },
2098
+ "model.layers.18.self_attn.k_proj": {
2099
+ "group_size": 64,
2100
+ "bits": 8,
2101
+ "mode": "affine"
2102
+ },
2103
+ "model.layers.18.self_attn.v_proj": {
2104
+ "group_size": 64,
2105
+ "bits": 8,
2106
+ "mode": "affine"
2107
+ },
2108
+ "model.layers.18.self_attn.o_proj": {
2109
+ "group_size": 64,
2110
+ "bits": 8,
2111
+ "mode": "affine"
2112
+ },
2113
+ "model.layers.19.self_attn.q_proj": {
2114
+ "group_size": 64,
2115
+ "bits": 8,
2116
+ "mode": "affine"
2117
+ },
2118
+ "model.layers.19.self_attn.k_proj": {
2119
+ "group_size": 64,
2120
+ "bits": 8,
2121
+ "mode": "affine"
2122
+ },
2123
+ "model.layers.19.self_attn.v_proj": {
2124
+ "group_size": 64,
2125
+ "bits": 8,
2126
+ "mode": "affine"
2127
+ },
2128
+ "model.layers.19.self_attn.o_proj": {
2129
+ "group_size": 64,
2130
+ "bits": 8,
2131
+ "mode": "affine"
2132
+ },
2133
+ "model.layers.20.self_attn.q_proj": {
2134
+ "group_size": 64,
2135
+ "bits": 8,
2136
+ "mode": "affine"
2137
+ },
2138
+ "model.layers.20.self_attn.k_proj": {
2139
+ "group_size": 64,
2140
+ "bits": 8,
2141
+ "mode": "affine"
2142
+ },
2143
+ "model.layers.20.self_attn.v_proj": {
2144
+ "group_size": 64,
2145
+ "bits": 8,
2146
+ "mode": "affine"
2147
+ },
2148
+ "model.layers.20.self_attn.o_proj": {
2149
+ "group_size": 64,
2150
+ "bits": 8,
2151
+ "mode": "affine"
2152
+ },
2153
+ "model.layers.21.self_attn.q_proj": {
2154
+ "group_size": 64,
2155
+ "bits": 8,
2156
+ "mode": "affine"
2157
+ },
2158
+ "model.layers.21.self_attn.k_proj": {
2159
+ "group_size": 64,
2160
+ "bits": 8,
2161
+ "mode": "affine"
2162
+ },
2163
+ "model.layers.21.self_attn.v_proj": {
2164
+ "group_size": 64,
2165
+ "bits": 8,
2166
+ "mode": "affine"
2167
+ },
2168
+ "model.layers.21.self_attn.o_proj": {
2169
+ "group_size": 64,
2170
+ "bits": 8,
2171
+ "mode": "affine"
2172
+ },
2173
+ "model.layers.22.self_attn.q_proj": {
2174
+ "group_size": 64,
2175
+ "bits": 8,
2176
+ "mode": "affine"
2177
+ },
2178
+ "model.layers.22.self_attn.k_proj": {
2179
+ "group_size": 64,
2180
+ "bits": 8,
2181
+ "mode": "affine"
2182
+ },
2183
+ "model.layers.22.self_attn.v_proj": {
2184
+ "group_size": 64,
2185
+ "bits": 8,
2186
+ "mode": "affine"
2187
+ },
2188
+ "model.layers.22.self_attn.o_proj": {
2189
+ "group_size": 64,
2190
+ "bits": 8,
2191
+ "mode": "affine"
2192
+ },
2193
+ "model.layers.23.self_attn.q_proj": {
2194
+ "group_size": 64,
2195
+ "bits": 8,
2196
+ "mode": "affine"
2197
+ },
2198
+ "model.layers.23.self_attn.k_proj": {
2199
+ "group_size": 64,
2200
+ "bits": 8,
2201
+ "mode": "affine"
2202
+ },
2203
+ "model.layers.23.self_attn.v_proj": {
2204
+ "group_size": 64,
2205
+ "bits": 8,
2206
+ "mode": "affine"
2207
+ },
2208
+ "model.layers.23.self_attn.o_proj": {
2209
+ "group_size": 64,
2210
+ "bits": 8,
2211
+ "mode": "affine"
2212
+ },
2213
+ "model.layers.24.self_attn.q_proj": {
2214
+ "group_size": 64,
2215
+ "bits": 8,
2216
+ "mode": "affine"
2217
+ },
2218
+ "model.layers.24.self_attn.k_proj": {
2219
+ "group_size": 64,
2220
+ "bits": 8,
2221
+ "mode": "affine"
2222
+ },
2223
+ "model.layers.24.self_attn.v_proj": {
2224
+ "group_size": 64,
2225
+ "bits": 8,
2226
+ "mode": "affine"
2227
+ },
2228
+ "model.layers.24.self_attn.o_proj": {
2229
+ "group_size": 64,
2230
+ "bits": 8,
2231
+ "mode": "affine"
2232
+ },
2233
+ "model.layers.25.self_attn.q_proj": {
2234
+ "group_size": 64,
2235
+ "bits": 8,
2236
+ "mode": "affine"
2237
+ },
2238
+ "model.layers.25.self_attn.k_proj": {
2239
+ "group_size": 64,
2240
+ "bits": 8,
2241
+ "mode": "affine"
2242
+ },
2243
+ "model.layers.25.self_attn.v_proj": {
2244
+ "group_size": 64,
2245
+ "bits": 8,
2246
+ "mode": "affine"
2247
+ },
2248
+ "model.layers.25.self_attn.o_proj": {
2249
+ "group_size": 64,
2250
+ "bits": 8,
2251
+ "mode": "affine"
2252
+ },
2253
+ "model.layers.26.self_attn.q_proj": {
2254
+ "group_size": 64,
2255
+ "bits": 8,
2256
+ "mode": "affine"
2257
+ },
2258
+ "model.layers.26.self_attn.k_proj": {
2259
+ "group_size": 64,
2260
+ "bits": 8,
2261
+ "mode": "affine"
2262
+ },
2263
+ "model.layers.26.self_attn.v_proj": {
2264
+ "group_size": 64,
2265
+ "bits": 8,
2266
+ "mode": "affine"
2267
+ },
2268
+ "model.layers.26.self_attn.o_proj": {
2269
+ "group_size": 64,
2270
+ "bits": 8,
2271
+ "mode": "affine"
2272
+ },
2273
+ "model.layers.27.self_attn.q_proj": {
2274
+ "group_size": 64,
2275
+ "bits": 8,
2276
+ "mode": "affine"
2277
+ },
2278
+ "model.layers.27.self_attn.k_proj": {
2279
+ "group_size": 64,
2280
+ "bits": 8,
2281
+ "mode": "affine"
2282
+ },
2283
+ "model.layers.27.self_attn.v_proj": {
2284
+ "group_size": 64,
2285
+ "bits": 8,
2286
+ "mode": "affine"
2287
+ },
2288
+ "model.layers.27.self_attn.o_proj": {
2289
+ "group_size": 64,
2290
+ "bits": 8,
2291
+ "mode": "affine"
2292
+ },
2293
+ "model.layers.28.self_attn.q_proj": {
2294
+ "group_size": 64,
2295
+ "bits": 8,
2296
+ "mode": "affine"
2297
+ },
2298
+ "model.layers.28.self_attn.k_proj": {
2299
+ "group_size": 64,
2300
+ "bits": 8,
2301
+ "mode": "affine"
2302
+ },
2303
+ "model.layers.28.self_attn.v_proj": {
2304
+ "group_size": 64,
2305
+ "bits": 8,
2306
+ "mode": "affine"
2307
+ },
2308
+ "model.layers.28.self_attn.o_proj": {
2309
+ "group_size": 64,
2310
+ "bits": 8,
2311
+ "mode": "affine"
2312
+ },
2313
+ "model.layers.29.self_attn.q_proj": {
2314
+ "group_size": 64,
2315
+ "bits": 8,
2316
+ "mode": "affine"
2317
+ },
2318
+ "model.layers.29.self_attn.k_proj": {
2319
+ "group_size": 64,
2320
+ "bits": 8,
2321
+ "mode": "affine"
2322
+ },
2323
+ "model.layers.29.self_attn.v_proj": {
2324
+ "group_size": 64,
2325
+ "bits": 8,
2326
+ "mode": "affine"
2327
+ },
2328
+ "model.layers.29.self_attn.o_proj": {
2329
+ "group_size": 64,
2330
+ "bits": 8,
2331
+ "mode": "affine"
2332
+ },
2333
+ "model.layers.30.self_attn.q_proj": {
2334
+ "group_size": 64,
2335
+ "bits": 8,
2336
+ "mode": "affine"
2337
+ },
2338
+ "model.layers.30.self_attn.k_proj": {
2339
+ "group_size": 64,
2340
+ "bits": 8,
2341
+ "mode": "affine"
2342
+ },
2343
+ "model.layers.30.self_attn.v_proj": {
2344
+ "group_size": 64,
2345
+ "bits": 8,
2346
+ "mode": "affine"
2347
+ },
2348
+ "model.layers.30.self_attn.o_proj": {
2349
+ "group_size": 64,
2350
+ "bits": 8,
2351
+ "mode": "affine"
2352
+ },
2353
+ "model.layers.31.self_attn.q_proj": {
2354
+ "group_size": 64,
2355
+ "bits": 8,
2356
+ "mode": "affine"
2357
+ },
2358
+ "model.layers.31.self_attn.k_proj": {
2359
+ "group_size": 64,
2360
+ "bits": 8,
2361
+ "mode": "affine"
2362
+ },
2363
+ "model.layers.31.self_attn.v_proj": {
2364
+ "group_size": 64,
2365
+ "bits": 8,
2366
+ "mode": "affine"
2367
+ },
2368
+ "model.layers.31.self_attn.o_proj": {
2369
+ "group_size": 64,
2370
+ "bits": 8,
2371
+ "mode": "affine"
2372
+ },
2373
+ "model.layers.32.self_attn.q_proj": {
2374
+ "group_size": 64,
2375
+ "bits": 8,
2376
+ "mode": "affine"
2377
+ },
2378
+ "model.layers.32.self_attn.k_proj": {
2379
+ "group_size": 64,
2380
+ "bits": 8,
2381
+ "mode": "affine"
2382
+ },
2383
+ "model.layers.32.self_attn.v_proj": {
2384
+ "group_size": 64,
2385
+ "bits": 8,
2386
+ "mode": "affine"
2387
+ },
2388
+ "model.layers.32.self_attn.o_proj": {
2389
+ "group_size": 64,
2390
+ "bits": 8,
2391
+ "mode": "affine"
2392
+ },
2393
+ "model.layers.33.self_attn.q_proj": {
2394
+ "group_size": 64,
2395
+ "bits": 8,
2396
+ "mode": "affine"
2397
+ },
2398
+ "model.layers.33.self_attn.k_proj": {
2399
+ "group_size": 64,
2400
+ "bits": 8,
2401
+ "mode": "affine"
2402
+ },
2403
+ "model.layers.33.self_attn.v_proj": {
2404
+ "group_size": 64,
2405
+ "bits": 8,
2406
+ "mode": "affine"
2407
+ },
2408
+ "model.layers.33.self_attn.o_proj": {
2409
+ "group_size": 64,
2410
+ "bits": 8,
2411
+ "mode": "affine"
2412
+ },
2413
+ "model.layers.34.self_attn.q_proj": {
2414
+ "group_size": 64,
2415
+ "bits": 8,
2416
+ "mode": "affine"
2417
+ },
2418
+ "model.layers.34.self_attn.k_proj": {
2419
+ "group_size": 64,
2420
+ "bits": 8,
2421
+ "mode": "affine"
2422
+ },
2423
+ "model.layers.34.self_attn.v_proj": {
2424
+ "group_size": 64,
2425
+ "bits": 8,
2426
+ "mode": "affine"
2427
+ },
2428
+ "model.layers.34.self_attn.o_proj": {
2429
+ "group_size": 64,
2430
+ "bits": 8,
2431
+ "mode": "affine"
2432
+ },
2433
+ "model.layers.35.self_attn.q_proj": {
2434
+ "group_size": 64,
2435
+ "bits": 8,
2436
+ "mode": "affine"
2437
+ },
2438
+ "model.layers.35.self_attn.k_proj": {
2439
+ "group_size": 64,
2440
+ "bits": 8,
2441
+ "mode": "affine"
2442
+ },
2443
+ "model.layers.35.self_attn.v_proj": {
2444
+ "group_size": 64,
2445
+ "bits": 8,
2446
+ "mode": "affine"
2447
+ },
2448
+ "model.layers.35.self_attn.o_proj": {
2449
+ "group_size": 64,
2450
+ "bits": 8,
2451
+ "mode": "affine"
2452
+ },
2453
+ "model.layers.36.self_attn.q_proj": {
2454
+ "group_size": 64,
2455
+ "bits": 8,
2456
+ "mode": "affine"
2457
+ },
2458
+ "model.layers.36.self_attn.k_proj": {
2459
+ "group_size": 64,
2460
+ "bits": 8,
2461
+ "mode": "affine"
2462
+ },
2463
+ "model.layers.36.self_attn.v_proj": {
2464
+ "group_size": 64,
2465
+ "bits": 8,
2466
+ "mode": "affine"
2467
+ },
2468
+ "model.layers.36.self_attn.o_proj": {
2469
+ "group_size": 64,
2470
+ "bits": 8,
2471
+ "mode": "affine"
2472
+ },
2473
+ "model.layers.37.self_attn.q_proj": {
2474
+ "group_size": 64,
2475
+ "bits": 8,
2476
+ "mode": "affine"
2477
+ },
2478
+ "model.layers.37.self_attn.k_proj": {
2479
+ "group_size": 64,
2480
+ "bits": 8,
2481
+ "mode": "affine"
2482
+ },
2483
+ "model.layers.37.self_attn.v_proj": {
2484
+ "group_size": 64,
2485
+ "bits": 8,
2486
+ "mode": "affine"
2487
+ },
2488
+ "model.layers.37.self_attn.o_proj": {
2489
+ "group_size": 64,
2490
+ "bits": 8,
2491
+ "mode": "affine"
2492
+ },
2493
+ "model.layers.38.self_attn.q_proj": {
2494
+ "group_size": 64,
2495
+ "bits": 8,
2496
+ "mode": "affine"
2497
+ },
2498
+ "model.layers.38.self_attn.k_proj": {
2499
+ "group_size": 64,
2500
+ "bits": 8,
2501
+ "mode": "affine"
2502
+ },
2503
+ "model.layers.38.self_attn.v_proj": {
2504
+ "group_size": 64,
2505
+ "bits": 8,
2506
+ "mode": "affine"
2507
+ },
2508
+ "model.layers.38.self_attn.o_proj": {
2509
+ "group_size": 64,
2510
+ "bits": 8,
2511
+ "mode": "affine"
2512
+ },
2513
+ "model.layers.39.self_attn.q_proj": {
2514
+ "group_size": 64,
2515
+ "bits": 8,
2516
+ "mode": "affine"
2517
+ },
2518
+ "model.layers.39.self_attn.k_proj": {
2519
+ "group_size": 64,
2520
+ "bits": 8,
2521
+ "mode": "affine"
2522
+ },
2523
+ "model.layers.39.self_attn.v_proj": {
2524
+ "group_size": 64,
2525
+ "bits": 8,
2526
+ "mode": "affine"
2527
+ },
2528
+ "model.layers.39.self_attn.o_proj": {
2529
+ "group_size": 64,
2530
+ "bits": 8,
2531
+ "mode": "affine"
2532
+ },
2533
+ "model.layers.40.self_attn.q_proj": {
2534
+ "group_size": 64,
2535
+ "bits": 8,
2536
+ "mode": "affine"
2537
+ },
2538
+ "model.layers.40.self_attn.k_proj": {
2539
+ "group_size": 64,
2540
+ "bits": 8,
2541
+ "mode": "affine"
2542
+ },
2543
+ "model.layers.40.self_attn.v_proj": {
2544
+ "group_size": 64,
2545
+ "bits": 8,
2546
+ "mode": "affine"
2547
+ },
2548
+ "model.layers.40.self_attn.o_proj": {
2549
+ "group_size": 64,
2550
+ "bits": 8,
2551
+ "mode": "affine"
2552
+ },
2553
+ "model.layers.41.self_attn.q_proj": {
2554
+ "group_size": 64,
2555
+ "bits": 8,
2556
+ "mode": "affine"
2557
+ },
2558
+ "model.layers.41.self_attn.k_proj": {
2559
+ "group_size": 64,
2560
+ "bits": 8,
2561
+ "mode": "affine"
2562
+ },
2563
+ "model.layers.41.self_attn.v_proj": {
2564
+ "group_size": 64,
2565
+ "bits": 8,
2566
+ "mode": "affine"
2567
+ },
2568
+ "model.layers.41.self_attn.o_proj": {
2569
+ "group_size": 64,
2570
+ "bits": 8,
2571
+ "mode": "affine"
2572
+ },
2573
+ "model.layers.42.self_attn.q_proj": {
2574
+ "group_size": 64,
2575
+ "bits": 8,
2576
+ "mode": "affine"
2577
+ },
2578
+ "model.layers.42.self_attn.k_proj": {
2579
+ "group_size": 64,
2580
+ "bits": 8,
2581
+ "mode": "affine"
2582
+ },
2583
+ "model.layers.42.self_attn.v_proj": {
2584
+ "group_size": 64,
2585
+ "bits": 8,
2586
+ "mode": "affine"
2587
+ },
2588
+ "model.layers.42.self_attn.o_proj": {
2589
+ "group_size": 64,
2590
+ "bits": 8,
2591
+ "mode": "affine"
2592
+ },
2593
+ "model.layers.43.self_attn.q_proj": {
2594
+ "group_size": 64,
2595
+ "bits": 8,
2596
+ "mode": "affine"
2597
+ },
2598
+ "model.layers.43.self_attn.k_proj": {
2599
+ "group_size": 64,
2600
+ "bits": 8,
2601
+ "mode": "affine"
2602
+ },
2603
+ "model.layers.43.self_attn.v_proj": {
2604
+ "group_size": 64,
2605
+ "bits": 8,
2606
+ "mode": "affine"
2607
+ },
2608
+ "model.layers.43.self_attn.o_proj": {
2609
+ "group_size": 64,
2610
+ "bits": 8,
2611
+ "mode": "affine"
2612
+ },
2613
+ "model.layers.44.self_attn.q_proj": {
2614
+ "group_size": 64,
2615
+ "bits": 8,
2616
+ "mode": "affine"
2617
+ },
2618
+ "model.layers.44.self_attn.k_proj": {
2619
+ "group_size": 64,
2620
+ "bits": 8,
2621
+ "mode": "affine"
2622
+ },
2623
+ "model.layers.44.self_attn.v_proj": {
2624
+ "group_size": 64,
2625
+ "bits": 8,
2626
+ "mode": "affine"
2627
+ },
2628
+ "model.layers.44.self_attn.o_proj": {
2629
+ "group_size": 64,
2630
+ "bits": 8,
2631
+ "mode": "affine"
2632
+ },
2633
+ "model.layers.45.self_attn.q_proj": {
2634
+ "group_size": 64,
2635
+ "bits": 8,
2636
+ "mode": "affine"
2637
+ },
2638
+ "model.layers.45.self_attn.k_proj": {
2639
+ "group_size": 64,
2640
+ "bits": 8,
2641
+ "mode": "affine"
2642
+ },
2643
+ "model.layers.45.self_attn.v_proj": {
2644
+ "group_size": 64,
2645
+ "bits": 8,
2646
+ "mode": "affine"
2647
+ },
2648
+ "model.layers.45.self_attn.o_proj": {
2649
+ "group_size": 64,
2650
+ "bits": 8,
2651
+ "mode": "affine"
2652
+ },
2653
+ "model.layers.46.self_attn.q_proj": {
2654
+ "group_size": 64,
2655
+ "bits": 8,
2656
+ "mode": "affine"
2657
+ },
2658
+ "model.layers.46.self_attn.k_proj": {
2659
+ "group_size": 64,
2660
+ "bits": 8,
2661
+ "mode": "affine"
2662
+ },
2663
+ "model.layers.46.self_attn.v_proj": {
2664
+ "group_size": 64,
2665
+ "bits": 8,
2666
+ "mode": "affine"
2667
+ },
2668
+ "model.layers.46.self_attn.o_proj": {
2669
+ "group_size": 64,
2670
+ "bits": 8,
2671
+ "mode": "affine"
2672
+ },
2673
+ "model.layers.47.self_attn.q_proj": {
2674
+ "group_size": 64,
2675
+ "bits": 8,
2676
+ "mode": "affine"
2677
+ },
2678
+ "model.layers.47.self_attn.k_proj": {
2679
+ "group_size": 64,
2680
+ "bits": 8,
2681
+ "mode": "affine"
2682
+ },
2683
+ "model.layers.47.self_attn.v_proj": {
2684
+ "group_size": 64,
2685
+ "bits": 8,
2686
+ "mode": "affine"
2687
+ },
2688
+ "model.layers.47.self_attn.o_proj": {
2689
+ "group_size": 64,
2690
+ "bits": 8,
2691
+ "mode": "affine"
2692
+ },
2693
+ "model.layers.48.self_attn.q_proj": {
2694
+ "group_size": 64,
2695
+ "bits": 8,
2696
+ "mode": "affine"
2697
+ },
2698
+ "model.layers.48.self_attn.k_proj": {
2699
+ "group_size": 64,
2700
+ "bits": 8,
2701
+ "mode": "affine"
2702
+ },
2703
+ "model.layers.48.self_attn.v_proj": {
2704
+ "group_size": 64,
2705
+ "bits": 8,
2706
+ "mode": "affine"
2707
+ },
2708
+ "model.layers.48.self_attn.o_proj": {
2709
+ "group_size": 64,
2710
+ "bits": 8,
2711
+ "mode": "affine"
2712
+ },
2713
+ "model.layers.49.self_attn.q_proj": {
2714
+ "group_size": 64,
2715
+ "bits": 8,
2716
+ "mode": "affine"
2717
+ },
2718
+ "model.layers.49.self_attn.k_proj": {
2719
+ "group_size": 64,
2720
+ "bits": 8,
2721
+ "mode": "affine"
2722
+ },
2723
+ "model.layers.49.self_attn.v_proj": {
2724
+ "group_size": 64,
2725
+ "bits": 8,
2726
+ "mode": "affine"
2727
+ },
2728
+ "model.layers.49.self_attn.o_proj": {
2729
+ "group_size": 64,
2730
+ "bits": 8,
2731
+ "mode": "affine"
2732
+ },
2733
+ "model.layers.50.self_attn.q_proj": {
2734
+ "group_size": 64,
2735
+ "bits": 8,
2736
+ "mode": "affine"
2737
+ },
2738
+ "model.layers.50.self_attn.k_proj": {
2739
+ "group_size": 64,
2740
+ "bits": 8,
2741
+ "mode": "affine"
2742
+ },
2743
+ "model.layers.50.self_attn.v_proj": {
2744
+ "group_size": 64,
2745
+ "bits": 8,
2746
+ "mode": "affine"
2747
+ },
2748
+ "model.layers.50.self_attn.o_proj": {
2749
+ "group_size": 64,
2750
+ "bits": 8,
2751
+ "mode": "affine"
2752
+ },
2753
+ "model.layers.51.self_attn.q_proj": {
2754
+ "group_size": 64,
2755
+ "bits": 8,
2756
+ "mode": "affine"
2757
+ },
2758
+ "model.layers.51.self_attn.k_proj": {
2759
+ "group_size": 64,
2760
+ "bits": 8,
2761
+ "mode": "affine"
2762
+ },
2763
+ "model.layers.51.self_attn.v_proj": {
2764
+ "group_size": 64,
2765
+ "bits": 8,
2766
+ "mode": "affine"
2767
+ },
2768
+ "model.layers.51.self_attn.o_proj": {
2769
+ "group_size": 64,
2770
+ "bits": 8,
2771
+ "mode": "affine"
2772
+ },
2773
+ "model.layers.52.self_attn.q_proj": {
2774
+ "group_size": 64,
2775
+ "bits": 8,
2776
+ "mode": "affine"
2777
+ },
2778
+ "model.layers.52.self_attn.k_proj": {
2779
+ "group_size": 64,
2780
+ "bits": 8,
2781
+ "mode": "affine"
2782
+ },
2783
+ "model.layers.52.self_attn.v_proj": {
2784
+ "group_size": 64,
2785
+ "bits": 8,
2786
+ "mode": "affine"
2787
+ },
2788
+ "model.layers.52.self_attn.o_proj": {
2789
+ "group_size": 64,
2790
+ "bits": 8,
2791
+ "mode": "affine"
2792
+ },
2793
+ "model.layers.53.self_attn.q_proj": {
2794
+ "group_size": 64,
2795
+ "bits": 8,
2796
+ "mode": "affine"
2797
+ },
2798
+ "model.layers.53.self_attn.k_proj": {
2799
+ "group_size": 64,
2800
+ "bits": 8,
2801
+ "mode": "affine"
2802
+ },
2803
+ "model.layers.53.self_attn.v_proj": {
2804
+ "group_size": 64,
2805
+ "bits": 8,
2806
+ "mode": "affine"
2807
+ },
2808
+ "model.layers.53.self_attn.o_proj": {
2809
+ "group_size": 64,
2810
+ "bits": 8,
2811
+ "mode": "affine"
2812
+ },
2813
+ "model.layers.54.self_attn.q_proj": {
2814
+ "group_size": 64,
2815
+ "bits": 8,
2816
+ "mode": "affine"
2817
+ },
2818
+ "model.layers.54.self_attn.k_proj": {
2819
+ "group_size": 64,
2820
+ "bits": 8,
2821
+ "mode": "affine"
2822
+ },
2823
+ "model.layers.54.self_attn.v_proj": {
2824
+ "group_size": 64,
2825
+ "bits": 8,
2826
+ "mode": "affine"
2827
+ },
2828
+ "model.layers.54.self_attn.o_proj": {
2829
+ "group_size": 64,
2830
+ "bits": 8,
2831
+ "mode": "affine"
2832
+ },
2833
+ "model.layers.55.self_attn.q_proj": {
2834
+ "group_size": 64,
2835
+ "bits": 8,
2836
+ "mode": "affine"
2837
+ },
2838
+ "model.layers.55.self_attn.k_proj": {
2839
+ "group_size": 64,
2840
+ "bits": 8,
2841
+ "mode": "affine"
2842
+ },
2843
+ "model.layers.55.self_attn.v_proj": {
2844
+ "group_size": 64,
2845
+ "bits": 8,
2846
+ "mode": "affine"
2847
+ },
2848
+ "model.layers.55.self_attn.o_proj": {
2849
+ "group_size": 64,
2850
+ "bits": 8,
2851
+ "mode": "affine"
2852
+ },
2853
+ "model.layers.56.self_attn.q_proj": {
2854
+ "group_size": 64,
2855
+ "bits": 8,
2856
+ "mode": "affine"
2857
+ },
2858
+ "model.layers.56.self_attn.k_proj": {
2859
+ "group_size": 64,
2860
+ "bits": 8,
2861
+ "mode": "affine"
2862
+ },
2863
+ "model.layers.56.self_attn.v_proj": {
2864
+ "group_size": 64,
2865
+ "bits": 8,
2866
+ "mode": "affine"
2867
+ },
2868
+ "model.layers.56.self_attn.o_proj": {
2869
+ "group_size": 64,
2870
+ "bits": 8,
2871
+ "mode": "affine"
2872
+ },
2873
+ "model.layers.57.self_attn.q_proj": {
2874
+ "group_size": 64,
2875
+ "bits": 8,
2876
+ "mode": "affine"
2877
+ },
2878
+ "model.layers.57.self_attn.k_proj": {
2879
+ "group_size": 64,
2880
+ "bits": 8,
2881
+ "mode": "affine"
2882
+ },
2883
+ "model.layers.57.self_attn.v_proj": {
2884
+ "group_size": 64,
2885
+ "bits": 8,
2886
+ "mode": "affine"
2887
+ },
2888
+ "model.layers.57.self_attn.o_proj": {
2889
+ "group_size": 64,
2890
+ "bits": 8,
2891
+ "mode": "affine"
2892
+ },
2893
+ "model.layers.58.self_attn.q_proj": {
2894
+ "group_size": 64,
2895
+ "bits": 8,
2896
+ "mode": "affine"
2897
+ },
2898
+ "model.layers.58.self_attn.k_proj": {
2899
+ "group_size": 64,
2900
+ "bits": 8,
2901
+ "mode": "affine"
2902
+ },
2903
+ "model.layers.58.self_attn.v_proj": {
2904
+ "group_size": 64,
2905
+ "bits": 8,
2906
+ "mode": "affine"
2907
+ },
2908
+ "model.layers.58.self_attn.o_proj": {
2909
+ "group_size": 64,
2910
+ "bits": 8,
2911
+ "mode": "affine"
2912
+ },
2913
+ "model.layers.59.self_attn.q_proj": {
2914
+ "group_size": 64,
2915
+ "bits": 8,
2916
+ "mode": "affine"
2917
+ },
2918
+ "model.layers.59.self_attn.k_proj": {
2919
+ "group_size": 64,
2920
+ "bits": 8,
2921
+ "mode": "affine"
2922
+ },
2923
+ "model.layers.59.self_attn.v_proj": {
2924
+ "group_size": 64,
2925
+ "bits": 8,
2926
+ "mode": "affine"
2927
+ },
2928
+ "model.layers.59.self_attn.o_proj": {
2929
+ "group_size": 64,
2930
+ "bits": 8,
2931
+ "mode": "affine"
2932
+ },
2933
+ "model.layers.60.self_attn.q_proj": {
2934
+ "group_size": 64,
2935
+ "bits": 8,
2936
+ "mode": "affine"
2937
+ },
2938
+ "model.layers.60.self_attn.k_proj": {
2939
+ "group_size": 64,
2940
+ "bits": 8,
2941
+ "mode": "affine"
2942
+ },
2943
+ "model.layers.60.self_attn.v_proj": {
2944
+ "group_size": 64,
2945
+ "bits": 8,
2946
+ "mode": "affine"
2947
+ },
2948
+ "model.layers.60.self_attn.o_proj": {
2949
+ "group_size": 64,
2950
+ "bits": 8,
2951
+ "mode": "affine"
2952
+ },
2953
+ "model.layers.61.self_attn.q_proj": {
2954
+ "group_size": 64,
2955
+ "bits": 8,
2956
+ "mode": "affine"
2957
+ },
2958
+ "model.layers.61.self_attn.k_proj": {
2959
+ "group_size": 64,
2960
+ "bits": 8,
2961
+ "mode": "affine"
2962
+ },
2963
+ "model.layers.61.self_attn.v_proj": {
2964
+ "group_size": 64,
2965
+ "bits": 8,
2966
+ "mode": "affine"
2967
+ },
2968
+ "model.layers.61.self_attn.o_proj": {
2969
+ "group_size": 64,
2970
+ "bits": 8,
2971
+ "mode": "affine"
2972
+ },
2973
+ "model.layers.62.self_attn.q_proj": {
2974
+ "group_size": 64,
2975
+ "bits": 8,
2976
+ "mode": "affine"
2977
+ },
2978
+ "model.layers.62.self_attn.k_proj": {
2979
+ "group_size": 64,
2980
+ "bits": 8,
2981
+ "mode": "affine"
2982
+ },
2983
+ "model.layers.62.self_attn.v_proj": {
2984
+ "group_size": 64,
2985
+ "bits": 8,
2986
+ "mode": "affine"
2987
+ },
2988
+ "model.layers.62.self_attn.o_proj": {
2989
+ "group_size": 64,
2990
+ "bits": 8,
2991
+ "mode": "affine"
2992
+ },
2993
+ "model.layers.63.self_attn.q_proj": {
2994
+ "group_size": 64,
2995
+ "bits": 8,
2996
+ "mode": "affine"
2997
+ },
2998
+ "model.layers.63.self_attn.k_proj": {
2999
+ "group_size": 64,
3000
+ "bits": 8,
3001
+ "mode": "affine"
3002
+ },
3003
+ "model.layers.63.self_attn.v_proj": {
3004
+ "group_size": 64,
3005
+ "bits": 8,
3006
+ "mode": "affine"
3007
+ },
3008
+ "model.layers.63.self_attn.o_proj": {
3009
+ "group_size": 64,
3010
+ "bits": 8,
3011
+ "mode": "affine"
3012
+ },
3013
+ "model.layers.64.self_attn.q_proj": {
3014
+ "group_size": 64,
3015
+ "bits": 8,
3016
+ "mode": "affine"
3017
+ },
3018
+ "model.layers.64.self_attn.k_proj": {
3019
+ "group_size": 64,
3020
+ "bits": 8,
3021
+ "mode": "affine"
3022
+ },
3023
+ "model.layers.64.self_attn.v_proj": {
3024
+ "group_size": 64,
3025
+ "bits": 8,
3026
+ "mode": "affine"
3027
+ },
3028
+ "model.layers.64.self_attn.o_proj": {
3029
+ "group_size": 64,
3030
+ "bits": 8,
3031
+ "mode": "affine"
3032
+ },
3033
+ "model.layers.65.self_attn.q_proj": {
3034
+ "group_size": 64,
3035
+ "bits": 8,
3036
+ "mode": "affine"
3037
+ },
3038
+ "model.layers.65.self_attn.k_proj": {
3039
+ "group_size": 64,
3040
+ "bits": 8,
3041
+ "mode": "affine"
3042
+ },
3043
+ "model.layers.65.self_attn.v_proj": {
3044
+ "group_size": 64,
3045
+ "bits": 8,
3046
+ "mode": "affine"
3047
+ },
3048
+ "model.layers.65.self_attn.o_proj": {
3049
+ "group_size": 64,
3050
+ "bits": 8,
3051
+ "mode": "affine"
3052
+ },
3053
+ "model.layers.66.self_attn.q_proj": {
3054
+ "group_size": 64,
3055
+ "bits": 8,
3056
+ "mode": "affine"
3057
+ },
3058
+ "model.layers.66.self_attn.k_proj": {
3059
+ "group_size": 64,
3060
+ "bits": 8,
3061
+ "mode": "affine"
3062
+ },
3063
+ "model.layers.66.self_attn.v_proj": {
3064
+ "group_size": 64,
3065
+ "bits": 8,
3066
+ "mode": "affine"
3067
+ },
3068
+ "model.layers.66.self_attn.o_proj": {
3069
+ "group_size": 64,
3070
+ "bits": 8,
3071
+ "mode": "affine"
3072
+ },
3073
+ "model.layers.67.self_attn.q_proj": {
3074
+ "group_size": 64,
3075
+ "bits": 8,
3076
+ "mode": "affine"
3077
+ },
3078
+ "model.layers.67.self_attn.k_proj": {
3079
+ "group_size": 64,
3080
+ "bits": 8,
3081
+ "mode": "affine"
3082
+ },
3083
+ "model.layers.67.self_attn.v_proj": {
3084
+ "group_size": 64,
3085
+ "bits": 8,
3086
+ "mode": "affine"
3087
+ },
3088
+ "model.layers.67.self_attn.o_proj": {
3089
+ "group_size": 64,
3090
+ "bits": 8,
3091
+ "mode": "affine"
3092
+ },
3093
+ "model.layers.68.self_attn.q_proj": {
3094
+ "group_size": 64,
3095
+ "bits": 8,
3096
+ "mode": "affine"
3097
+ },
3098
+ "model.layers.68.self_attn.k_proj": {
3099
+ "group_size": 64,
3100
+ "bits": 8,
3101
+ "mode": "affine"
3102
+ },
3103
+ "model.layers.68.self_attn.v_proj": {
3104
+ "group_size": 64,
3105
+ "bits": 8,
3106
+ "mode": "affine"
3107
+ },
3108
+ "model.layers.68.self_attn.o_proj": {
3109
+ "group_size": 64,
3110
+ "bits": 8,
3111
+ "mode": "affine"
3112
+ },
3113
+ "model.layers.69.self_attn.q_proj": {
3114
+ "group_size": 64,
3115
+ "bits": 8,
3116
+ "mode": "affine"
3117
+ },
3118
+ "model.layers.69.self_attn.k_proj": {
3119
+ "group_size": 64,
3120
+ "bits": 8,
3121
+ "mode": "affine"
3122
+ },
3123
+ "model.layers.69.self_attn.v_proj": {
3124
+ "group_size": 64,
3125
+ "bits": 8,
3126
+ "mode": "affine"
3127
+ },
3128
+ "model.layers.69.self_attn.o_proj": {
3129
+ "group_size": 64,
3130
+ "bits": 8,
3131
+ "mode": "affine"
3132
+ },
3133
+ "lm_head": {
3134
+ "group_size": 64,
3135
+ "bits": 8,
3136
+ "mode": "affine"
3137
+ }
3138
+ },
3139
+ "rope_parameters": {
3140
+ "partial_rotary_factor": 0.334,
3141
+ "rope_theta": 10000000,
3142
+ "rope_type": "default",
3143
+ "type": "default"
3144
+ },
3145
+ "rope_theta": 10000000,
3146
+ "routed_scaling_factor": null,
3147
+ "scoring_func": "sigmoid",
3148
+ "sliding_window": 128,
3149
+ "sliding_window_size": 128,
3150
+ "swa_head_dim": 192,
3151
+ "swa_num_attention_heads": 128,
3152
+ "swa_num_key_value_heads": 8,
3153
+ "swa_rope_theta": 10000,
3154
+ "swa_v_head_dim": 128,
3155
+ "tie_word_embeddings": false,
3156
+ "topk_group": 1,
3157
+ "topk_method": "noaux_tc",
3158
+ "transformers_version": "5.3.0",
3159
+ "use_cache": true,
3160
+ "v_head_dim": 128,
3161
+ "video_token_id": 151656,
3162
+ "vision_end_token_id": 151653,
3163
+ "vision_model_type": "mimovl",
3164
+ "vision_start_token_id": 151652,
3165
+ "vocab_size": 152576
3166
+ }
configuration_mimo_v2.py ADDED
@@ -0,0 +1,247 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # coding=utf-8
2
+ #
3
+ # Copyright 2026 Xiaomi Corporation.
4
+ # Copyright 2026 The HuggingFace Inc. team.
5
+ #
6
+ # Licensed under the Apache License, Version 2.0 (the "License");
7
+ # you may not use this file except in compliance with the License.
8
+ # You may obtain a copy of the License at
9
+ #
10
+ # http://www.apache.org/licenses/LICENSE-2.0
11
+ #
12
+ # Unless required by applicable law or agreed to in writing, software
13
+ # distributed under the License is distributed on an "AS IS" BASIS,
14
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
15
+ # See the License for the specific language governing permissions and
16
+ # limitations under the License.
17
+
18
+ from copy import deepcopy
19
+
20
+ from transformers.configuration_utils import PretrainedConfig
21
+ from transformers.modeling_rope_utils import rope_config_validation
22
+ from transformers.utils import logging
23
+
24
+
25
+ logger = logging.get_logger(__name__)
26
+
27
+
28
+ _MIMOV2_ATTENTION_PROJECTION_LAYOUTS = {"split", "fused_qkv"}
29
+
30
+ _MIMOV2_SPLIT_TP_PLAN = {
31
+ "layers.*.self_attn.q_proj": "colwise",
32
+ "layers.*.self_attn.k_proj": "colwise",
33
+ "layers.*.self_attn.v_proj": "colwise",
34
+ "layers.*.self_attn.o_proj": "rowwise",
35
+ "layers.*.mlp.gate_proj": "colwise",
36
+ "layers.*.mlp.up_proj": "colwise",
37
+ "layers.*.mlp.down_proj": "rowwise",
38
+ }
39
+
40
+ _MIMOV2_FUSED_QKV_TP_PLAN = {
41
+ "layers.*.self_attn.qkv_proj": "colwise",
42
+ "layers.*.self_attn.o_proj": "rowwise",
43
+ "layers.*.mlp.gate_proj": "colwise",
44
+ "layers.*.mlp.up_proj": "colwise",
45
+ "layers.*.mlp.down_proj": "rowwise",
46
+ }
47
+
48
+ _MIMOV2_PP_PLAN = {
49
+ "embed_tokens": (["input_ids"], ["inputs_embeds"]),
50
+ "layers": (["hidden_states", "attention_mask"], ["hidden_states"]),
51
+ "norm": (["hidden_states"], ["hidden_states"]),
52
+ }
53
+
54
+
55
+ def _to_plain_dict(value):
56
+ if value is None:
57
+ return {}
58
+ if isinstance(value, dict):
59
+ return deepcopy(value)
60
+ if hasattr(value, "to_dict"):
61
+ return deepcopy(value.to_dict())
62
+ if hasattr(value, "__dict__"):
63
+ return deepcopy(vars(value))
64
+ raise TypeError(f"Unsupported config value type: {type(value)!r}")
65
+
66
+
67
+ class MiMoV2Config(PretrainedConfig):
68
+
69
+ model_type = "mimo_v2"
70
+ keys_to_ignore_at_inference = ["past_key_values"]
71
+
72
+ base_model_tp_plan = _MIMOV2_SPLIT_TP_PLAN
73
+ base_model_pp_plan = _MIMOV2_PP_PLAN
74
+
75
+ attribute_map = {
76
+ "num_local_experts": "n_routed_experts",
77
+ }
78
+
79
+ def __init__(
80
+ self,
81
+ vocab_size=151936,
82
+ hidden_size=4096,
83
+ intermediate_size=22016,
84
+ num_hidden_layers=32,
85
+ num_attention_heads=32,
86
+ num_key_value_heads=32,
87
+ hidden_act="silu",
88
+ max_position_embeddings=32768,
89
+ initializer_range=0.02,
90
+ layernorm_epsilon=1e-6,
91
+ use_cache=True,
92
+ tie_word_embeddings=False,
93
+ rope_theta=10000.0,
94
+ rope_scaling=None,
95
+ attention_dropout=0.0,
96
+ attention_bias=False,
97
+ attention_value_scale=None,
98
+ head_dim=None,
99
+ v_head_dim=None,
100
+ swa_num_attention_heads=None,
101
+ swa_num_key_value_heads=None,
102
+ swa_head_dim=None,
103
+ swa_v_head_dim=None,
104
+ swa_rope_theta=None,
105
+ sliding_window=None,
106
+ sliding_window_size=None,
107
+ add_full_attention_sink_bias=False,
108
+ add_swa_attention_sink_bias=False,
109
+ hybrid_block_size=None,
110
+ hybrid_layer_pattern=None,
111
+ partial_rotary_factor=1.0,
112
+ n_routed_experts=None,
113
+ moe_intermediate_size=None,
114
+ num_experts_per_tok=None,
115
+ routed_scaling_factor=None,
116
+ scoring_func="sigmoid",
117
+ topk_method="noaux_tc",
118
+ n_group=None,
119
+ topk_group=None,
120
+ norm_topk_prob=True,
121
+ moe_layer_freq=None,
122
+ attention_projection_layout="split",
123
+ vision_config=None,
124
+ audio_config=None,
125
+ processor_config=None,
126
+ image_token_id=None,
127
+ video_token_id=None,
128
+ vision_start_token_id=None,
129
+ vision_end_token_id=None,
130
+ vision_model_type=None,
131
+ **kwargs,
132
+ ):
133
+ rope_parameters = kwargs.pop("rope_parameters", None)
134
+ if rope_scaling is None and rope_parameters is not None:
135
+ rope_scaling = rope_parameters
136
+
137
+ if attention_projection_layout is None:
138
+ attention_projection_layout = "split"
139
+ if attention_projection_layout not in _MIMOV2_ATTENTION_PROJECTION_LAYOUTS:
140
+ raise ValueError(f"Unsupported MiMoV2 attention projection layout: {attention_projection_layout}")
141
+
142
+ self.attention_projection_layout = attention_projection_layout
143
+ self.base_model_tp_plan = (
144
+ _MIMOV2_FUSED_QKV_TP_PLAN.copy()
145
+ if attention_projection_layout == "fused_qkv"
146
+ else _MIMOV2_SPLIT_TP_PLAN.copy()
147
+ )
148
+ self.base_model_pp_plan = _MIMOV2_PP_PLAN.copy()
149
+
150
+ self.vocab_size = vocab_size
151
+ self.max_position_embeddings = max_position_embeddings
152
+ self.hidden_size = hidden_size
153
+ self.intermediate_size = intermediate_size
154
+ self.num_hidden_layers = num_hidden_layers
155
+ self.num_attention_heads = num_attention_heads
156
+
157
+ if num_key_value_heads is None:
158
+ num_key_value_heads = num_attention_heads
159
+ if num_attention_heads % num_key_value_heads != 0:
160
+ raise ValueError("num_attention_heads must be divisible by num_key_value_heads")
161
+
162
+ self.num_key_value_heads = num_key_value_heads
163
+ self.hidden_act = hidden_act
164
+ self.initializer_range = initializer_range
165
+ self.layernorm_epsilon = layernorm_epsilon
166
+ self.use_cache = use_cache
167
+ self.rope_theta = rope_theta
168
+ self.rope_scaling = rope_scaling
169
+ self.attention_dropout = attention_dropout
170
+ self.attention_bias = attention_bias
171
+ self.attention_value_scale = attention_value_scale
172
+
173
+ self.head_dim = head_dim if head_dim is not None else hidden_size // num_attention_heads
174
+ self.v_head_dim = v_head_dim if v_head_dim is not None else self.head_dim
175
+ self.swa_num_attention_heads = (
176
+ swa_num_attention_heads if swa_num_attention_heads is not None else num_attention_heads
177
+ )
178
+ self.swa_num_key_value_heads = (
179
+ swa_num_key_value_heads if swa_num_key_value_heads is not None else num_key_value_heads
180
+ )
181
+ if self.swa_num_attention_heads % self.swa_num_key_value_heads != 0:
182
+ raise ValueError("swa_num_attention_heads must be divisible by swa_num_key_value_heads")
183
+ self.swa_head_dim = swa_head_dim if swa_head_dim is not None else self.head_dim
184
+ self.swa_v_head_dim = swa_v_head_dim if swa_v_head_dim is not None else self.swa_head_dim
185
+ self.swa_rope_theta = swa_rope_theta if swa_rope_theta is not None else rope_theta
186
+
187
+ if sliding_window is None:
188
+ sliding_window = sliding_window_size
189
+ self.sliding_window = sliding_window
190
+ self.sliding_window_size = sliding_window_size if sliding_window_size is not None else sliding_window
191
+ self.add_full_attention_sink_bias = add_full_attention_sink_bias
192
+ self.add_swa_attention_sink_bias = add_swa_attention_sink_bias
193
+
194
+ if hybrid_block_size is not None and hybrid_layer_pattern is None:
195
+ hybrid_layer_pattern = [0 if ((i + 1) % hybrid_block_size == 0) else 1 for i in range(num_hidden_layers)]
196
+ elif hybrid_layer_pattern is None:
197
+ hybrid_layer_pattern = [0] * num_hidden_layers
198
+ if len(hybrid_layer_pattern) != num_hidden_layers:
199
+ raise ValueError("hybrid_layer_pattern length must match num_hidden_layers")
200
+ self.hybrid_block_size = hybrid_block_size
201
+ self.hybrid_layer_pattern = hybrid_layer_pattern
202
+
203
+ self.partial_rotary_factor = partial_rotary_factor
204
+
205
+ self.n_routed_experts = n_routed_experts
206
+ self.moe_intermediate_size = moe_intermediate_size if moe_intermediate_size is not None else intermediate_size
207
+ self.num_experts_per_tok = num_experts_per_tok
208
+ self.routed_scaling_factor = routed_scaling_factor
209
+ self.scoring_func = scoring_func
210
+ self.topk_method = topk_method
211
+ self.n_group = n_group
212
+ self.topk_group = topk_group
213
+ self.norm_topk_prob = norm_topk_prob
214
+ if isinstance(moe_layer_freq, int):
215
+ moe_layer_freq = [moe_layer_freq > 0 and i % moe_layer_freq == 0 for i in range(num_hidden_layers)]
216
+ elif moe_layer_freq is None:
217
+ moe_layer_freq = [False] * num_hidden_layers
218
+ if len(moe_layer_freq) != num_hidden_layers:
219
+ raise ValueError("moe_layer_freq length must match num_hidden_layers")
220
+ self.moe_layer_freq = moe_layer_freq
221
+
222
+ self.vision_config = _to_plain_dict(vision_config)
223
+ self.audio_config = _to_plain_dict(audio_config)
224
+ self.processor_config = _to_plain_dict(processor_config)
225
+ self.image_token_id = image_token_id
226
+ self.video_token_id = video_token_id
227
+ self.vision_start_token_id = vision_start_token_id
228
+ self.vision_end_token_id = vision_end_token_id
229
+ self.vision_model_type = vision_model_type
230
+ self.audio_token_id = self.processor_config.get("audio_token_id", None) if self.processor_config else None
231
+ self.audio_start_token_id = (
232
+ self.processor_config.get("audio_start_token_id", None) if self.processor_config else None
233
+ )
234
+ self.audio_end_token_id = (
235
+ self.processor_config.get("audio_end_token_id", None) if self.processor_config else None
236
+ )
237
+
238
+ if self.rope_scaling is not None and "type" in self.rope_scaling:
239
+ self.rope_scaling["rope_type"] = self.rope_scaling["type"]
240
+ rope_config_validation(self)
241
+
242
+ super().__init__(
243
+ tie_word_embeddings=tie_word_embeddings,
244
+ **kwargs,
245
+ )
246
+
247
+ __all__ = ["MiMoV2Config"]
generation_config.json ADDED
@@ -0,0 +1,9 @@
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "bos_token_id": 151643,
3
+ "do_sample": false,
4
+ "eos_token_id": [151643, 151645, 151672],
5
+ "temperature": 1.0,
6
+ "top_p": 0.95,
7
+ "max_new_tokens": 2048,
8
+ "transformers_version": "4.37.0"
9
+ }
model-00001-of-00116.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:b71c68549bb25b67bdcd38763814e7864566a512af88f1e7a0d11734a8626527
3
+ size 4452021053
model-00004-of-00116.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:3c59b09d4dae9cfb06b8a941058f1de40e4d1022b7ceccb97add942cc4296e9a
3
+ size 5271679021
model-00005-of-00116.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:fc7597df22a91d15a223acfa917d867166ee55c240254169cde9ec715524ac72
3
+ size 5289568827
model-00013-of-00116.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:a123c926745f62746458aab3c4fa1e1d4d1b3d3c3016038aac3eea7d1020de84
3
+ size 5133828628
model-00016-of-00116.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:881e6c67456511abb5b88732f80ea68a7e4c79b48cc32b64198166eae592fe6e
3
+ size 5266933900
model-00019-of-00116.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:05c6e1a6e898e14fb007394a593c24ad376bf92f191ff237f180e64c8af24a82
3
+ size 5271679045
model-00021-of-00116.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:103f261876b2bc0a656d39f6c4eaa889727bfee1f6477c24da66804d70185544
3
+ size 5266933900
model-00024-of-00116.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:4e452118488d37ce6fa1767db4b30485e9813ed60c65e6dbbf3f373398d34f91
3
+ size 5271678693
model-00032-of-00116.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:b5af6fa54c198d8ed6bf78d29512681742a893ae7070f0a41ac031e39d0df565
3
+ size 3006754852
model-00033-of-00116.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:102d5140458c7c67700ffdf4533c4d66cfb6aaa6d170b5b4fb07e2b6b1ece1f0
3
+ size 5133828632
model-00036-of-00116.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:dae201e5e3e334f8098744addf644cc659827185bdc1953acaaeba608db674dd
3
+ size 5266933868
model-00037-of-00116.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:cd85e856c77e613b83bd6d8ae03a8a955bd49061e1cd5f6d071d3d688bc924c0
3
+ size 3006754496
model-00038-of-00116.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:a8d35a5dbbf16f8b9271f4a599b15f9a0c525b2298ebef0e9f7b849589ff2359
3
+ size 5133828630
model-00039-of-00116.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:7153296d58501f7329ce53eef90627290b81e54adfd848a0e0419cce4c5309f2
3
+ size 5271679067
model-00040-of-00116.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:ab15fb04453759e8fd81f106fb9f892070192d625c6101c36d8921d4be769598
3
+ size 5289568846
model-00044-of-00116.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:d9e370cbcbd604a8b8a61da3e0808d14fc20db6e5a3155d202594e1ebc4890f7
3
+ size 5271679067
model-00045-of-00116.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:ec698aeee4c9df24dc9a50be739edd0e2132f1377cfb570659198f9c83857eed
3
+ size 5289568852
model-00046-of-00116.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:97c6be437aa31da73dffd93a01e93b3c6d2da92190af92c52d94797a7c0442df
3
+ size 5266933900
model-00052-of-00116.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:74c12a954081adf1a3d574c76fe355294913cc23e25c5fe6d348aa812fafaa00
3
+ size 3006754888
model-00057-of-00116.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:b25d58a12e4bd358d7f3962fa0f559158d7fd11c603ade53e7a683203e1a29fe
3
+ size 3006754874
model-00058-of-00116.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:565db83ddb5b6e9444c921926d0e48c6f2e5ebf584dc2fe3d17655014fbb4af0
3
+ size 5133828632
model-00060-of-00116.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:94832707dcd5df966c82b8b94b7691b3f0353c86e3ff5924e302d8d136468d6e
3
+ size 5289568856
model-00065-of-00116.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:79efdf70c2672bb48f31804b65804402d6739fc56608391faec1a11e86d281b2
3
+ size 5289568834
model-00072-of-00116.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:3354ac0d6c447eae2d0cdf4ed08816d99336cd89278c653c0af6561d991c0753
3
+ size 3006754836
model-00073-of-00116.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:48ae583c062f19829ec6aa4f87e33e68aa52e75bc11a3b95ea5582dc25ed2ed8
3
+ size 5133828634
model-00076-of-00116.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:ff7196ae46bad526575d2eee45ac3b2321b9012d38ce44c0c6c824c23f53f669
3
+ size 5266933898
model-00077-of-00116.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:f11fbf8ca5940c08a5fb8d5e1ff484d11b3508cec095dd81f67f80ab081d0c4f
3
+ size 3006754480
model-00078-of-00116.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:3e32d5f4bebcbe88fbcce1ee75418335b1dadafc60b5162c359ea47bbddcd17a
3
+ size 5133828634
model-00079-of-00116.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:0ef133d10805a923f0c8e2b78e54d98c3aff20fdada72cb1076dc6fdf5af90b4
3
+ size 5271679067
model-00080-of-00116.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:047e48127f16798388ff4501ad3d3e5a8425eaf60d2a3c9021113d5b43224b41
3
+ size 5289568854
model-00081-of-00116.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:1fd6e433655294b40a86dff36d9b7125cd1e729defb44c4baaf187c7ce5e5d47
3
+ size 5266933900
model-00084-of-00116.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:23f3e51534bcdccca7d26c66c53cf74b18b5c984addedd1cc76c3bba60360a1f
3
+ size 5271679035
model-00085-of-00116.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:eae9b61aab37cb1ed9aec8325fd5cd59dd7b78d5cc0347e1487c1ba1d3568644
3
+ size 5289568832
model-00093-of-00116.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:6c6508d8482c967da8ecaf8ad9a190ea9967566ee307cf87fabfb306db09253b
3
+ size 5133828632
model-00096-of-00116.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:adb483ced1ee9162756831ea5d3be352ec28aeb704a308ea1f3d9fc176eb05e7
3
+ size 5266933900
model-00099-of-00116.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:42d9a4996aaf8046ce3c6b8ed8de0becfe0cfdec906d49a5de03cd1381408587
3
+ size 5271679043
model-00100-of-00116.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:301780d494e4afcb52dc253524b61bada14db29e637164e4c20cbd2bbbaddd51
3
+ size 5289568844
model-00105-of-00116.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:60a7fb03426f879656d199e8e495e7a42f351ec31a81caee889ae9acafe1a316
3
+ size 5289568828
model-00112-of-00116.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:d4fdd523e0e694c91d3176fcb00231851c56baa579d464b92f0cce6cf5866a19
3
+ size 3006754834
model-00113-of-00116.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:3cafdca088e5867cd22cf75a1c1c3c4a12f6f080cbc73429dc35b6a032a11725
3
+ size 5133828630
model-00116-of-00116.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:3d378afa5c1911dadcff8ead242a599d8dc9606b9c26f8edf13ef0621a7053e1
3
+ size 996016431
model.safetensors.index.json ADDED
The diff for this file is too large to render. See raw diff
 
modeling_mimo_v2.py ADDED
@@ -0,0 +1,1878 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # coding=utf-8
2
+ #
3
+ # Copyright 2026 Xiaomi Corporation.
4
+ # Copyright 2026 The HuggingFace Inc. team.
5
+ #
6
+ # Licensed under the Apache License, Version 2.0 (the "License");
7
+ # you may not use this file except in compliance with the License.
8
+ # You may obtain a copy of the License at
9
+ #
10
+ # http://www.apache.org/licenses/LICENSE-2.0
11
+ #
12
+ # Unless required by applicable law or agreed to in writing, software
13
+ # distributed under the License is distributed on an "AS IS" BASIS,
14
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
15
+ # See the License for the specific language governing permissions and
16
+ # limitations under the License.
17
+
18
+ import math
19
+ from copy import copy
20
+ from types import SimpleNamespace
21
+ from typing import Callable, Optional, Union
22
+
23
+ import torch
24
+ import torch.nn as nn
25
+ import torch.nn.functional as F
26
+
27
+ from transformers.activations import ACT2FN
28
+ from transformers.cache_utils import Cache, DynamicCache
29
+ from transformers.configuration_utils import PretrainedConfig
30
+ from transformers.generation import GenerationMixin
31
+ from transformers.integrations import use_kernel_forward_from_hub
32
+ from transformers.masking_utils import create_causal_mask, create_sliding_window_causal_mask
33
+ from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast
34
+ from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update
35
+ from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel
36
+ from transformers.models.qwen2.configuration_qwen2 import Qwen2Config
37
+ from transformers.models.qwen2.modeling_qwen2 import Qwen2Model
38
+ from transformers.processing_utils import Unpack
39
+ from transformers.utils import TransformersKwargs, can_return_tuple, logging
40
+
41
+ from .configuration_mimo_v2 import MiMoV2Config
42
+
43
+
44
+ logger = logging.get_logger(__name__)
45
+
46
+
47
+ def rotate_half(x):
48
+ """Rotates half the hidden dims of the input."""
49
+ x1 = x[..., : x.shape[-1] // 2]
50
+ x2 = x[..., x.shape[-1] // 2 :]
51
+ return torch.cat((-x2, x1), dim=-1)
52
+
53
+
54
+ def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1):
55
+ """Applies rotary position embedding to query and key tensors."""
56
+ cos = cos.unsqueeze(unsqueeze_dim)
57
+ sin = sin.unsqueeze(unsqueeze_dim)
58
+ q_embed = (q * cos) + (rotate_half(q) * sin)
59
+ k_embed = (k * cos) + (rotate_half(k) * sin)
60
+ return q_embed, k_embed
61
+
62
+
63
+ def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
64
+ batch, num_key_value_heads, slen, head_dim = hidden_states.shape
65
+ if n_rep == 1:
66
+ return hidden_states
67
+ hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim)
68
+ return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)
69
+
70
+
71
+ def eager_attention_forward(
72
+ module: nn.Module,
73
+ query: torch.Tensor,
74
+ key: torch.Tensor,
75
+ value: torch.Tensor,
76
+ attention_mask: Optional[torch.Tensor],
77
+ scaling: float,
78
+ dropout: float = 0.0,
79
+ sinks: Optional[torch.Tensor] = None,
80
+ **kwargs,
81
+ ):
82
+ key_states = repeat_kv(key, module.num_key_value_groups)
83
+ value_states = repeat_kv(value, module.num_key_value_groups)
84
+ attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling
85
+ if attention_mask is not None:
86
+ causal_mask = attention_mask[:, :, :, : key_states.shape[-2]]
87
+ attn_weights = attn_weights + causal_mask
88
+
89
+ if sinks is not None:
90
+ sinks = module.attention_sink_bias.reshape(1, -1, 1, 1).expand(query.shape[0], -1, query.shape[-2], -1)
91
+ attn_weights = torch.cat([attn_weights, sinks], dim=-1)
92
+
93
+ attn_weights = attn_weights - attn_weights.max(dim=-1, keepdim=True).values
94
+ probs = F.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype)
95
+
96
+ if sinks is not None:
97
+ probs = probs[..., :-1]
98
+
99
+ attn_weights = nn.functional.dropout(probs, p=dropout, training=module.training)
100
+ attn_output = torch.matmul(attn_weights, value_states)
101
+ attn_output = attn_output.transpose(1, 2).contiguous()
102
+ return attn_output, attn_weights
103
+
104
+
105
+ @use_kernel_forward_from_hub("RMSNorm")
106
+ class MiMoV2RMSNorm(nn.Module):
107
+ def __init__(self, hidden_size, eps=1e-6):
108
+ super().__init__()
109
+ self.weight = nn.Parameter(torch.ones(hidden_size))
110
+ self.variance_epsilon = eps
111
+
112
+ def forward(self, hidden_states):
113
+ input_dtype = hidden_states.dtype
114
+ hidden_states = hidden_states.to(torch.float32)
115
+ variance = hidden_states.pow(2).mean(-1, keepdim=True)
116
+ hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
117
+ return self.weight * hidden_states.to(input_dtype)
118
+
119
+
120
+ class MiMoV2MLP(nn.Module):
121
+ def __init__(self, config, intermediate_size=None):
122
+ super().__init__()
123
+ self.config = config
124
+ self.hidden_size = config.hidden_size
125
+ self.intermediate_size = config.intermediate_size if intermediate_size is None else intermediate_size
126
+ self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
127
+ self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
128
+ self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)
129
+ self.act_fn = ACT2FN[config.hidden_act]
130
+
131
+ def forward(self, hidden_states):
132
+ return self.down_proj(self.act_fn(self.gate_proj(hidden_states)) * self.up_proj(hidden_states))
133
+
134
+
135
+ class MiMoV2MoEGate(nn.Module):
136
+ def __init__(self, config):
137
+ super().__init__()
138
+ self.config = config
139
+ self.top_k = config.num_experts_per_tok
140
+ self.n_routed_experts = config.n_routed_experts
141
+ self.routed_scaling_factor = config.routed_scaling_factor if config.routed_scaling_factor is not None else 1.0
142
+ self.scoring_func = config.scoring_func
143
+ self.topk_method = config.topk_method
144
+ self.n_group = config.n_group
145
+ self.topk_group = config.topk_group
146
+ self.norm_topk_prob = config.norm_topk_prob
147
+ self.gating_dim = config.hidden_size
148
+ self.weight = nn.Parameter(torch.empty((self.n_routed_experts, self.gating_dim)))
149
+ if self.topk_method == "noaux_tc":
150
+ self.e_score_correction_bias = nn.Parameter(torch.empty((self.n_routed_experts)))
151
+
152
+ def forward(self, hidden_states):
153
+ bsz, seq_len, h = hidden_states.shape
154
+ hidden_states = hidden_states.view(-1, h)
155
+ logits = F.linear(hidden_states.type(torch.float32), self.weight.type(torch.float32), None)
156
+ if self.scoring_func == "sigmoid":
157
+ scores = logits.sigmoid()
158
+ else:
159
+ raise NotImplementedError(f"Unsupported scoring function for MoE gating: {self.scoring_func}")
160
+
161
+ if self.topk_method == "noaux_tc":
162
+ if self.training:
163
+ raise ValueError("MiMoV2 noaux_tc routing is only implemented for inference.")
164
+ scores_for_choice = scores.view(bsz * seq_len, -1) + self.e_score_correction_bias.unsqueeze(0)
165
+ group_scores = scores_for_choice.view(bsz * seq_len, self.n_group, -1).topk(2, dim=-1)[0].sum(dim=-1)
166
+ group_idx = torch.topk(group_scores, k=self.topk_group, dim=-1, sorted=False)[1]
167
+ group_mask = torch.zeros_like(group_scores)
168
+ group_mask.scatter_(1, group_idx, 1)
169
+ score_mask = (
170
+ group_mask.unsqueeze(-1)
171
+ .expand(bsz * seq_len, self.n_group, self.n_routed_experts // self.n_group)
172
+ .reshape(bsz * seq_len, -1)
173
+ )
174
+ tmp_scores = scores_for_choice.masked_fill(~score_mask.bool(), float("-inf"))
175
+ _, topk_idx = torch.topk(tmp_scores, k=self.top_k, dim=-1, sorted=False)
176
+ topk_weight = scores.gather(1, topk_idx)
177
+ else:
178
+ raise NotImplementedError(f"Unsupported TopK function for MoE gating: {self.topk_method}")
179
+
180
+ if self.top_k > 1 and self.norm_topk_prob:
181
+ denominator = topk_weight.sum(dim=-1, keepdim=True) + 1e-20
182
+ topk_weight = topk_weight / denominator
183
+ topk_weight = topk_weight * self.routed_scaling_factor
184
+ return topk_idx, topk_weight
185
+
186
+
187
+ class MiMoV2MoE(nn.Module):
188
+ def __init__(self, config):
189
+ super().__init__()
190
+ self.config = config
191
+ self.experts = nn.ModuleList(
192
+ [MiMoV2MLP(config, intermediate_size=config.moe_intermediate_size) for _ in range(config.n_routed_experts)]
193
+ )
194
+ self.gate = MiMoV2MoEGate(config)
195
+
196
+ def moe(self, hidden_states: torch.Tensor, topk_indices: torch.Tensor, topk_weights: torch.Tensor):
197
+ final_hidden_states = torch.zeros_like(hidden_states, dtype=topk_weights.dtype)
198
+ expert_mask = torch.nn.functional.one_hot(topk_indices, num_classes=len(self.experts))
199
+ expert_mask = expert_mask.permute(2, 0, 1)
200
+
201
+ for expert_idx, expert in enumerate(self.experts):
202
+ mask = expert_mask[expert_idx]
203
+ token_indices, weight_indices = torch.where(mask)
204
+ if token_indices.numel() > 0:
205
+ expert_weights = topk_weights[token_indices, weight_indices]
206
+ expert_input = hidden_states[token_indices]
207
+ expert_output = expert(expert_input)
208
+ final_hidden_states.index_add_(0, token_indices, expert_output * expert_weights.unsqueeze(-1))
209
+
210
+ return final_hidden_states.type(hidden_states.dtype)
211
+
212
+ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
213
+ orig_shape = hidden_states.shape
214
+ topk_indices, topk_weights = self.gate(hidden_states)
215
+ hidden_states = hidden_states.view(-1, hidden_states.shape[-1])
216
+ hidden_states = self.moe(hidden_states, topk_indices, topk_weights).view(*orig_shape)
217
+ return hidden_states
218
+
219
+
220
+ class MiMoV2Attention(nn.Module):
221
+ """MiMoV2 attention.
222
+
223
+ `projection_layout` only controls how checkpoint weights are named and
224
+ stored: Flash uses separate q/k/v projections, while Pro uses fused qkv.
225
+ The attention computation after projection is shared.
226
+ """
227
+
228
+ def __init__(self, config, is_swa: bool, layer_idx: int, projection_layout: str = "split"):
229
+ super().__init__()
230
+ if projection_layout not in {"split", "fused_qkv"}:
231
+ raise ValueError(f"Unsupported MiMoV2 attention projection layout: {projection_layout}")
232
+
233
+ self.config = config
234
+ self.layer_idx = layer_idx
235
+ self.is_swa = is_swa
236
+ self.is_causal = True
237
+ self.projection_layout = projection_layout
238
+
239
+ default_head_dim = config.hidden_size // config.num_attention_heads
240
+ default_v_head_dim = getattr(config, "v_head_dim", default_head_dim)
241
+
242
+ if is_swa:
243
+ self.head_dim = getattr(config, "swa_head_dim", getattr(config, "head_dim", default_head_dim))
244
+ self.v_head_dim = getattr(config, "swa_v_head_dim", default_v_head_dim)
245
+ self.num_attention_heads = getattr(config, "swa_num_attention_heads", config.num_attention_heads)
246
+ self.num_key_value_heads = getattr(config, "swa_num_key_value_heads", config.num_key_value_heads)
247
+ else:
248
+ self.head_dim = getattr(config, "head_dim", default_head_dim)
249
+ self.v_head_dim = getattr(config, "v_head_dim", self.head_dim)
250
+ self.num_attention_heads = config.num_attention_heads
251
+ self.num_key_value_heads = config.num_key_value_heads
252
+
253
+ self.rope_dim = int(self.head_dim * getattr(config, "partial_rotary_factor", 1.0))
254
+ if self.rope_dim % 2 != 0:
255
+ raise ValueError(
256
+ f"MiMoV2 rotary dimension must be even, got {self.rope_dim} from "
257
+ f"head_dim={self.head_dim} and partial_rotary_factor={getattr(config, 'partial_rotary_factor', 1.0)}"
258
+ )
259
+ self.num_key_value_groups = self.num_attention_heads // self.num_key_value_heads
260
+ self.attention_dropout = getattr(config, "attention_dropout", 0.0)
261
+ self.scaling = self.head_dim**-0.5
262
+ self.sliding_window = getattr(config, "sliding_window", None) if is_swa else None
263
+ self.q_size = self.num_attention_heads * self.head_dim
264
+ self.k_size = self.num_key_value_heads * self.head_dim
265
+ self.v_size = self.num_key_value_heads * self.v_head_dim
266
+ self.o_hidden_size = self.num_attention_heads * self.v_head_dim
267
+ self.v_scale = getattr(config, "attention_value_scale", None)
268
+ self.attention_sink_bias = (
269
+ nn.Parameter(torch.empty(self.num_attention_heads), requires_grad=False)
270
+ if (
271
+ (getattr(config, "add_full_attention_sink_bias", False) and not is_swa)
272
+ or (getattr(config, "add_swa_attention_sink_bias", False) and is_swa)
273
+ )
274
+ else None
275
+ )
276
+
277
+ attention_bias = getattr(config, "attention_bias", False)
278
+ if self.projection_layout == "fused_qkv":
279
+ self.qkv_proj = nn.Linear(
280
+ config.hidden_size,
281
+ self.q_size + self.k_size + self.v_size,
282
+ bias=attention_bias,
283
+ )
284
+ else:
285
+ self.q_proj = nn.Linear(config.hidden_size, self.q_size, bias=attention_bias)
286
+ self.k_proj = nn.Linear(config.hidden_size, self.k_size, bias=attention_bias)
287
+ self.v_proj = nn.Linear(config.hidden_size, self.v_size, bias=attention_bias)
288
+ self.o_proj = nn.Linear(self.o_hidden_size, config.hidden_size, bias=False)
289
+
290
+ def _forward_attention(
291
+ self,
292
+ query_states: torch.Tensor,
293
+ key_states: torch.Tensor,
294
+ value_states: torch.Tensor,
295
+ input_shape: torch.Size,
296
+ position_embeddings: tuple[torch.Tensor, torch.Tensor],
297
+ attention_mask: Optional[torch.Tensor],
298
+ past_key_values: Optional[Cache] = None,
299
+ cache_position: Optional[torch.LongTensor] = None,
300
+ position_ids: Optional[torch.LongTensor] = None,
301
+ ) -> tuple[torch.Tensor, torch.Tensor]:
302
+ if self.v_scale is not None:
303
+ value_states = value_states * self.v_scale
304
+
305
+ cos, sin = position_embeddings
306
+ query_rope, query_nope = query_states.split([self.rope_dim, self.head_dim - self.rope_dim], dim=-1)
307
+ key_rope, key_nope = key_states.split([self.rope_dim, self.head_dim - self.rope_dim], dim=-1)
308
+ query_rope, key_rope = apply_rotary_pos_emb(query_rope, key_rope, cos, sin)
309
+ query_states = torch.cat([query_rope, query_nope], dim=-1)
310
+ key_states = torch.cat([key_rope, key_nope], dim=-1)
311
+
312
+ if past_key_values is not None:
313
+ cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position}
314
+ key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx, cache_kwargs)
315
+
316
+ attn_implementation = self.config._attn_implementation
317
+ if attn_implementation is not None and attn_implementation.startswith("paged|"):
318
+ raise ValueError(
319
+ "MiMoV2 remote code does not support paged attention cache. "
320
+ "Please use eager, sdpa, flex_attention, or flash_attention_2."
321
+ )
322
+
323
+ attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface(
324
+ attn_implementation, eager_attention_forward
325
+ )
326
+ if self.attention_sink_bias is not None and attn_implementation == "sdpa":
327
+ logger.warning_once(
328
+ "MiMoV2 attention sink bias is not supported by SDPA; falling back to eager attention for correctness."
329
+ )
330
+ attention_interface = eager_attention_forward
331
+
332
+ attention_kwargs = {
333
+ "dropout": 0.0 if not self.training else self.attention_dropout,
334
+ "scaling": self.scaling,
335
+ "position_ids": position_ids,
336
+ "is_causal": self.is_causal,
337
+ }
338
+ if attention_interface is eager_attention_forward:
339
+ attention_kwargs["sinks"] = self.attention_sink_bias
340
+ else:
341
+ if self.attention_sink_bias is not None:
342
+ attention_kwargs["s_aux"] = self.attention_sink_bias
343
+ if self.sliding_window is not None:
344
+ attention_kwargs["sliding_window"] = self.sliding_window
345
+
346
+ attn_output, attn_weights = attention_interface(
347
+ self,
348
+ query_states,
349
+ key_states,
350
+ value_states,
351
+ attention_mask,
352
+ **attention_kwargs,
353
+ )
354
+ attn_output = attn_output.reshape(*input_shape, -1).contiguous()
355
+ attn_output = self.o_proj(attn_output)
356
+ return attn_output, attn_weights
357
+
358
+ def forward(
359
+ self,
360
+ hidden_states: torch.Tensor,
361
+ position_embeddings: tuple[torch.Tensor, torch.Tensor],
362
+ attention_mask: Optional[torch.Tensor],
363
+ past_key_values: Optional[Cache] = None,
364
+ cache_position: Optional[torch.LongTensor] = None,
365
+ position_ids: Optional[torch.LongTensor] = None,
366
+ **kwargs: Unpack[TransformersKwargs],
367
+ ) -> tuple[torch.Tensor, torch.Tensor]:
368
+ input_shape = hidden_states.shape[:-1]
369
+
370
+ if self.projection_layout == "fused_qkv":
371
+ qkv_states = self.qkv_proj(hidden_states)
372
+ query_states, key_states, value_states = qkv_states.split([self.q_size, self.k_size, self.v_size], dim=-1)
373
+ else:
374
+ query_states = self.q_proj(hidden_states)
375
+ key_states = self.k_proj(hidden_states)
376
+ value_states = self.v_proj(hidden_states)
377
+
378
+ query_states = query_states.view(*input_shape, self.num_attention_heads, self.head_dim).transpose(1, 2)
379
+ key_states = key_states.view(*input_shape, self.num_key_value_heads, self.head_dim).transpose(1, 2)
380
+ value_states = value_states.view(*input_shape, self.num_key_value_heads, self.v_head_dim).transpose(1, 2)
381
+ return self._forward_attention(
382
+ query_states,
383
+ key_states,
384
+ value_states,
385
+ input_shape,
386
+ position_embeddings,
387
+ attention_mask,
388
+ past_key_values=past_key_values,
389
+ cache_position=cache_position,
390
+ position_ids=position_ids,
391
+ )
392
+
393
+
394
+ class MiMoV2DecoderLayer(nn.Module):
395
+ attention_projection_layout = "split"
396
+
397
+ def __init__(self, config, layer_idx: int, attention_projection_layout: Optional[str] = None):
398
+ super().__init__()
399
+ attention_projection_layout = attention_projection_layout or self.attention_projection_layout
400
+ is_swa_layer = config.hybrid_layer_pattern[layer_idx] == 1
401
+ self.attention_type = "sliding_window_attention" if is_swa_layer else "full_attention"
402
+ self.self_attn = MiMoV2Attention(
403
+ config, is_swa_layer, layer_idx, projection_layout=attention_projection_layout
404
+ )
405
+ self.mlp = (
406
+ MiMoV2MoE(config)
407
+ if getattr(config, "n_routed_experts", None) is not None and config.moe_layer_freq[layer_idx]
408
+ else MiMoV2MLP(config)
409
+ )
410
+ self.input_layernorm = MiMoV2RMSNorm(config.hidden_size, eps=config.layernorm_epsilon)
411
+ self.post_attention_layernorm = MiMoV2RMSNorm(config.hidden_size, eps=config.layernorm_epsilon)
412
+
413
+ def forward(
414
+ self,
415
+ hidden_states: torch.Tensor,
416
+ attention_mask: Optional[torch.Tensor] = None,
417
+ position_ids: Optional[torch.LongTensor] = None,
418
+ past_key_values: Optional[Cache] = None,
419
+ use_cache: Optional[bool] = False,
420
+ cache_position: Optional[torch.LongTensor] = None,
421
+ position_embeddings: Optional[tuple[torch.Tensor, torch.Tensor]] = None,
422
+ **kwargs: Unpack[TransformersKwargs],
423
+ ) -> torch.Tensor:
424
+ residual = hidden_states
425
+ hidden_states = self.input_layernorm(hidden_states)
426
+ hidden_states, _ = self.self_attn(
427
+ hidden_states=hidden_states,
428
+ attention_mask=attention_mask,
429
+ position_ids=position_ids,
430
+ past_key_values=past_key_values,
431
+ use_cache=use_cache,
432
+ cache_position=cache_position,
433
+ position_embeddings=position_embeddings,
434
+ **kwargs,
435
+ )
436
+ hidden_states = residual + hidden_states
437
+
438
+ residual = hidden_states
439
+ hidden_states = self.post_attention_layernorm(hidden_states)
440
+ hidden_states = self.mlp(hidden_states)
441
+ hidden_states = residual + hidden_states
442
+ return hidden_states
443
+
444
+
445
+ class MiMoV2RotaryEmbedding(nn.Module):
446
+ inv_freq: torch.Tensor
447
+
448
+ def __init__(self, config, is_swa: bool, device=None):
449
+ super().__init__()
450
+ if hasattr(config, "rope_scaling") and isinstance(config.rope_scaling, dict):
451
+ self.rope_type = config.rope_scaling.get("rope_type", config.rope_scaling.get("type", "default"))
452
+ else:
453
+ self.rope_type = "default"
454
+ self.max_seq_len_cached = config.max_position_embeddings
455
+ self.original_max_seq_len = config.max_position_embeddings
456
+
457
+ self.config = copy(config)
458
+ self.config.rope_parameters = copy(getattr(config, "rope_parameters", None) or {})
459
+ if is_swa:
460
+ self.config.rope_theta = getattr(config, "swa_rope_theta", config.rope_theta)
461
+ self.config.head_dim = getattr(config, "swa_head_dim", getattr(config, "head_dim", None))
462
+ if self.config.rope_parameters:
463
+ self.config.rope_parameters["rope_theta"] = self.config.rope_theta
464
+ self.rope_init_fn = (
465
+ self.compute_default_rope_parameters
466
+ if self.rope_type == "default"
467
+ else ROPE_INIT_FUNCTIONS[self.rope_type]
468
+ )
469
+
470
+ inv_freq, self.attention_scaling = self.rope_init_fn(self.config, device)
471
+ self.register_buffer("inv_freq", inv_freq, persistent=False)
472
+ self.original_inv_freq = self.inv_freq
473
+
474
+ @staticmethod
475
+ def compute_default_rope_parameters(config, device=None, seq_len=None, layer_type=None):
476
+ config.standardize_rope_params()
477
+ rope_parameters = config.rope_parameters[layer_type] if layer_type is not None else config.rope_parameters
478
+ base = rope_parameters["rope_theta"]
479
+ partial_rotary_factor = rope_parameters.get("partial_rotary_factor", 1.0)
480
+ head_dim = getattr(config, "head_dim", None) or config.hidden_size // config.num_attention_heads
481
+ dim = int(head_dim * partial_rotary_factor)
482
+ if dim % 2 != 0:
483
+ raise ValueError(
484
+ f"MiMoV2 rotary dimension must be even, got {dim} from "
485
+ f"head_dim={head_dim} and partial_rotary_factor={partial_rotary_factor}"
486
+ )
487
+ inv_freq = 1.0 / (
488
+ base ** (torch.arange(0, dim, 2, dtype=torch.int64).to(device=device, dtype=torch.float) / dim)
489
+ )
490
+ return inv_freq, 1.0
491
+
492
+ @torch.no_grad()
493
+ @dynamic_rope_update
494
+ def forward(self, x, position_ids):
495
+ inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1).to(x.device)
496
+ position_ids_expanded = position_ids[:, None, :].float()
497
+
498
+ device_type = x.device.type if isinstance(x.device.type, str) and x.device.type != "mps" else "cpu"
499
+ with torch.autocast(device_type=device_type, enabled=False):
500
+ freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2)
501
+ emb = torch.cat((freqs, freqs), dim=-1)
502
+ cos = emb.cos() * self.attention_scaling
503
+ sin = emb.sin() * self.attention_scaling
504
+
505
+ return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)
506
+
507
+
508
+ # ---------------------------------------------------------------------------
509
+ # Multimodal helpers
510
+ # ---------------------------------------------------------------------------
511
+
512
+
513
+ def _as_namespace(config_like):
514
+ if config_like is None:
515
+ return SimpleNamespace()
516
+ if isinstance(config_like, dict):
517
+ return SimpleNamespace(**config_like)
518
+ return config_like
519
+
520
+
521
+ def _parse_maybe_list(value: str | int, length: int) -> list[int]:
522
+ if isinstance(value, str) and "-" in value:
523
+ return [int(x) for x in value.split("-")]
524
+ return [int(value)] * length
525
+
526
+
527
+ def _build_speech_embeddings(config) -> nn.ModuleList:
528
+ audio_channels = getattr(config, "audio_channels")
529
+ input_local_dim = getattr(config, "input_local_dim")
530
+ speech_empty_ids = _parse_maybe_list(getattr(config, "speech_zeroemb_idx"), audio_channels)
531
+ speech_vocab_sizes = _parse_maybe_list(getattr(config, "speech_vocab_size"), audio_channels)
532
+ return nn.ModuleList(
533
+ [
534
+ nn.Embedding(speech_vocab_sizes[i], input_local_dim, padding_idx=speech_empty_ids[i])
535
+ for i in range(audio_channels)
536
+ ]
537
+ )
538
+
539
+
540
+ def _pad_and_group_audio_codes(
541
+ audio_codes: torch.Tensor, audio_channels: int, group_size: int
542
+ ) -> torch.Tensor:
543
+ """Slice to `audio_channels`, pad to `group_size` boundary, reshape to [G, group_size, C]."""
544
+ if audio_codes.dim() != 2:
545
+ raise ValueError(f"`audio_codes` must be 2D [T, C], got shape={tuple(audio_codes.shape)}")
546
+ audio_codes = audio_codes[:, :audio_channels]
547
+ T = audio_codes.shape[0]
548
+ padded_T = ((T + group_size - 1) // group_size) * group_size
549
+ if padded_T > T:
550
+ audio_codes = torch.cat([audio_codes, audio_codes[-1:].expand(padded_T - T, -1)], dim=0)
551
+ return audio_codes.reshape(padded_T // group_size, group_size, audio_channels)
552
+
553
+
554
+ def _replace_modal_embeddings_inplace(
555
+ input_ids: torch.Tensor,
556
+ inputs_embeds: torch.Tensor,
557
+ token_id: int | None,
558
+ modal_embeds: torch.Tensor | None,
559
+ ) -> None:
560
+ if token_id is None or modal_embeds is None:
561
+ return
562
+
563
+ if modal_embeds.dim() != 2:
564
+ raise ValueError(f"`modal_embeds` must be 2D [N, H], got shape={tuple(modal_embeds.shape)}")
565
+
566
+ mask = input_ids.eq(token_id)
567
+ num_slots = int(mask.sum().item())
568
+ if num_slots == 0:
569
+ return
570
+
571
+ if modal_embeds.shape[0] != num_slots:
572
+ raise ValueError(
573
+ f"Modal embedding count mismatch for token_id={token_id}: "
574
+ f"found {num_slots} placeholders but got {modal_embeds.shape[0]} embeddings."
575
+ )
576
+
577
+ inputs_embeds[mask] = modal_embeds.to(device=inputs_embeds.device, dtype=inputs_embeds.dtype)
578
+
579
+
580
+ # ---------------------------------------------------------------------------
581
+ # Vision encoder
582
+ # ---------------------------------------------------------------------------
583
+
584
+
585
+ def _rotate_half_vision(x: torch.Tensor) -> torch.Tensor:
586
+ x1 = x[..., : x.shape[-1] // 2]
587
+ x2 = x[..., x.shape[-1] // 2 :]
588
+ return torch.cat((-x2, x1), dim=-1)
589
+
590
+
591
+ def _apply_rotary_pos_emb_vision(
592
+ q: torch.Tensor, k: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor
593
+ ) -> tuple[torch.Tensor, torch.Tensor]:
594
+ orig_q_dtype, orig_k_dtype = q.dtype, k.dtype
595
+ q, k = q.float(), k.float()
596
+ cos, sin = cos.unsqueeze(-2).float(), sin.unsqueeze(-2).float()
597
+ q_embed = (q * cos) + (_rotate_half_vision(q) * sin)
598
+ k_embed = (k * cos) + (_rotate_half_vision(k) * sin)
599
+ return q_embed.to(orig_q_dtype), k_embed.to(orig_k_dtype)
600
+
601
+
602
+ class MiMoVisionRotaryEmbedding(nn.Module):
603
+ def __init__(self, dim: int, theta: float = 10000.0) -> None:
604
+ super().__init__()
605
+ inv_freq = 1.0 / (theta ** (torch.arange(0, dim, 2, dtype=torch.float) / dim))
606
+ self.register_buffer("inv_freq", inv_freq, persistent=False)
607
+
608
+ def forward(self, seqlen: int) -> torch.Tensor:
609
+ seq = torch.arange(seqlen, device=self.inv_freq.device, dtype=self.inv_freq.dtype)
610
+ return torch.outer(seq, self.inv_freq)
611
+
612
+
613
+ class MiMoVisionPatchEmbed(nn.Module):
614
+ def __init__(
615
+ self, patch_size: int = 16, temporal_patch_size: int = 2, in_channels: int = 3, embed_dim: int = 1280
616
+ ):
617
+ super().__init__()
618
+ self.patch_size = patch_size
619
+ self.temporal_patch_size = temporal_patch_size
620
+ self.in_channels = in_channels
621
+ self.embed_dim = embed_dim
622
+ kernel_size = [temporal_patch_size, patch_size, patch_size]
623
+ self.proj = nn.Conv3d(in_channels, embed_dim, kernel_size=kernel_size, stride=kernel_size, bias=False)
624
+
625
+ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
626
+ target_dtype = self.proj.weight.dtype
627
+ hidden_states = hidden_states.view(
628
+ -1, self.in_channels, self.temporal_patch_size, self.patch_size, self.patch_size
629
+ )
630
+ return self.proj(hidden_states.to(dtype=target_dtype)).view(-1, self.embed_dim)
631
+
632
+
633
+ class MiMoVisionSwiGLUMLP(nn.Module):
634
+ def __init__(self, dim: int, intermediate_dim: int, hidden_act: str = "silu"):
635
+ super().__init__()
636
+ self.gate_proj = nn.Linear(dim, intermediate_dim, bias=True)
637
+ self.up_proj = nn.Linear(dim, intermediate_dim, bias=True)
638
+ self.down_proj = nn.Linear(intermediate_dim, dim, bias=True)
639
+ self.act_fn = ACT2FN[hidden_act]
640
+
641
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
642
+ return self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
643
+
644
+
645
+ class MiMoVisionAttention(nn.Module):
646
+ def __init__(
647
+ self,
648
+ dim: int,
649
+ num_heads: int,
650
+ num_kv_heads: int | None = None,
651
+ head_dim: int | None = None,
652
+ use_sinks: bool = False,
653
+ window_size: int = -1,
654
+ ):
655
+ super().__init__()
656
+ self.dim = dim
657
+ self.num_heads = num_heads
658
+ self.num_kv_heads = num_kv_heads if num_kv_heads is not None else num_heads
659
+ self.head_dim = head_dim if head_dim is not None else dim // num_heads
660
+ self.num_kv_groups = self.num_heads // self.num_kv_heads
661
+ self.scaling = self.head_dim**-0.5
662
+ self.window_size = window_size
663
+
664
+ qkv_dim = (self.num_heads + 2 * self.num_kv_heads) * self.head_dim
665
+ self.qkv = nn.Linear(dim, qkv_dim, bias=True)
666
+ self.proj = nn.Linear(self.num_heads * self.head_dim, dim, bias=True)
667
+ self.sinks = nn.Parameter(torch.zeros(self.num_heads)) if use_sinks else None
668
+
669
+ def _build_window_mask(self, seq_len: int, device: torch.device, dtype: torch.dtype) -> torch.Tensor | None:
670
+ if self.window_size <= 0:
671
+ return None
672
+ row_idx = torch.arange(seq_len, device=device).unsqueeze(1)
673
+ col_idx = torch.arange(seq_len, device=device).unsqueeze(0)
674
+ mask = torch.zeros(seq_len, seq_len, device=device, dtype=dtype)
675
+ mask = mask.masked_fill((row_idx - col_idx).abs() > self.window_size, float("-inf"))
676
+ return mask
677
+
678
+ def forward(
679
+ self,
680
+ hidden_states: torch.Tensor,
681
+ cu_seqlens: torch.Tensor,
682
+ position_embeddings: tuple[torch.Tensor, torch.Tensor],
683
+ full_attn: bool = False,
684
+ ) -> torch.Tensor:
685
+ seq_len = hidden_states.shape[0]
686
+ qkv = self.qkv(hidden_states)
687
+
688
+ q_dim = self.num_heads * self.head_dim
689
+ kv_dim = self.num_kv_heads * self.head_dim
690
+ q = qkv[:, :q_dim].view(seq_len, self.num_heads, self.head_dim)
691
+ k = qkv[:, q_dim : q_dim + kv_dim].view(seq_len, self.num_kv_heads, self.head_dim)
692
+ v = qkv[:, q_dim + kv_dim :].view(seq_len, self.num_kv_heads, self.head_dim)
693
+
694
+ cos, sin = position_embeddings
695
+ q, k = _apply_rotary_pos_emb_vision(q, k, cos, sin)
696
+
697
+ lengths = cu_seqlens[1:] - cu_seqlens[:-1]
698
+ q_chunks = torch.split(q, lengths.tolist(), dim=0)
699
+ k_chunks = torch.split(k, lengths.tolist(), dim=0)
700
+ v_chunks = torch.split(v, lengths.tolist(), dim=0)
701
+
702
+ outputs = []
703
+ for q_c, k_c, v_c in zip(q_chunks, k_chunks, v_chunks):
704
+ q_c = q_c.unsqueeze(0).transpose(1, 2)
705
+ k_c = k_c.unsqueeze(0).transpose(1, 2)
706
+ v_c = v_c.unsqueeze(0).transpose(1, 2)
707
+
708
+ if self.num_kv_groups > 1:
709
+ k_c = k_c.repeat_interleave(self.num_kv_groups, dim=1)
710
+ v_c = v_c.repeat_interleave(self.num_kv_groups, dim=1)
711
+
712
+ attn_mask = None
713
+ if not full_attn:
714
+ attn_mask = self._build_window_mask(q_c.shape[2], q_c.device, q_c.dtype)
715
+
716
+ if self.sinks is not None:
717
+ sink_bias = torch.zeros(
718
+ 1, self.num_heads, q_c.shape[2], k_c.shape[2], device=q_c.device, dtype=q_c.dtype
719
+ )
720
+ sink_bias[..., 0] = self.sinks.view(1, self.num_heads, 1)
721
+ attn_mask = sink_bias if attn_mask is None else attn_mask + sink_bias
722
+
723
+ attn_out = F.scaled_dot_product_attention(q_c, k_c, v_c, attn_mask=attn_mask, scale=self.scaling)
724
+ outputs.append(attn_out.squeeze(0).transpose(0, 1))
725
+
726
+ attn_output = torch.cat(outputs, dim=0)
727
+ attn_output = attn_output.reshape(seq_len, -1)
728
+ return self.proj(attn_output)
729
+
730
+
731
+ class MiMoVisionBlock(nn.Module):
732
+ def __init__(
733
+ self,
734
+ dim: int,
735
+ intermediate_dim: int,
736
+ num_heads: int,
737
+ num_kv_heads: int | None = None,
738
+ head_dim: int | None = None,
739
+ hidden_act: str = "silu",
740
+ rms_norm_eps: float = 1e-6,
741
+ use_sinks: bool = False,
742
+ window_size: int = -1,
743
+ ):
744
+ super().__init__()
745
+ self.norm1 = nn.RMSNorm(dim, eps=rms_norm_eps)
746
+ self.norm2 = nn.RMSNorm(dim, eps=rms_norm_eps)
747
+ self.attn = MiMoVisionAttention(
748
+ dim=dim, num_heads=num_heads, num_kv_heads=num_kv_heads, head_dim=head_dim,
749
+ use_sinks=use_sinks, window_size=window_size,
750
+ )
751
+ self.mlp = MiMoVisionSwiGLUMLP(dim=dim, intermediate_dim=intermediate_dim, hidden_act=hidden_act)
752
+
753
+ def forward(
754
+ self,
755
+ hidden_states: torch.Tensor,
756
+ cu_seqlens: torch.Tensor,
757
+ position_embeddings: tuple[torch.Tensor, torch.Tensor],
758
+ full_attn: bool = False,
759
+ ) -> torch.Tensor:
760
+ hidden_states = hidden_states + self.attn(
761
+ self.norm1(hidden_states), cu_seqlens=cu_seqlens,
762
+ position_embeddings=position_embeddings, full_attn=full_attn,
763
+ )
764
+ hidden_states = hidden_states + self.mlp(self.norm2(hidden_states))
765
+ return hidden_states
766
+
767
+
768
+ class MiMoVisionPatchMerger(nn.Module):
769
+ def __init__(self, dim: int, context_dim: int, spatial_merge_size: int = 2):
770
+ super().__init__()
771
+ self.hidden_size = context_dim * (spatial_merge_size**2)
772
+ self.ln_q = nn.LayerNorm(context_dim, eps=1e-6)
773
+ self.mlp = nn.Sequential(
774
+ nn.Linear(self.hidden_size, self.hidden_size),
775
+ nn.GELU(),
776
+ nn.Linear(self.hidden_size, dim),
777
+ )
778
+
779
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
780
+ return self.mlp(self.ln_q(x).view(-1, self.hidden_size))
781
+
782
+
783
+ class MiMoVisionTransformer(nn.Module):
784
+ def __init__(self, config):
785
+ super().__init__()
786
+ self.config = config
787
+ hidden_size = config.hidden_size
788
+ depth = config.depth
789
+ num_heads = config.num_heads
790
+ num_kv_heads = getattr(config, "num_key_value_heads", num_heads)
791
+ head_dim = getattr(config, "qk_channels", 64)
792
+ spatial_merge_size = getattr(config, "spatial_merge_size", 2)
793
+ rms_norm_eps = getattr(config, "rms_norm_eps", 1e-6)
794
+ self.fullatt_block_indexes = getattr(config, "fullatt_block_indexes", [])
795
+ use_sink = getattr(config, "use_sink", False)
796
+ visual_token_window_size = getattr(config, "visual_token_window_size", -1)
797
+ self.vit_window_attn_types = getattr(config, "vit_window_attn_types", None) or [-1] * depth
798
+
799
+ self.spatial_merge_size = spatial_merge_size
800
+ self.spatial_merge_unit = spatial_merge_size * spatial_merge_size
801
+
802
+ self.patch_embed = MiMoVisionPatchEmbed(
803
+ patch_size=config.patch_size,
804
+ temporal_patch_size=config.temporal_patch_size,
805
+ in_channels=getattr(config, "in_channels", None) or getattr(config, "in_chans", 3),
806
+ embed_dim=hidden_size,
807
+ )
808
+
809
+ self.rotary_pos_emb = MiMoVisionRotaryEmbedding(head_dim // 2)
810
+
811
+ self.blocks = nn.ModuleList(
812
+ [
813
+ MiMoVisionBlock(
814
+ dim=hidden_size,
815
+ intermediate_dim=config.intermediate_size,
816
+ num_heads=num_heads,
817
+ num_kv_heads=num_kv_heads,
818
+ head_dim=head_dim,
819
+ hidden_act=config.hidden_act,
820
+ rms_norm_eps=rms_norm_eps,
821
+ use_sinks=use_sink and (i not in self.fullatt_block_indexes),
822
+ window_size=visual_token_window_size,
823
+ )
824
+ for i in range(depth)
825
+ ]
826
+ )
827
+
828
+ self.merger = MiMoVisionPatchMerger(
829
+ dim=config.out_hidden_size,
830
+ context_dim=hidden_size,
831
+ spatial_merge_size=spatial_merge_size,
832
+ )
833
+
834
+ @property
835
+ def dtype(self) -> torch.dtype:
836
+ return self.patch_embed.proj.weight.dtype
837
+
838
+ def apply_index(self, tensor: torch.Tensor, index: torch.Tensor) -> torch.Tensor:
839
+ tensor = tensor.unflatten(0, (-1, self.spatial_merge_unit))
840
+ tensor = tensor[index]
841
+ return tensor.flatten(0, 1)
842
+
843
+ def get_window_index_1d(self, grid_thw: torch.Tensor, col: bool = True) -> torch.Tensor:
844
+ window_index = []
845
+ window_index_id = 0
846
+ for grid_t, grid_h, grid_w in grid_thw:
847
+ llm_grid_h = grid_h // self.spatial_merge_size
848
+ llm_grid_w = grid_w // self.spatial_merge_size
849
+ index = torch.arange(grid_t * llm_grid_h * llm_grid_w).reshape(grid_t, llm_grid_h, llm_grid_w)
850
+ index_new = index.transpose(1, 2).reshape(-1) if col else index.reshape(-1)
851
+ window_index.append(index_new + window_index_id)
852
+ window_index_id += (grid_t * llm_grid_h * llm_grid_w).item()
853
+ return torch.cat(window_index, dim=0)
854
+
855
+ def rot_pos_emb(self, grid_thw: torch.Tensor) -> torch.Tensor:
856
+ pos_ids = []
857
+ for t, h, w in grid_thw:
858
+ hpos_ids = torch.arange(h).unsqueeze(1).expand(-1, w)
859
+ hpos_ids = hpos_ids.reshape(
860
+ h // self.spatial_merge_size, self.spatial_merge_size,
861
+ w // self.spatial_merge_size, self.spatial_merge_size,
862
+ )
863
+ hpos_ids = hpos_ids.permute(0, 2, 1, 3).flatten()
864
+
865
+ wpos_ids = torch.arange(w).unsqueeze(0).expand(h, -1)
866
+ wpos_ids = wpos_ids.reshape(
867
+ h // self.spatial_merge_size, self.spatial_merge_size,
868
+ w // self.spatial_merge_size, self.spatial_merge_size,
869
+ )
870
+ wpos_ids = wpos_ids.permute(0, 2, 1, 3).flatten()
871
+
872
+ pos_ids.append(torch.stack([hpos_ids, wpos_ids], dim=-1).repeat(t, 1))
873
+ pos_ids = torch.cat(pos_ids, dim=0)
874
+ max_grid_size = grid_thw[:, 1:].max()
875
+ rotary_pos_emb_full = self.rotary_pos_emb(max_grid_size)
876
+ return rotary_pos_emb_full[pos_ids].flatten(1)
877
+
878
+ def forward(self, pixel_values: torch.Tensor, grid_thw: torch.Tensor) -> torch.Tensor:
879
+ x = pixel_values.to(device=self.patch_embed.proj.weight.device, dtype=self.dtype)
880
+ x = self.patch_embed(x)
881
+
882
+ rotary_emb = self.rot_pos_emb(grid_thw)
883
+ rotary_emb = rotary_emb.to(device=x.device)
884
+ emb = torch.cat((rotary_emb, rotary_emb), dim=-1)
885
+
886
+ window_index_1d_col = self.get_window_index_1d(grid_thw, col=True).to(device=x.device)
887
+ reverse_window_index_1d_col = torch.argsort(window_index_1d_col).to(device=x.device)
888
+
889
+ row_based_embeddings = (emb.cos(), emb.sin())
890
+ col_emb = self.apply_index(emb, window_index_1d_col)
891
+ col_based_embeddings = (col_emb.cos(), col_emb.sin())
892
+
893
+ cu_seqlens = torch.repeat_interleave(grid_thw[:, 1] * grid_thw[:, 2], grid_thw[:, 0]).cumsum(
894
+ dim=0, dtype=torch.int32
895
+ )
896
+ cu_seqlens = F.pad(cu_seqlens, (1, 0), value=0).to(device=x.device)
897
+
898
+ for i, blk in enumerate(self.blocks):
899
+ window_attn_type = self.vit_window_attn_types[i]
900
+
901
+ if window_attn_type == 1 and (i == 0 or self.vit_window_attn_types[i - 1] != 1):
902
+ x = self.apply_index(x, window_index_1d_col)
903
+
904
+ if i > 0 and window_attn_type != 1 and self.vit_window_attn_types[i - 1] == 1:
905
+ x = self.apply_index(x, reverse_window_index_1d_col)
906
+
907
+ position_embeddings = col_based_embeddings if window_attn_type == 1 else row_based_embeddings
908
+ full_attn = i in self.fullatt_block_indexes
909
+ x = blk(x, cu_seqlens=cu_seqlens, position_embeddings=position_embeddings, full_attn=full_attn)
910
+
911
+ return self.merger(x)
912
+
913
+
914
+ # ---------------------------------------------------------------------------
915
+ # Audio encoder
916
+ # ---------------------------------------------------------------------------
917
+
918
+
919
+ class AudioProjection(nn.Module):
920
+ def __init__(self, input_size: int, hidden_size: int, output_size: int):
921
+ super().__init__()
922
+ self.mlp = nn.Sequential(
923
+ nn.Linear(input_size, hidden_size, bias=False),
924
+ nn.GELU(),
925
+ nn.Linear(hidden_size, output_size, bias=False),
926
+ )
927
+
928
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
929
+ return self.mlp(x)
930
+
931
+
932
+ class MiMoAudioEncoder(nn.Module):
933
+ def __init__(self, config):
934
+ super().__init__()
935
+ self.config = config
936
+
937
+ self.audio_channels = getattr(config, "audio_channels")
938
+ self.group_size = getattr(config, "group_size")
939
+ self.input_local_dim = getattr(config, "input_local_dim")
940
+ self.out_hidden_size = getattr(config, "out_hidden_size")
941
+ self.input_full_attention = getattr(config, "input_full_attention", True)
942
+ self.audio_segment_size = getattr(config, "audio_segment_size", 6000)
943
+
944
+ input_local_config = Qwen2Config(
945
+ hidden_size=getattr(config, "input_local_dim"),
946
+ num_hidden_layers=getattr(config, "input_local_layers"),
947
+ num_attention_heads=getattr(config, "input_local_attn_heads"),
948
+ num_key_value_heads=getattr(config, "input_local_attn_heads"),
949
+ intermediate_size=getattr(config, "input_local_intermediate_size"),
950
+ attention_dropout=getattr(config, "input_local_hidden_dropout", 0.0),
951
+ rope_theta=getattr(config, "rope_theta", 640000.0),
952
+ partial_rotary_factor=getattr(config, "partial_rotary_factor", 1.0),
953
+ )
954
+ self.input_local_transformer = Qwen2Model(input_local_config)
955
+
956
+ if not getattr(config, "add_post_norm", True):
957
+ self.input_local_transformer.norm = nn.Identity()
958
+
959
+ proj_in = self.input_local_dim * self.group_size
960
+ projection_layers = getattr(config, "projection_layers", 2)
961
+ if projection_layers == 1:
962
+ self.projection = nn.Linear(proj_in, self.out_hidden_size, bias=False)
963
+ elif projection_layers == 2:
964
+ self.projection = AudioProjection(proj_in, proj_in * 4, self.out_hidden_size)
965
+ else:
966
+ raise ValueError(f"Unsupported projection_layers={projection_layers}, expected 1 or 2.")
967
+
968
+ def _apply_speech_embeddings(self, audio_codes: torch.Tensor, speech_embeddings: nn.ModuleList) -> torch.Tensor:
969
+ num_segments = audio_codes.shape[0]
970
+ out = torch.zeros(
971
+ (num_segments, self.group_size, self.input_local_dim),
972
+ dtype=speech_embeddings[0].weight.dtype,
973
+ device=audio_codes.device,
974
+ )
975
+ for i in range(self.audio_channels):
976
+ out.add_(speech_embeddings[i](audio_codes[:, :, i].long()))
977
+ return out
978
+
979
+ def _apply_input_local_transformer(self, speech_embeddings: torch.Tensor) -> torch.Tensor:
980
+ output = self.input_local_transformer(
981
+ inputs_embeds=speech_embeddings, return_dict=True, use_cache=False,
982
+ is_causal=not self.input_full_attention,
983
+ )
984
+ return output.last_hidden_state
985
+
986
+ def _process_audio_codes(self, audio_codes: torch.Tensor, speech_embeddings: nn.ModuleList) -> torch.Tensor:
987
+ audio_codes = _pad_and_group_audio_codes(audio_codes, self.audio_channels, self.group_size)
988
+ audio_embs = self._apply_speech_embeddings(audio_codes, speech_embeddings)
989
+ audio_hidden = self._apply_input_local_transformer(audio_embs)
990
+ return self.projection(audio_hidden.reshape(audio_hidden.shape[0], -1))
991
+
992
+ def get_audio_feature(
993
+ self,
994
+ mels: list[torch.Tensor],
995
+ speech_embeddings: nn.ModuleList,
996
+ audio_tokenizer_encoder,
997
+ ) -> torch.Tensor:
998
+ """Full pipeline: mel spectrograms → tokenize → codes → embed → project."""
999
+ if not mels:
1000
+ device = next(self.projection.parameters()).device
1001
+ dtype = next(self.projection.parameters()).dtype
1002
+ return torch.empty(0, self.out_hidden_size, device=device, dtype=dtype)
1003
+
1004
+ device = next(audio_tokenizer_encoder.parameters()).device
1005
+ code_list = tokenize_audio_batch(
1006
+ mels, audio_tokenizer_encoder, segment_size=self.audio_segment_size, device=device,
1007
+ )
1008
+
1009
+ codecs_to_concat = []
1010
+ for codecs in code_list:
1011
+ codecs_to_concat.append(_pad_and_group_audio_codes(codecs, self.audio_channels, self.group_size))
1012
+ audio_codes = torch.cat(codecs_to_concat, dim=0)
1013
+
1014
+ audio_embs = self._apply_speech_embeddings(audio_codes, speech_embeddings)
1015
+ audio_hidden = self._apply_input_local_transformer(audio_embs)
1016
+ return self.projection(audio_hidden.reshape(audio_hidden.shape[0], -1))
1017
+
1018
+ def forward(
1019
+ self,
1020
+ speech_embeddings: nn.ModuleList,
1021
+ audio_codes: torch.Tensor | None = None,
1022
+ audio_embeds: torch.Tensor | None = None,
1023
+ ) -> torch.Tensor:
1024
+ if audio_embeds is not None:
1025
+ if audio_embeds.dim() != 2:
1026
+ raise ValueError(f"`audio_embeds` must be 2D [N, H], got shape={tuple(audio_embeds.shape)}")
1027
+ if audio_embeds.shape[-1] != self.out_hidden_size:
1028
+ raise ValueError(
1029
+ f"Unexpected audio_embeds hidden size {audio_embeds.shape[-1]}, expected {self.out_hidden_size}"
1030
+ )
1031
+ return audio_embeds
1032
+
1033
+ if audio_codes is None:
1034
+ raise ValueError("Either `audio_codes` or `audio_embeds` must be provided.")
1035
+
1036
+ return self._process_audio_codes(audio_codes, speech_embeddings)
1037
+
1038
+
1039
+ # ---------------------------------------------------------------------------
1040
+ # Audio tokenizer (codec: mel → encoder → VQ → codes)
1041
+ # Adapted from https://github.com/XiaomiMiMo/MiMo-Audio-Tokenizer.git
1042
+ # ---------------------------------------------------------------------------
1043
+
1044
+
1045
+ class MiMoAudioTokenizerConfig(PretrainedConfig):
1046
+ model_type = "mimo_audio_tokenizer"
1047
+
1048
+ def __init__(
1049
+ self,
1050
+ max_audio_seconds: int = 1800,
1051
+ stride_size: int = 2,
1052
+ avg_pooler: int = 1,
1053
+ d_model: int = 768,
1054
+ scale_embedding: bool = True,
1055
+ kernel_size: int = 3,
1056
+ activation_function: str = "gelu",
1057
+ encoder_layers: int = 8,
1058
+ encoder_skip_layer_id: int = None,
1059
+ encoder_attention_heads: int = 12,
1060
+ encoder_ffn_dim: int = 3072,
1061
+ encoder_causal: bool = False,
1062
+ encoder_attn_window_size: list = None,
1063
+ decoder_layers: int = 8,
1064
+ decoder_attention_heads: int = 12,
1065
+ decoder_ffn_dim: int = 3072,
1066
+ decoder_kernel_size: int = 3,
1067
+ decoder_stride_size: int = 2,
1068
+ decoder_causal: bool = True,
1069
+ decoder_attn_window_size: list = None,
1070
+ nfft: int = 1024,
1071
+ vocoder_dim: int = 512,
1072
+ vocoder_intermediate_dim: int = 4096,
1073
+ vocoder_num_layers: int = 30,
1074
+ n_mels: int = 80,
1075
+ sampling_rate: int = 24000,
1076
+ hop_length: int = 240,
1077
+ window_size: int = 1024,
1078
+ vocoder_padding: str = "same",
1079
+ fmin: int = 0,
1080
+ fmax: int = None,
1081
+ num_quantizers: int = 12,
1082
+ codebook_size: list = None,
1083
+ threshold_ema_dead_code: int = 10,
1084
+ position_embedding_type: str = "rope",
1085
+ rope_theta: int = 10000,
1086
+ rope_type: str = "default",
1087
+ ln_type: str = "LayerNorm",
1088
+ vocoder_attention_heads: int = 4,
1089
+ vocoder_attn_window_size: list = None,
1090
+ use_istft_only: bool = False,
1091
+ hybrid_attention: bool = False,
1092
+ hybrid_block_size: int = 8,
1093
+ swa_per_block: int = 2,
1094
+ **kwargs,
1095
+ ):
1096
+ super().__init__(**kwargs)
1097
+ self.max_audio_seconds = max_audio_seconds
1098
+ self.stride_size = stride_size
1099
+ self.avg_pooler = avg_pooler
1100
+ self.d_model = d_model
1101
+ self.scale_embedding = scale_embedding
1102
+ self.kernel_size = kernel_size
1103
+ self.activation_function = activation_function
1104
+ self.encoder_layers = encoder_layers
1105
+ self.encoder_skip_layer_id = encoder_skip_layer_id
1106
+ self.encoder_attention_heads = encoder_attention_heads
1107
+ self.encoder_ffn_dim = encoder_ffn_dim
1108
+ self.encoder_causal = encoder_causal
1109
+ self.encoder_attn_window_size = encoder_attn_window_size if encoder_attn_window_size is not None else [-1, -1]
1110
+ self.decoder_layers = decoder_layers
1111
+ self.decoder_attention_heads = decoder_attention_heads
1112
+ self.decoder_ffn_dim = decoder_ffn_dim
1113
+ self.decoder_kernel_size = decoder_kernel_size
1114
+ self.decoder_stride_size = decoder_stride_size
1115
+ self.decoder_causal = decoder_causal
1116
+ self.decoder_attn_window_size = decoder_attn_window_size if decoder_attn_window_size is not None else [-1, -1]
1117
+ self.nfft = nfft
1118
+ self.vocoder_dim = vocoder_dim
1119
+ self.vocoder_intermediate_dim = vocoder_intermediate_dim
1120
+ self.vocoder_num_layers = vocoder_num_layers
1121
+ self.n_mels = n_mels
1122
+ self.sampling_rate = sampling_rate
1123
+ self.hop_length = hop_length
1124
+ self.window_size = window_size
1125
+ self.vocoder_padding = vocoder_padding
1126
+ self.fmin = fmin
1127
+ self.fmax = fmax
1128
+ self.num_quantizers = num_quantizers
1129
+ self.codebook_size = codebook_size if codebook_size is not None else [1024]
1130
+ self.threshold_ema_dead_code = threshold_ema_dead_code
1131
+ self.position_embedding_type = position_embedding_type
1132
+ self.rope_theta = rope_theta
1133
+ self.rope_type = rope_type
1134
+ self.ln_type = ln_type
1135
+ self.vocoder_attention_heads = vocoder_attention_heads
1136
+ self.vocoder_attn_window_size = vocoder_attn_window_size if vocoder_attn_window_size is not None else [40, 10]
1137
+ self.use_istft_only = use_istft_only
1138
+ self.hybrid_attention = hybrid_attention
1139
+ self.hybrid_block_size = hybrid_block_size
1140
+ self.swa_per_block = swa_per_block
1141
+
1142
+
1143
+ class EuclideanCodebook(nn.Module):
1144
+ def __init__(self, dim: int, codebook_size: int, kmeans_init: bool = False, **kwargs):
1145
+ super().__init__()
1146
+ init_fn = torch.zeros if kmeans_init else self._uniform_init
1147
+ embed = init_fn(codebook_size, dim)
1148
+ self.codebook_size = codebook_size
1149
+ self.register_buffer("inited", torch.Tensor([not kmeans_init]))
1150
+ self.register_buffer("cluster_size", torch.zeros(codebook_size))
1151
+ self.register_buffer("embed", embed)
1152
+ self.register_buffer("embed_avg", embed.clone())
1153
+
1154
+ def quantize(self, x):
1155
+ embed = self.embed.t()
1156
+ dist = -(x.pow(2).sum(1, keepdim=True) - 2 * x @ embed + embed.pow(2).sum(0, keepdim=True))
1157
+ return dist.max(dim=-1).indices
1158
+
1159
+ def encode(self, x):
1160
+ shape = x.shape
1161
+ x = x.reshape(-1, x.shape[-1])
1162
+ embed_ind = self.quantize(x)
1163
+ return embed_ind.view(*shape[:-1])
1164
+
1165
+ def decode(self, embed_ind):
1166
+ return F.embedding(embed_ind, self.embed)
1167
+
1168
+ @staticmethod
1169
+ def _uniform_init(*shape: int):
1170
+ t = torch.empty(shape)
1171
+ nn.init.kaiming_uniform_(t)
1172
+ return t
1173
+
1174
+
1175
+ class VectorQuantization(nn.Module):
1176
+ def __init__(self, dim: int, codebook_size: int, codebook_dim: Optional[int] = None, kmeans_init: bool = True, **kwargs):
1177
+ super().__init__()
1178
+ _codebook_dim = codebook_dim if codebook_dim is not None else dim
1179
+ requires_projection = _codebook_dim != dim
1180
+ self.project_in = nn.Linear(dim, _codebook_dim) if requires_projection else nn.Identity()
1181
+ self.project_out = nn.Linear(_codebook_dim, dim) if requires_projection else nn.Identity()
1182
+ self._codebook = EuclideanCodebook(dim=_codebook_dim, codebook_size=codebook_size, kmeans_init=kmeans_init)
1183
+ self.codebook_size = codebook_size
1184
+
1185
+ def encode(self, x):
1186
+ return self._codebook.encode(self.project_in(x))
1187
+
1188
+ def decode(self, embed_ind):
1189
+ return self.project_out(self._codebook.decode(embed_ind))
1190
+
1191
+
1192
+ class ResidualVectorQuantization(nn.Module):
1193
+ def __init__(self, *, num_quantizers, codebook_size, **kwargs):
1194
+ super().__init__()
1195
+ if isinstance(codebook_size, int):
1196
+ codebook_size = [codebook_size] * num_quantizers
1197
+ elif len(codebook_size) < num_quantizers:
1198
+ codebook_size += [codebook_size[-1]] * (num_quantizers - len(codebook_size))
1199
+ self.layers = nn.ModuleList(
1200
+ [VectorQuantization(codebook_size=codebook_size[i], **kwargs) for i in range(num_quantizers)]
1201
+ )
1202
+
1203
+ def encode(self, x: torch.Tensor, n_q: Optional[int] = None, st: Optional[int] = None) -> torch.Tensor:
1204
+ residual = x
1205
+ all_indices = []
1206
+ n_q = len(self.layers) if n_q is None else n_q
1207
+ st = 0 if st is None else st
1208
+ for layer in self.layers[st:n_q]:
1209
+ indices = layer.encode(residual)
1210
+ quantized = layer.decode(indices)
1211
+ residual = residual - quantized
1212
+ all_indices.append(indices)
1213
+ return torch.stack(all_indices)
1214
+
1215
+ def decode(self, q_indices: torch.Tensor, st: int = 0) -> torch.Tensor:
1216
+ quantized_out = self.layers[st].decode(q_indices[0])
1217
+ for i in range(1, len(q_indices)):
1218
+ quantized_out = quantized_out + self.layers[st + i].decode(q_indices[i])
1219
+ return quantized_out
1220
+
1221
+
1222
+ class ResidualVectorQuantizer(nn.Module):
1223
+ def __init__(self, dimension: int = 256, n_q: int = 8, bins: int | list = 1024, kmeans_init: bool = True, **kwargs):
1224
+ super().__init__()
1225
+ self.n_q = n_q
1226
+ self.vq = ResidualVectorQuantization(dim=dimension, codebook_size=bins, num_quantizers=n_q, kmeans_init=kmeans_init)
1227
+
1228
+ def encode(self, x: torch.Tensor, n_q: Optional[int] = None, st: Optional[int] = None) -> torch.Tensor:
1229
+ return self.vq.encode(x, n_q=n_q or self.n_q, st=st or 0)
1230
+
1231
+ def decode(self, codes: torch.Tensor, st: int = 0) -> torch.Tensor:
1232
+ return self.vq.decode(codes, st=st)
1233
+
1234
+
1235
+ class AudioTokenizerRotaryEmbedding(nn.Module):
1236
+ def __init__(self, base, dim, max_seq_len, rope_type="default", device=None):
1237
+ super().__init__()
1238
+ self.attention_scaling = 1.0
1239
+ inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2, dtype=torch.float, device=device) / dim))
1240
+ self.register_buffer("inv_freq", inv_freq, persistent=False)
1241
+
1242
+ @torch.no_grad()
1243
+ def forward(self, x, position_ids):
1244
+ inv_freq_expanded = self.inv_freq[:, None].float().expand(-1, 1).to(x.device)
1245
+ position_ids_expanded = position_ids[None, :].float()
1246
+ with torch.autocast(device_type="cpu", enabled=False):
1247
+ freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(0, 1)
1248
+ emb = torch.cat((freqs, freqs), dim=-1)
1249
+ cos = emb.cos() * self.attention_scaling
1250
+ sin = emb.sin() * self.attention_scaling
1251
+ return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)
1252
+
1253
+
1254
+ def _at_get_position_ids(lengths):
1255
+ total_len = lengths.sum()
1256
+ offset = torch.cat([torch.zeros(1, device=lengths.device, dtype=lengths.dtype), lengths[:-1].cumsum(dim=0)])
1257
+ offset = torch.repeat_interleave(offset, lengths)
1258
+ return torch.arange(0, total_len, device=lengths.device) - offset
1259
+
1260
+
1261
+ def _at_get_sequence_mask(inputs, inputs_length):
1262
+ if inputs.dim() == 3:
1263
+ bsz, tgt_len, _ = inputs.size()
1264
+ else:
1265
+ bsz, tgt_len = inputs_length.shape[0], torch.max(inputs_length)
1266
+ sequence_mask = torch.arange(0, tgt_len, device=inputs.device)
1267
+ sequence_mask = torch.lt(sequence_mask, inputs_length.reshape(bsz, 1)).view(bsz, tgt_len, 1)
1268
+ unpacking_index = torch.cumsum(sequence_mask.to(torch.int64).view(-1), dim=0) - 1
1269
+ return sequence_mask, unpacking_index
1270
+
1271
+
1272
+ def _at_unpack_hidden_states(hidden_states, lengths, sequence_mask=None, unpacking_index=None):
1273
+ bsz = lengths.shape[0]
1274
+ if sequence_mask is None or unpacking_index is None:
1275
+ sequence_mask, unpacking_index = _at_get_sequence_mask(hidden_states, lengths)
1276
+ hidden_states = torch.index_select(hidden_states, 0, unpacking_index).view(
1277
+ bsz, torch.max(lengths), hidden_states.shape[-1]
1278
+ )
1279
+ return torch.where(sequence_mask, hidden_states, 0)
1280
+
1281
+
1282
+ def _at_rotate_half(x):
1283
+ x1 = x[..., : x.shape[-1] // 2]
1284
+ x2 = x[..., x.shape[-1] // 2 :]
1285
+ return torch.cat((-x2, x1), dim=-1)
1286
+
1287
+
1288
+ def _at_apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1):
1289
+ cos = cos.unsqueeze(unsqueeze_dim)
1290
+ sin = sin.unsqueeze(unsqueeze_dim)
1291
+ return (q * cos) + (_at_rotate_half(q) * sin), (k * cos) + (_at_rotate_half(k) * sin)
1292
+
1293
+
1294
+ _AT_LAYER_NORM = {"LayerNorm": nn.LayerNorm}
1295
+
1296
+
1297
+ class AudioTokenizerAttention(nn.Module):
1298
+ def __init__(self, embed_dim: int, num_heads: int, window_size: tuple[int, int] = (-1, -1), causal: bool = False):
1299
+ super().__init__()
1300
+ self.embed_dim = embed_dim
1301
+ self.num_heads = num_heads
1302
+ self.head_dim = embed_dim // num_heads
1303
+ self.window_size = window_size
1304
+ self.causal = causal
1305
+ self.scaling = self.head_dim**-0.5
1306
+
1307
+ self.k_proj = nn.Linear(embed_dim, embed_dim, bias=False)
1308
+ self.v_proj = nn.Linear(embed_dim, embed_dim, bias=True)
1309
+ self.q_proj = nn.Linear(embed_dim, embed_dim, bias=True)
1310
+ self.out_proj = nn.Linear(embed_dim, embed_dim, bias=True)
1311
+
1312
+ def _build_attn_mask(self, seq_len: int, device: torch.device, dtype: torch.dtype) -> torch.Tensor | None:
1313
+ has_window = self.window_size[0] > 0
1314
+ if not self.causal and not has_window:
1315
+ return None
1316
+ mask = torch.zeros(seq_len, seq_len, device=device, dtype=dtype)
1317
+ if self.causal:
1318
+ mask = mask + torch.triu(torch.full((seq_len, seq_len), float("-inf"), device=device, dtype=dtype), diagonal=1)
1319
+ if has_window:
1320
+ row_idx = torch.arange(seq_len, device=device).unsqueeze(1)
1321
+ col_idx = torch.arange(seq_len, device=device).unsqueeze(0)
1322
+ mask = mask.masked_fill((row_idx - col_idx).abs() > self.window_size[0], float("-inf"))
1323
+ return mask
1324
+
1325
+ def forward(self, hidden_states, cu_seqlens, max_seqlen, rope_position_embeddings=None):
1326
+ total_len = hidden_states.shape[0]
1327
+ q = self.q_proj(hidden_states).view(total_len, self.num_heads, self.head_dim)
1328
+ k = self.k_proj(hidden_states).view(total_len, self.num_heads, self.head_dim)
1329
+ v = self.v_proj(hidden_states).view(total_len, self.num_heads, self.head_dim)
1330
+ if rope_position_embeddings is not None:
1331
+ cos, sin = rope_position_embeddings
1332
+ q, k = _at_apply_rotary_pos_emb(q, k, cos, sin)
1333
+ num_seqs = cu_seqlens.shape[0] - 1
1334
+ outputs = []
1335
+ for i in range(num_seqs):
1336
+ start, end = cu_seqlens[i].item(), cu_seqlens[i + 1].item()
1337
+ seq_len = end - start
1338
+ q_seq = q[start:end].transpose(0, 1).unsqueeze(0)
1339
+ k_seq = k[start:end].transpose(0, 1).unsqueeze(0)
1340
+ v_seq = v[start:end].transpose(0, 1).unsqueeze(0)
1341
+ attn_mask = self._build_attn_mask(seq_len, q_seq.device, q_seq.dtype)
1342
+ out = F.scaled_dot_product_attention(q_seq, k_seq, v_seq, attn_mask=attn_mask, scale=self.scaling)
1343
+ outputs.append(out.squeeze(0).transpose(0, 1))
1344
+ return self.out_proj(torch.cat(outputs, dim=0).reshape(total_len, self.embed_dim))
1345
+
1346
+
1347
+ class AudioTokenizerTransformerLayer(nn.Module):
1348
+ def __init__(self, config: MiMoAudioTokenizerConfig, causal: bool, attn_window_size: tuple[int, int] = (-1, -1)):
1349
+ super().__init__()
1350
+ self.embed_dim = config.d_model
1351
+ self.self_attn = AudioTokenizerAttention(
1352
+ embed_dim=self.embed_dim, num_heads=config.encoder_attention_heads,
1353
+ window_size=attn_window_size, causal=causal,
1354
+ )
1355
+ self.self_attn_layer_norm = _AT_LAYER_NORM[config.ln_type](self.embed_dim)
1356
+ self.activation_fn = ACT2FN[config.activation_function]
1357
+ self.fc1 = nn.Linear(self.embed_dim, config.encoder_ffn_dim)
1358
+ self.fc2 = nn.Linear(config.encoder_ffn_dim, self.embed_dim)
1359
+ self.final_layer_norm = _AT_LAYER_NORM[config.ln_type](self.embed_dim)
1360
+
1361
+ def forward(self, hidden_states, cu_seqlens, max_seqlen, rope_position_embeddings):
1362
+ residual = hidden_states
1363
+ hidden_states = self.self_attn_layer_norm(hidden_states)
1364
+ hidden_states = self.self_attn(hidden_states, cu_seqlens, max_seqlen, rope_position_embeddings=rope_position_embeddings)
1365
+ hidden_states = residual + hidden_states
1366
+ residual = hidden_states
1367
+ hidden_states = self.final_layer_norm(hidden_states)
1368
+ hidden_states = self.activation_fn(self.fc1(hidden_states))
1369
+ hidden_states = self.fc2(hidden_states)
1370
+ hidden_states = residual + hidden_states
1371
+ return hidden_states
1372
+
1373
+
1374
+ class AudioTokenizerEncoder(nn.Module):
1375
+ def __init__(self, config: MiMoAudioTokenizerConfig):
1376
+ super().__init__()
1377
+ self.config = config
1378
+ self.max_source_positions = (config.max_audio_seconds * config.sampling_rate // config.hop_length) // config.stride_size
1379
+ self.embed_scale = math.sqrt(config.d_model) if config.scale_embedding else 1.0
1380
+ self.skip_layer_idx = config.encoder_skip_layer_id
1381
+
1382
+ self.conv1 = nn.Conv1d(config.n_mels, config.d_model, kernel_size=config.kernel_size, padding=1)
1383
+ self.conv2 = nn.Conv1d(config.d_model, config.d_model, kernel_size=config.kernel_size, stride=config.stride_size, padding=1)
1384
+
1385
+ self.position_embedding = AudioTokenizerRotaryEmbedding(
1386
+ config.rope_theta, config.d_model // config.encoder_attention_heads,
1387
+ self.max_source_positions, config.rope_type,
1388
+ )
1389
+
1390
+ attn_window_sizes = []
1391
+ if config.hybrid_attention:
1392
+ for i in range(config.encoder_layers):
1393
+ if i % config.swa_per_block < config.swa_per_block - 1:
1394
+ attn_window_sizes.append(tuple(config.encoder_attn_window_size))
1395
+ else:
1396
+ attn_window_sizes.append((-1, -1))
1397
+ else:
1398
+ attn_window_sizes = [tuple(config.encoder_attn_window_size)] * config.encoder_layers
1399
+
1400
+ self.layers = nn.ModuleList([
1401
+ AudioTokenizerTransformerLayer(config=config, causal=config.encoder_causal, attn_window_size=attn_window_sizes[i])
1402
+ for i in range(config.encoder_layers)
1403
+ ])
1404
+
1405
+ self.layer_norm = _AT_LAYER_NORM[config.ln_type](config.d_model)
1406
+
1407
+ if config.avg_pooler != 1:
1408
+ self.down_sample_layer = nn.Sequential(
1409
+ nn.Conv1d(config.d_model, config.d_model, config.avg_pooler, config.avg_pooler, bias=False),
1410
+ nn.GELU(),
1411
+ )
1412
+ self.down_sample_norm = _AT_LAYER_NORM[config.ln_type](config.d_model)
1413
+ else:
1414
+ self.down_sample_layer = None
1415
+
1416
+ if config.num_quantizers != 0:
1417
+ self.quantizer = ResidualVectorQuantizer(
1418
+ dimension=config.d_model, n_q=config.num_quantizers,
1419
+ bins=config.codebook_size,
1420
+ threshold_ema_dead_code=config.threshold_ema_dead_code,
1421
+ )
1422
+ else:
1423
+ self.quantizer = None
1424
+
1425
+ def get_output_length(self, mel_len):
1426
+ tgt_len = mel_len + 3 - self.config.kernel_size
1427
+ return (tgt_len + 2 - self.config.kernel_size) // self.config.stride_size + 1
1428
+
1429
+ def get_features(self, input_features, output_length):
1430
+ input_features = input_features.to(self.conv1.weight)
1431
+ inputs_embeds = F.gelu(self.conv1(input_features))
1432
+ inputs_embeds = F.gelu(self.conv2(inputs_embeds))
1433
+ inputs_embeds = inputs_embeds.permute(0, 2, 1)
1434
+ bsz, tgt_len, _ = inputs_embeds.size()
1435
+
1436
+ position_ids = _at_get_position_ids(output_length).long().to(input_features.device)
1437
+ rope_position_embeddings = self.position_embedding(input_features, position_ids)
1438
+
1439
+ attention_mask, unpacking_index = _at_get_sequence_mask(inputs_embeds, output_length)
1440
+ hidden_states = torch.masked_select(inputs_embeds, attention_mask).view(
1441
+ torch.sum(output_length), self.config.d_model
1442
+ )
1443
+
1444
+ cu_seqlens = F.pad(torch.cumsum(output_length, dim=0), (1, 0), "constant", 0).to(
1445
+ device=hidden_states.device, dtype=torch.int32
1446
+ )
1447
+ max_seqlen = torch.max(output_length).to(torch.int32).item()
1448
+
1449
+ skip_connect_hidden_states = 0.0
1450
+ for idx, encoder_layer in enumerate(self.layers):
1451
+ hidden_states = encoder_layer(hidden_states, cu_seqlens, max_seqlen, rope_position_embeddings=rope_position_embeddings)
1452
+ if self.skip_layer_idx is not None and idx == self.skip_layer_idx - 1:
1453
+ skip_connect_hidden_states = hidden_states.clone()
1454
+
1455
+ hidden_states += skip_connect_hidden_states
1456
+ hidden_states = self.layer_norm(hidden_states)
1457
+
1458
+ if self.down_sample_layer is not None:
1459
+ hidden_states = torch.index_select(hidden_states, 0, unpacking_index).view(bsz, tgt_len, self.config.d_model)
1460
+ if hidden_states.size(1) % self.config.avg_pooler:
1461
+ pad_len = self.config.avg_pooler - hidden_states.size(1) % self.config.avg_pooler
1462
+ hidden_states = F.pad(hidden_states, (0, 0, 0, pad_len), mode="constant", value=0.0)
1463
+ tgt_len += pad_len
1464
+ tgt_len = tgt_len // self.config.avg_pooler
1465
+ hidden_states = self.down_sample_layer(hidden_states.transpose(1, 2))
1466
+ output_length = output_length // self.config.avg_pooler + (output_length % self.config.avg_pooler != 0).int()
1467
+ hidden_states = hidden_states.transpose(1, 2)
1468
+ attention_mask, unpacking_index = _at_get_sequence_mask(hidden_states, output_length)
1469
+ hidden_states = torch.masked_select(hidden_states, attention_mask).view(
1470
+ torch.sum(output_length), self.config.d_model
1471
+ )
1472
+ hidden_states = self.down_sample_norm(hidden_states)
1473
+
1474
+ return hidden_states, output_length, attention_mask, unpacking_index, tgt_len, bsz
1475
+
1476
+ @torch.no_grad()
1477
+ def encode(self, input_features, input_lens=None, output_length=None, return_codes_only=False, n_q=None, use_quantizer=True):
1478
+ if output_length is None:
1479
+ output_length = self.get_output_length(input_lens)
1480
+ input_features = _at_unpack_hidden_states(input_features, input_lens)
1481
+ hidden_states, output_length, attention_mask, unpacking_index, tgt_len, bsz = self.get_features(
1482
+ input_features=input_features.transpose(1, 2), output_length=output_length,
1483
+ )
1484
+ dtype = hidden_states.dtype
1485
+ if use_quantizer and self.quantizer is not None:
1486
+ self.quantizer.float()
1487
+ codes = self.quantizer.encode(hidden_states.float(), n_q=n_q)
1488
+ if return_codes_only:
1489
+ return codes, output_length
1490
+ hidden_states = self.quantizer.decode(codes)
1491
+ hidden_states = hidden_states.to(dtype)
1492
+ else:
1493
+ codes = None
1494
+ hidden_states_packed = hidden_states.clone()
1495
+ hidden_states = torch.index_select(hidden_states, 0, unpacking_index).view(bsz, tgt_len, self.config.d_model)
1496
+ hidden_states = torch.where(attention_mask, hidden_states, 0)
1497
+ return hidden_states, hidden_states_packed, output_length, codes
1498
+
1499
+
1500
+ class MiMoAudioTokenizer(PreTrainedModel):
1501
+ config_class = MiMoAudioTokenizerConfig
1502
+
1503
+ def __init__(self, config: MiMoAudioTokenizerConfig):
1504
+ super().__init__(config)
1505
+ self.config = config
1506
+ self.sampling_rate = config.sampling_rate
1507
+ self.encoder = AudioTokenizerEncoder(config=config)
1508
+ self.downsample_rate = int(config.hop_length * 2 * config.avg_pooler)
1509
+
1510
+ def get_output_length(self, mel_len):
1511
+ return self.encoder.get_output_length(mel_len)
1512
+
1513
+ @torch.no_grad()
1514
+ def encode(self, mels, input_lens, use_quantizer=True):
1515
+ return self.encoder.encode(mels, input_lens=input_lens, use_quantizer=use_quantizer)
1516
+
1517
+
1518
+ def _at_group_by_length(features, lengths, max_length):
1519
+ split_points, current_sum = [], 0
1520
+ for i, seq_len in enumerate(lengths):
1521
+ if current_sum + seq_len > max_length and current_sum > 0:
1522
+ split_points.append(i)
1523
+ current_sum = seq_len.item()
1524
+ else:
1525
+ current_sum += seq_len.item()
1526
+ group_sizes, prev = [], 0
1527
+ for point in split_points:
1528
+ group_sizes.append(point - prev)
1529
+ prev = point
1530
+ if prev < len(lengths):
1531
+ group_sizes.append(len(lengths) - prev)
1532
+ len_groups = torch.split(lengths, group_sizes)
1533
+ feature_groups = torch.split(features, [g.sum().item() for g in len_groups])
1534
+ return feature_groups, len_groups
1535
+
1536
+
1537
+ @torch.no_grad()
1538
+ def tokenize_audio_batch(mels, audio_tokenizer_encoder, segment_size=6000, device=None):
1539
+ if not mels:
1540
+ return []
1541
+ if device is None:
1542
+ device = next(audio_tokenizer_encoder.parameters()).device
1543
+ input_len_seg_per_mel = []
1544
+ for m in mels:
1545
+ input_len = m.size(0)
1546
+ segs = [segment_size] * (input_len // segment_size)
1547
+ if input_len % segment_size > 0:
1548
+ segs.append(input_len % segment_size)
1549
+ input_len_seg_per_mel.append(segs)
1550
+ input_lens_flat = [s for segs in input_len_seg_per_mel for s in segs]
1551
+ input_features = torch.cat([m.to(device) for m in mels], dim=0)
1552
+ input_lens_t = torch.tensor(input_lens_flat, dtype=torch.long, device=device)
1553
+ feature_groups, len_groups = _at_group_by_length(input_features, input_lens_t, 256000)
1554
+ encoded_parts = []
1555
+ for features, lengths in zip(feature_groups, len_groups):
1556
+ codes, _ = audio_tokenizer_encoder.encode(input_features=features, input_lens=lengths, return_codes_only=True)
1557
+ encoded_parts.append(codes)
1558
+ codes = torch.cat(encoded_parts, dim=-1).transpose(0, 1).detach()
1559
+ code_lengths = []
1560
+ for segs in input_len_seg_per_mel:
1561
+ out_len = audio_tokenizer_encoder.get_output_length(torch.tensor(segs, dtype=torch.long, device=device))
1562
+ if getattr(audio_tokenizer_encoder, "down_sample_layer", None) is not None:
1563
+ avg = audio_tokenizer_encoder.config.avg_pooler
1564
+ out_len = out_len // avg + (out_len % avg != 0).long()
1565
+ code_lengths.append(out_len.sum().item())
1566
+ return list(torch.split(codes, code_lengths))
1567
+
1568
+
1569
+ # ---------------------------------------------------------------------------
1570
+ # LLM backbone
1571
+ # ---------------------------------------------------------------------------
1572
+
1573
+
1574
+ class MiMoV2Model(PreTrainedModel):
1575
+ config_class = MiMoV2Config
1576
+ attention_projection_layout = "split"
1577
+
1578
+ def __init__(self, config):
1579
+ super().__init__(config)
1580
+ self.attention_projection_layout = getattr(
1581
+ config, "attention_projection_layout", self.attention_projection_layout
1582
+ )
1583
+ self.vocab_size = config.vocab_size
1584
+ self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size)
1585
+ self.layers = nn.ModuleList(
1586
+ [
1587
+ MiMoV2DecoderLayer(
1588
+ config,
1589
+ layer_idx,
1590
+ attention_projection_layout=self.attention_projection_layout,
1591
+ )
1592
+ for layer_idx in range(config.num_hidden_layers)
1593
+ ]
1594
+ )
1595
+ self.norm = MiMoV2RMSNorm(config.hidden_size, eps=config.layernorm_epsilon)
1596
+ self.rotary_emb = MiMoV2RotaryEmbedding(config=config, is_swa=False)
1597
+ self.swa_rotary_emb = MiMoV2RotaryEmbedding(config=config, is_swa=True)
1598
+ self.has_sliding_layers = any(pattern == 1 for pattern in config.hybrid_layer_pattern)
1599
+ self.config.layer_types = [
1600
+ "sliding_attention" if config.hybrid_layer_pattern[i] == 1 else "full_attention"
1601
+ for i in range(config.num_hidden_layers)
1602
+ ]
1603
+ self.post_init()
1604
+
1605
+ def get_input_embeddings(self):
1606
+ return self.embed_tokens
1607
+
1608
+ def set_input_embeddings(self, value):
1609
+ self.embed_tokens = value
1610
+
1611
+ def forward(
1612
+ self,
1613
+ input_ids: Optional[torch.LongTensor] = None,
1614
+ attention_mask: Optional[torch.Tensor] = None,
1615
+ position_ids: Optional[torch.LongTensor] = None,
1616
+ past_key_values: Optional[Cache] = None,
1617
+ inputs_embeds: Optional[torch.FloatTensor] = None,
1618
+ use_cache: Optional[bool] = None,
1619
+ cache_position: Optional[torch.LongTensor] = None,
1620
+ **kwargs: Unpack[TransformersKwargs],
1621
+ ) -> BaseModelOutputWithPast:
1622
+ if (input_ids is None) ^ (inputs_embeds is not None):
1623
+ raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
1624
+
1625
+ use_cache = use_cache if use_cache is not None else self.config.use_cache
1626
+
1627
+ if inputs_embeds is None:
1628
+ inputs_embeds = self.embed_tokens(input_ids)
1629
+
1630
+ if use_cache and past_key_values is None:
1631
+ past_key_values = DynamicCache(config=self.config)
1632
+
1633
+ if cache_position is None:
1634
+ past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
1635
+ cache_position = torch.arange(
1636
+ past_seen_tokens, past_seen_tokens + inputs_embeds.shape[1], device=inputs_embeds.device
1637
+ )
1638
+
1639
+ if position_ids is None:
1640
+ position_ids = cache_position.unsqueeze(0)
1641
+
1642
+ if not isinstance(causal_mask_mapping := attention_mask, dict):
1643
+ mask_kwargs = {
1644
+ "config": self.config,
1645
+ "input_embeds": inputs_embeds,
1646
+ "attention_mask": attention_mask,
1647
+ "cache_position": cache_position,
1648
+ "past_key_values": past_key_values,
1649
+ "position_ids": position_ids,
1650
+ }
1651
+ causal_mask_mapping = {
1652
+ "full_attention": create_causal_mask(**mask_kwargs),
1653
+ }
1654
+ if self.has_sliding_layers:
1655
+ if getattr(self.config, "sliding_window", None) is None:
1656
+ raise ValueError("MiMoV2 config `sliding_window` must be set when hybrid_layer_pattern uses SWA.")
1657
+ causal_mask_mapping["sliding_window_attention"] = create_sliding_window_causal_mask(**mask_kwargs)
1658
+
1659
+ hidden_states = inputs_embeds
1660
+ position_embeddings = self.rotary_emb(hidden_states, position_ids)
1661
+ swa_position_embeddings = self.swa_rotary_emb(hidden_states, position_ids)
1662
+
1663
+ for decoder_layer in self.layers[: self.config.num_hidden_layers]:
1664
+ hidden_states = decoder_layer(
1665
+ hidden_states,
1666
+ attention_mask=causal_mask_mapping[decoder_layer.attention_type],
1667
+ position_embeddings=position_embeddings
1668
+ if decoder_layer.attention_type == "full_attention"
1669
+ else swa_position_embeddings,
1670
+ position_ids=position_ids,
1671
+ past_key_values=past_key_values,
1672
+ use_cache=use_cache,
1673
+ cache_position=cache_position,
1674
+ **kwargs,
1675
+ )
1676
+
1677
+ hidden_states = self.norm(hidden_states)
1678
+ return BaseModelOutputWithPast(
1679
+ last_hidden_state=hidden_states,
1680
+ past_key_values=past_key_values if use_cache else None,
1681
+ )
1682
+
1683
+
1684
+ class MiMoV2ForCausalLM(PreTrainedModel, GenerationMixin):
1685
+ config_class = MiMoV2Config
1686
+ model_class = MiMoV2Model
1687
+ _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"}
1688
+ _tp_plan = {"lm_head": "colwise_rep"}
1689
+ _pp_plan = {"lm_head": (["hidden_states"], ["logits"])}
1690
+ _keys_to_ignore_on_load_unexpected = [
1691
+ r"model\.(swa_)?rotary_emb\.inv_freq",
1692
+ r"model\.layers\.\d+\.self_attn\.rotary_emb\.inv_freq",
1693
+ r"model\.layers\.\d+\.self_attn\.rotary_emb\.(cos_cached|sin_cached)",
1694
+ r"model\.mtp\..*",
1695
+ ]
1696
+ _keys_to_ignore_on_load_missing = [
1697
+ r"audio_encoder\.input_local_transformer\.embed_tokens\.weight",
1698
+ ]
1699
+
1700
+ def __init__(self, config):
1701
+ super().__init__(config)
1702
+ self.model = self.model_class(config)
1703
+ self.vocab_size = config.vocab_size
1704
+ self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
1705
+
1706
+ if config.vision_config:
1707
+ self.visual = MiMoVisionTransformer(_as_namespace(config.vision_config))
1708
+ if config.audio_config:
1709
+ audio_cfg = _as_namespace(config.audio_config)
1710
+ self.speech_embeddings = _build_speech_embeddings(audio_cfg)
1711
+ self.audio_encoder = MiMoAudioEncoder(audio_cfg)
1712
+
1713
+ self.audio_tokenizer = None
1714
+ self.post_init()
1715
+
1716
+ def load_audio_tokenizer(self, path: str, device: torch.device | str | None = None, dtype: torch.dtype = torch.bfloat16):
1717
+ """Load the audio tokenizer from a directory containing config.json and model.safetensors."""
1718
+ import json
1719
+ import os
1720
+
1721
+ from safetensors.torch import load_file
1722
+
1723
+ config_path = os.path.join(path, "config.json")
1724
+ with open(config_path) as f:
1725
+ config_dict = json.load(f)
1726
+ tokenizer_config = MiMoAudioTokenizerConfig(**config_dict)
1727
+ tokenizer_model = MiMoAudioTokenizer(tokenizer_config)
1728
+
1729
+ safetensors_path = os.path.join(path, "model.safetensors")
1730
+ bin_path = os.path.join(path, "pytorch_model.bin")
1731
+ if os.path.exists(safetensors_path):
1732
+ state_dict = load_file(safetensors_path, device="cpu")
1733
+ elif os.path.exists(bin_path):
1734
+ state_dict = torch.load(bin_path, map_location="cpu", weights_only=True)
1735
+ else:
1736
+ raise FileNotFoundError(f"No model weights found in {path}")
1737
+ tokenizer_model.load_state_dict(state_dict, strict=False)
1738
+
1739
+ if device is None:
1740
+ device = next(self.parameters()).device
1741
+ tokenizer_model = tokenizer_model.to(device=device, dtype=dtype)
1742
+ tokenizer_model.eval()
1743
+ tokenizer_model.requires_grad_(False)
1744
+ self.audio_tokenizer = tokenizer_model
1745
+
1746
+ def get_input_embeddings(self):
1747
+ return self.model.embed_tokens
1748
+
1749
+ def set_input_embeddings(self, value):
1750
+ self.model.embed_tokens = value
1751
+
1752
+ def get_output_embeddings(self):
1753
+ return self.lm_head
1754
+
1755
+ def set_output_embeddings(self, new_embeddings):
1756
+ self.lm_head = new_embeddings
1757
+
1758
+ def _get_multimodal_embeds(
1759
+ self,
1760
+ input_ids: torch.Tensor,
1761
+ inputs_embeds: torch.Tensor,
1762
+ pixel_values: Optional[torch.Tensor] = None,
1763
+ image_grid_thw: Optional[torch.Tensor] = None,
1764
+ image_embeds: Optional[torch.Tensor] = None,
1765
+ video_pixel_values: Optional[torch.Tensor] = None,
1766
+ video_grid_thw: Optional[torch.Tensor] = None,
1767
+ video_embeds: Optional[torch.Tensor] = None,
1768
+ audio_codes: Optional[torch.Tensor] = None,
1769
+ audio_embeds: Optional[torch.Tensor] = None,
1770
+ ) -> torch.Tensor:
1771
+ has_image = image_embeds is not None or pixel_values is not None
1772
+ has_video = video_embeds is not None or video_pixel_values is not None
1773
+ has_audio = audio_embeds is not None or audio_codes is not None
1774
+
1775
+ if not (has_image or has_video or has_audio):
1776
+ return inputs_embeds
1777
+
1778
+ inputs_embeds = inputs_embeds.clone()
1779
+
1780
+ if has_image:
1781
+ cur_image_embeds = image_embeds if image_embeds is not None else self.visual(pixel_values=pixel_values, grid_thw=image_grid_thw)
1782
+ _replace_modal_embeddings_inplace(
1783
+ input_ids=input_ids, inputs_embeds=inputs_embeds,
1784
+ token_id=getattr(self.config, "image_token_id", None), modal_embeds=cur_image_embeds,
1785
+ )
1786
+
1787
+ if has_video:
1788
+ cur_video_embeds = video_embeds if video_embeds is not None else self.visual(pixel_values=video_pixel_values, grid_thw=video_grid_thw)
1789
+ _replace_modal_embeddings_inplace(
1790
+ input_ids=input_ids, inputs_embeds=inputs_embeds,
1791
+ token_id=getattr(self.config, "video_token_id", None), modal_embeds=cur_video_embeds,
1792
+ )
1793
+
1794
+ if has_audio:
1795
+ _replace_modal_embeddings_inplace(
1796
+ input_ids=input_ids, inputs_embeds=inputs_embeds,
1797
+ token_id=getattr(self.config, "audio_token_id", None),
1798
+ modal_embeds=self.audio_encoder(
1799
+ speech_embeddings=self.speech_embeddings, audio_codes=audio_codes, audio_embeds=audio_embeds,
1800
+ ),
1801
+ )
1802
+
1803
+ return inputs_embeds
1804
+
1805
+ @can_return_tuple
1806
+ def forward(
1807
+ self,
1808
+ input_ids: Optional[torch.LongTensor] = None,
1809
+ attention_mask: Optional[torch.Tensor] = None,
1810
+ position_ids: Optional[torch.LongTensor] = None,
1811
+ past_key_values: Optional[Cache] = None,
1812
+ inputs_embeds: Optional[torch.FloatTensor] = None,
1813
+ labels: Optional[torch.LongTensor] = None,
1814
+ use_cache: Optional[bool] = None,
1815
+ cache_position: Optional[torch.LongTensor] = None,
1816
+ logits_to_keep: Union[int, torch.Tensor] = 0,
1817
+ pixel_values: Optional[torch.Tensor] = None,
1818
+ image_grid_thw: Optional[torch.Tensor] = None,
1819
+ image_embeds: Optional[torch.Tensor] = None,
1820
+ video_pixel_values: Optional[torch.Tensor] = None,
1821
+ video_grid_thw: Optional[torch.Tensor] = None,
1822
+ video_embeds: Optional[torch.Tensor] = None,
1823
+ audio_codes: Optional[torch.Tensor] = None,
1824
+ audio_embeds: Optional[torch.Tensor] = None,
1825
+ **kwargs: Unpack[TransformersKwargs],
1826
+ ) -> CausalLMOutputWithPast:
1827
+ if inputs_embeds is None and input_ids is not None:
1828
+ inputs_embeds = self.model.get_input_embeddings()(input_ids)
1829
+ if any(x is not None for x in [pixel_values, image_embeds, video_pixel_values, video_embeds, audio_codes, audio_embeds]):
1830
+ inputs_embeds = self._get_multimodal_embeds(
1831
+ input_ids=input_ids, inputs_embeds=inputs_embeds,
1832
+ pixel_values=pixel_values, image_grid_thw=image_grid_thw, image_embeds=image_embeds,
1833
+ video_pixel_values=video_pixel_values, video_grid_thw=video_grid_thw, video_embeds=video_embeds,
1834
+ audio_codes=audio_codes, audio_embeds=audio_embeds,
1835
+ )
1836
+ input_ids = None
1837
+
1838
+ outputs: BaseModelOutputWithPast = self.model(
1839
+ input_ids=input_ids,
1840
+ attention_mask=attention_mask,
1841
+ position_ids=position_ids,
1842
+ past_key_values=past_key_values,
1843
+ inputs_embeds=inputs_embeds,
1844
+ use_cache=use_cache,
1845
+ cache_position=cache_position,
1846
+ **kwargs,
1847
+ )
1848
+
1849
+ hidden_states = outputs.last_hidden_state
1850
+ slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
1851
+ logits = self.lm_head(hidden_states[:, slice_indices, :])
1852
+
1853
+ loss = None
1854
+ if labels is not None:
1855
+ loss = self.loss_function(logits=logits, labels=labels, vocab_size=self.config.vocab_size, **kwargs)
1856
+
1857
+ return CausalLMOutputWithPast(
1858
+ loss=loss,
1859
+ logits=logits,
1860
+ past_key_values=outputs.past_key_values,
1861
+ hidden_states=outputs.hidden_states,
1862
+ attentions=outputs.attentions,
1863
+ )
1864
+
1865
+
1866
+ __all__ = [
1867
+ "MiMoAudioTokenizer",
1868
+ "MiMoAudioTokenizerConfig",
1869
+ "MiMoV2Attention",
1870
+ "MiMoV2DecoderLayer",
1871
+ "MiMoV2ForCausalLM",
1872
+ "MiMoV2MLP",
1873
+ "MiMoV2MoE",
1874
+ "MiMoV2MoEGate",
1875
+ "MiMoV2Model",
1876
+ "MiMoV2RMSNorm",
1877
+ "MiMoV2RotaryEmbedding",
1878
+ ]
tokenizer.json ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:8d4b7746684abeca0980d9b740e761c8e135cbf927feab55b377d07a6cca83e8
3
+ size 11423819