Raullen commited on
Commit
3b16649
·
verified ·
1 Parent(s): d402cb0

Upload folder using huggingface_hub

Browse files
.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,58 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ ---
2
+ license: mit
3
+ base_model: inclusionAI/Ling-3.0-tiny
4
+ library_name: mlx
5
+ tags:
6
+ - mlx
7
+ - ling
8
+ - bailing_hybrid
9
+ - moe
10
+ - apple-silicon
11
+ pipeline_tag: text-generation
12
+ ---
13
+
14
+ # Ling-3.0-tiny — MLX 4-bit
15
+
16
+ The first MLX conversion of [inclusionAI/Ling-3.0-tiny](https://huggingface.co/inclusionAI/Ling-3.0-tiny):
17
+ a 7.9B-total / **1.3B-active** sparse-MoE reasoner (128 experts, top-8 + 1 shared)
18
+ with a KDA + MLA hybrid attention stack and 131K context, MIT licensed.
19
+
20
+ **4.2 GB at 4.507 bits/weight** — it fits and runs on an 8 GB Apple Silicon Mac.
21
+
22
+ | | |
23
+ |---|---|
24
+ | Quantization | 4-bit, group size 64 (router kept 8-bit, short-conv weights fp) |
25
+ | Size on disk | 4.2 GB |
26
+ | Context | 131,072 tokens |
27
+ | Active parameters | 1.3B per token |
28
+ | License | MIT (inherited from the base model) |
29
+
30
+ ## Serve it
31
+
32
+ The `bailing_hybrid` architecture is not in upstream `mlx-lm` yet — this
33
+ checkpoint is served by [rapid-mlx](https://rapidmlx.com), which ships a
34
+ verified native implementation (reference parity 1.5e-6 against the
35
+ official modeling code):
36
+
37
+ ```bash
38
+ pip install -U rapid-mlx
39
+ rapid-mlx serve rapid-mlx/Ling-3.0-tiny-MLX-4bit
40
+ ```
41
+
42
+ You get an OpenAI-compatible server on `localhost:8000` with reasoning
43
+ (`reasoning_content`) and tool calling parsed natively — thinking is
44
+ controlled with `chat_template_kwargs: {"enable_thinking": true}` or the
45
+ model's `detailed thinking on/off` system-prompt switch.
46
+
47
+ Once `mlx-lm` gains native `bailing_hybrid` support, this checkpoint will
48
+ load there unchanged.
49
+
50
+ ## Conversion provenance
51
+
52
+ Converted with `mlx_lm.convert` (quantize=True, q_bits=4, q_group_size=64)
53
+ running rapid-mlx's vendored `bailing_hybrid` implementation
54
+ ([PR #1817](https://github.com/raullenchai/Rapid-MLX/pull/1817)), which was
55
+ verified against the official `modeling_bailing_moe_v3.py` on identical
56
+ random weights to a max logits deviation of 1.5e-6 (full prefill) /
57
+ 1.9e-6 (token-by-token incremental) before conversion. End-to-end
58
+ chat / reasoning / tool-call behaviour validated on an M2 Pro Mac mini.
chat_template.jinja ADDED
@@ -0,0 +1,130 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {#- Bailing V3 chat template -#}
2
+ {#- Supports: thinking option, tool calling -#}
3
+
4
+ {#- ==================== thinking option normalization ==================== -#}
5
+ {%- if enable_thinking is defined %}
6
+ {%- if enable_thinking %}
7
+ {%- set thinking_option = 'on' %}
8
+ {%- else %}
9
+ {%- set thinking_option = 'off' %}
10
+ {%- endif %}
11
+ {%- elif thinking_option is not defined %}
12
+ {%- set thinking_option = 'on' %}
13
+ {%- endif %}
14
+
15
+ {#- ==================== preserved thinking ==================== -#}
16
+ {% set preserved_thinking = true %}
17
+
18
+ {#- ==================== system message ==================== -#}
19
+ {{- '<role>SYSTEM</role>' }}
20
+ {%- if tools %}
21
+ {%- if messages[0].role == 'system' %}
22
+ {{- messages[0].content + '\n' }}
23
+ {%- endif %}
24
+ {{- "# Tools\n\nYou may call one or more functions to assist with the user query.\n\nYou are provided with function signatures within <tools></tools> XML tags:\n<tools>" }}
25
+ {%- for tool in tools %}
26
+ {{- "\n" }}
27
+ {{- tool | tojson }}
28
+ {%- endfor %}
29
+ {{- "\n</tools>\n\nIf none of the functions can be used, point it out. If the given question lacks the parameters required by the function, also point it out.\nIf you need to use a function, for each function call, output the function name and arguments within the following XML format:\n<tool_call>{function-name}\n<arg_key>{arg-key-1}</arg_key>\n<arg_value>{arg-value-1}</arg_value>\n<arg_key>{arg-key-2}</arg_key>\n<arg_value>{arg-value-2}</arg_value>\n...\n</tool_call>\n" }}
30
+ {%- if messages[0].role == 'system' and messages[0].content is string and ('detailed thinking on' in messages[0].content or 'detailed thinking off' in messages[0].content) %}
31
+ {{- '<|role_end|>' }}
32
+ {%- else %}
33
+ {{- 'detailed thinking ' + thinking_option + '<|role_end|>' }}
34
+ {%- endif %}
35
+ {%- else %}
36
+ {%- if messages[0].role == 'system' %}
37
+ {%- if 'detailed thinking on' in messages[0].content or 'detailed thinking off' in messages[0].content %}
38
+ {{- messages[0].content + '<|role_end|>' }}
39
+ {%- else %}
40
+ {{- messages[0].content + '\n' }}
41
+ {{- 'detailed thinking ' + thinking_option + '<|role_end|>' }}
42
+ {%- endif %}
43
+ {% else %}
44
+ {{- 'detailed thinking ' + thinking_option + '<|role_end|>' }}
45
+ {%- endif %}
46
+ {%- endif %}
47
+ {%- set ns = namespace(multi_step_tool=true, last_query_index=messages|length - 1) %}
48
+ {%- for message in messages[::-1] %}
49
+ {%- set index = (messages|length - 1) - loop.index0 %}
50
+ {%- if ns.multi_step_tool and message.role == "user" and message.content is string and not(message.content.startswith('<tool_response>') and message.content.endswith('</tool_response>')) %}
51
+ {%- set ns.multi_step_tool = false %}
52
+ {%- set ns.last_query_index = index %}
53
+ {%- endif %}
54
+ {%- endfor %}
55
+ {%- for message in messages %}
56
+ {%- if message.content is string %}
57
+ {%- set content = message.content %}
58
+ {%- else %}
59
+ {%- set content = '' %}
60
+ {%- endif %}
61
+ {%- if message.role == "user" %}
62
+ {{- '<role>HUMAN</role>' + message.content + '<|role_end|>' }}
63
+ {%- elif message.role == "system" and not loop.first %}
64
+ {{- '<role>SYSTEM</role>' + message.content + '<|role_end|>' }}
65
+ {%- elif message.role == "assistant" %}
66
+ {%- set reasoning_content = '' %}
67
+ {%- if message.reasoning_content is string and message.reasoning_content != '' %}
68
+ {%- set reasoning_content = message.reasoning_content %}
69
+ {%- else %}
70
+ {%- if '</think>' in content %}
71
+ {%- set reasoning_content = content.split('</think>')[0].rstrip('\n').split('<think>')[-1].lstrip('\n') %}
72
+ {%- set content = content.split('</think>')[-1].lstrip('\n') %}
73
+ {%- endif %}
74
+ {%- endif %}
75
+ {%- if preserved_thinking or loop.index0 > ns.last_query_index %}
76
+ {%- if reasoning_content != '' %}
77
+ {{- '<role>ASSISTANT</role>' + '\n<think>' + reasoning_content.strip('\n') + '</think>' + content.lstrip('\n') }}
78
+ {%- else %}
79
+ {{- '<role>ASSISTANT</role>\n<think></think>' + content }}
80
+ {%- endif %}
81
+ {%- else %}
82
+ {{- '<role>ASSISTANT</role>\n<think></think>' + content }}
83
+ {%- endif %}
84
+ {%- if message.tool_calls %}
85
+ {%- for tool_call in message.tool_calls %}
86
+ {%- if (loop.first and content) or (not loop.first) %}
87
+ {{- '\n' }}
88
+ {%- endif %}
89
+ {%- set tc = tool_call %}
90
+ {%- if tool_call.function %}
91
+ {%- set tc = tool_call.function %}
92
+ {%- endif %}
93
+ {{- '<tool_call>' + tc.name }}
94
+ {% set _args = tc.arguments %}
95
+ {%- for k, v in _args.items() %}
96
+ {{- '<arg_key>' + k + '</arg_key>' }}
97
+ {{- '\n<arg_value>' }}
98
+ {%- if v is string %}
99
+ {{- v }}
100
+ {%- else %}
101
+ {{- v | tojson(ensure_ascii=False) }}
102
+ {%- endif %}
103
+ {{- '</arg_value>' }}
104
+ {%- endfor %}
105
+ {{- '\n</tool_call>' }}
106
+ {%- endfor %}
107
+ {%- endif %}
108
+ {{- '<|role_end|>' }}
109
+ {%- elif message.role == "tool" %}
110
+ {%- if loop.first or (messages[loop.index0 - 1].role != "tool") %}
111
+ {{- '<role>OBSERVATION</role>' }}
112
+ {%- endif %}
113
+ {{- '\n<tool_response>\n' }}
114
+ {{- content }}
115
+ {{- '\n</tool_response>' }}
116
+ {%- if loop.last or (messages[loop.index0 + 1].role != "tool") %}
117
+ {{- '<|role_end|>' }}
118
+ {%- endif %}
119
+ {%- endif %}
120
+ {%- endfor %}
121
+
122
+ {#- ==================== generation prompt ==================== -#}
123
+ {%- if add_generation_prompt %}
124
+ {{- '<role>ASSISTANT</role>' }}
125
+ {%- if thinking_option == 'on' %}
126
+ {{- '\n<think>' }}
127
+ {%- elif thinking_option == 'off' %}
128
+ {{- '\n<think></think>' }}
129
+ {%- endif %}
130
+ {%- endif %}
config.json ADDED
@@ -0,0 +1,93 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "architectures": [
3
+ "BailingMoeV3ForCausalLM"
4
+ ],
5
+ "attention_dropout": 0.0,
6
+ "auto_map": {
7
+ "AutoConfig": "configuration_bailing_moe_v3.BailingMoeV3Config",
8
+ "AutoModel": "modeling_bailing_moe_v3.BailingMoeV3Model",
9
+ "AutoModelForCausalLM": "modeling_bailing_moe_v3.BailingMoeV3ForCausalLM"
10
+ },
11
+ "embedding_dropout": 0.0,
12
+ "eos_token_id": 156895,
13
+ "expert_swiglu_limit_list": null,
14
+ "first_k_dense_replace": 1,
15
+ "gated_attention_proj_granularity_type": "head_wise",
16
+ "group_norm_size": 1,
17
+ "head_dim": 128,
18
+ "hidden_act": "silu",
19
+ "hidden_size": 1536,
20
+ "initializer_range": 0.02,
21
+ "intermediate_size": 4608,
22
+ "kda_lower_bound": -5,
23
+ "kda_safe_gate": true,
24
+ "kv_lora_rank": 512,
25
+ "layer_group_size": 4,
26
+ "linear_silu": true,
27
+ "max_position_embeddings": 131072,
28
+ "max_window_layers": 20,
29
+ "model_type": "bailing_hybrid",
30
+ "moe_intermediate_size": 512,
31
+ "moe_router_enable_expert_bias": true,
32
+ "moe_shared_expert_intermediate_size": 512,
33
+ "mtp_loss_scaling_factor": 0,
34
+ "mtp_use_kda": false,
35
+ "n_group": 8,
36
+ "no_kda_lora": true,
37
+ "norm_topk_prob": true,
38
+ "num_attention_heads": 16,
39
+ "num_experts": 128,
40
+ "num_experts_per_tok": 8,
41
+ "num_hidden_layers": 24,
42
+ "num_key_value_heads": 16,
43
+ "num_kv_heads_for_linear_attn": 0,
44
+ "num_nextn_predict_layers": 0,
45
+ "num_shared_experts": 1,
46
+ "output_dropout": 0.0,
47
+ "output_router_logits": false,
48
+ "pad_token_id": 156892,
49
+ "partial_rotary_factor": 0.5,
50
+ "q_lora_rank": 256,
51
+ "qk_head_dim": 192,
52
+ "qk_nope_head_dim": 128,
53
+ "qk_rope_head_dim": 64,
54
+ "quantization": {
55
+ "group_size": 64,
56
+ "bits": 4,
57
+ "mode": "affine"
58
+ },
59
+ "quantization_config": {
60
+ "group_size": 64,
61
+ "bits": 4,
62
+ "mode": "affine"
63
+ },
64
+ "rms_norm_eps": 1e-06,
65
+ "rope_interleave": true,
66
+ "rope_scaling": null,
67
+ "rope_theta": 6000000,
68
+ "rotary_dim": 64,
69
+ "routed_scaling_factor": 2.5,
70
+ "router_dtype": "fp32",
71
+ "scale_router_input": false,
72
+ "score_function": "sigmoid",
73
+ "scoring_func": "sigmoid",
74
+ "seq_aux": true,
75
+ "share_expert_swiglu_limit_list": null,
76
+ "short_conv_kernel_size": 4,
77
+ "tie_word_embeddings": false,
78
+ "topk_group": 4,
79
+ "topk_method": "noaux_tc",
80
+ "torch_dtype": "bfloat16",
81
+ "transformers_version": "4.45.0",
82
+ "up_proj_norm": false,
83
+ "use_bias": false,
84
+ "use_cache": true,
85
+ "use_kda_lora": false,
86
+ "use_mla_nope": false,
87
+ "use_nGPT": false,
88
+ "use_qk_norm": true,
89
+ "use_qkv_bias": false,
90
+ "v_head_dim": 128,
91
+ "value_norm": false,
92
+ "vocab_size": 157184
93
+ }
configuration_bailing_moe_v3.py ADDED
@@ -0,0 +1,126 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Bailing MoE V2 model configuration"""
2
+
3
+ from transformers.configuration_utils import PretrainedConfig
4
+
5
+
6
+ class BailingMoeV3Config(PretrainedConfig):
7
+
8
+ def __init__(
9
+ self,
10
+ vocab_size=157184,
11
+ hidden_size=2048,
12
+ intermediate_size=5120,
13
+ num_hidden_layers=20,
14
+ num_attention_heads=16,
15
+ num_key_value_heads=4,
16
+ hidden_act="silu",
17
+ use_qkv_bias=False, # bailing only
18
+ use_bias=False, # bailing only
19
+ rms_norm_eps=1e-06,
20
+ tie_word_embeddings=False, # PretrainedConfig key, here change default value.
21
+ embedding_dropout=0.0,
22
+ attention_dropout=0.0,
23
+ output_dropout=0.0,
24
+ initializer_range=0.02,
25
+ max_position_embeddings=32768,
26
+ rope_theta=600000.0,
27
+ use_cache=True,
28
+ max_window_layers=20,
29
+ rope_scaling=None,
30
+ pad_token_id=156892,
31
+ eos_token_id=156892,
32
+ num_experts=256,
33
+ num_shared_experts=1,
34
+ num_experts_per_tok=8,
35
+ n_group=8,
36
+ topk_group=4,
37
+ moe_intermediate_size=512,
38
+ moe_shared_expert_intermediate_size=512,
39
+ first_k_dense_replace=1,
40
+ head_dim=128,
41
+ output_router_logits=False,
42
+ use_qk_norm=True,
43
+ num_nextn_predict_layers=0,
44
+ mtp_loss_scaling_factor=0,
45
+ moe_router_enable_expert_bias=True,
46
+ routed_scaling_factor=1.0,
47
+ layer_group_size=5,
48
+ kv_lora_rank=512,
49
+ q_lora_rank=None,
50
+ qk_rope_head_dim=64,
51
+ v_head_dim=128,
52
+ qk_nope_head_dim=128,
53
+ rope_interleave=True,
54
+ score_function="sigmoid",
55
+ scoring_func="sigmoid",
56
+ seq_aux=True,
57
+ topk_method="noaux_tc",
58
+ router_dtype="fp32",
59
+ gated_attention_proj_granularity_type=None,
60
+ no_kda_lora=False,
61
+ kda_safe_gate=False,
62
+ kda_lower_bound=None,
63
+ short_conv_kernel_size=4,
64
+ **kwargs,
65
+ ):
66
+ self.num_hidden_layers = num_hidden_layers
67
+ self.vocab_size = vocab_size
68
+ self.hidden_size = hidden_size
69
+ self.intermediate_size = intermediate_size
70
+ self.num_attention_heads = num_attention_heads
71
+ self.num_key_value_heads = num_key_value_heads
72
+ self.hidden_act = hidden_act
73
+ self.use_qkv_bias = use_qkv_bias
74
+ self.use_bias = use_bias
75
+ self.rms_norm_eps = rms_norm_eps
76
+ self.embedding_dropout = embedding_dropout
77
+ self.attention_dropout = attention_dropout
78
+ self.output_dropout = output_dropout
79
+ self.num_nextn_predict_layers = num_nextn_predict_layers
80
+ self.mtp_loss_scaling_factor = mtp_loss_scaling_factor
81
+ self.initializer_range = initializer_range
82
+ self.max_position_embeddings = max_position_embeddings
83
+ self.rope_theta = rope_theta
84
+ self.use_cache = use_cache
85
+ self.max_window_layers = max_window_layers
86
+ self.head_dim = head_dim or self.hidden_size // self.num_attention_heads
87
+ self.rope_scaling = rope_scaling
88
+ self.use_qk_norm = use_qk_norm
89
+ self.moe_router_enable_expert_bias = moe_router_enable_expert_bias
90
+ self.routed_scaling_factor = routed_scaling_factor
91
+
92
+ # MoE configs
93
+ self.num_experts = num_experts
94
+ self.num_shared_experts = num_shared_experts
95
+ self.num_experts_per_tok = num_experts_per_tok
96
+ self.n_group = n_group
97
+ self.topk_group = topk_group
98
+ self.moe_intermediate_size = moe_intermediate_size
99
+ self.moe_shared_expert_intermediate_size = moe_shared_expert_intermediate_size
100
+ self.first_k_dense_replace = first_k_dense_replace
101
+ self.output_router_logits = output_router_logits
102
+
103
+ # Linear configs
104
+ self.layer_group_size = layer_group_size
105
+ # mla
106
+ self.kv_lora_rank = kv_lora_rank
107
+ self.q_lora_rank = q_lora_rank
108
+ self.qk_rope_head_dim = qk_rope_head_dim
109
+
110
+ self.score_function = score_function
111
+ self.scoring_func = scoring_func
112
+ self.seq_aux = seq_aux
113
+ self.topk_method = topk_method
114
+ self.v_head_dim = v_head_dim
115
+ self.qk_nope_head_dim = qk_nope_head_dim
116
+ self.qk_head_dim = qk_nope_head_dim + qk_rope_head_dim
117
+ self.rope_interleave = rope_interleave
118
+ self.router_dtype = router_dtype
119
+ self.gated_attention_proj_granularity_type = gated_attention_proj_granularity_type
120
+ self.no_kda_lora = no_kda_lora
121
+ self.kda_safe_gate = kda_safe_gate
122
+ self.kda_lower_bound = kda_lower_bound
123
+ self.short_conv_kernel_size = short_conv_kernel_size
124
+ super().__init__(
125
+ pad_token_id=pad_token_id, eos_token_id=eos_token_id, tie_word_embeddings=tie_word_embeddings, **kwargs
126
+ )
generation_config.json ADDED
@@ -0,0 +1,5 @@
 
 
 
 
 
 
1
+ {
2
+ "temperature": 1.0,
3
+ "top_p": 0.95,
4
+ "top_k": 20
5
+ }
model.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:be523d3d361ad38d85e48f33bdbd942b9512493d0a2746ec4e962e0a3183243b
3
+ size 4447478988
model.safetensors.index.json ADDED
@@ -0,0 +1,1138 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "metadata": {
3
+ "total_size": 4447345472,
4
+ "total_parameters": 7893355648
5
+ },
6
+ "weight_map": {
7
+ "lm_head.biases": "model.safetensors",
8
+ "lm_head.scales": "model.safetensors",
9
+ "lm_head.weight": "model.safetensors",
10
+ "model.layers.0.attention.A_log": "model.safetensors",
11
+ "model.layers.0.attention.b_proj.biases": "model.safetensors",
12
+ "model.layers.0.attention.b_proj.scales": "model.safetensors",
13
+ "model.layers.0.attention.b_proj.weight": "model.safetensors",
14
+ "model.layers.0.attention.dt_bias": "model.safetensors",
15
+ "model.layers.0.attention.f_proj.biases": "model.safetensors",
16
+ "model.layers.0.attention.f_proj.scales": "model.safetensors",
17
+ "model.layers.0.attention.f_proj.weight": "model.safetensors",
18
+ "model.layers.0.attention.g_proj.biases": "model.safetensors",
19
+ "model.layers.0.attention.g_proj.scales": "model.safetensors",
20
+ "model.layers.0.attention.g_proj.weight": "model.safetensors",
21
+ "model.layers.0.attention.k_conv1d.conv.weight": "model.safetensors",
22
+ "model.layers.0.attention.k_proj.biases": "model.safetensors",
23
+ "model.layers.0.attention.k_proj.scales": "model.safetensors",
24
+ "model.layers.0.attention.k_proj.weight": "model.safetensors",
25
+ "model.layers.0.attention.o_norm.weight": "model.safetensors",
26
+ "model.layers.0.attention.o_proj.biases": "model.safetensors",
27
+ "model.layers.0.attention.o_proj.scales": "model.safetensors",
28
+ "model.layers.0.attention.o_proj.weight": "model.safetensors",
29
+ "model.layers.0.attention.q_conv1d.conv.weight": "model.safetensors",
30
+ "model.layers.0.attention.q_proj.biases": "model.safetensors",
31
+ "model.layers.0.attention.q_proj.scales": "model.safetensors",
32
+ "model.layers.0.attention.q_proj.weight": "model.safetensors",
33
+ "model.layers.0.attention.v_conv1d.conv.weight": "model.safetensors",
34
+ "model.layers.0.attention.v_proj.biases": "model.safetensors",
35
+ "model.layers.0.attention.v_proj.scales": "model.safetensors",
36
+ "model.layers.0.attention.v_proj.weight": "model.safetensors",
37
+ "model.layers.0.input_layernorm.weight": "model.safetensors",
38
+ "model.layers.0.mlp.down_proj.biases": "model.safetensors",
39
+ "model.layers.0.mlp.down_proj.scales": "model.safetensors",
40
+ "model.layers.0.mlp.down_proj.weight": "model.safetensors",
41
+ "model.layers.0.mlp.gate_proj.biases": "model.safetensors",
42
+ "model.layers.0.mlp.gate_proj.scales": "model.safetensors",
43
+ "model.layers.0.mlp.gate_proj.weight": "model.safetensors",
44
+ "model.layers.0.mlp.up_proj.biases": "model.safetensors",
45
+ "model.layers.0.mlp.up_proj.scales": "model.safetensors",
46
+ "model.layers.0.mlp.up_proj.weight": "model.safetensors",
47
+ "model.layers.0.post_attention_layernorm.weight": "model.safetensors",
48
+ "model.layers.1.attention.A_log": "model.safetensors",
49
+ "model.layers.1.attention.b_proj.biases": "model.safetensors",
50
+ "model.layers.1.attention.b_proj.scales": "model.safetensors",
51
+ "model.layers.1.attention.b_proj.weight": "model.safetensors",
52
+ "model.layers.1.attention.dt_bias": "model.safetensors",
53
+ "model.layers.1.attention.f_proj.biases": "model.safetensors",
54
+ "model.layers.1.attention.f_proj.scales": "model.safetensors",
55
+ "model.layers.1.attention.f_proj.weight": "model.safetensors",
56
+ "model.layers.1.attention.g_proj.biases": "model.safetensors",
57
+ "model.layers.1.attention.g_proj.scales": "model.safetensors",
58
+ "model.layers.1.attention.g_proj.weight": "model.safetensors",
59
+ "model.layers.1.attention.k_conv1d.conv.weight": "model.safetensors",
60
+ "model.layers.1.attention.k_proj.biases": "model.safetensors",
61
+ "model.layers.1.attention.k_proj.scales": "model.safetensors",
62
+ "model.layers.1.attention.k_proj.weight": "model.safetensors",
63
+ "model.layers.1.attention.o_norm.weight": "model.safetensors",
64
+ "model.layers.1.attention.o_proj.biases": "model.safetensors",
65
+ "model.layers.1.attention.o_proj.scales": "model.safetensors",
66
+ "model.layers.1.attention.o_proj.weight": "model.safetensors",
67
+ "model.layers.1.attention.q_conv1d.conv.weight": "model.safetensors",
68
+ "model.layers.1.attention.q_proj.biases": "model.safetensors",
69
+ "model.layers.1.attention.q_proj.scales": "model.safetensors",
70
+ "model.layers.1.attention.q_proj.weight": "model.safetensors",
71
+ "model.layers.1.attention.v_conv1d.conv.weight": "model.safetensors",
72
+ "model.layers.1.attention.v_proj.biases": "model.safetensors",
73
+ "model.layers.1.attention.v_proj.scales": "model.safetensors",
74
+ "model.layers.1.attention.v_proj.weight": "model.safetensors",
75
+ "model.layers.1.input_layernorm.weight": "model.safetensors",
76
+ "model.layers.1.mlp.experts.down_proj.biases": "model.safetensors",
77
+ "model.layers.1.mlp.experts.down_proj.scales": "model.safetensors",
78
+ "model.layers.1.mlp.experts.down_proj.weight": "model.safetensors",
79
+ "model.layers.1.mlp.experts.gate_proj.biases": "model.safetensors",
80
+ "model.layers.1.mlp.experts.gate_proj.scales": "model.safetensors",
81
+ "model.layers.1.mlp.experts.gate_proj.weight": "model.safetensors",
82
+ "model.layers.1.mlp.experts.up_proj.biases": "model.safetensors",
83
+ "model.layers.1.mlp.experts.up_proj.scales": "model.safetensors",
84
+ "model.layers.1.mlp.experts.up_proj.weight": "model.safetensors",
85
+ "model.layers.1.mlp.gate.expert_bias": "model.safetensors",
86
+ "model.layers.1.mlp.gate.weight": "model.safetensors",
87
+ "model.layers.1.mlp.shared_experts.down_proj.biases": "model.safetensors",
88
+ "model.layers.1.mlp.shared_experts.down_proj.scales": "model.safetensors",
89
+ "model.layers.1.mlp.shared_experts.down_proj.weight": "model.safetensors",
90
+ "model.layers.1.mlp.shared_experts.gate_proj.biases": "model.safetensors",
91
+ "model.layers.1.mlp.shared_experts.gate_proj.scales": "model.safetensors",
92
+ "model.layers.1.mlp.shared_experts.gate_proj.weight": "model.safetensors",
93
+ "model.layers.1.mlp.shared_experts.up_proj.biases": "model.safetensors",
94
+ "model.layers.1.mlp.shared_experts.up_proj.scales": "model.safetensors",
95
+ "model.layers.1.mlp.shared_experts.up_proj.weight": "model.safetensors",
96
+ "model.layers.1.post_attention_layernorm.weight": "model.safetensors",
97
+ "model.layers.10.attention.A_log": "model.safetensors",
98
+ "model.layers.10.attention.b_proj.biases": "model.safetensors",
99
+ "model.layers.10.attention.b_proj.scales": "model.safetensors",
100
+ "model.layers.10.attention.b_proj.weight": "model.safetensors",
101
+ "model.layers.10.attention.dt_bias": "model.safetensors",
102
+ "model.layers.10.attention.f_proj.biases": "model.safetensors",
103
+ "model.layers.10.attention.f_proj.scales": "model.safetensors",
104
+ "model.layers.10.attention.f_proj.weight": "model.safetensors",
105
+ "model.layers.10.attention.g_proj.biases": "model.safetensors",
106
+ "model.layers.10.attention.g_proj.scales": "model.safetensors",
107
+ "model.layers.10.attention.g_proj.weight": "model.safetensors",
108
+ "model.layers.10.attention.k_conv1d.conv.weight": "model.safetensors",
109
+ "model.layers.10.attention.k_proj.biases": "model.safetensors",
110
+ "model.layers.10.attention.k_proj.scales": "model.safetensors",
111
+ "model.layers.10.attention.k_proj.weight": "model.safetensors",
112
+ "model.layers.10.attention.o_norm.weight": "model.safetensors",
113
+ "model.layers.10.attention.o_proj.biases": "model.safetensors",
114
+ "model.layers.10.attention.o_proj.scales": "model.safetensors",
115
+ "model.layers.10.attention.o_proj.weight": "model.safetensors",
116
+ "model.layers.10.attention.q_conv1d.conv.weight": "model.safetensors",
117
+ "model.layers.10.attention.q_proj.biases": "model.safetensors",
118
+ "model.layers.10.attention.q_proj.scales": "model.safetensors",
119
+ "model.layers.10.attention.q_proj.weight": "model.safetensors",
120
+ "model.layers.10.attention.v_conv1d.conv.weight": "model.safetensors",
121
+ "model.layers.10.attention.v_proj.biases": "model.safetensors",
122
+ "model.layers.10.attention.v_proj.scales": "model.safetensors",
123
+ "model.layers.10.attention.v_proj.weight": "model.safetensors",
124
+ "model.layers.10.input_layernorm.weight": "model.safetensors",
125
+ "model.layers.10.mlp.experts.down_proj.biases": "model.safetensors",
126
+ "model.layers.10.mlp.experts.down_proj.scales": "model.safetensors",
127
+ "model.layers.10.mlp.experts.down_proj.weight": "model.safetensors",
128
+ "model.layers.10.mlp.experts.gate_proj.biases": "model.safetensors",
129
+ "model.layers.10.mlp.experts.gate_proj.scales": "model.safetensors",
130
+ "model.layers.10.mlp.experts.gate_proj.weight": "model.safetensors",
131
+ "model.layers.10.mlp.experts.up_proj.biases": "model.safetensors",
132
+ "model.layers.10.mlp.experts.up_proj.scales": "model.safetensors",
133
+ "model.layers.10.mlp.experts.up_proj.weight": "model.safetensors",
134
+ "model.layers.10.mlp.gate.expert_bias": "model.safetensors",
135
+ "model.layers.10.mlp.gate.weight": "model.safetensors",
136
+ "model.layers.10.mlp.shared_experts.down_proj.biases": "model.safetensors",
137
+ "model.layers.10.mlp.shared_experts.down_proj.scales": "model.safetensors",
138
+ "model.layers.10.mlp.shared_experts.down_proj.weight": "model.safetensors",
139
+ "model.layers.10.mlp.shared_experts.gate_proj.biases": "model.safetensors",
140
+ "model.layers.10.mlp.shared_experts.gate_proj.scales": "model.safetensors",
141
+ "model.layers.10.mlp.shared_experts.gate_proj.weight": "model.safetensors",
142
+ "model.layers.10.mlp.shared_experts.up_proj.biases": "model.safetensors",
143
+ "model.layers.10.mlp.shared_experts.up_proj.scales": "model.safetensors",
144
+ "model.layers.10.mlp.shared_experts.up_proj.weight": "model.safetensors",
145
+ "model.layers.10.post_attention_layernorm.weight": "model.safetensors",
146
+ "model.layers.11.attention.dense.biases": "model.safetensors",
147
+ "model.layers.11.attention.dense.scales": "model.safetensors",
148
+ "model.layers.11.attention.dense.weight": "model.safetensors",
149
+ "model.layers.11.attention.g_proj.biases": "model.safetensors",
150
+ "model.layers.11.attention.g_proj.scales": "model.safetensors",
151
+ "model.layers.11.attention.g_proj.weight": "model.safetensors",
152
+ "model.layers.11.attention.kv_a_layernorm.weight": "model.safetensors",
153
+ "model.layers.11.attention.kv_a_proj_with_mqa.biases": "model.safetensors",
154
+ "model.layers.11.attention.kv_a_proj_with_mqa.scales": "model.safetensors",
155
+ "model.layers.11.attention.kv_a_proj_with_mqa.weight": "model.safetensors",
156
+ "model.layers.11.attention.kv_b_proj.biases": "model.safetensors",
157
+ "model.layers.11.attention.kv_b_proj.scales": "model.safetensors",
158
+ "model.layers.11.attention.kv_b_proj.weight": "model.safetensors",
159
+ "model.layers.11.attention.q_a_layernorm.weight": "model.safetensors",
160
+ "model.layers.11.attention.q_a_proj.biases": "model.safetensors",
161
+ "model.layers.11.attention.q_a_proj.scales": "model.safetensors",
162
+ "model.layers.11.attention.q_a_proj.weight": "model.safetensors",
163
+ "model.layers.11.attention.q_b_proj.biases": "model.safetensors",
164
+ "model.layers.11.attention.q_b_proj.scales": "model.safetensors",
165
+ "model.layers.11.attention.q_b_proj.weight": "model.safetensors",
166
+ "model.layers.11.input_layernorm.weight": "model.safetensors",
167
+ "model.layers.11.mlp.experts.down_proj.biases": "model.safetensors",
168
+ "model.layers.11.mlp.experts.down_proj.scales": "model.safetensors",
169
+ "model.layers.11.mlp.experts.down_proj.weight": "model.safetensors",
170
+ "model.layers.11.mlp.experts.gate_proj.biases": "model.safetensors",
171
+ "model.layers.11.mlp.experts.gate_proj.scales": "model.safetensors",
172
+ "model.layers.11.mlp.experts.gate_proj.weight": "model.safetensors",
173
+ "model.layers.11.mlp.experts.up_proj.biases": "model.safetensors",
174
+ "model.layers.11.mlp.experts.up_proj.scales": "model.safetensors",
175
+ "model.layers.11.mlp.experts.up_proj.weight": "model.safetensors",
176
+ "model.layers.11.mlp.gate.expert_bias": "model.safetensors",
177
+ "model.layers.11.mlp.gate.weight": "model.safetensors",
178
+ "model.layers.11.mlp.shared_experts.down_proj.biases": "model.safetensors",
179
+ "model.layers.11.mlp.shared_experts.down_proj.scales": "model.safetensors",
180
+ "model.layers.11.mlp.shared_experts.down_proj.weight": "model.safetensors",
181
+ "model.layers.11.mlp.shared_experts.gate_proj.biases": "model.safetensors",
182
+ "model.layers.11.mlp.shared_experts.gate_proj.scales": "model.safetensors",
183
+ "model.layers.11.mlp.shared_experts.gate_proj.weight": "model.safetensors",
184
+ "model.layers.11.mlp.shared_experts.up_proj.biases": "model.safetensors",
185
+ "model.layers.11.mlp.shared_experts.up_proj.scales": "model.safetensors",
186
+ "model.layers.11.mlp.shared_experts.up_proj.weight": "model.safetensors",
187
+ "model.layers.11.post_attention_layernorm.weight": "model.safetensors",
188
+ "model.layers.12.attention.A_log": "model.safetensors",
189
+ "model.layers.12.attention.b_proj.biases": "model.safetensors",
190
+ "model.layers.12.attention.b_proj.scales": "model.safetensors",
191
+ "model.layers.12.attention.b_proj.weight": "model.safetensors",
192
+ "model.layers.12.attention.dt_bias": "model.safetensors",
193
+ "model.layers.12.attention.f_proj.biases": "model.safetensors",
194
+ "model.layers.12.attention.f_proj.scales": "model.safetensors",
195
+ "model.layers.12.attention.f_proj.weight": "model.safetensors",
196
+ "model.layers.12.attention.g_proj.biases": "model.safetensors",
197
+ "model.layers.12.attention.g_proj.scales": "model.safetensors",
198
+ "model.layers.12.attention.g_proj.weight": "model.safetensors",
199
+ "model.layers.12.attention.k_conv1d.conv.weight": "model.safetensors",
200
+ "model.layers.12.attention.k_proj.biases": "model.safetensors",
201
+ "model.layers.12.attention.k_proj.scales": "model.safetensors",
202
+ "model.layers.12.attention.k_proj.weight": "model.safetensors",
203
+ "model.layers.12.attention.o_norm.weight": "model.safetensors",
204
+ "model.layers.12.attention.o_proj.biases": "model.safetensors",
205
+ "model.layers.12.attention.o_proj.scales": "model.safetensors",
206
+ "model.layers.12.attention.o_proj.weight": "model.safetensors",
207
+ "model.layers.12.attention.q_conv1d.conv.weight": "model.safetensors",
208
+ "model.layers.12.attention.q_proj.biases": "model.safetensors",
209
+ "model.layers.12.attention.q_proj.scales": "model.safetensors",
210
+ "model.layers.12.attention.q_proj.weight": "model.safetensors",
211
+ "model.layers.12.attention.v_conv1d.conv.weight": "model.safetensors",
212
+ "model.layers.12.attention.v_proj.biases": "model.safetensors",
213
+ "model.layers.12.attention.v_proj.scales": "model.safetensors",
214
+ "model.layers.12.attention.v_proj.weight": "model.safetensors",
215
+ "model.layers.12.input_layernorm.weight": "model.safetensors",
216
+ "model.layers.12.mlp.experts.down_proj.biases": "model.safetensors",
217
+ "model.layers.12.mlp.experts.down_proj.scales": "model.safetensors",
218
+ "model.layers.12.mlp.experts.down_proj.weight": "model.safetensors",
219
+ "model.layers.12.mlp.experts.gate_proj.biases": "model.safetensors",
220
+ "model.layers.12.mlp.experts.gate_proj.scales": "model.safetensors",
221
+ "model.layers.12.mlp.experts.gate_proj.weight": "model.safetensors",
222
+ "model.layers.12.mlp.experts.up_proj.biases": "model.safetensors",
223
+ "model.layers.12.mlp.experts.up_proj.scales": "model.safetensors",
224
+ "model.layers.12.mlp.experts.up_proj.weight": "model.safetensors",
225
+ "model.layers.12.mlp.gate.expert_bias": "model.safetensors",
226
+ "model.layers.12.mlp.gate.weight": "model.safetensors",
227
+ "model.layers.12.mlp.shared_experts.down_proj.biases": "model.safetensors",
228
+ "model.layers.12.mlp.shared_experts.down_proj.scales": "model.safetensors",
229
+ "model.layers.12.mlp.shared_experts.down_proj.weight": "model.safetensors",
230
+ "model.layers.12.mlp.shared_experts.gate_proj.biases": "model.safetensors",
231
+ "model.layers.12.mlp.shared_experts.gate_proj.scales": "model.safetensors",
232
+ "model.layers.12.mlp.shared_experts.gate_proj.weight": "model.safetensors",
233
+ "model.layers.12.mlp.shared_experts.up_proj.biases": "model.safetensors",
234
+ "model.layers.12.mlp.shared_experts.up_proj.scales": "model.safetensors",
235
+ "model.layers.12.mlp.shared_experts.up_proj.weight": "model.safetensors",
236
+ "model.layers.12.post_attention_layernorm.weight": "model.safetensors",
237
+ "model.layers.13.attention.A_log": "model.safetensors",
238
+ "model.layers.13.attention.b_proj.biases": "model.safetensors",
239
+ "model.layers.13.attention.b_proj.scales": "model.safetensors",
240
+ "model.layers.13.attention.b_proj.weight": "model.safetensors",
241
+ "model.layers.13.attention.dt_bias": "model.safetensors",
242
+ "model.layers.13.attention.f_proj.biases": "model.safetensors",
243
+ "model.layers.13.attention.f_proj.scales": "model.safetensors",
244
+ "model.layers.13.attention.f_proj.weight": "model.safetensors",
245
+ "model.layers.13.attention.g_proj.biases": "model.safetensors",
246
+ "model.layers.13.attention.g_proj.scales": "model.safetensors",
247
+ "model.layers.13.attention.g_proj.weight": "model.safetensors",
248
+ "model.layers.13.attention.k_conv1d.conv.weight": "model.safetensors",
249
+ "model.layers.13.attention.k_proj.biases": "model.safetensors",
250
+ "model.layers.13.attention.k_proj.scales": "model.safetensors",
251
+ "model.layers.13.attention.k_proj.weight": "model.safetensors",
252
+ "model.layers.13.attention.o_norm.weight": "model.safetensors",
253
+ "model.layers.13.attention.o_proj.biases": "model.safetensors",
254
+ "model.layers.13.attention.o_proj.scales": "model.safetensors",
255
+ "model.layers.13.attention.o_proj.weight": "model.safetensors",
256
+ "model.layers.13.attention.q_conv1d.conv.weight": "model.safetensors",
257
+ "model.layers.13.attention.q_proj.biases": "model.safetensors",
258
+ "model.layers.13.attention.q_proj.scales": "model.safetensors",
259
+ "model.layers.13.attention.q_proj.weight": "model.safetensors",
260
+ "model.layers.13.attention.v_conv1d.conv.weight": "model.safetensors",
261
+ "model.layers.13.attention.v_proj.biases": "model.safetensors",
262
+ "model.layers.13.attention.v_proj.scales": "model.safetensors",
263
+ "model.layers.13.attention.v_proj.weight": "model.safetensors",
264
+ "model.layers.13.input_layernorm.weight": "model.safetensors",
265
+ "model.layers.13.mlp.experts.down_proj.biases": "model.safetensors",
266
+ "model.layers.13.mlp.experts.down_proj.scales": "model.safetensors",
267
+ "model.layers.13.mlp.experts.down_proj.weight": "model.safetensors",
268
+ "model.layers.13.mlp.experts.gate_proj.biases": "model.safetensors",
269
+ "model.layers.13.mlp.experts.gate_proj.scales": "model.safetensors",
270
+ "model.layers.13.mlp.experts.gate_proj.weight": "model.safetensors",
271
+ "model.layers.13.mlp.experts.up_proj.biases": "model.safetensors",
272
+ "model.layers.13.mlp.experts.up_proj.scales": "model.safetensors",
273
+ "model.layers.13.mlp.experts.up_proj.weight": "model.safetensors",
274
+ "model.layers.13.mlp.gate.expert_bias": "model.safetensors",
275
+ "model.layers.13.mlp.gate.weight": "model.safetensors",
276
+ "model.layers.13.mlp.shared_experts.down_proj.biases": "model.safetensors",
277
+ "model.layers.13.mlp.shared_experts.down_proj.scales": "model.safetensors",
278
+ "model.layers.13.mlp.shared_experts.down_proj.weight": "model.safetensors",
279
+ "model.layers.13.mlp.shared_experts.gate_proj.biases": "model.safetensors",
280
+ "model.layers.13.mlp.shared_experts.gate_proj.scales": "model.safetensors",
281
+ "model.layers.13.mlp.shared_experts.gate_proj.weight": "model.safetensors",
282
+ "model.layers.13.mlp.shared_experts.up_proj.biases": "model.safetensors",
283
+ "model.layers.13.mlp.shared_experts.up_proj.scales": "model.safetensors",
284
+ "model.layers.13.mlp.shared_experts.up_proj.weight": "model.safetensors",
285
+ "model.layers.13.post_attention_layernorm.weight": "model.safetensors",
286
+ "model.layers.14.attention.A_log": "model.safetensors",
287
+ "model.layers.14.attention.b_proj.biases": "model.safetensors",
288
+ "model.layers.14.attention.b_proj.scales": "model.safetensors",
289
+ "model.layers.14.attention.b_proj.weight": "model.safetensors",
290
+ "model.layers.14.attention.dt_bias": "model.safetensors",
291
+ "model.layers.14.attention.f_proj.biases": "model.safetensors",
292
+ "model.layers.14.attention.f_proj.scales": "model.safetensors",
293
+ "model.layers.14.attention.f_proj.weight": "model.safetensors",
294
+ "model.layers.14.attention.g_proj.biases": "model.safetensors",
295
+ "model.layers.14.attention.g_proj.scales": "model.safetensors",
296
+ "model.layers.14.attention.g_proj.weight": "model.safetensors",
297
+ "model.layers.14.attention.k_conv1d.conv.weight": "model.safetensors",
298
+ "model.layers.14.attention.k_proj.biases": "model.safetensors",
299
+ "model.layers.14.attention.k_proj.scales": "model.safetensors",
300
+ "model.layers.14.attention.k_proj.weight": "model.safetensors",
301
+ "model.layers.14.attention.o_norm.weight": "model.safetensors",
302
+ "model.layers.14.attention.o_proj.biases": "model.safetensors",
303
+ "model.layers.14.attention.o_proj.scales": "model.safetensors",
304
+ "model.layers.14.attention.o_proj.weight": "model.safetensors",
305
+ "model.layers.14.attention.q_conv1d.conv.weight": "model.safetensors",
306
+ "model.layers.14.attention.q_proj.biases": "model.safetensors",
307
+ "model.layers.14.attention.q_proj.scales": "model.safetensors",
308
+ "model.layers.14.attention.q_proj.weight": "model.safetensors",
309
+ "model.layers.14.attention.v_conv1d.conv.weight": "model.safetensors",
310
+ "model.layers.14.attention.v_proj.biases": "model.safetensors",
311
+ "model.layers.14.attention.v_proj.scales": "model.safetensors",
312
+ "model.layers.14.attention.v_proj.weight": "model.safetensors",
313
+ "model.layers.14.input_layernorm.weight": "model.safetensors",
314
+ "model.layers.14.mlp.experts.down_proj.biases": "model.safetensors",
315
+ "model.layers.14.mlp.experts.down_proj.scales": "model.safetensors",
316
+ "model.layers.14.mlp.experts.down_proj.weight": "model.safetensors",
317
+ "model.layers.14.mlp.experts.gate_proj.biases": "model.safetensors",
318
+ "model.layers.14.mlp.experts.gate_proj.scales": "model.safetensors",
319
+ "model.layers.14.mlp.experts.gate_proj.weight": "model.safetensors",
320
+ "model.layers.14.mlp.experts.up_proj.biases": "model.safetensors",
321
+ "model.layers.14.mlp.experts.up_proj.scales": "model.safetensors",
322
+ "model.layers.14.mlp.experts.up_proj.weight": "model.safetensors",
323
+ "model.layers.14.mlp.gate.expert_bias": "model.safetensors",
324
+ "model.layers.14.mlp.gate.weight": "model.safetensors",
325
+ "model.layers.14.mlp.shared_experts.down_proj.biases": "model.safetensors",
326
+ "model.layers.14.mlp.shared_experts.down_proj.scales": "model.safetensors",
327
+ "model.layers.14.mlp.shared_experts.down_proj.weight": "model.safetensors",
328
+ "model.layers.14.mlp.shared_experts.gate_proj.biases": "model.safetensors",
329
+ "model.layers.14.mlp.shared_experts.gate_proj.scales": "model.safetensors",
330
+ "model.layers.14.mlp.shared_experts.gate_proj.weight": "model.safetensors",
331
+ "model.layers.14.mlp.shared_experts.up_proj.biases": "model.safetensors",
332
+ "model.layers.14.mlp.shared_experts.up_proj.scales": "model.safetensors",
333
+ "model.layers.14.mlp.shared_experts.up_proj.weight": "model.safetensors",
334
+ "model.layers.14.post_attention_layernorm.weight": "model.safetensors",
335
+ "model.layers.15.attention.dense.biases": "model.safetensors",
336
+ "model.layers.15.attention.dense.scales": "model.safetensors",
337
+ "model.layers.15.attention.dense.weight": "model.safetensors",
338
+ "model.layers.15.attention.g_proj.biases": "model.safetensors",
339
+ "model.layers.15.attention.g_proj.scales": "model.safetensors",
340
+ "model.layers.15.attention.g_proj.weight": "model.safetensors",
341
+ "model.layers.15.attention.kv_a_layernorm.weight": "model.safetensors",
342
+ "model.layers.15.attention.kv_a_proj_with_mqa.biases": "model.safetensors",
343
+ "model.layers.15.attention.kv_a_proj_with_mqa.scales": "model.safetensors",
344
+ "model.layers.15.attention.kv_a_proj_with_mqa.weight": "model.safetensors",
345
+ "model.layers.15.attention.kv_b_proj.biases": "model.safetensors",
346
+ "model.layers.15.attention.kv_b_proj.scales": "model.safetensors",
347
+ "model.layers.15.attention.kv_b_proj.weight": "model.safetensors",
348
+ "model.layers.15.attention.q_a_layernorm.weight": "model.safetensors",
349
+ "model.layers.15.attention.q_a_proj.biases": "model.safetensors",
350
+ "model.layers.15.attention.q_a_proj.scales": "model.safetensors",
351
+ "model.layers.15.attention.q_a_proj.weight": "model.safetensors",
352
+ "model.layers.15.attention.q_b_proj.biases": "model.safetensors",
353
+ "model.layers.15.attention.q_b_proj.scales": "model.safetensors",
354
+ "model.layers.15.attention.q_b_proj.weight": "model.safetensors",
355
+ "model.layers.15.input_layernorm.weight": "model.safetensors",
356
+ "model.layers.15.mlp.experts.down_proj.biases": "model.safetensors",
357
+ "model.layers.15.mlp.experts.down_proj.scales": "model.safetensors",
358
+ "model.layers.15.mlp.experts.down_proj.weight": "model.safetensors",
359
+ "model.layers.15.mlp.experts.gate_proj.biases": "model.safetensors",
360
+ "model.layers.15.mlp.experts.gate_proj.scales": "model.safetensors",
361
+ "model.layers.15.mlp.experts.gate_proj.weight": "model.safetensors",
362
+ "model.layers.15.mlp.experts.up_proj.biases": "model.safetensors",
363
+ "model.layers.15.mlp.experts.up_proj.scales": "model.safetensors",
364
+ "model.layers.15.mlp.experts.up_proj.weight": "model.safetensors",
365
+ "model.layers.15.mlp.gate.expert_bias": "model.safetensors",
366
+ "model.layers.15.mlp.gate.weight": "model.safetensors",
367
+ "model.layers.15.mlp.shared_experts.down_proj.biases": "model.safetensors",
368
+ "model.layers.15.mlp.shared_experts.down_proj.scales": "model.safetensors",
369
+ "model.layers.15.mlp.shared_experts.down_proj.weight": "model.safetensors",
370
+ "model.layers.15.mlp.shared_experts.gate_proj.biases": "model.safetensors",
371
+ "model.layers.15.mlp.shared_experts.gate_proj.scales": "model.safetensors",
372
+ "model.layers.15.mlp.shared_experts.gate_proj.weight": "model.safetensors",
373
+ "model.layers.15.mlp.shared_experts.up_proj.biases": "model.safetensors",
374
+ "model.layers.15.mlp.shared_experts.up_proj.scales": "model.safetensors",
375
+ "model.layers.15.mlp.shared_experts.up_proj.weight": "model.safetensors",
376
+ "model.layers.15.post_attention_layernorm.weight": "model.safetensors",
377
+ "model.layers.16.attention.A_log": "model.safetensors",
378
+ "model.layers.16.attention.b_proj.biases": "model.safetensors",
379
+ "model.layers.16.attention.b_proj.scales": "model.safetensors",
380
+ "model.layers.16.attention.b_proj.weight": "model.safetensors",
381
+ "model.layers.16.attention.dt_bias": "model.safetensors",
382
+ "model.layers.16.attention.f_proj.biases": "model.safetensors",
383
+ "model.layers.16.attention.f_proj.scales": "model.safetensors",
384
+ "model.layers.16.attention.f_proj.weight": "model.safetensors",
385
+ "model.layers.16.attention.g_proj.biases": "model.safetensors",
386
+ "model.layers.16.attention.g_proj.scales": "model.safetensors",
387
+ "model.layers.16.attention.g_proj.weight": "model.safetensors",
388
+ "model.layers.16.attention.k_conv1d.conv.weight": "model.safetensors",
389
+ "model.layers.16.attention.k_proj.biases": "model.safetensors",
390
+ "model.layers.16.attention.k_proj.scales": "model.safetensors",
391
+ "model.layers.16.attention.k_proj.weight": "model.safetensors",
392
+ "model.layers.16.attention.o_norm.weight": "model.safetensors",
393
+ "model.layers.16.attention.o_proj.biases": "model.safetensors",
394
+ "model.layers.16.attention.o_proj.scales": "model.safetensors",
395
+ "model.layers.16.attention.o_proj.weight": "model.safetensors",
396
+ "model.layers.16.attention.q_conv1d.conv.weight": "model.safetensors",
397
+ "model.layers.16.attention.q_proj.biases": "model.safetensors",
398
+ "model.layers.16.attention.q_proj.scales": "model.safetensors",
399
+ "model.layers.16.attention.q_proj.weight": "model.safetensors",
400
+ "model.layers.16.attention.v_conv1d.conv.weight": "model.safetensors",
401
+ "model.layers.16.attention.v_proj.biases": "model.safetensors",
402
+ "model.layers.16.attention.v_proj.scales": "model.safetensors",
403
+ "model.layers.16.attention.v_proj.weight": "model.safetensors",
404
+ "model.layers.16.input_layernorm.weight": "model.safetensors",
405
+ "model.layers.16.mlp.experts.down_proj.biases": "model.safetensors",
406
+ "model.layers.16.mlp.experts.down_proj.scales": "model.safetensors",
407
+ "model.layers.16.mlp.experts.down_proj.weight": "model.safetensors",
408
+ "model.layers.16.mlp.experts.gate_proj.biases": "model.safetensors",
409
+ "model.layers.16.mlp.experts.gate_proj.scales": "model.safetensors",
410
+ "model.layers.16.mlp.experts.gate_proj.weight": "model.safetensors",
411
+ "model.layers.16.mlp.experts.up_proj.biases": "model.safetensors",
412
+ "model.layers.16.mlp.experts.up_proj.scales": "model.safetensors",
413
+ "model.layers.16.mlp.experts.up_proj.weight": "model.safetensors",
414
+ "model.layers.16.mlp.gate.expert_bias": "model.safetensors",
415
+ "model.layers.16.mlp.gate.weight": "model.safetensors",
416
+ "model.layers.16.mlp.shared_experts.down_proj.biases": "model.safetensors",
417
+ "model.layers.16.mlp.shared_experts.down_proj.scales": "model.safetensors",
418
+ "model.layers.16.mlp.shared_experts.down_proj.weight": "model.safetensors",
419
+ "model.layers.16.mlp.shared_experts.gate_proj.biases": "model.safetensors",
420
+ "model.layers.16.mlp.shared_experts.gate_proj.scales": "model.safetensors",
421
+ "model.layers.16.mlp.shared_experts.gate_proj.weight": "model.safetensors",
422
+ "model.layers.16.mlp.shared_experts.up_proj.biases": "model.safetensors",
423
+ "model.layers.16.mlp.shared_experts.up_proj.scales": "model.safetensors",
424
+ "model.layers.16.mlp.shared_experts.up_proj.weight": "model.safetensors",
425
+ "model.layers.16.post_attention_layernorm.weight": "model.safetensors",
426
+ "model.layers.17.attention.A_log": "model.safetensors",
427
+ "model.layers.17.attention.b_proj.biases": "model.safetensors",
428
+ "model.layers.17.attention.b_proj.scales": "model.safetensors",
429
+ "model.layers.17.attention.b_proj.weight": "model.safetensors",
430
+ "model.layers.17.attention.dt_bias": "model.safetensors",
431
+ "model.layers.17.attention.f_proj.biases": "model.safetensors",
432
+ "model.layers.17.attention.f_proj.scales": "model.safetensors",
433
+ "model.layers.17.attention.f_proj.weight": "model.safetensors",
434
+ "model.layers.17.attention.g_proj.biases": "model.safetensors",
435
+ "model.layers.17.attention.g_proj.scales": "model.safetensors",
436
+ "model.layers.17.attention.g_proj.weight": "model.safetensors",
437
+ "model.layers.17.attention.k_conv1d.conv.weight": "model.safetensors",
438
+ "model.layers.17.attention.k_proj.biases": "model.safetensors",
439
+ "model.layers.17.attention.k_proj.scales": "model.safetensors",
440
+ "model.layers.17.attention.k_proj.weight": "model.safetensors",
441
+ "model.layers.17.attention.o_norm.weight": "model.safetensors",
442
+ "model.layers.17.attention.o_proj.biases": "model.safetensors",
443
+ "model.layers.17.attention.o_proj.scales": "model.safetensors",
444
+ "model.layers.17.attention.o_proj.weight": "model.safetensors",
445
+ "model.layers.17.attention.q_conv1d.conv.weight": "model.safetensors",
446
+ "model.layers.17.attention.q_proj.biases": "model.safetensors",
447
+ "model.layers.17.attention.q_proj.scales": "model.safetensors",
448
+ "model.layers.17.attention.q_proj.weight": "model.safetensors",
449
+ "model.layers.17.attention.v_conv1d.conv.weight": "model.safetensors",
450
+ "model.layers.17.attention.v_proj.biases": "model.safetensors",
451
+ "model.layers.17.attention.v_proj.scales": "model.safetensors",
452
+ "model.layers.17.attention.v_proj.weight": "model.safetensors",
453
+ "model.layers.17.input_layernorm.weight": "model.safetensors",
454
+ "model.layers.17.mlp.experts.down_proj.biases": "model.safetensors",
455
+ "model.layers.17.mlp.experts.down_proj.scales": "model.safetensors",
456
+ "model.layers.17.mlp.experts.down_proj.weight": "model.safetensors",
457
+ "model.layers.17.mlp.experts.gate_proj.biases": "model.safetensors",
458
+ "model.layers.17.mlp.experts.gate_proj.scales": "model.safetensors",
459
+ "model.layers.17.mlp.experts.gate_proj.weight": "model.safetensors",
460
+ "model.layers.17.mlp.experts.up_proj.biases": "model.safetensors",
461
+ "model.layers.17.mlp.experts.up_proj.scales": "model.safetensors",
462
+ "model.layers.17.mlp.experts.up_proj.weight": "model.safetensors",
463
+ "model.layers.17.mlp.gate.expert_bias": "model.safetensors",
464
+ "model.layers.17.mlp.gate.weight": "model.safetensors",
465
+ "model.layers.17.mlp.shared_experts.down_proj.biases": "model.safetensors",
466
+ "model.layers.17.mlp.shared_experts.down_proj.scales": "model.safetensors",
467
+ "model.layers.17.mlp.shared_experts.down_proj.weight": "model.safetensors",
468
+ "model.layers.17.mlp.shared_experts.gate_proj.biases": "model.safetensors",
469
+ "model.layers.17.mlp.shared_experts.gate_proj.scales": "model.safetensors",
470
+ "model.layers.17.mlp.shared_experts.gate_proj.weight": "model.safetensors",
471
+ "model.layers.17.mlp.shared_experts.up_proj.biases": "model.safetensors",
472
+ "model.layers.17.mlp.shared_experts.up_proj.scales": "model.safetensors",
473
+ "model.layers.17.mlp.shared_experts.up_proj.weight": "model.safetensors",
474
+ "model.layers.17.post_attention_layernorm.weight": "model.safetensors",
475
+ "model.layers.18.attention.A_log": "model.safetensors",
476
+ "model.layers.18.attention.b_proj.biases": "model.safetensors",
477
+ "model.layers.18.attention.b_proj.scales": "model.safetensors",
478
+ "model.layers.18.attention.b_proj.weight": "model.safetensors",
479
+ "model.layers.18.attention.dt_bias": "model.safetensors",
480
+ "model.layers.18.attention.f_proj.biases": "model.safetensors",
481
+ "model.layers.18.attention.f_proj.scales": "model.safetensors",
482
+ "model.layers.18.attention.f_proj.weight": "model.safetensors",
483
+ "model.layers.18.attention.g_proj.biases": "model.safetensors",
484
+ "model.layers.18.attention.g_proj.scales": "model.safetensors",
485
+ "model.layers.18.attention.g_proj.weight": "model.safetensors",
486
+ "model.layers.18.attention.k_conv1d.conv.weight": "model.safetensors",
487
+ "model.layers.18.attention.k_proj.biases": "model.safetensors",
488
+ "model.layers.18.attention.k_proj.scales": "model.safetensors",
489
+ "model.layers.18.attention.k_proj.weight": "model.safetensors",
490
+ "model.layers.18.attention.o_norm.weight": "model.safetensors",
491
+ "model.layers.18.attention.o_proj.biases": "model.safetensors",
492
+ "model.layers.18.attention.o_proj.scales": "model.safetensors",
493
+ "model.layers.18.attention.o_proj.weight": "model.safetensors",
494
+ "model.layers.18.attention.q_conv1d.conv.weight": "model.safetensors",
495
+ "model.layers.18.attention.q_proj.biases": "model.safetensors",
496
+ "model.layers.18.attention.q_proj.scales": "model.safetensors",
497
+ "model.layers.18.attention.q_proj.weight": "model.safetensors",
498
+ "model.layers.18.attention.v_conv1d.conv.weight": "model.safetensors",
499
+ "model.layers.18.attention.v_proj.biases": "model.safetensors",
500
+ "model.layers.18.attention.v_proj.scales": "model.safetensors",
501
+ "model.layers.18.attention.v_proj.weight": "model.safetensors",
502
+ "model.layers.18.input_layernorm.weight": "model.safetensors",
503
+ "model.layers.18.mlp.experts.down_proj.biases": "model.safetensors",
504
+ "model.layers.18.mlp.experts.down_proj.scales": "model.safetensors",
505
+ "model.layers.18.mlp.experts.down_proj.weight": "model.safetensors",
506
+ "model.layers.18.mlp.experts.gate_proj.biases": "model.safetensors",
507
+ "model.layers.18.mlp.experts.gate_proj.scales": "model.safetensors",
508
+ "model.layers.18.mlp.experts.gate_proj.weight": "model.safetensors",
509
+ "model.layers.18.mlp.experts.up_proj.biases": "model.safetensors",
510
+ "model.layers.18.mlp.experts.up_proj.scales": "model.safetensors",
511
+ "model.layers.18.mlp.experts.up_proj.weight": "model.safetensors",
512
+ "model.layers.18.mlp.gate.expert_bias": "model.safetensors",
513
+ "model.layers.18.mlp.gate.weight": "model.safetensors",
514
+ "model.layers.18.mlp.shared_experts.down_proj.biases": "model.safetensors",
515
+ "model.layers.18.mlp.shared_experts.down_proj.scales": "model.safetensors",
516
+ "model.layers.18.mlp.shared_experts.down_proj.weight": "model.safetensors",
517
+ "model.layers.18.mlp.shared_experts.gate_proj.biases": "model.safetensors",
518
+ "model.layers.18.mlp.shared_experts.gate_proj.scales": "model.safetensors",
519
+ "model.layers.18.mlp.shared_experts.gate_proj.weight": "model.safetensors",
520
+ "model.layers.18.mlp.shared_experts.up_proj.biases": "model.safetensors",
521
+ "model.layers.18.mlp.shared_experts.up_proj.scales": "model.safetensors",
522
+ "model.layers.18.mlp.shared_experts.up_proj.weight": "model.safetensors",
523
+ "model.layers.18.post_attention_layernorm.weight": "model.safetensors",
524
+ "model.layers.19.attention.dense.biases": "model.safetensors",
525
+ "model.layers.19.attention.dense.scales": "model.safetensors",
526
+ "model.layers.19.attention.dense.weight": "model.safetensors",
527
+ "model.layers.19.attention.g_proj.biases": "model.safetensors",
528
+ "model.layers.19.attention.g_proj.scales": "model.safetensors",
529
+ "model.layers.19.attention.g_proj.weight": "model.safetensors",
530
+ "model.layers.19.attention.kv_a_layernorm.weight": "model.safetensors",
531
+ "model.layers.19.attention.kv_a_proj_with_mqa.biases": "model.safetensors",
532
+ "model.layers.19.attention.kv_a_proj_with_mqa.scales": "model.safetensors",
533
+ "model.layers.19.attention.kv_a_proj_with_mqa.weight": "model.safetensors",
534
+ "model.layers.19.attention.kv_b_proj.biases": "model.safetensors",
535
+ "model.layers.19.attention.kv_b_proj.scales": "model.safetensors",
536
+ "model.layers.19.attention.kv_b_proj.weight": "model.safetensors",
537
+ "model.layers.19.attention.q_a_layernorm.weight": "model.safetensors",
538
+ "model.layers.19.attention.q_a_proj.biases": "model.safetensors",
539
+ "model.layers.19.attention.q_a_proj.scales": "model.safetensors",
540
+ "model.layers.19.attention.q_a_proj.weight": "model.safetensors",
541
+ "model.layers.19.attention.q_b_proj.biases": "model.safetensors",
542
+ "model.layers.19.attention.q_b_proj.scales": "model.safetensors",
543
+ "model.layers.19.attention.q_b_proj.weight": "model.safetensors",
544
+ "model.layers.19.input_layernorm.weight": "model.safetensors",
545
+ "model.layers.19.mlp.experts.down_proj.biases": "model.safetensors",
546
+ "model.layers.19.mlp.experts.down_proj.scales": "model.safetensors",
547
+ "model.layers.19.mlp.experts.down_proj.weight": "model.safetensors",
548
+ "model.layers.19.mlp.experts.gate_proj.biases": "model.safetensors",
549
+ "model.layers.19.mlp.experts.gate_proj.scales": "model.safetensors",
550
+ "model.layers.19.mlp.experts.gate_proj.weight": "model.safetensors",
551
+ "model.layers.19.mlp.experts.up_proj.biases": "model.safetensors",
552
+ "model.layers.19.mlp.experts.up_proj.scales": "model.safetensors",
553
+ "model.layers.19.mlp.experts.up_proj.weight": "model.safetensors",
554
+ "model.layers.19.mlp.gate.expert_bias": "model.safetensors",
555
+ "model.layers.19.mlp.gate.weight": "model.safetensors",
556
+ "model.layers.19.mlp.shared_experts.down_proj.biases": "model.safetensors",
557
+ "model.layers.19.mlp.shared_experts.down_proj.scales": "model.safetensors",
558
+ "model.layers.19.mlp.shared_experts.down_proj.weight": "model.safetensors",
559
+ "model.layers.19.mlp.shared_experts.gate_proj.biases": "model.safetensors",
560
+ "model.layers.19.mlp.shared_experts.gate_proj.scales": "model.safetensors",
561
+ "model.layers.19.mlp.shared_experts.gate_proj.weight": "model.safetensors",
562
+ "model.layers.19.mlp.shared_experts.up_proj.biases": "model.safetensors",
563
+ "model.layers.19.mlp.shared_experts.up_proj.scales": "model.safetensors",
564
+ "model.layers.19.mlp.shared_experts.up_proj.weight": "model.safetensors",
565
+ "model.layers.19.post_attention_layernorm.weight": "model.safetensors",
566
+ "model.layers.2.attention.A_log": "model.safetensors",
567
+ "model.layers.2.attention.b_proj.biases": "model.safetensors",
568
+ "model.layers.2.attention.b_proj.scales": "model.safetensors",
569
+ "model.layers.2.attention.b_proj.weight": "model.safetensors",
570
+ "model.layers.2.attention.dt_bias": "model.safetensors",
571
+ "model.layers.2.attention.f_proj.biases": "model.safetensors",
572
+ "model.layers.2.attention.f_proj.scales": "model.safetensors",
573
+ "model.layers.2.attention.f_proj.weight": "model.safetensors",
574
+ "model.layers.2.attention.g_proj.biases": "model.safetensors",
575
+ "model.layers.2.attention.g_proj.scales": "model.safetensors",
576
+ "model.layers.2.attention.g_proj.weight": "model.safetensors",
577
+ "model.layers.2.attention.k_conv1d.conv.weight": "model.safetensors",
578
+ "model.layers.2.attention.k_proj.biases": "model.safetensors",
579
+ "model.layers.2.attention.k_proj.scales": "model.safetensors",
580
+ "model.layers.2.attention.k_proj.weight": "model.safetensors",
581
+ "model.layers.2.attention.o_norm.weight": "model.safetensors",
582
+ "model.layers.2.attention.o_proj.biases": "model.safetensors",
583
+ "model.layers.2.attention.o_proj.scales": "model.safetensors",
584
+ "model.layers.2.attention.o_proj.weight": "model.safetensors",
585
+ "model.layers.2.attention.q_conv1d.conv.weight": "model.safetensors",
586
+ "model.layers.2.attention.q_proj.biases": "model.safetensors",
587
+ "model.layers.2.attention.q_proj.scales": "model.safetensors",
588
+ "model.layers.2.attention.q_proj.weight": "model.safetensors",
589
+ "model.layers.2.attention.v_conv1d.conv.weight": "model.safetensors",
590
+ "model.layers.2.attention.v_proj.biases": "model.safetensors",
591
+ "model.layers.2.attention.v_proj.scales": "model.safetensors",
592
+ "model.layers.2.attention.v_proj.weight": "model.safetensors",
593
+ "model.layers.2.input_layernorm.weight": "model.safetensors",
594
+ "model.layers.2.mlp.experts.down_proj.biases": "model.safetensors",
595
+ "model.layers.2.mlp.experts.down_proj.scales": "model.safetensors",
596
+ "model.layers.2.mlp.experts.down_proj.weight": "model.safetensors",
597
+ "model.layers.2.mlp.experts.gate_proj.biases": "model.safetensors",
598
+ "model.layers.2.mlp.experts.gate_proj.scales": "model.safetensors",
599
+ "model.layers.2.mlp.experts.gate_proj.weight": "model.safetensors",
600
+ "model.layers.2.mlp.experts.up_proj.biases": "model.safetensors",
601
+ "model.layers.2.mlp.experts.up_proj.scales": "model.safetensors",
602
+ "model.layers.2.mlp.experts.up_proj.weight": "model.safetensors",
603
+ "model.layers.2.mlp.gate.expert_bias": "model.safetensors",
604
+ "model.layers.2.mlp.gate.weight": "model.safetensors",
605
+ "model.layers.2.mlp.shared_experts.down_proj.biases": "model.safetensors",
606
+ "model.layers.2.mlp.shared_experts.down_proj.scales": "model.safetensors",
607
+ "model.layers.2.mlp.shared_experts.down_proj.weight": "model.safetensors",
608
+ "model.layers.2.mlp.shared_experts.gate_proj.biases": "model.safetensors",
609
+ "model.layers.2.mlp.shared_experts.gate_proj.scales": "model.safetensors",
610
+ "model.layers.2.mlp.shared_experts.gate_proj.weight": "model.safetensors",
611
+ "model.layers.2.mlp.shared_experts.up_proj.biases": "model.safetensors",
612
+ "model.layers.2.mlp.shared_experts.up_proj.scales": "model.safetensors",
613
+ "model.layers.2.mlp.shared_experts.up_proj.weight": "model.safetensors",
614
+ "model.layers.2.post_attention_layernorm.weight": "model.safetensors",
615
+ "model.layers.20.attention.A_log": "model.safetensors",
616
+ "model.layers.20.attention.b_proj.biases": "model.safetensors",
617
+ "model.layers.20.attention.b_proj.scales": "model.safetensors",
618
+ "model.layers.20.attention.b_proj.weight": "model.safetensors",
619
+ "model.layers.20.attention.dt_bias": "model.safetensors",
620
+ "model.layers.20.attention.f_proj.biases": "model.safetensors",
621
+ "model.layers.20.attention.f_proj.scales": "model.safetensors",
622
+ "model.layers.20.attention.f_proj.weight": "model.safetensors",
623
+ "model.layers.20.attention.g_proj.biases": "model.safetensors",
624
+ "model.layers.20.attention.g_proj.scales": "model.safetensors",
625
+ "model.layers.20.attention.g_proj.weight": "model.safetensors",
626
+ "model.layers.20.attention.k_conv1d.conv.weight": "model.safetensors",
627
+ "model.layers.20.attention.k_proj.biases": "model.safetensors",
628
+ "model.layers.20.attention.k_proj.scales": "model.safetensors",
629
+ "model.layers.20.attention.k_proj.weight": "model.safetensors",
630
+ "model.layers.20.attention.o_norm.weight": "model.safetensors",
631
+ "model.layers.20.attention.o_proj.biases": "model.safetensors",
632
+ "model.layers.20.attention.o_proj.scales": "model.safetensors",
633
+ "model.layers.20.attention.o_proj.weight": "model.safetensors",
634
+ "model.layers.20.attention.q_conv1d.conv.weight": "model.safetensors",
635
+ "model.layers.20.attention.q_proj.biases": "model.safetensors",
636
+ "model.layers.20.attention.q_proj.scales": "model.safetensors",
637
+ "model.layers.20.attention.q_proj.weight": "model.safetensors",
638
+ "model.layers.20.attention.v_conv1d.conv.weight": "model.safetensors",
639
+ "model.layers.20.attention.v_proj.biases": "model.safetensors",
640
+ "model.layers.20.attention.v_proj.scales": "model.safetensors",
641
+ "model.layers.20.attention.v_proj.weight": "model.safetensors",
642
+ "model.layers.20.input_layernorm.weight": "model.safetensors",
643
+ "model.layers.20.mlp.experts.down_proj.biases": "model.safetensors",
644
+ "model.layers.20.mlp.experts.down_proj.scales": "model.safetensors",
645
+ "model.layers.20.mlp.experts.down_proj.weight": "model.safetensors",
646
+ "model.layers.20.mlp.experts.gate_proj.biases": "model.safetensors",
647
+ "model.layers.20.mlp.experts.gate_proj.scales": "model.safetensors",
648
+ "model.layers.20.mlp.experts.gate_proj.weight": "model.safetensors",
649
+ "model.layers.20.mlp.experts.up_proj.biases": "model.safetensors",
650
+ "model.layers.20.mlp.experts.up_proj.scales": "model.safetensors",
651
+ "model.layers.20.mlp.experts.up_proj.weight": "model.safetensors",
652
+ "model.layers.20.mlp.gate.expert_bias": "model.safetensors",
653
+ "model.layers.20.mlp.gate.weight": "model.safetensors",
654
+ "model.layers.20.mlp.shared_experts.down_proj.biases": "model.safetensors",
655
+ "model.layers.20.mlp.shared_experts.down_proj.scales": "model.safetensors",
656
+ "model.layers.20.mlp.shared_experts.down_proj.weight": "model.safetensors",
657
+ "model.layers.20.mlp.shared_experts.gate_proj.biases": "model.safetensors",
658
+ "model.layers.20.mlp.shared_experts.gate_proj.scales": "model.safetensors",
659
+ "model.layers.20.mlp.shared_experts.gate_proj.weight": "model.safetensors",
660
+ "model.layers.20.mlp.shared_experts.up_proj.biases": "model.safetensors",
661
+ "model.layers.20.mlp.shared_experts.up_proj.scales": "model.safetensors",
662
+ "model.layers.20.mlp.shared_experts.up_proj.weight": "model.safetensors",
663
+ "model.layers.20.post_attention_layernorm.weight": "model.safetensors",
664
+ "model.layers.21.attention.A_log": "model.safetensors",
665
+ "model.layers.21.attention.b_proj.biases": "model.safetensors",
666
+ "model.layers.21.attention.b_proj.scales": "model.safetensors",
667
+ "model.layers.21.attention.b_proj.weight": "model.safetensors",
668
+ "model.layers.21.attention.dt_bias": "model.safetensors",
669
+ "model.layers.21.attention.f_proj.biases": "model.safetensors",
670
+ "model.layers.21.attention.f_proj.scales": "model.safetensors",
671
+ "model.layers.21.attention.f_proj.weight": "model.safetensors",
672
+ "model.layers.21.attention.g_proj.biases": "model.safetensors",
673
+ "model.layers.21.attention.g_proj.scales": "model.safetensors",
674
+ "model.layers.21.attention.g_proj.weight": "model.safetensors",
675
+ "model.layers.21.attention.k_conv1d.conv.weight": "model.safetensors",
676
+ "model.layers.21.attention.k_proj.biases": "model.safetensors",
677
+ "model.layers.21.attention.k_proj.scales": "model.safetensors",
678
+ "model.layers.21.attention.k_proj.weight": "model.safetensors",
679
+ "model.layers.21.attention.o_norm.weight": "model.safetensors",
680
+ "model.layers.21.attention.o_proj.biases": "model.safetensors",
681
+ "model.layers.21.attention.o_proj.scales": "model.safetensors",
682
+ "model.layers.21.attention.o_proj.weight": "model.safetensors",
683
+ "model.layers.21.attention.q_conv1d.conv.weight": "model.safetensors",
684
+ "model.layers.21.attention.q_proj.biases": "model.safetensors",
685
+ "model.layers.21.attention.q_proj.scales": "model.safetensors",
686
+ "model.layers.21.attention.q_proj.weight": "model.safetensors",
687
+ "model.layers.21.attention.v_conv1d.conv.weight": "model.safetensors",
688
+ "model.layers.21.attention.v_proj.biases": "model.safetensors",
689
+ "model.layers.21.attention.v_proj.scales": "model.safetensors",
690
+ "model.layers.21.attention.v_proj.weight": "model.safetensors",
691
+ "model.layers.21.input_layernorm.weight": "model.safetensors",
692
+ "model.layers.21.mlp.experts.down_proj.biases": "model.safetensors",
693
+ "model.layers.21.mlp.experts.down_proj.scales": "model.safetensors",
694
+ "model.layers.21.mlp.experts.down_proj.weight": "model.safetensors",
695
+ "model.layers.21.mlp.experts.gate_proj.biases": "model.safetensors",
696
+ "model.layers.21.mlp.experts.gate_proj.scales": "model.safetensors",
697
+ "model.layers.21.mlp.experts.gate_proj.weight": "model.safetensors",
698
+ "model.layers.21.mlp.experts.up_proj.biases": "model.safetensors",
699
+ "model.layers.21.mlp.experts.up_proj.scales": "model.safetensors",
700
+ "model.layers.21.mlp.experts.up_proj.weight": "model.safetensors",
701
+ "model.layers.21.mlp.gate.expert_bias": "model.safetensors",
702
+ "model.layers.21.mlp.gate.weight": "model.safetensors",
703
+ "model.layers.21.mlp.shared_experts.down_proj.biases": "model.safetensors",
704
+ "model.layers.21.mlp.shared_experts.down_proj.scales": "model.safetensors",
705
+ "model.layers.21.mlp.shared_experts.down_proj.weight": "model.safetensors",
706
+ "model.layers.21.mlp.shared_experts.gate_proj.biases": "model.safetensors",
707
+ "model.layers.21.mlp.shared_experts.gate_proj.scales": "model.safetensors",
708
+ "model.layers.21.mlp.shared_experts.gate_proj.weight": "model.safetensors",
709
+ "model.layers.21.mlp.shared_experts.up_proj.biases": "model.safetensors",
710
+ "model.layers.21.mlp.shared_experts.up_proj.scales": "model.safetensors",
711
+ "model.layers.21.mlp.shared_experts.up_proj.weight": "model.safetensors",
712
+ "model.layers.21.post_attention_layernorm.weight": "model.safetensors",
713
+ "model.layers.22.attention.A_log": "model.safetensors",
714
+ "model.layers.22.attention.b_proj.biases": "model.safetensors",
715
+ "model.layers.22.attention.b_proj.scales": "model.safetensors",
716
+ "model.layers.22.attention.b_proj.weight": "model.safetensors",
717
+ "model.layers.22.attention.dt_bias": "model.safetensors",
718
+ "model.layers.22.attention.f_proj.biases": "model.safetensors",
719
+ "model.layers.22.attention.f_proj.scales": "model.safetensors",
720
+ "model.layers.22.attention.f_proj.weight": "model.safetensors",
721
+ "model.layers.22.attention.g_proj.biases": "model.safetensors",
722
+ "model.layers.22.attention.g_proj.scales": "model.safetensors",
723
+ "model.layers.22.attention.g_proj.weight": "model.safetensors",
724
+ "model.layers.22.attention.k_conv1d.conv.weight": "model.safetensors",
725
+ "model.layers.22.attention.k_proj.biases": "model.safetensors",
726
+ "model.layers.22.attention.k_proj.scales": "model.safetensors",
727
+ "model.layers.22.attention.k_proj.weight": "model.safetensors",
728
+ "model.layers.22.attention.o_norm.weight": "model.safetensors",
729
+ "model.layers.22.attention.o_proj.biases": "model.safetensors",
730
+ "model.layers.22.attention.o_proj.scales": "model.safetensors",
731
+ "model.layers.22.attention.o_proj.weight": "model.safetensors",
732
+ "model.layers.22.attention.q_conv1d.conv.weight": "model.safetensors",
733
+ "model.layers.22.attention.q_proj.biases": "model.safetensors",
734
+ "model.layers.22.attention.q_proj.scales": "model.safetensors",
735
+ "model.layers.22.attention.q_proj.weight": "model.safetensors",
736
+ "model.layers.22.attention.v_conv1d.conv.weight": "model.safetensors",
737
+ "model.layers.22.attention.v_proj.biases": "model.safetensors",
738
+ "model.layers.22.attention.v_proj.scales": "model.safetensors",
739
+ "model.layers.22.attention.v_proj.weight": "model.safetensors",
740
+ "model.layers.22.input_layernorm.weight": "model.safetensors",
741
+ "model.layers.22.mlp.experts.down_proj.biases": "model.safetensors",
742
+ "model.layers.22.mlp.experts.down_proj.scales": "model.safetensors",
743
+ "model.layers.22.mlp.experts.down_proj.weight": "model.safetensors",
744
+ "model.layers.22.mlp.experts.gate_proj.biases": "model.safetensors",
745
+ "model.layers.22.mlp.experts.gate_proj.scales": "model.safetensors",
746
+ "model.layers.22.mlp.experts.gate_proj.weight": "model.safetensors",
747
+ "model.layers.22.mlp.experts.up_proj.biases": "model.safetensors",
748
+ "model.layers.22.mlp.experts.up_proj.scales": "model.safetensors",
749
+ "model.layers.22.mlp.experts.up_proj.weight": "model.safetensors",
750
+ "model.layers.22.mlp.gate.expert_bias": "model.safetensors",
751
+ "model.layers.22.mlp.gate.weight": "model.safetensors",
752
+ "model.layers.22.mlp.shared_experts.down_proj.biases": "model.safetensors",
753
+ "model.layers.22.mlp.shared_experts.down_proj.scales": "model.safetensors",
754
+ "model.layers.22.mlp.shared_experts.down_proj.weight": "model.safetensors",
755
+ "model.layers.22.mlp.shared_experts.gate_proj.biases": "model.safetensors",
756
+ "model.layers.22.mlp.shared_experts.gate_proj.scales": "model.safetensors",
757
+ "model.layers.22.mlp.shared_experts.gate_proj.weight": "model.safetensors",
758
+ "model.layers.22.mlp.shared_experts.up_proj.biases": "model.safetensors",
759
+ "model.layers.22.mlp.shared_experts.up_proj.scales": "model.safetensors",
760
+ "model.layers.22.mlp.shared_experts.up_proj.weight": "model.safetensors",
761
+ "model.layers.22.post_attention_layernorm.weight": "model.safetensors",
762
+ "model.layers.23.attention.dense.biases": "model.safetensors",
763
+ "model.layers.23.attention.dense.scales": "model.safetensors",
764
+ "model.layers.23.attention.dense.weight": "model.safetensors",
765
+ "model.layers.23.attention.g_proj.biases": "model.safetensors",
766
+ "model.layers.23.attention.g_proj.scales": "model.safetensors",
767
+ "model.layers.23.attention.g_proj.weight": "model.safetensors",
768
+ "model.layers.23.attention.kv_a_layernorm.weight": "model.safetensors",
769
+ "model.layers.23.attention.kv_a_proj_with_mqa.biases": "model.safetensors",
770
+ "model.layers.23.attention.kv_a_proj_with_mqa.scales": "model.safetensors",
771
+ "model.layers.23.attention.kv_a_proj_with_mqa.weight": "model.safetensors",
772
+ "model.layers.23.attention.kv_b_proj.biases": "model.safetensors",
773
+ "model.layers.23.attention.kv_b_proj.scales": "model.safetensors",
774
+ "model.layers.23.attention.kv_b_proj.weight": "model.safetensors",
775
+ "model.layers.23.attention.q_a_layernorm.weight": "model.safetensors",
776
+ "model.layers.23.attention.q_a_proj.biases": "model.safetensors",
777
+ "model.layers.23.attention.q_a_proj.scales": "model.safetensors",
778
+ "model.layers.23.attention.q_a_proj.weight": "model.safetensors",
779
+ "model.layers.23.attention.q_b_proj.biases": "model.safetensors",
780
+ "model.layers.23.attention.q_b_proj.scales": "model.safetensors",
781
+ "model.layers.23.attention.q_b_proj.weight": "model.safetensors",
782
+ "model.layers.23.input_layernorm.weight": "model.safetensors",
783
+ "model.layers.23.mlp.experts.down_proj.biases": "model.safetensors",
784
+ "model.layers.23.mlp.experts.down_proj.scales": "model.safetensors",
785
+ "model.layers.23.mlp.experts.down_proj.weight": "model.safetensors",
786
+ "model.layers.23.mlp.experts.gate_proj.biases": "model.safetensors",
787
+ "model.layers.23.mlp.experts.gate_proj.scales": "model.safetensors",
788
+ "model.layers.23.mlp.experts.gate_proj.weight": "model.safetensors",
789
+ "model.layers.23.mlp.experts.up_proj.biases": "model.safetensors",
790
+ "model.layers.23.mlp.experts.up_proj.scales": "model.safetensors",
791
+ "model.layers.23.mlp.experts.up_proj.weight": "model.safetensors",
792
+ "model.layers.23.mlp.gate.expert_bias": "model.safetensors",
793
+ "model.layers.23.mlp.gate.weight": "model.safetensors",
794
+ "model.layers.23.mlp.shared_experts.down_proj.biases": "model.safetensors",
795
+ "model.layers.23.mlp.shared_experts.down_proj.scales": "model.safetensors",
796
+ "model.layers.23.mlp.shared_experts.down_proj.weight": "model.safetensors",
797
+ "model.layers.23.mlp.shared_experts.gate_proj.biases": "model.safetensors",
798
+ "model.layers.23.mlp.shared_experts.gate_proj.scales": "model.safetensors",
799
+ "model.layers.23.mlp.shared_experts.gate_proj.weight": "model.safetensors",
800
+ "model.layers.23.mlp.shared_experts.up_proj.biases": "model.safetensors",
801
+ "model.layers.23.mlp.shared_experts.up_proj.scales": "model.safetensors",
802
+ "model.layers.23.mlp.shared_experts.up_proj.weight": "model.safetensors",
803
+ "model.layers.23.post_attention_layernorm.weight": "model.safetensors",
804
+ "model.layers.3.attention.dense.biases": "model.safetensors",
805
+ "model.layers.3.attention.dense.scales": "model.safetensors",
806
+ "model.layers.3.attention.dense.weight": "model.safetensors",
807
+ "model.layers.3.attention.g_proj.biases": "model.safetensors",
808
+ "model.layers.3.attention.g_proj.scales": "model.safetensors",
809
+ "model.layers.3.attention.g_proj.weight": "model.safetensors",
810
+ "model.layers.3.attention.kv_a_layernorm.weight": "model.safetensors",
811
+ "model.layers.3.attention.kv_a_proj_with_mqa.biases": "model.safetensors",
812
+ "model.layers.3.attention.kv_a_proj_with_mqa.scales": "model.safetensors",
813
+ "model.layers.3.attention.kv_a_proj_with_mqa.weight": "model.safetensors",
814
+ "model.layers.3.attention.kv_b_proj.biases": "model.safetensors",
815
+ "model.layers.3.attention.kv_b_proj.scales": "model.safetensors",
816
+ "model.layers.3.attention.kv_b_proj.weight": "model.safetensors",
817
+ "model.layers.3.attention.q_a_layernorm.weight": "model.safetensors",
818
+ "model.layers.3.attention.q_a_proj.biases": "model.safetensors",
819
+ "model.layers.3.attention.q_a_proj.scales": "model.safetensors",
820
+ "model.layers.3.attention.q_a_proj.weight": "model.safetensors",
821
+ "model.layers.3.attention.q_b_proj.biases": "model.safetensors",
822
+ "model.layers.3.attention.q_b_proj.scales": "model.safetensors",
823
+ "model.layers.3.attention.q_b_proj.weight": "model.safetensors",
824
+ "model.layers.3.input_layernorm.weight": "model.safetensors",
825
+ "model.layers.3.mlp.experts.down_proj.biases": "model.safetensors",
826
+ "model.layers.3.mlp.experts.down_proj.scales": "model.safetensors",
827
+ "model.layers.3.mlp.experts.down_proj.weight": "model.safetensors",
828
+ "model.layers.3.mlp.experts.gate_proj.biases": "model.safetensors",
829
+ "model.layers.3.mlp.experts.gate_proj.scales": "model.safetensors",
830
+ "model.layers.3.mlp.experts.gate_proj.weight": "model.safetensors",
831
+ "model.layers.3.mlp.experts.up_proj.biases": "model.safetensors",
832
+ "model.layers.3.mlp.experts.up_proj.scales": "model.safetensors",
833
+ "model.layers.3.mlp.experts.up_proj.weight": "model.safetensors",
834
+ "model.layers.3.mlp.gate.expert_bias": "model.safetensors",
835
+ "model.layers.3.mlp.gate.weight": "model.safetensors",
836
+ "model.layers.3.mlp.shared_experts.down_proj.biases": "model.safetensors",
837
+ "model.layers.3.mlp.shared_experts.down_proj.scales": "model.safetensors",
838
+ "model.layers.3.mlp.shared_experts.down_proj.weight": "model.safetensors",
839
+ "model.layers.3.mlp.shared_experts.gate_proj.biases": "model.safetensors",
840
+ "model.layers.3.mlp.shared_experts.gate_proj.scales": "model.safetensors",
841
+ "model.layers.3.mlp.shared_experts.gate_proj.weight": "model.safetensors",
842
+ "model.layers.3.mlp.shared_experts.up_proj.biases": "model.safetensors",
843
+ "model.layers.3.mlp.shared_experts.up_proj.scales": "model.safetensors",
844
+ "model.layers.3.mlp.shared_experts.up_proj.weight": "model.safetensors",
845
+ "model.layers.3.post_attention_layernorm.weight": "model.safetensors",
846
+ "model.layers.4.attention.A_log": "model.safetensors",
847
+ "model.layers.4.attention.b_proj.biases": "model.safetensors",
848
+ "model.layers.4.attention.b_proj.scales": "model.safetensors",
849
+ "model.layers.4.attention.b_proj.weight": "model.safetensors",
850
+ "model.layers.4.attention.dt_bias": "model.safetensors",
851
+ "model.layers.4.attention.f_proj.biases": "model.safetensors",
852
+ "model.layers.4.attention.f_proj.scales": "model.safetensors",
853
+ "model.layers.4.attention.f_proj.weight": "model.safetensors",
854
+ "model.layers.4.attention.g_proj.biases": "model.safetensors",
855
+ "model.layers.4.attention.g_proj.scales": "model.safetensors",
856
+ "model.layers.4.attention.g_proj.weight": "model.safetensors",
857
+ "model.layers.4.attention.k_conv1d.conv.weight": "model.safetensors",
858
+ "model.layers.4.attention.k_proj.biases": "model.safetensors",
859
+ "model.layers.4.attention.k_proj.scales": "model.safetensors",
860
+ "model.layers.4.attention.k_proj.weight": "model.safetensors",
861
+ "model.layers.4.attention.o_norm.weight": "model.safetensors",
862
+ "model.layers.4.attention.o_proj.biases": "model.safetensors",
863
+ "model.layers.4.attention.o_proj.scales": "model.safetensors",
864
+ "model.layers.4.attention.o_proj.weight": "model.safetensors",
865
+ "model.layers.4.attention.q_conv1d.conv.weight": "model.safetensors",
866
+ "model.layers.4.attention.q_proj.biases": "model.safetensors",
867
+ "model.layers.4.attention.q_proj.scales": "model.safetensors",
868
+ "model.layers.4.attention.q_proj.weight": "model.safetensors",
869
+ "model.layers.4.attention.v_conv1d.conv.weight": "model.safetensors",
870
+ "model.layers.4.attention.v_proj.biases": "model.safetensors",
871
+ "model.layers.4.attention.v_proj.scales": "model.safetensors",
872
+ "model.layers.4.attention.v_proj.weight": "model.safetensors",
873
+ "model.layers.4.input_layernorm.weight": "model.safetensors",
874
+ "model.layers.4.mlp.experts.down_proj.biases": "model.safetensors",
875
+ "model.layers.4.mlp.experts.down_proj.scales": "model.safetensors",
876
+ "model.layers.4.mlp.experts.down_proj.weight": "model.safetensors",
877
+ "model.layers.4.mlp.experts.gate_proj.biases": "model.safetensors",
878
+ "model.layers.4.mlp.experts.gate_proj.scales": "model.safetensors",
879
+ "model.layers.4.mlp.experts.gate_proj.weight": "model.safetensors",
880
+ "model.layers.4.mlp.experts.up_proj.biases": "model.safetensors",
881
+ "model.layers.4.mlp.experts.up_proj.scales": "model.safetensors",
882
+ "model.layers.4.mlp.experts.up_proj.weight": "model.safetensors",
883
+ "model.layers.4.mlp.gate.expert_bias": "model.safetensors",
884
+ "model.layers.4.mlp.gate.weight": "model.safetensors",
885
+ "model.layers.4.mlp.shared_experts.down_proj.biases": "model.safetensors",
886
+ "model.layers.4.mlp.shared_experts.down_proj.scales": "model.safetensors",
887
+ "model.layers.4.mlp.shared_experts.down_proj.weight": "model.safetensors",
888
+ "model.layers.4.mlp.shared_experts.gate_proj.biases": "model.safetensors",
889
+ "model.layers.4.mlp.shared_experts.gate_proj.scales": "model.safetensors",
890
+ "model.layers.4.mlp.shared_experts.gate_proj.weight": "model.safetensors",
891
+ "model.layers.4.mlp.shared_experts.up_proj.biases": "model.safetensors",
892
+ "model.layers.4.mlp.shared_experts.up_proj.scales": "model.safetensors",
893
+ "model.layers.4.mlp.shared_experts.up_proj.weight": "model.safetensors",
894
+ "model.layers.4.post_attention_layernorm.weight": "model.safetensors",
895
+ "model.layers.5.attention.A_log": "model.safetensors",
896
+ "model.layers.5.attention.b_proj.biases": "model.safetensors",
897
+ "model.layers.5.attention.b_proj.scales": "model.safetensors",
898
+ "model.layers.5.attention.b_proj.weight": "model.safetensors",
899
+ "model.layers.5.attention.dt_bias": "model.safetensors",
900
+ "model.layers.5.attention.f_proj.biases": "model.safetensors",
901
+ "model.layers.5.attention.f_proj.scales": "model.safetensors",
902
+ "model.layers.5.attention.f_proj.weight": "model.safetensors",
903
+ "model.layers.5.attention.g_proj.biases": "model.safetensors",
904
+ "model.layers.5.attention.g_proj.scales": "model.safetensors",
905
+ "model.layers.5.attention.g_proj.weight": "model.safetensors",
906
+ "model.layers.5.attention.k_conv1d.conv.weight": "model.safetensors",
907
+ "model.layers.5.attention.k_proj.biases": "model.safetensors",
908
+ "model.layers.5.attention.k_proj.scales": "model.safetensors",
909
+ "model.layers.5.attention.k_proj.weight": "model.safetensors",
910
+ "model.layers.5.attention.o_norm.weight": "model.safetensors",
911
+ "model.layers.5.attention.o_proj.biases": "model.safetensors",
912
+ "model.layers.5.attention.o_proj.scales": "model.safetensors",
913
+ "model.layers.5.attention.o_proj.weight": "model.safetensors",
914
+ "model.layers.5.attention.q_conv1d.conv.weight": "model.safetensors",
915
+ "model.layers.5.attention.q_proj.biases": "model.safetensors",
916
+ "model.layers.5.attention.q_proj.scales": "model.safetensors",
917
+ "model.layers.5.attention.q_proj.weight": "model.safetensors",
918
+ "model.layers.5.attention.v_conv1d.conv.weight": "model.safetensors",
919
+ "model.layers.5.attention.v_proj.biases": "model.safetensors",
920
+ "model.layers.5.attention.v_proj.scales": "model.safetensors",
921
+ "model.layers.5.attention.v_proj.weight": "model.safetensors",
922
+ "model.layers.5.input_layernorm.weight": "model.safetensors",
923
+ "model.layers.5.mlp.experts.down_proj.biases": "model.safetensors",
924
+ "model.layers.5.mlp.experts.down_proj.scales": "model.safetensors",
925
+ "model.layers.5.mlp.experts.down_proj.weight": "model.safetensors",
926
+ "model.layers.5.mlp.experts.gate_proj.biases": "model.safetensors",
927
+ "model.layers.5.mlp.experts.gate_proj.scales": "model.safetensors",
928
+ "model.layers.5.mlp.experts.gate_proj.weight": "model.safetensors",
929
+ "model.layers.5.mlp.experts.up_proj.biases": "model.safetensors",
930
+ "model.layers.5.mlp.experts.up_proj.scales": "model.safetensors",
931
+ "model.layers.5.mlp.experts.up_proj.weight": "model.safetensors",
932
+ "model.layers.5.mlp.gate.expert_bias": "model.safetensors",
933
+ "model.layers.5.mlp.gate.weight": "model.safetensors",
934
+ "model.layers.5.mlp.shared_experts.down_proj.biases": "model.safetensors",
935
+ "model.layers.5.mlp.shared_experts.down_proj.scales": "model.safetensors",
936
+ "model.layers.5.mlp.shared_experts.down_proj.weight": "model.safetensors",
937
+ "model.layers.5.mlp.shared_experts.gate_proj.biases": "model.safetensors",
938
+ "model.layers.5.mlp.shared_experts.gate_proj.scales": "model.safetensors",
939
+ "model.layers.5.mlp.shared_experts.gate_proj.weight": "model.safetensors",
940
+ "model.layers.5.mlp.shared_experts.up_proj.biases": "model.safetensors",
941
+ "model.layers.5.mlp.shared_experts.up_proj.scales": "model.safetensors",
942
+ "model.layers.5.mlp.shared_experts.up_proj.weight": "model.safetensors",
943
+ "model.layers.5.post_attention_layernorm.weight": "model.safetensors",
944
+ "model.layers.6.attention.A_log": "model.safetensors",
945
+ "model.layers.6.attention.b_proj.biases": "model.safetensors",
946
+ "model.layers.6.attention.b_proj.scales": "model.safetensors",
947
+ "model.layers.6.attention.b_proj.weight": "model.safetensors",
948
+ "model.layers.6.attention.dt_bias": "model.safetensors",
949
+ "model.layers.6.attention.f_proj.biases": "model.safetensors",
950
+ "model.layers.6.attention.f_proj.scales": "model.safetensors",
951
+ "model.layers.6.attention.f_proj.weight": "model.safetensors",
952
+ "model.layers.6.attention.g_proj.biases": "model.safetensors",
953
+ "model.layers.6.attention.g_proj.scales": "model.safetensors",
954
+ "model.layers.6.attention.g_proj.weight": "model.safetensors",
955
+ "model.layers.6.attention.k_conv1d.conv.weight": "model.safetensors",
956
+ "model.layers.6.attention.k_proj.biases": "model.safetensors",
957
+ "model.layers.6.attention.k_proj.scales": "model.safetensors",
958
+ "model.layers.6.attention.k_proj.weight": "model.safetensors",
959
+ "model.layers.6.attention.o_norm.weight": "model.safetensors",
960
+ "model.layers.6.attention.o_proj.biases": "model.safetensors",
961
+ "model.layers.6.attention.o_proj.scales": "model.safetensors",
962
+ "model.layers.6.attention.o_proj.weight": "model.safetensors",
963
+ "model.layers.6.attention.q_conv1d.conv.weight": "model.safetensors",
964
+ "model.layers.6.attention.q_proj.biases": "model.safetensors",
965
+ "model.layers.6.attention.q_proj.scales": "model.safetensors",
966
+ "model.layers.6.attention.q_proj.weight": "model.safetensors",
967
+ "model.layers.6.attention.v_conv1d.conv.weight": "model.safetensors",
968
+ "model.layers.6.attention.v_proj.biases": "model.safetensors",
969
+ "model.layers.6.attention.v_proj.scales": "model.safetensors",
970
+ "model.layers.6.attention.v_proj.weight": "model.safetensors",
971
+ "model.layers.6.input_layernorm.weight": "model.safetensors",
972
+ "model.layers.6.mlp.experts.down_proj.biases": "model.safetensors",
973
+ "model.layers.6.mlp.experts.down_proj.scales": "model.safetensors",
974
+ "model.layers.6.mlp.experts.down_proj.weight": "model.safetensors",
975
+ "model.layers.6.mlp.experts.gate_proj.biases": "model.safetensors",
976
+ "model.layers.6.mlp.experts.gate_proj.scales": "model.safetensors",
977
+ "model.layers.6.mlp.experts.gate_proj.weight": "model.safetensors",
978
+ "model.layers.6.mlp.experts.up_proj.biases": "model.safetensors",
979
+ "model.layers.6.mlp.experts.up_proj.scales": "model.safetensors",
980
+ "model.layers.6.mlp.experts.up_proj.weight": "model.safetensors",
981
+ "model.layers.6.mlp.gate.expert_bias": "model.safetensors",
982
+ "model.layers.6.mlp.gate.weight": "model.safetensors",
983
+ "model.layers.6.mlp.shared_experts.down_proj.biases": "model.safetensors",
984
+ "model.layers.6.mlp.shared_experts.down_proj.scales": "model.safetensors",
985
+ "model.layers.6.mlp.shared_experts.down_proj.weight": "model.safetensors",
986
+ "model.layers.6.mlp.shared_experts.gate_proj.biases": "model.safetensors",
987
+ "model.layers.6.mlp.shared_experts.gate_proj.scales": "model.safetensors",
988
+ "model.layers.6.mlp.shared_experts.gate_proj.weight": "model.safetensors",
989
+ "model.layers.6.mlp.shared_experts.up_proj.biases": "model.safetensors",
990
+ "model.layers.6.mlp.shared_experts.up_proj.scales": "model.safetensors",
991
+ "model.layers.6.mlp.shared_experts.up_proj.weight": "model.safetensors",
992
+ "model.layers.6.post_attention_layernorm.weight": "model.safetensors",
993
+ "model.layers.7.attention.dense.biases": "model.safetensors",
994
+ "model.layers.7.attention.dense.scales": "model.safetensors",
995
+ "model.layers.7.attention.dense.weight": "model.safetensors",
996
+ "model.layers.7.attention.g_proj.biases": "model.safetensors",
997
+ "model.layers.7.attention.g_proj.scales": "model.safetensors",
998
+ "model.layers.7.attention.g_proj.weight": "model.safetensors",
999
+ "model.layers.7.attention.kv_a_layernorm.weight": "model.safetensors",
1000
+ "model.layers.7.attention.kv_a_proj_with_mqa.biases": "model.safetensors",
1001
+ "model.layers.7.attention.kv_a_proj_with_mqa.scales": "model.safetensors",
1002
+ "model.layers.7.attention.kv_a_proj_with_mqa.weight": "model.safetensors",
1003
+ "model.layers.7.attention.kv_b_proj.biases": "model.safetensors",
1004
+ "model.layers.7.attention.kv_b_proj.scales": "model.safetensors",
1005
+ "model.layers.7.attention.kv_b_proj.weight": "model.safetensors",
1006
+ "model.layers.7.attention.q_a_layernorm.weight": "model.safetensors",
1007
+ "model.layers.7.attention.q_a_proj.biases": "model.safetensors",
1008
+ "model.layers.7.attention.q_a_proj.scales": "model.safetensors",
1009
+ "model.layers.7.attention.q_a_proj.weight": "model.safetensors",
1010
+ "model.layers.7.attention.q_b_proj.biases": "model.safetensors",
1011
+ "model.layers.7.attention.q_b_proj.scales": "model.safetensors",
1012
+ "model.layers.7.attention.q_b_proj.weight": "model.safetensors",
1013
+ "model.layers.7.input_layernorm.weight": "model.safetensors",
1014
+ "model.layers.7.mlp.experts.down_proj.biases": "model.safetensors",
1015
+ "model.layers.7.mlp.experts.down_proj.scales": "model.safetensors",
1016
+ "model.layers.7.mlp.experts.down_proj.weight": "model.safetensors",
1017
+ "model.layers.7.mlp.experts.gate_proj.biases": "model.safetensors",
1018
+ "model.layers.7.mlp.experts.gate_proj.scales": "model.safetensors",
1019
+ "model.layers.7.mlp.experts.gate_proj.weight": "model.safetensors",
1020
+ "model.layers.7.mlp.experts.up_proj.biases": "model.safetensors",
1021
+ "model.layers.7.mlp.experts.up_proj.scales": "model.safetensors",
1022
+ "model.layers.7.mlp.experts.up_proj.weight": "model.safetensors",
1023
+ "model.layers.7.mlp.gate.expert_bias": "model.safetensors",
1024
+ "model.layers.7.mlp.gate.weight": "model.safetensors",
1025
+ "model.layers.7.mlp.shared_experts.down_proj.biases": "model.safetensors",
1026
+ "model.layers.7.mlp.shared_experts.down_proj.scales": "model.safetensors",
1027
+ "model.layers.7.mlp.shared_experts.down_proj.weight": "model.safetensors",
1028
+ "model.layers.7.mlp.shared_experts.gate_proj.biases": "model.safetensors",
1029
+ "model.layers.7.mlp.shared_experts.gate_proj.scales": "model.safetensors",
1030
+ "model.layers.7.mlp.shared_experts.gate_proj.weight": "model.safetensors",
1031
+ "model.layers.7.mlp.shared_experts.up_proj.biases": "model.safetensors",
1032
+ "model.layers.7.mlp.shared_experts.up_proj.scales": "model.safetensors",
1033
+ "model.layers.7.mlp.shared_experts.up_proj.weight": "model.safetensors",
1034
+ "model.layers.7.post_attention_layernorm.weight": "model.safetensors",
1035
+ "model.layers.8.attention.A_log": "model.safetensors",
1036
+ "model.layers.8.attention.b_proj.biases": "model.safetensors",
1037
+ "model.layers.8.attention.b_proj.scales": "model.safetensors",
1038
+ "model.layers.8.attention.b_proj.weight": "model.safetensors",
1039
+ "model.layers.8.attention.dt_bias": "model.safetensors",
1040
+ "model.layers.8.attention.f_proj.biases": "model.safetensors",
1041
+ "model.layers.8.attention.f_proj.scales": "model.safetensors",
1042
+ "model.layers.8.attention.f_proj.weight": "model.safetensors",
1043
+ "model.layers.8.attention.g_proj.biases": "model.safetensors",
1044
+ "model.layers.8.attention.g_proj.scales": "model.safetensors",
1045
+ "model.layers.8.attention.g_proj.weight": "model.safetensors",
1046
+ "model.layers.8.attention.k_conv1d.conv.weight": "model.safetensors",
1047
+ "model.layers.8.attention.k_proj.biases": "model.safetensors",
1048
+ "model.layers.8.attention.k_proj.scales": "model.safetensors",
1049
+ "model.layers.8.attention.k_proj.weight": "model.safetensors",
1050
+ "model.layers.8.attention.o_norm.weight": "model.safetensors",
1051
+ "model.layers.8.attention.o_proj.biases": "model.safetensors",
1052
+ "model.layers.8.attention.o_proj.scales": "model.safetensors",
1053
+ "model.layers.8.attention.o_proj.weight": "model.safetensors",
1054
+ "model.layers.8.attention.q_conv1d.conv.weight": "model.safetensors",
1055
+ "model.layers.8.attention.q_proj.biases": "model.safetensors",
1056
+ "model.layers.8.attention.q_proj.scales": "model.safetensors",
1057
+ "model.layers.8.attention.q_proj.weight": "model.safetensors",
1058
+ "model.layers.8.attention.v_conv1d.conv.weight": "model.safetensors",
1059
+ "model.layers.8.attention.v_proj.biases": "model.safetensors",
1060
+ "model.layers.8.attention.v_proj.scales": "model.safetensors",
1061
+ "model.layers.8.attention.v_proj.weight": "model.safetensors",
1062
+ "model.layers.8.input_layernorm.weight": "model.safetensors",
1063
+ "model.layers.8.mlp.experts.down_proj.biases": "model.safetensors",
1064
+ "model.layers.8.mlp.experts.down_proj.scales": "model.safetensors",
1065
+ "model.layers.8.mlp.experts.down_proj.weight": "model.safetensors",
1066
+ "model.layers.8.mlp.experts.gate_proj.biases": "model.safetensors",
1067
+ "model.layers.8.mlp.experts.gate_proj.scales": "model.safetensors",
1068
+ "model.layers.8.mlp.experts.gate_proj.weight": "model.safetensors",
1069
+ "model.layers.8.mlp.experts.up_proj.biases": "model.safetensors",
1070
+ "model.layers.8.mlp.experts.up_proj.scales": "model.safetensors",
1071
+ "model.layers.8.mlp.experts.up_proj.weight": "model.safetensors",
1072
+ "model.layers.8.mlp.gate.expert_bias": "model.safetensors",
1073
+ "model.layers.8.mlp.gate.weight": "model.safetensors",
1074
+ "model.layers.8.mlp.shared_experts.down_proj.biases": "model.safetensors",
1075
+ "model.layers.8.mlp.shared_experts.down_proj.scales": "model.safetensors",
1076
+ "model.layers.8.mlp.shared_experts.down_proj.weight": "model.safetensors",
1077
+ "model.layers.8.mlp.shared_experts.gate_proj.biases": "model.safetensors",
1078
+ "model.layers.8.mlp.shared_experts.gate_proj.scales": "model.safetensors",
1079
+ "model.layers.8.mlp.shared_experts.gate_proj.weight": "model.safetensors",
1080
+ "model.layers.8.mlp.shared_experts.up_proj.biases": "model.safetensors",
1081
+ "model.layers.8.mlp.shared_experts.up_proj.scales": "model.safetensors",
1082
+ "model.layers.8.mlp.shared_experts.up_proj.weight": "model.safetensors",
1083
+ "model.layers.8.post_attention_layernorm.weight": "model.safetensors",
1084
+ "model.layers.9.attention.A_log": "model.safetensors",
1085
+ "model.layers.9.attention.b_proj.biases": "model.safetensors",
1086
+ "model.layers.9.attention.b_proj.scales": "model.safetensors",
1087
+ "model.layers.9.attention.b_proj.weight": "model.safetensors",
1088
+ "model.layers.9.attention.dt_bias": "model.safetensors",
1089
+ "model.layers.9.attention.f_proj.biases": "model.safetensors",
1090
+ "model.layers.9.attention.f_proj.scales": "model.safetensors",
1091
+ "model.layers.9.attention.f_proj.weight": "model.safetensors",
1092
+ "model.layers.9.attention.g_proj.biases": "model.safetensors",
1093
+ "model.layers.9.attention.g_proj.scales": "model.safetensors",
1094
+ "model.layers.9.attention.g_proj.weight": "model.safetensors",
1095
+ "model.layers.9.attention.k_conv1d.conv.weight": "model.safetensors",
1096
+ "model.layers.9.attention.k_proj.biases": "model.safetensors",
1097
+ "model.layers.9.attention.k_proj.scales": "model.safetensors",
1098
+ "model.layers.9.attention.k_proj.weight": "model.safetensors",
1099
+ "model.layers.9.attention.o_norm.weight": "model.safetensors",
1100
+ "model.layers.9.attention.o_proj.biases": "model.safetensors",
1101
+ "model.layers.9.attention.o_proj.scales": "model.safetensors",
1102
+ "model.layers.9.attention.o_proj.weight": "model.safetensors",
1103
+ "model.layers.9.attention.q_conv1d.conv.weight": "model.safetensors",
1104
+ "model.layers.9.attention.q_proj.biases": "model.safetensors",
1105
+ "model.layers.9.attention.q_proj.scales": "model.safetensors",
1106
+ "model.layers.9.attention.q_proj.weight": "model.safetensors",
1107
+ "model.layers.9.attention.v_conv1d.conv.weight": "model.safetensors",
1108
+ "model.layers.9.attention.v_proj.biases": "model.safetensors",
1109
+ "model.layers.9.attention.v_proj.scales": "model.safetensors",
1110
+ "model.layers.9.attention.v_proj.weight": "model.safetensors",
1111
+ "model.layers.9.input_layernorm.weight": "model.safetensors",
1112
+ "model.layers.9.mlp.experts.down_proj.biases": "model.safetensors",
1113
+ "model.layers.9.mlp.experts.down_proj.scales": "model.safetensors",
1114
+ "model.layers.9.mlp.experts.down_proj.weight": "model.safetensors",
1115
+ "model.layers.9.mlp.experts.gate_proj.biases": "model.safetensors",
1116
+ "model.layers.9.mlp.experts.gate_proj.scales": "model.safetensors",
1117
+ "model.layers.9.mlp.experts.gate_proj.weight": "model.safetensors",
1118
+ "model.layers.9.mlp.experts.up_proj.biases": "model.safetensors",
1119
+ "model.layers.9.mlp.experts.up_proj.scales": "model.safetensors",
1120
+ "model.layers.9.mlp.experts.up_proj.weight": "model.safetensors",
1121
+ "model.layers.9.mlp.gate.expert_bias": "model.safetensors",
1122
+ "model.layers.9.mlp.gate.weight": "model.safetensors",
1123
+ "model.layers.9.mlp.shared_experts.down_proj.biases": "model.safetensors",
1124
+ "model.layers.9.mlp.shared_experts.down_proj.scales": "model.safetensors",
1125
+ "model.layers.9.mlp.shared_experts.down_proj.weight": "model.safetensors",
1126
+ "model.layers.9.mlp.shared_experts.gate_proj.biases": "model.safetensors",
1127
+ "model.layers.9.mlp.shared_experts.gate_proj.scales": "model.safetensors",
1128
+ "model.layers.9.mlp.shared_experts.gate_proj.weight": "model.safetensors",
1129
+ "model.layers.9.mlp.shared_experts.up_proj.biases": "model.safetensors",
1130
+ "model.layers.9.mlp.shared_experts.up_proj.scales": "model.safetensors",
1131
+ "model.layers.9.mlp.shared_experts.up_proj.weight": "model.safetensors",
1132
+ "model.layers.9.post_attention_layernorm.weight": "model.safetensors",
1133
+ "model.norm.weight": "model.safetensors",
1134
+ "model.word_embeddings.biases": "model.safetensors",
1135
+ "model.word_embeddings.scales": "model.safetensors",
1136
+ "model.word_embeddings.weight": "model.safetensors"
1137
+ }
1138
+ }
modeling_bailing_moe_v3.py ADDED
@@ -0,0 +1,1625 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # coding=utf-8
2
+ # Copyright 2025 Antgroup and The HuggingFace Inc. team. All rights reserved.
3
+ #
4
+ # This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX
5
+ # and OPT implementations in this library. It has been modified from its
6
+ # original forms to accommodate minor architectural differences compared
7
+ # to GPT-NeoX and OPT used by the Meta AI team that trained the model.
8
+ #
9
+ # Licensed under the Apache License, Version 2.0 (the "License");
10
+ # you may not use this file except in compliance with the License.
11
+ # You may obtain a copy of the License at
12
+ #
13
+ # http://www.apache.org/licenses/LICENSE-2.0
14
+ #
15
+ # Unless required by applicable law or agreed to in writing, software
16
+ # distributed under the License is distributed on an "AS IS" BASIS,
17
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
18
+ # See the License for the specific language governing permissions and
19
+ # limitations under the License.
20
+ """PyTorch BailingMoE model."""
21
+
22
+ import math
23
+ import warnings
24
+ from typing import List, Optional, Tuple, Union, Callable
25
+ from copy import deepcopy
26
+
27
+ import torch
28
+ import torch.nn.functional as F
29
+ from torch import nn
30
+
31
+ from transformers.activations import ACT2FN
32
+ from transformers.cache_utils import Cache, DynamicCache
33
+ from transformers.modeling_attn_mask_utils import (
34
+ AttentionMaskConverter,
35
+ _prepare_4d_attention_mask,
36
+ _prepare_4d_causal_attention_mask,
37
+ _prepare_4d_causal_attention_mask_for_sdpa,
38
+ )
39
+ from transformers.modeling_outputs import MoeModelOutputWithPast
40
+ from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update
41
+ from transformers.modeling_utils import PreTrainedModel
42
+ from transformers.pytorch_utils import ALL_LAYERNORM_LAYERS, is_torch_greater_or_equal_than_1_13
43
+ from transformers.utils import (
44
+ add_start_docstrings,
45
+ add_start_docstrings_to_model_forward,
46
+ logging,
47
+ replace_return_docstrings,
48
+ )
49
+ from transformers.utils.import_utils import is_torch_fx_available
50
+ from .configuration_bailing_moe_v3 import BailingMoeV3Config
51
+ from transformers.generation.utils import GenerationMixin
52
+ from dataclasses import dataclass
53
+ from transformers.utils import ModelOutput
54
+ from transformers import DynamicLayer
55
+ from transformers.processing_utils import Unpack
56
+ from transformers.utils import TransformersKwargs
57
+ from transformers.utils.deprecation import deprecate_kwarg
58
+ from transformers.modeling_flash_attention_utils import FlashAttentionKwargs
59
+
60
+ from fla.ops.simple_gla.fused_recurrent import fused_recurrent_simple_gla
61
+ from fla.ops.simple_gla.chunk import chunk_simple_gla
62
+ from einops import rearrange, repeat
63
+
64
+
65
+ try:
66
+ from fla.modules import FusedRMSNormGated, ShortConvolution
67
+ from fla.ops.kda import chunk_kda, fused_recurrent_kda
68
+
69
+ from fla.ops.utils.index import prepare_cu_seqlens_from_mask, prepare_lens_from_mask
70
+ from fla.utils import tensor_cache
71
+ except ImportError:
72
+ raise ImportError("Plese run `pip install -U fla-core`")
73
+
74
+ # This makes `_prepare_4d_causal_attention_mask` a leaf function in the FX graph.
75
+ # It means that the function will not be traced through and simply appear as a node in the graph.
76
+ if is_torch_fx_available():
77
+ if not is_torch_greater_or_equal_than_1_13:
78
+ import torch.fx
79
+
80
+ _prepare_4d_causal_attention_mask = torch.fx.wrap(_prepare_4d_causal_attention_mask)
81
+
82
+
83
+ logger = logging.get_logger(__name__)
84
+
85
+ _CONFIG_FOR_DOC = "BailingMoeV3Config"
86
+
87
+
88
+ def roll_tensor(tensor, shifts=-1, dims=-1, fill_value=0):
89
+ """Roll the tensor input along the given dimension(s).
90
+ Inserted elements are set to be 0.0.
91
+ """
92
+ rolled_tensor = torch.roll(tensor, shifts=shifts, dims=dims)
93
+ rolled_tensor.select(dims, shifts).fill_(fill_value)
94
+ return rolled_tensor, rolled_tensor.sum()
95
+
96
+
97
+ @dataclass
98
+ class MoEV3CausalLMOutputWithPast(ModelOutput):
99
+ """
100
+ Base class for causal language model (or autoregressive) outputs as well as Mixture of Expert's router hidden
101
+ states terms, to train a MoE model.
102
+ Args:
103
+ loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` is provided):
104
+ Language modeling loss (for next-token prediction).
105
+ logits (`torch.FloatTensor` of shape `(batch_size, sequence_length, config.vocab_size)`):
106
+ Prediction scores of the language modeling head (scores for each vocabulary token before SoftMax).
107
+ past_key_values (`Cache`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`):
108
+ It is a [`~cache_utils.Cache`] instance. For more details, see our [kv cache guide](https://huggingface.co/docs/transformers/en/kv_cache).
109
+ Contains pre-computed hidden-states (key and values in the self-attention blocks) that can be used (see
110
+ `past_key_values` input) to speed up sequential decoding.
111
+ hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):
112
+ Tuple of `torch.FloatTensor` (one for the output of the embeddings, if the model has an embedding layer, +
113
+ one for the output of each layer) of shape `(batch_size, sequence_length, hidden_size)`.
114
+ Hidden-states of the model at the output of each layer plus the optional initial embedding outputs.
115
+ attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):
116
+ Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length,
117
+ sequence_length)`.
118
+ Attentions weights after the attention softmax, used to compute the weighted average in the self-attention
119
+ heads.
120
+ z_loss (`torch.FloatTensor`, *optional*, returned when `labels` is provided):
121
+ z_loss for the sparse modules.
122
+ aux_loss (`torch.FloatTensor`, *optional*, returned when `labels` is provided):
123
+ aux_loss for the sparse modules.
124
+ router_logits (`tuple(torch.FloatTensor)`, *optional*, returned when `output_router_logits=True` is passed or when `config.add_router_probs=True`):
125
+ Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, sequence_length, num_experts)`.
126
+ Router logits of the encoder model, useful to compute the auxiliary loss and the z_loss for the sparse
127
+ modules.
128
+ """
129
+
130
+ loss: Optional[torch.FloatTensor] = None
131
+ logits: Optional[torch.FloatTensor] = None
132
+ past_key_values: Optional[Cache] = None
133
+ hidden_states: Optional[tuple[torch.FloatTensor, ...]] = None
134
+ attentions: Optional[tuple[torch.FloatTensor, ...]] = None
135
+ z_loss: Optional[torch.FloatTensor] = None
136
+ aux_loss: Optional[torch.FloatTensor] = None
137
+ router_logits: Optional[tuple[torch.FloatTensor]] = None
138
+ mtp_loss: Optional[torch.FloatTensor] = None
139
+ mtp_logits: Optional[tuple[torch.FloatTensor, ...]] = None
140
+
141
+
142
+ class MoeV3ModelOutputWithPast(MoeModelOutputWithPast):
143
+
144
+ def __init__(self, mtp_hidden_states=None, **kwargs):
145
+ super().__init__(**kwargs)
146
+ self.mtp_hidden_states = mtp_hidden_states
147
+
148
+
149
+ def index_first_axis(x, indices):
150
+ other_shape = x.shape[1:]
151
+ second_dim = other_shape.numel()
152
+ return torch.gather(
153
+ rearrange(x, "b ... -> b (...)"),
154
+ 0,
155
+ repeat(indices, "z -> z d", d=second_dim),
156
+ ).reshape(-1, *other_shape)
157
+
158
+
159
+ def index_put_first_axis(x, indices, first_axis_dim):
160
+ y = torch.zeros(first_axis_dim, *x.shape[1:], device=x.device, dtype=x.dtype)
161
+ y[indices] = x
162
+ # y.scatter_(0, repeat(indices, 'z -> z d', d=x.shape[1]), x)
163
+ return y
164
+
165
+
166
+ def pad_input(
167
+ hidden_states: torch.Tensor,
168
+ indices: torch.LongTensor,
169
+ batch_size: int,
170
+ seq_len: int,
171
+ ) -> torch.Tensor:
172
+ output = index_put_first_axis(hidden_states, indices, batch_size * seq_len)
173
+ return rearrange(output, "(b s) ... -> b s ...", b=batch_size)
174
+
175
+
176
+ @tensor_cache
177
+ def _get_unpad_data(attention_mask):
178
+ seqlens_in_batch = attention_mask.sum(dim=-1, dtype=torch.int32)
179
+ indices = torch.nonzero(attention_mask.flatten(), as_tuple=False).flatten()
180
+ max_seqlen_in_batch = seqlens_in_batch.max().item()
181
+ cu_seqlens = F.pad(torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.torch.int32), (1, 0))
182
+ return (
183
+ indices,
184
+ cu_seqlens,
185
+ max_seqlen_in_batch,
186
+ )
187
+
188
+
189
+ def _expand_mask(mask: torch.Tensor, dtype: torch.dtype, tgt_len: Optional[int] = None):
190
+ warnings.warn(
191
+ "Calling `transformers.models.BailingMoeV3.modeling_BailingMoeV3._prepare_4d_attention_mask` is deprecated and will be removed in v4.37. Use `transformers.modeling_attn_mask_utils._prepare_4d_attention_mask"
192
+ )
193
+ return _prepare_4d_attention_mask(mask=mask, dtype=dtype, tgt_len=tgt_len)
194
+
195
+
196
+ def _make_causal_mask(
197
+ input_ids_shape: torch.Size, dtype: torch.dtype, device: torch.device, past_key_values_length: int = 0
198
+ ):
199
+ warnings.warn(
200
+ "Calling `transformers.models.BailingMoeV3.modeling_BailingMoeV3._make_causal_mask` is deprecated and will be removed in v4.37. Use `transformers.models.BailingMoeV3.modeling_BailingMoeV3.AttentionMaskConverter._make_causal_mask"
201
+ )
202
+ return AttentionMaskConverter._make_causal_mask(
203
+ input_ids_shape=input_ids_shape, dtype=dtype, device=device, past_key_values_length=past_key_values_length
204
+ )
205
+
206
+
207
+ class BailingMoeV3RMSNorm(nn.Module):
208
+ def __init__(self, hidden_size, eps=1e-6):
209
+ """
210
+ BailingMoeV3RMSNorm is equivalent to T5LayerNorm
211
+ """
212
+ super().__init__()
213
+ self.weight = nn.Parameter(torch.ones(hidden_size))
214
+ self.variance_epsilon = eps
215
+
216
+ def forward(self, hidden_states):
217
+ input_dtype = hidden_states.dtype
218
+ hidden_states = hidden_states.to(torch.float32)
219
+ variance = hidden_states.pow(2).mean(-1, keepdim=True)
220
+ hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
221
+ return self.weight * hidden_states.to(input_dtype)
222
+
223
+
224
+ class BailingMoeV3GroupRMSNorm(nn.Module):
225
+ def __init__(self, hidden_size, group_norm_size, eps=1e-6):
226
+ """
227
+ BailingMoeV3RMSNorm is equivalent to T5LayerNorm
228
+ """
229
+ super().__init__()
230
+ self.weight = nn.Parameter(torch.ones(hidden_size))
231
+ self.group_norm_size = group_norm_size
232
+ assert hidden_size % group_norm_size == 0, "hidden_size must be divisible by group_norm_size"
233
+ self.variance_epsilon = eps
234
+
235
+ def forward(self, hidden_states):
236
+ input_dtype = hidden_states.dtype
237
+ input_shape = hidden_states.size()
238
+ group_input_shape = input_shape[:-1] + (self.group_norm_size, input_shape[-1] // self.group_norm_size)
239
+ hidden_states = hidden_states.view(group_input_shape)
240
+ hidden_states = hidden_states.to(torch.float32)
241
+ variance = hidden_states.pow(2).mean(-1, keepdim=True)
242
+ hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
243
+ return self.weight * hidden_states.to(input_dtype).view(input_shape)
244
+
245
+
246
+ ALL_LAYERNORM_LAYERS.append(BailingMoeV3RMSNorm)
247
+
248
+
249
+ class BailingMoeV3RotaryEmbedding(nn.Module):
250
+ def __init__(self, config: BailingMoeV3Config, device=None):
251
+ super().__init__()
252
+ # BC: "rope_type" was originally "type"
253
+ if hasattr(config, "rope_scaling") and config.rope_scaling is not None:
254
+ self.rope_type = config.rope_scaling.get("rope_type", config.rope_scaling.get("type"))
255
+ else:
256
+ self.rope_type = "default"
257
+ self.max_seq_len_cached = config.max_position_embeddings
258
+ self.original_max_seq_len = config.max_position_embeddings
259
+
260
+ self.config = deepcopy(config)
261
+ self.config.head_dim = config.qk_rope_head_dim
262
+ self.config.partial_rotary_factor = 1.0
263
+ self.rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type]
264
+
265
+ inv_freq, self.attention_scaling = self.rope_init_fn(self.config, device)
266
+ self.register_buffer("inv_freq", inv_freq, persistent=False)
267
+ self.original_inv_freq = self.inv_freq
268
+
269
+ @torch.no_grad()
270
+ @dynamic_rope_update # power user: used with advanced RoPE types (e.g. dynamic rope)
271
+ def forward(self, x, position_ids):
272
+ inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1).to(x.device)
273
+ position_ids_expanded = position_ids[:, None, :].float()
274
+
275
+ device_type = x.device.type if isinstance(x.device.type, str) and x.device.type != "mps" else "cpu"
276
+ with torch.autocast(device_type=device_type, enabled=False): # Force float32
277
+ freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2)
278
+ emb = torch.cat((freqs, freqs), dim=-1)
279
+ cos = emb.cos() * self.attention_scaling
280
+ sin = emb.sin() * self.attention_scaling
281
+
282
+ return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)
283
+
284
+
285
+ # Copied from transformers.models.llama.modeling_llama.rotate_half
286
+ def rotate_half(x):
287
+ """Rotates half the hidden dims of the input."""
288
+ x1 = x[..., : x.shape[-1] // 2]
289
+ x2 = x[..., x.shape[-1] // 2 :]
290
+ return torch.cat((-x2, x1), dim=-1)
291
+
292
+
293
+ # Copied from transformers.models.llama.modeling_llama.apply_rotary_pos_emb
294
+ def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1):
295
+ """Applies Rotary Position Embedding to the query and key tensors.
296
+ Args:
297
+ q (`torch.Tensor`): The query tensor.
298
+ k (`torch.Tensor`): The key tensor.
299
+ cos (`torch.Tensor`): The cosine part of the rotary embedding.
300
+ sin (`torch.Tensor`): The sine part of the rotary embedding.
301
+ unsqueeze_dim (`int`, *optional*, defaults to 1):
302
+ The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and
303
+ sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note
304
+ that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and
305
+ k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes
306
+ cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have
307
+ the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2.
308
+ Returns:
309
+ `tuple(torch.Tensor)` comprising the query and key tensors rotated using the Rotary Position Embedding.
310
+ """
311
+ cos = cos.unsqueeze(unsqueeze_dim)
312
+ sin = sin.unsqueeze(unsqueeze_dim)
313
+
314
+ # Keep half or full tensor for later concatenation
315
+ rotary_dim = cos.shape[-1]
316
+ q_rot, q_pass = q[..., :rotary_dim], q[..., rotary_dim:]
317
+ k_rot, k_pass = k[..., :rotary_dim], k[..., rotary_dim:]
318
+
319
+ # Apply rotary embeddings on the first half or full tensor
320
+ q_embed = (q_rot * cos) + (rotate_half(q_rot) * sin)
321
+ k_embed = (k_rot * cos) + (rotate_half(k_rot) * sin)
322
+
323
+ # Concatenate back to full shape
324
+ q_embed = torch.cat([q_embed, q_pass], dim=-1)
325
+ k_embed = torch.cat([k_embed, k_pass], dim=-1)
326
+ return q_embed, k_embed
327
+
328
+
329
+ class BailingMoeV3MLP(nn.Module):
330
+ def __init__(self, config: BailingMoeV3Config, intermediate_size: int):
331
+ super().__init__()
332
+ self.config = config
333
+ self.hidden_size = config.hidden_size
334
+ self.intermediate_size = intermediate_size
335
+
336
+ self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
337
+ self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
338
+ self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)
339
+ self.act_fn = ACT2FN[config.hidden_act]
340
+
341
+ def forward(self, x):
342
+ return self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
343
+
344
+
345
+ class BailingMoeV3Gate(nn.Module):
346
+ def __init__(self, config):
347
+ super().__init__()
348
+ self.config = config
349
+ self.top_k = config.num_experts_per_tok
350
+ self.num_experts = config.num_experts
351
+
352
+ self.n_group = config.n_group
353
+ self.topk_group = config.topk_group
354
+
355
+ # topk selection algorithm
356
+ self.gating_dim = config.hidden_size
357
+ self.weight = nn.Parameter(torch.empty((self.num_experts, self.gating_dim)))
358
+ self.routed_scaling_factor = config.routed_scaling_factor
359
+
360
+ self.register_buffer("expert_bias", torch.zeros((self.num_experts)))
361
+ self.reset_parameters()
362
+
363
+ def reset_parameters(self) -> None:
364
+ import torch.nn.init as init
365
+
366
+ init.kaiming_uniform_(self.weight, a=math.sqrt(5))
367
+
368
+ def group_limited_topk(
369
+ self,
370
+ scores: torch.Tensor,
371
+ ):
372
+ num_tokens, _ = scores.size()
373
+ # Organize the experts into groups
374
+ group_scores = scores.view(num_tokens, self.n_group, -1).topk(2, dim=-1)[0].sum(dim=-1)
375
+ group_idx = torch.topk(group_scores, k=self.topk_group, dim=-1, sorted=False)[1]
376
+ group_mask = torch.zeros_like(group_scores)
377
+ group_mask.scatter_(1, group_idx, 1)
378
+
379
+ # Mask the experts based on selection groups
380
+ score_mask = (
381
+ group_mask.unsqueeze(-1)
382
+ .expand(num_tokens, self.n_group, self.num_experts // self.n_group)
383
+ .reshape(num_tokens, -1)
384
+ )
385
+
386
+ masked_scores = scores.masked_fill(~score_mask.bool(), float('-inf'))
387
+ probs, top_indices = torch.topk(masked_scores, k=self.top_k, dim=-1)
388
+
389
+ return probs, top_indices
390
+
391
+ def forward(self, hidden_states):
392
+ # compute gating score
393
+ hidden_states = hidden_states.view(-1, hidden_states.shape[-1])
394
+ logits = F.linear(hidden_states.type(torch.float32), self.weight.type(torch.float32))
395
+
396
+ scores = torch.sigmoid(logits.float()).type_as(logits)
397
+
398
+ scores_for_routing = scores + self.expert_bias
399
+ _, topk_idx = self.group_limited_topk(scores_for_routing)
400
+
401
+ scores = torch.gather(scores, dim=1, index=topk_idx).type_as(logits)
402
+
403
+ topk_weight = scores / (scores.sum(dim=-1, keepdim=True) + 1e-20) if self.top_k > 1 else scores
404
+ topk_weight = topk_weight * self.routed_scaling_factor
405
+
406
+ return topk_idx, topk_weight, logits
407
+
408
+
409
+ class BailingMoeV3SparseMoeBlock(nn.Module):
410
+ """
411
+ A mixed expert module containing shared experts.
412
+ """
413
+
414
+ def __init__(self, config: BailingMoeV3Config):
415
+ super().__init__()
416
+ self.config = config
417
+ self.num_experts_per_tok = config.num_experts_per_tok
418
+ self._setup_experts()
419
+ self.gate = BailingMoeV3Gate(config)
420
+ if config.num_shared_experts is not None:
421
+ self.shared_experts = BailingMoeV3MLP(
422
+ config=config, intermediate_size=config.moe_shared_expert_intermediate_size * config.num_shared_experts
423
+ )
424
+
425
+ def _setup_experts(self):
426
+ self.experts = nn.ModuleList(
427
+ [
428
+ BailingMoeV3MLP(config=self.config, intermediate_size=self.config.moe_intermediate_size)
429
+ for _ in range(self.config.num_experts)
430
+ ]
431
+ )
432
+
433
+ def forward(self, hidden_states):
434
+ identity = hidden_states
435
+ bsz, seq_len, h = hidden_states.shape
436
+ topk_idx, topk_weight, router_logits = self.gate(hidden_states)
437
+ hidden_states = hidden_states.view(-1, hidden_states.shape[-1])
438
+ flat_topk_idx = topk_idx.view(-1)
439
+ if self.training:
440
+ hidden_states = hidden_states.repeat_interleave(self.num_experts_per_tok, dim=0)
441
+ y = torch.empty_like(hidden_states)
442
+ for i, expert in enumerate(self.experts):
443
+ y[flat_topk_idx == i] = expert(hidden_states[flat_topk_idx == i])
444
+ y = (y.view(*topk_weight.shape, -1) * topk_weight.unsqueeze(-1)).sum(dim=1)
445
+ y = y.to(hidden_states.dtype).view(bsz, seq_len, h)
446
+ else:
447
+ y = self.moe_infer(hidden_states, topk_idx, topk_weight).view(bsz, seq_len, h)
448
+ if self.config.num_shared_experts is not None:
449
+ y = y + self.shared_experts(identity)
450
+ return y, (router_logits.view(bsz, seq_len, -1), topk_idx.view(bsz, seq_len, -1))
451
+
452
+ @torch.no_grad()
453
+ def moe_infer(self, x, topk_ids, topk_weight):
454
+ cnts = topk_ids.new_zeros((topk_ids.shape[0], len(self.experts)))
455
+ cnts.scatter_(1, topk_ids, 1)
456
+ tokens_per_expert = cnts.sum(dim=0)
457
+ idxs = topk_ids.view(-1).argsort()
458
+ sorted_tokens = x[idxs // topk_ids.shape[1]]
459
+ tokens_per_expert = tokens_per_expert.cpu().numpy()
460
+ outputs = []
461
+ start_idx = 0
462
+ for i, num_tokens in enumerate(tokens_per_expert):
463
+ end_idx = start_idx + num_tokens
464
+ if num_tokens == 0:
465
+ continue
466
+ expert = self.experts[i]
467
+ tokens_for_this_expert = sorted_tokens[start_idx:end_idx]
468
+ expert_out = expert(tokens_for_this_expert)
469
+ outputs.append(expert_out.to(x.device))
470
+ start_idx = end_idx
471
+
472
+ outs = torch.cat(outputs, dim=0) if len(outputs) else sorted_tokens.new_empty(0)
473
+ new_x = torch.empty_like(outs)
474
+ new_x[idxs] = outs
475
+ final_out = (
476
+ new_x.view(*topk_ids.shape, -1)
477
+ .type(topk_weight.dtype)
478
+ .mul_(topk_weight.unsqueeze(dim=-1))
479
+ .sum(dim=1)
480
+ .type(new_x.dtype)
481
+ )
482
+ return final_out
483
+
484
+
485
+ # Copied from transformers.models.llama.modeling_llama.repeat_kv
486
+ def repeat_kv(hidden_states: torch.Tensor, n_rep: int, head_first: bool = True) -> torch.Tensor:
487
+ """
488
+ This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). If head_first is True, the hidden states go from (batch,
489
+ num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim)
490
+ """
491
+ if n_rep == 1:
492
+ return hidden_states
493
+ if head_first:
494
+ batch, num_key_value_heads, slen, head_dim = hidden_states.shape
495
+ hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim)
496
+ return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)
497
+ else:
498
+ batch, slen, num_key_value_heads, head_dim = hidden_states.shape
499
+ hidden_states = hidden_states[:, :, :, None, :].expand(batch, slen, num_key_value_heads, n_rep, head_dim)
500
+ return hidden_states.reshape(batch, slen, num_key_value_heads * n_rep, head_dim)
501
+
502
+
503
+ def repeat_kv2(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
504
+ """
505
+ This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch,
506
+ num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim)
507
+ """
508
+ batch, num_key_value_heads, slen, head_dim = hidden_states.shape
509
+ if n_rep == 1:
510
+ return hidden_states
511
+ hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim)
512
+ return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)
513
+
514
+
515
+ def eager_attention_forward(
516
+ module: nn.Module,
517
+ query: torch.Tensor,
518
+ key: torch.Tensor,
519
+ value: torch.Tensor,
520
+ attention_mask: Optional[torch.Tensor],
521
+ scaling: float,
522
+ dropout: float = 0.0,
523
+ **kwargs: Unpack[TransformersKwargs],
524
+ ):
525
+ key_states = repeat_kv2(key, module.num_key_value_groups)
526
+ value_states = repeat_kv2(value, module.num_key_value_groups)
527
+
528
+ attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling
529
+ if attention_mask is not None:
530
+ causal_mask = attention_mask[:, :, :, : key_states.shape[-2]]
531
+ attn_weights = attn_weights + causal_mask
532
+
533
+ attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype)
534
+ attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training)
535
+ attn_output = torch.matmul(attn_weights, value_states)
536
+ attn_output = attn_output.transpose(1, 2).contiguous()
537
+
538
+ return attn_output, attn_weights
539
+
540
+
541
+ def apply_rotary_pos_emb_interleave(q, k, cos, sin, position_ids=None, unsqueeze_dim=1):
542
+ r"""
543
+ TODO let's just use the original freqcis computation to not have the view
544
+ transpose + reshape! This is not optimized!
545
+ Applies Rotary Position Embedding to the query and key tensors.
546
+
547
+ Args:
548
+ q (`torch.Tensor`): The query tensor.
549
+ k (`torch.Tensor`): The key tensor.
550
+ cos (`torch.Tensor`): The cosine part of the rotary embedding.
551
+ sin (`torch.Tensor`): The sine part of the rotary embedding.
552
+ position_ids (`torch.Tensor`):
553
+ The position indices of the tokens corresponding to the query and key tensors. For example, this can be
554
+ used to pass offsetted position ids when working with a KV-cache.
555
+ unsqueeze_dim (`int`, *optional*, defaults to 1):
556
+ The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and
557
+ sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note
558
+ that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and
559
+ k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes
560
+ cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have
561
+ the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2.
562
+ Returns:
563
+ `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding.
564
+ """
565
+ cos = cos.unsqueeze(unsqueeze_dim)
566
+ sin = sin.unsqueeze(unsqueeze_dim)
567
+
568
+ b, h, s, d = q.shape
569
+ q = q.view(b, h, s, d // 2, 2).transpose(4, 3).reshape(b, h, s, d)
570
+
571
+ b, h, s, d = k.shape
572
+ k = k.view(b, h, s, d // 2, 2).transpose(4, 3).reshape(b, h, s, d)
573
+
574
+ q_embed = (q * cos) + (rotate_half(q) * sin)
575
+ k_embed = (k * cos) + (rotate_half(k) * sin)
576
+ return q_embed, k_embed
577
+
578
+
579
+ def yarn_get_mscale(scale=1, mscale=1):
580
+ if scale <= 1:
581
+ return 1.0
582
+ return 0.1 * mscale * math.log(scale) + 1.0
583
+
584
+
585
+ class BailingMoeV3MultiLatentAttention(nn.Module):
586
+ """Multi-headed attention from 'Attention Is All You Need' paper"""
587
+
588
+ def __init__(self, config: BailingMoeV3Config, layer_idx: int):
589
+ super().__init__()
590
+ self.config = config
591
+ self.layer_idx = layer_idx
592
+ self.num_key_value_groups = config.num_attention_heads // config.num_key_value_heads
593
+ self.attention_dropout = config.attention_dropout
594
+ self.num_heads = config.num_attention_heads
595
+ self.rope_theta = config.rope_theta
596
+ self.q_lora_rank = config.q_lora_rank
597
+ self.qk_rope_head_dim = config.qk_rope_head_dim
598
+ self.kv_lora_rank = config.kv_lora_rank
599
+ self.v_head_dim = config.v_head_dim
600
+ self.qk_nope_head_dim = config.qk_nope_head_dim
601
+ self.qk_head_dim = config.qk_head_dim
602
+ self.gated_attention_proj_granularity_type = config.gated_attention_proj_granularity_type
603
+
604
+ self.is_causal = True
605
+ if self.q_lora_rank is None:
606
+ self.q_proj = nn.Linear(config.hidden_size, self.num_heads * self.qk_head_dim, bias=False)
607
+ else:
608
+ self.q_a_proj = nn.Linear(config.hidden_size, config.q_lora_rank, bias=config.use_qkv_bias)
609
+ self.q_a_layernorm = BailingMoeV3RMSNorm(config.q_lora_rank)
610
+ self.q_b_proj = nn.Linear(config.q_lora_rank, self.num_heads * self.qk_head_dim, bias=False)
611
+
612
+ self.kv_a_proj_with_mqa = nn.Linear(
613
+ config.hidden_size,
614
+ self.kv_lora_rank + self.qk_rope_head_dim,
615
+ bias=config.use_qkv_bias,
616
+ )
617
+ self.kv_a_layernorm = BailingMoeV3RMSNorm(self.kv_lora_rank)
618
+ self.kv_b_proj = nn.Linear(
619
+ self.kv_lora_rank,
620
+ self.num_heads * (self.qk_nope_head_dim + self.v_head_dim),
621
+ bias=False,
622
+ )
623
+ if self.gated_attention_proj_granularity_type is None:
624
+ self.g_proj = None
625
+ elif self.gated_attention_proj_granularity_type == "head_wise":
626
+ self.g_proj = nn.Linear(config.hidden_size, self.num_heads, bias=False)
627
+ elif self.gated_attention_proj_granularity_type == "element_wise":
628
+ self.g_proj = nn.Linear(config.hidden_size, self.num_heads * self.v_head_dim, bias=False)
629
+
630
+ self.dense = nn.Linear(
631
+ self.num_heads * self.v_head_dim,
632
+ config.hidden_size,
633
+ bias=config.use_qkv_bias,
634
+ )
635
+
636
+ self.scaling = self.qk_head_dim ** (-0.5)
637
+ if self.config.rope_scaling is not None:
638
+ mscale_all_dim = self.config.rope_scaling.get("mscale_all_dim", 0)
639
+ scaling_factor = self.config.rope_scaling["factor"]
640
+ if mscale_all_dim:
641
+ mscale = yarn_get_mscale(scaling_factor, mscale_all_dim)
642
+ self.scaling = self.scaling * mscale * mscale
643
+
644
+ @deprecate_kwarg("past_key_value", new_name="past_key_values", version="4.58")
645
+ def forward(
646
+ self,
647
+ hidden_states: torch.Tensor,
648
+ position_embeddings: tuple[torch.Tensor, torch.Tensor],
649
+ attention_mask: Optional[torch.Tensor],
650
+ past_key_values: Optional[Cache] = None,
651
+ cache_position: Optional[torch.LongTensor] = None,
652
+ **kwargs: Unpack[FlashAttentionKwargs],
653
+ ) -> tuple[torch.Tensor, Optional[torch.Tensor], Optional[tuple[torch.Tensor]]]:
654
+
655
+ batch_size, seq_length = hidden_states.shape[:-1]
656
+ query_shape = (batch_size, seq_length, -1, self.qk_head_dim)
657
+ key_shape = (batch_size, seq_length, -1, self.qk_nope_head_dim + self.v_head_dim)
658
+
659
+ if self.q_lora_rank is None:
660
+ q_states = self.q_proj(hidden_states)
661
+ else:
662
+ q_states = self.q_b_proj(self.q_a_layernorm(self.q_a_proj(hidden_states)))
663
+ q_states = q_states.view(query_shape).transpose(1, 2)
664
+ q_pass, q_rot = torch.split(q_states, [self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1)
665
+
666
+ compressed_kv = self.kv_a_proj_with_mqa(hidden_states)
667
+ k_pass, k_rot = torch.split(compressed_kv, [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1)
668
+
669
+ k_pass = self.kv_b_proj(self.kv_a_layernorm(k_pass)).view(key_shape).transpose(1, 2)
670
+ k_pass, value_states = torch.split(k_pass, [self.qk_nope_head_dim, self.v_head_dim], dim=-1)
671
+
672
+ k_rot = k_rot.view(batch_size, 1, seq_length, self.qk_rope_head_dim)
673
+
674
+ cos, sin = position_embeddings # tptest
675
+ if self.config.rope_interleave: # support using interleaved weights for efficiency
676
+ q_rot, k_rot = apply_rotary_pos_emb_interleave(q_rot, k_rot, cos, sin)
677
+ else:
678
+ x = 1 / 0
679
+ q_rot, k_rot = apply_rotary_pos_emb(q_rot, k_rot, cos, sin)
680
+ k_rot = k_rot.expand(*k_pass.shape[:-1], -1)
681
+
682
+ query_states = torch.cat((q_pass, q_rot), dim=-1)
683
+ key_states = torch.cat((k_pass, k_rot), dim=-1)
684
+
685
+ if past_key_values is not None:
686
+ # sin and cos are specific to RoPE models; cache_position needed for the static cache
687
+ cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position}
688
+ key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx, cache_kwargs)
689
+
690
+ if self.config._attn_implementation == "flash_attention_2" and self.qk_head_dim != self.v_head_dim:
691
+ value_states = F.pad(value_states, [0, self.qk_head_dim - self.v_head_dim])
692
+
693
+ attention_interface: Callable = eager_attention_forward
694
+
695
+ attn_output, attn_weights = attention_interface(
696
+ self,
697
+ query_states,
698
+ key_states,
699
+ value_states,
700
+ attention_mask,
701
+ dropout=0.0 if not self.training else self.attention_dropout,
702
+ scaling=self.scaling,
703
+ **kwargs,
704
+ )
705
+
706
+ if self.config._attn_implementation == "flash_attention_2" and self.qk_head_dim != self.v_head_dim:
707
+ attn_output = attn_output[:, :, :, : self.v_head_dim]
708
+
709
+ if self.g_proj is not None:
710
+ gate = self.g_proj(hidden_states)
711
+ gate = F.sigmoid(gate.float()).type_as(hidden_states)
712
+ if self.gated_attention_proj_granularity_type == "head_wise":
713
+ attn_output = attn_output * gate[:, :, :, None]
714
+ else:
715
+ attn_output = attn_output * gate.view(batch_size, seq_length, self.num_heads, self.v_head_dim)
716
+
717
+ attn_output = attn_output.reshape(batch_size, seq_length, -1).contiguous()
718
+ attn_output = self.dense(attn_output)
719
+ return attn_output, attn_weights, past_key_values
720
+
721
+
722
+ class BailingMoeV3KimiDeltaAttention(nn.Module):
723
+ def __init__(self, config: BailingMoeV3Config, layer_idx: int):
724
+ super().__init__()
725
+ self.config = config
726
+ self.mode = "chunk"
727
+
728
+ self.hidden_size = config.hidden_size
729
+ self.conv_size = config.short_conv_kernel_size
730
+ self.head_dim = config.head_dim
731
+ self.num_heads = config.num_attention_heads
732
+ self.head_k_dim = self.head_dim
733
+ self.num_k_heads = self.num_heads
734
+ self.no_kda_lora = config.no_kda_lora
735
+ self.safe_gate = config.kda_safe_gate
736
+ self.lower_bound = config.kda_lower_bound
737
+
738
+ self.layer_idx = layer_idx
739
+
740
+ assert self.mode in ['chunk', 'fused_recurrent'], f"Not suppoerted mode `{self.mode}`."
741
+
742
+ projection_k_size = self.head_k_dim * self.num_k_heads
743
+ projection_size = self.head_dim * self.num_heads
744
+
745
+ self.q_proj = nn.Linear(self.hidden_size, projection_k_size, bias=False)
746
+ self.k_proj = nn.Linear(self.hidden_size, projection_k_size, bias=False)
747
+ self.v_proj = nn.Linear(self.hidden_size, projection_size, bias=False)
748
+
749
+ self.q_conv1d = ShortConvolution(
750
+ hidden_size=projection_k_size,
751
+ kernel_size=self.conv_size,
752
+ activation='silu',
753
+ )
754
+ self.k_conv1d = ShortConvolution(
755
+ hidden_size=projection_k_size,
756
+ kernel_size=self.conv_size,
757
+ activation='silu',
758
+ )
759
+ self.v_conv1d = ShortConvolution(
760
+ hidden_size=projection_size,
761
+ kernel_size=self.conv_size,
762
+ activation='silu',
763
+ )
764
+
765
+ self.A_log = torch.nn.Parameter(torch.log(torch.empty(self.num_heads, dtype=torch.float32).uniform_(1, 16)))
766
+
767
+ if self.no_kda_lora:
768
+ self.f_proj = nn.Linear(self.hidden_size, projection_size, bias=False)
769
+ else:
770
+ self.f_a_proj = nn.Linear(self.hidden_size, self.head_dim, bias=False)
771
+ self.f_b_proj = nn.Linear(self.head_dim, projection_size, bias=False)
772
+
773
+ self.dt_bias = nn.Parameter(torch.empty(projection_size, dtype=torch.float32))
774
+
775
+ self.b_proj = nn.Linear(self.hidden_size, self.num_heads, bias=False)
776
+
777
+ if self.no_kda_lora:
778
+ self.g_proj = nn.Linear(self.hidden_size, projection_size, bias=False)
779
+ else:
780
+ self.g_a_proj = nn.Linear(self.hidden_size, self.head_dim, bias=False)
781
+ self.g_b_proj = nn.Linear(self.head_dim, projection_size, bias=False)
782
+
783
+ self.o_norm = FusedRMSNormGated(self.head_dim, eps=config.rms_norm_eps, activation='sigmoid')
784
+ self.o_proj = nn.Linear(projection_size, self.hidden_size, bias=False)
785
+
786
+ def forward(
787
+ self,
788
+ hidden_states: torch.Tensor,
789
+ attention_mask: torch.Tensor | None = None,
790
+ past_key_value=None,
791
+ **kwargs: Unpack[dict],
792
+ ) -> tuple[torch.Tensor, torch.Tensor | None, Cache | None]:
793
+ attention_mask = None
794
+ if attention_mask is not None:
795
+ if attention_mask.dim() != 2:
796
+ attention_mask = kwargs.get("padding_mask")
797
+
798
+ if attention_mask is not None and attention_mask.dim() != 2:
799
+ raise ValueError(
800
+ "attention_mask must be a 0-1 matrix of shape [batch_size, seq_len] "
801
+ "(0 = padding). 3D masks are not supported here.",
802
+ )
803
+ use_cache = past_key_value is not None
804
+ batch_size, q_len, _ = hidden_states.shape
805
+ mode = 'fused_recurrent' if q_len <= 64 else self.mode
806
+ if self.training:
807
+ assert mode == 'chunk', "Only chunk mode is supported in training."
808
+
809
+ cu_seqlens = kwargs.get('cu_seqlens')
810
+ indices = None
811
+ if attention_mask is not None:
812
+ indices, cu_seqlens, _ = _get_unpad_data(attention_mask[:, -q_len:])
813
+ hidden_states = index_first_axis(rearrange(hidden_states, "b s ... -> (b s) ..."), indices).unsqueeze(0)
814
+
815
+ conv_state_q, conv_state_k, conv_state_v = None, None, None
816
+ recurrent_state = None
817
+
818
+ if past_key_value is not None and isinstance(past_key_value, Cache):
819
+ # ensure the cache list is long enough
820
+ while len(past_key_value.layers) <= self.layer_idx:
821
+ past_key_value.layers.append(DynamicLayer())
822
+
823
+ if past_key_value.layers[self.layer_idx].keys is not None:
824
+ recurrent_state = past_key_value.layers[self.layer_idx].keys
825
+ # ensure recurrent_state is on the same device as hidden_states
826
+ if recurrent_state.device != hidden_states.device:
827
+ recurrent_state = recurrent_state.to(hidden_states.device).contiguous()
828
+
829
+ if past_key_value.layers[self.layer_idx].values is not None:
830
+ conv_state_q, conv_state_k, conv_state_v = past_key_value.layers[self.layer_idx].values
831
+
832
+ q, conv_state_q = self.q_conv1d(
833
+ x=self.q_proj(hidden_states),
834
+ cache=conv_state_q,
835
+ output_final_state=use_cache,
836
+ cu_seqlens=cu_seqlens,
837
+ )
838
+ k, conv_state_k = self.k_conv1d(
839
+ x=self.k_proj(hidden_states),
840
+ cache=conv_state_k,
841
+ output_final_state=use_cache,
842
+ cu_seqlens=cu_seqlens,
843
+ )
844
+ v, conv_state_v = self.v_conv1d(
845
+ x=self.v_proj(hidden_states),
846
+ cache=conv_state_v,
847
+ output_final_state=use_cache,
848
+ cu_seqlens=cu_seqlens,
849
+ )
850
+
851
+ if self.no_kda_lora:
852
+ g = self.f_proj(hidden_states)
853
+ else:
854
+ g = self.f_b_proj(self.f_a_proj(hidden_states))
855
+
856
+ beta = self.b_proj(hidden_states).float().sigmoid()
857
+
858
+ q, k = map(lambda x: rearrange(x, '... (h d) -> ... h d', d=self.head_k_dim), (q, k))
859
+ v = rearrange(v, '... (h d) -> ... h d', d=self.head_dim)
860
+ g = rearrange(g, '... (h d) -> ... h d', d=self.head_dim)
861
+
862
+ if mode == 'chunk':
863
+ o, recurrent_state = chunk_kda(
864
+ q=q,
865
+ k=k,
866
+ v=v,
867
+ g=g,
868
+ beta=beta,
869
+ A_log=self.A_log,
870
+ dt_bias=self.dt_bias,
871
+ initial_state=recurrent_state,
872
+ output_final_state=True,
873
+ use_qk_l2norm_in_kernel=True,
874
+ use_gate_in_kernel=True,
875
+ safe_gate=self.safe_gate,
876
+ lower_bound=self.lower_bound,
877
+ cu_seqlens=cu_seqlens,
878
+ )
879
+ else:
880
+ o, recurrent_state = fused_recurrent_kda(
881
+ q=q,
882
+ k=k,
883
+ v=v,
884
+ g=g,
885
+ beta=beta,
886
+ A_log=self.A_log,
887
+ dt_bias=self.dt_bias,
888
+ initial_state=recurrent_state,
889
+ output_final_state=True,
890
+ use_qk_l2norm_in_kernel=True,
891
+ use_gate_in_kernel=True,
892
+ lower_bound=self.lower_bound,
893
+ cu_seqlens=cu_seqlens,
894
+ )
895
+
896
+ if use_cache and past_key_value is not None and isinstance(past_key_value, Cache):
897
+ target_device = None
898
+ for cache in past_key_value.layers:
899
+ if cache.keys is not None:
900
+ target_device = cache.keys.device
901
+ break
902
+ if target_device is None:
903
+ target_device = recurrent_state.device
904
+
905
+ # move to target device
906
+ if recurrent_state.device != target_device:
907
+ recurrent_state = recurrent_state.to(target_device)
908
+
909
+ past_key_value.layers[self.layer_idx].keys = recurrent_state
910
+ past_key_value.layers[self.layer_idx].values = (conv_state_q, conv_state_k, conv_state_v)
911
+
912
+ if self.no_kda_lora:
913
+ g = self.g_proj(hidden_states)
914
+ else:
915
+ g = self.g_b_proj(self.g_a_proj(hidden_states))
916
+ g = rearrange(g, '... (h d) -> ... h d', d=self.head_dim)
917
+ o = self.o_norm(o, g)
918
+
919
+ o = rearrange(o, 'b t h d -> b t (h d)')
920
+ o = self.o_proj(o)
921
+ if attention_mask is not None:
922
+ o = pad_input(o.squeeze(0), indices, batch_size, q_len)
923
+
924
+ return o, None, past_key_value
925
+
926
+
927
+ class BailingMoeV3MTPLayer(nn.Module):
928
+ def __init__(self, config: BailingMoeV3Config, layer_idx: int):
929
+ super().__init__()
930
+ self.layer_idx = layer_idx
931
+ self.input_layernorm = BailingMoeV3RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
932
+ self.enorm = BailingMoeV3RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
933
+
934
+ self.eh_proj = nn.Linear(config.hidden_size * 2, config.hidden_size, bias=False)
935
+ self.post_attention_layernorm = BailingMoeV3RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
936
+ self.attention = BailingMoeV3MultiLatentAttention(config=config, layer_idx=layer_idx)
937
+ self.mlp = BailingMoeV3SparseMoeBlock(config)
938
+
939
+ self.hnorm = BailingMoeV3RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
940
+ self.final_layernorm = BailingMoeV3RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
941
+
942
+ def forward(
943
+ self,
944
+ input_embeds,
945
+ hidden_states: torch.Tensor,
946
+ attention_mask: Optional[torch.Tensor] = None,
947
+ position_ids: Optional[torch.LongTensor] = None,
948
+ past_key_value: Optional[Tuple[torch.Tensor]] = None,
949
+ output_attentions: Optional[bool] = False,
950
+ output_router_logits: Optional[bool] = False,
951
+ use_cache: Optional[bool] = False,
952
+ position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, # necessary, but kept here for BC
953
+ **kwargs,
954
+ ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]:
955
+ input_embeds = self.enorm(input_embeds)
956
+ hidden_states = self.hnorm(hidden_states)
957
+ hidden_states = self.eh_proj(torch.cat([input_embeds, hidden_states], dim=-1))
958
+ residual = hidden_states
959
+
960
+ hidden_states = self.input_layernorm(hidden_states)
961
+
962
+ # Self Attention
963
+ hidden_states, self_attn_weights, present_key_value = self.attention(
964
+ hidden_states=hidden_states,
965
+ attention_mask=attention_mask,
966
+ position_ids=position_ids,
967
+ past_key_value=past_key_value,
968
+ output_attentions=output_attentions,
969
+ position_embeddings=position_embeddings,
970
+ use_cache=use_cache,
971
+ )
972
+ hidden_states = residual + hidden_states
973
+
974
+ # Fully Connected
975
+ residual = hidden_states
976
+ hidden_states = self.post_attention_layernorm(hidden_states)
977
+ hidden_states = self.mlp(hidden_states)
978
+ if isinstance(hidden_states, tuple):
979
+ hidden_states, router_logits = hidden_states
980
+ else:
981
+ router_logits = None
982
+ hidden_states = residual + hidden_states.to(residual.device)
983
+ hidden_states = self.final_layernorm(hidden_states)
984
+
985
+ outputs = (hidden_states,)
986
+
987
+ if output_attentions:
988
+ outputs += (self_attn_weights,)
989
+
990
+ if use_cache:
991
+ outputs += (present_key_value,)
992
+
993
+ if output_router_logits:
994
+ outputs += (router_logits,)
995
+
996
+ return outputs
997
+
998
+
999
+ class BailingMoeV3DecoderLayer(nn.Module):
1000
+ def __init__(self, config: BailingMoeV3Config, layer_idx: int):
1001
+ super().__init__()
1002
+ self.hidden_size = config.hidden_size
1003
+ self.layer_idx = layer_idx
1004
+ self.attention_layer_type = (
1005
+ "attention"
1006
+ if (layer_idx + 1) % config.layer_group_size == 0
1007
+ or layer_idx >= config.num_hidden_layers // config.layer_group_size * config.layer_group_size
1008
+ else "linear_attention"
1009
+ )
1010
+
1011
+ if self.attention_layer_type == "attention":
1012
+ self.attention = BailingMoeV3MultiLatentAttention(config=config, layer_idx=layer_idx)
1013
+ else:
1014
+ self.attention = BailingMoeV3KimiDeltaAttention(config=config, layer_idx=layer_idx)
1015
+
1016
+ self.mlp = (
1017
+ BailingMoeV3SparseMoeBlock(config)
1018
+ if (config.num_experts is not None and layer_idx >= config.first_k_dense_replace)
1019
+ else BailingMoeV3MLP(config=config, intermediate_size=config.intermediate_size)
1020
+ )
1021
+ self.input_layernorm = BailingMoeV3RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
1022
+ self.post_attention_layernorm = BailingMoeV3RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
1023
+
1024
+ def forward(
1025
+ self,
1026
+ hidden_states: torch.Tensor,
1027
+ attention_mask: Optional[torch.Tensor] = None,
1028
+ position_ids: Optional[torch.LongTensor] = None,
1029
+ past_key_value: Optional[Tuple[torch.Tensor]] = None,
1030
+ cache_position: Optional[torch.LongTensor] = None,
1031
+ output_attentions: Optional[bool] = False,
1032
+ output_router_logits: Optional[bool] = False,
1033
+ use_cache: Optional[bool] = False,
1034
+ position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, # necessary, but kept here for BC
1035
+ **kwargs,
1036
+ ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]:
1037
+ """
1038
+ Args:
1039
+ hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)`
1040
+ attention_mask (`torch.FloatTensor`, *optional*):
1041
+ attention mask of size `(batch_size, sequence_length)` if flash attention is used or `(batch_size, 1,
1042
+ query_sequence_length, key_sequence_length)` if default attention is used.
1043
+ position_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
1044
+ Indices of positions of each input sequence tokens in the position embeddings. Selected in the range `[0,
1045
+ config.n_positions - 1]`.
1046
+ past_key_value (`Tuple(torch.FloatTensor)`, *optional*):
1047
+ cached past key and value projection states
1048
+ output_attentions (`bool`, *optional*):
1049
+ Whether to return the attentions tensors of all attention layers. See `attentions` under
1050
+ returned tensors for more detail.
1051
+ output_router_logits (`bool`, *optional*):
1052
+ Whether or not to return the logits of all the routers. They are useful for computing the router loss,
1053
+ and should not be returned during inference.
1054
+ use_cache (`bool`, *optional*):
1055
+ If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding
1056
+ (see `past_key_values`).
1057
+ """
1058
+ residual = hidden_states
1059
+
1060
+ hidden_states = self.input_layernorm(hidden_states)
1061
+
1062
+ # Self Attention
1063
+ if self.attention_layer_type == "attention":
1064
+ hidden_states, self_attn_weights, present_key_value = self.attention(
1065
+ hidden_states=hidden_states,
1066
+ attention_mask=attention_mask,
1067
+ position_ids=position_ids,
1068
+ past_key_values=past_key_value,
1069
+ use_cache=use_cache,
1070
+ cache_position=cache_position, #
1071
+ position_embeddings=position_embeddings, #
1072
+ **kwargs,
1073
+ )
1074
+ else:
1075
+ batch_size, seq_len = hidden_states.shape[0], hidden_states.shape[1]
1076
+ device = hidden_states.device
1077
+
1078
+ if attention_mask is None:
1079
+ # if attention_mask is None, create a full mask
1080
+ attention_mask = torch.ones((batch_size, seq_len), dtype=torch.int32, device=device)
1081
+ elif attention_mask.dim() == 4 and attention_mask.shape[1] == 1:
1082
+ attention_mask = attention_mask[:, 0, -1, :].to(torch.int32)
1083
+ attention_mask = (attention_mask > -1e4).to(torch.int32)
1084
+ elif attention_mask.dim() == 2:
1085
+ attention_mask = attention_mask.to(torch.int32)
1086
+ else:
1087
+ raise ValueError(f"Unsupported mask dimension: {attention_mask.shape}")
1088
+
1089
+ hidden_states, self_attn_weights, present_key_value = self.attention(
1090
+ hidden_states=hidden_states,
1091
+ attention_mask=attention_mask,
1092
+ past_key_value=past_key_value,
1093
+ position_ids=position_ids,
1094
+ use_cache=use_cache,
1095
+ output_attentions=output_attentions,
1096
+ )
1097
+
1098
+ hidden_states = residual + hidden_states
1099
+
1100
+ # Fully Connected
1101
+ residual = hidden_states
1102
+ hidden_states = self.post_attention_layernorm(hidden_states)
1103
+ hidden_states = self.mlp(hidden_states)
1104
+ if isinstance(hidden_states, tuple):
1105
+ hidden_states, router_logits = hidden_states
1106
+ else:
1107
+ router_logits = None
1108
+ hidden_states = residual + hidden_states.to(residual.device)
1109
+
1110
+ outputs = (hidden_states,)
1111
+
1112
+ if output_attentions:
1113
+ outputs += (self_attn_weights,)
1114
+
1115
+ if use_cache:
1116
+ outputs += (present_key_value,)
1117
+
1118
+ if output_router_logits:
1119
+ outputs += (router_logits,)
1120
+
1121
+ return outputs
1122
+
1123
+
1124
+ BAILINGMOEV3_START_DOCSTRING = r"""
1125
+ This model inherits from [`PreTrainedModel`]. Check the superclass documentation for the generic methods the
1126
+ library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads
1127
+ etc.)
1128
+ This model is also a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass.
1129
+ Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage
1130
+ and behavior.
1131
+ Parameters:
1132
+ config ([`BailingMoeV3Config`]):
1133
+ Model configuration class with all the parameters of the model. Initializing with a config file does not
1134
+ load the weights associated with the model, only the configuration. Check out the
1135
+ [`~PreTrainedModel.from_pretrained`] method to load the model weights.
1136
+ """
1137
+
1138
+
1139
+ @add_start_docstrings(
1140
+ "The bare BailingMoeV3 Model outputting raw hidden-states without any specific head on top.",
1141
+ BAILINGMOEV3_START_DOCSTRING,
1142
+ )
1143
+ class BailingMoeV3PreTrainedModel(PreTrainedModel):
1144
+ config_class = BailingMoeV3Config
1145
+ base_model_prefix = "model"
1146
+ supports_gradient_checkpointing = True
1147
+ _no_split_modules = ["BailingMoeV3DecoderLayer"]
1148
+ _skip_keys_device_placement = "past_key_values"
1149
+ _supports_flash_attn_2 = True
1150
+ _supports_sdpa = True
1151
+ _supports_cache_class = True
1152
+
1153
+ def _init_weights(self, module):
1154
+ std = self.config.initializer_range
1155
+ if isinstance(module, nn.Linear):
1156
+ module.weight.data.normal_(mean=0.0, std=std)
1157
+ if module.bias is not None:
1158
+ module.bias.data.zero_()
1159
+ elif isinstance(module, nn.Embedding):
1160
+ module.weight.data.normal_(mean=0.0, std=std)
1161
+ if module.padding_idx is not None:
1162
+ module.weight.data[module.padding_idx].zero_()
1163
+
1164
+
1165
+ BAILINGMOEV3_INPUTS_DOCSTRING = r"""
1166
+ Args:
1167
+ input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`):
1168
+ Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you provide
1169
+ it.
1170
+ Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
1171
+ [`PreTrainedTokenizer.__call__`] for details.
1172
+ [What are input IDs?](../glossary#input-ids)
1173
+ attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*):
1174
+ Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:
1175
+ - 1 for tokens that are **not masked**,
1176
+ - 0 for tokens that are **masked**.
1177
+ [What are attention masks?](../glossary#attention-mask)
1178
+ Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
1179
+ [`PreTrainedTokenizer.__call__`] for details.
1180
+ If `past_key_values` is used, optionally only the last `input_ids` have to be input (see
1181
+ `past_key_values`).
1182
+ If you want to change padding behavior, you should read [`modeling_opt._prepare_decoder_attention_mask`]
1183
+ and modify to your needs. See diagram 1 in [the paper](https://arxiv.org/abs/1910.13461) for more
1184
+ information on the default strategy.
1185
+ - 1 indicates the head is **not masked**,
1186
+ - 0 indicates the head is **masked**.
1187
+ position_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
1188
+ Indices of positions of each input sequence tokens in the position embeddings. Selected in the range `[0,
1189
+ config.n_positions - 1]`.
1190
+ [What are position IDs?](../glossary#position-ids)
1191
+ past_key_values (`Cache` or `tuple(tuple(torch.FloatTensor))`, *optional*):
1192
+ Pre-computed hidden-states (key and values in the self-attention blocks and in the cross-attention
1193
+ blocks) that can be used to speed up sequential decoding. This typically consists in the `past_key_values`
1194
+ returned by the model at a previous stage of decoding, when `use_cache=True` or `config.use_cache=True`.
1195
+ Two formats are allowed:
1196
+ - a [`~cache_utils.Cache`] instance;
1197
+ - Tuple of `tuple(torch.FloatTensor)` of length `config.n_layers`, with each tuple having 2 tensors of
1198
+ shape `(batch_size, num_heads, sequence_length, embed_size_per_head)`). This is also known as the legacy
1199
+ cache format.
1200
+ The model will output the same cache format that is fed as input. If no `past_key_values` are passed, the
1201
+ legacy cache format will be returned.
1202
+ If `past_key_values` are used, the user can optionally input only the last `input_ids` (those that don't
1203
+ have their past key value states given to this model) of shape `(batch_size, 1)` instead of all `input_ids`
1204
+ of shape `(batch_size, sequence_length)`.
1205
+ inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*):
1206
+ Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation. This
1207
+ is useful if you want more control over how to convert `input_ids` indices into associated vectors than the
1208
+ model's internal embedding lookup matrix.
1209
+ use_cache (`bool`, *optional*):
1210
+ If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding (see
1211
+ `past_key_values`).
1212
+ output_attentions (`bool`, *optional*):
1213
+ Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned
1214
+ tensors for more detail.
1215
+ output_hidden_states (`bool`, *optional*):
1216
+ Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
1217
+ more detail.
1218
+ return_dict (`bool`, *optional*):
1219
+ Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
1220
+ """
1221
+
1222
+
1223
+ @add_start_docstrings(
1224
+ "The bare BailingMoeV3 Model outputting raw hidden-states without any specific head on top.",
1225
+ BAILINGMOEV3_START_DOCSTRING,
1226
+ )
1227
+ class BailingMoeV3Model(BailingMoeV3PreTrainedModel):
1228
+ """
1229
+ Transformer decoder consisting of *config.num_hidden_layers* layers. Each layer is a [`BailingMoeV3DecoderLayer`]
1230
+ Args:
1231
+ config: BailingMoeV3Config
1232
+ """
1233
+
1234
+ def __init__(self, config: BailingMoeV3Config):
1235
+ super().__init__(config)
1236
+ self.padding_idx = config.pad_token_id
1237
+ self.vocab_size = config.vocab_size
1238
+ self.num_nextn_predict_layers = config.num_nextn_predict_layers
1239
+
1240
+ self.word_embeddings = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)
1241
+ self.layers = []
1242
+ for layer_idx in range(config.num_hidden_layers + config.num_nextn_predict_layers):
1243
+ layer_cls = BailingMoeV3DecoderLayer if layer_idx < config.num_hidden_layers else BailingMoeV3MTPLayer
1244
+ self.layers.append(layer_cls(config, layer_idx))
1245
+
1246
+ self.layers = nn.ModuleList(self.layers)
1247
+
1248
+ self._use_sdpa = config._attn_implementation == "sdpa"
1249
+ self._use_flash_attention_2 = config._attn_implementation == "flash_attention_2"
1250
+ self.norm = BailingMoeV3RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
1251
+ self.rotary_emb = BailingMoeV3RotaryEmbedding(config=config)
1252
+ self.gradient_checkpointing = False
1253
+ # Initialize weights and apply final processing
1254
+ self.post_init()
1255
+
1256
+ def get_input_embeddings(self):
1257
+ return self.word_embeddings
1258
+
1259
+ def set_input_embeddings(self, value):
1260
+ self.word_embeddings = value
1261
+
1262
+ @add_start_docstrings_to_model_forward(BAILINGMOEV3_INPUTS_DOCSTRING)
1263
+ def forward(
1264
+ self,
1265
+ input_ids: torch.LongTensor = None,
1266
+ attention_mask: Optional[torch.Tensor] = None,
1267
+ position_ids: Optional[torch.LongTensor] = None,
1268
+ past_key_values: Optional[List[torch.FloatTensor]] = None,
1269
+ inputs_embeds: Optional[torch.FloatTensor] = None,
1270
+ cache_position: Optional[torch.LongTensor] = None,
1271
+ use_cache: Optional[bool] = None,
1272
+ output_attentions: Optional[bool] = None,
1273
+ output_hidden_states: Optional[bool] = None,
1274
+ output_router_logits: Optional[bool] = None,
1275
+ return_dict: Optional[bool] = None,
1276
+ **kwargs,
1277
+ ) -> Union[Tuple, MoeV3ModelOutputWithPast]:
1278
+ output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
1279
+ output_hidden_states = (
1280
+ output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
1281
+ )
1282
+ output_router_logits = (
1283
+ output_router_logits if output_router_logits is not None else self.config.output_router_logits
1284
+ )
1285
+ use_cache = use_cache if use_cache is not None else self.config.use_cache
1286
+
1287
+ return_dict = return_dict if return_dict is not None else self.config.use_return_dict
1288
+
1289
+ # retrieve input_ids and inputs_embeds
1290
+ if input_ids is not None and inputs_embeds is not None:
1291
+ raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
1292
+ elif input_ids is not None:
1293
+ batch_size, seq_length = input_ids.shape[:2]
1294
+ elif inputs_embeds is not None:
1295
+ batch_size, seq_length = inputs_embeds.shape[:2]
1296
+ else:
1297
+ raise ValueError("You have to specify either input_ids or inputs_embeds")
1298
+
1299
+ if self.gradient_checkpointing and self.training:
1300
+ if use_cache:
1301
+ logger.warning_once(
1302
+ "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`transformers."
1303
+ )
1304
+ use_cache = False
1305
+
1306
+ if use_cache and past_key_values is None:
1307
+ past_key_values = DynamicCache()
1308
+
1309
+ if inputs_embeds is None:
1310
+ inputs_embeds = self.word_embeddings(input_ids)
1311
+
1312
+ if cache_position is None:
1313
+ past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
1314
+ cache_position: torch.Tensor = torch.arange(
1315
+ past_seen_tokens, past_seen_tokens + inputs_embeds.shape[1], device=inputs_embeds.device
1316
+ )
1317
+
1318
+ if position_ids is None:
1319
+ position_ids = cache_position.unsqueeze(0)
1320
+
1321
+ softmax_attention_layer_id = self.config.layer_group_size - 1
1322
+ past_seen_tokens = (
1323
+ past_key_values.get_seq_length(layer_idx=softmax_attention_layer_id) if past_key_values is not None else 0
1324
+ )
1325
+
1326
+ if position_ids is None:
1327
+ position_ids = torch.arange(
1328
+ past_seen_tokens, past_seen_tokens + inputs_embeds.shape[1], device=inputs_embeds.device
1329
+ )
1330
+ position_ids = position_ids.unsqueeze(0)
1331
+
1332
+ if self._use_flash_attention_2:
1333
+ # 2d mask is passed through the layers
1334
+ attention_mask = attention_mask if (attention_mask is not None and 0 in attention_mask) else None
1335
+ elif self._use_sdpa and not output_attentions:
1336
+ # output_attentions=True can not be supported when using SDPA, and we fall back on
1337
+ # the manual implementation that requires a 4D causal mask in all cases.
1338
+ attention_mask = _prepare_4d_causal_attention_mask_for_sdpa(
1339
+ attention_mask,
1340
+ (batch_size, seq_length),
1341
+ inputs_embeds,
1342
+ past_seen_tokens,
1343
+ )
1344
+ else:
1345
+ # 4d mask is passed through the layers
1346
+ attention_mask = _prepare_4d_causal_attention_mask(
1347
+ attention_mask, (batch_size, seq_length), inputs_embeds, past_seen_tokens
1348
+ )
1349
+
1350
+ # embed positions
1351
+ hidden_states = inputs_embeds
1352
+
1353
+ # create position embeddings to be shared across the decoder layers
1354
+ position_embeddings = self.rotary_emb(hidden_states, position_ids)
1355
+
1356
+ # decoder layers
1357
+ all_hidden_states = () if output_hidden_states else None
1358
+ all_self_attns = () if output_attentions else None
1359
+ all_router_logits = () if output_router_logits else None
1360
+ next_decoder_cache = None
1361
+ layers = self.layers[: -self.num_nextn_predict_layers] if self.num_nextn_predict_layers > 0 else self.layers
1362
+ mtp_layers = self.layers[-self.num_nextn_predict_layers :] if self.num_nextn_predict_layers > 0 else None
1363
+
1364
+ # tptest miss causal_mask = create_causal_mask(
1365
+
1366
+ for decoder_layer in layers:
1367
+ if output_hidden_states:
1368
+ all_hidden_states += (hidden_states,)
1369
+
1370
+ if self.gradient_checkpointing and self.training:
1371
+ layer_outputs = self._gradient_checkpointing_func(
1372
+ decoder_layer.__call__,
1373
+ hidden_states,
1374
+ attention_mask,
1375
+ position_ids,
1376
+ past_key_values,
1377
+ cache_position,
1378
+ output_attentions,
1379
+ output_router_logits,
1380
+ use_cache,
1381
+ position_embeddings,
1382
+ )
1383
+ else:
1384
+ layer_outputs = decoder_layer(
1385
+ hidden_states,
1386
+ attention_mask=attention_mask,
1387
+ position_ids=position_ids,
1388
+ past_key_value=past_key_values,
1389
+ cache_position=cache_position,
1390
+ output_attentions=output_attentions,
1391
+ output_router_logits=output_router_logits,
1392
+ use_cache=use_cache,
1393
+ position_embeddings=position_embeddings,
1394
+ )
1395
+ hidden_states = layer_outputs[0]
1396
+
1397
+ if use_cache:
1398
+ next_decoder_cache = layer_outputs[2 if output_attentions else 1]
1399
+
1400
+ if output_attentions:
1401
+ all_self_attns += (layer_outputs[1],)
1402
+
1403
+ if output_router_logits and layer_outputs[-1] is not None:
1404
+ all_router_logits += (layer_outputs[-1],)
1405
+
1406
+ hidden_states = self.norm(hidden_states)
1407
+ main_hidden_states = hidden_states
1408
+
1409
+ # add hidden states from the last decoder layer
1410
+ if output_hidden_states:
1411
+ all_hidden_states += (main_hidden_states,)
1412
+
1413
+ mtp_hidden_states = None
1414
+
1415
+ if mtp_layers:
1416
+ for decoder_layer in mtp_layers:
1417
+ input_ids, _ = roll_tensor(input_ids, shifts=-1, dims=-1)
1418
+ inputs_embeds = self.word_embeddings(input_ids)
1419
+
1420
+ if self.gradient_checkpointing and self.training:
1421
+ layer_outputs = self._gradient_checkpointing_func(
1422
+ decoder_layer.__call__,
1423
+ inputs_embeds,
1424
+ hidden_states,
1425
+ attention_mask,
1426
+ position_ids,
1427
+ past_key_values,
1428
+ output_attentions,
1429
+ output_router_logits,
1430
+ use_cache,
1431
+ position_embeddings,
1432
+ )
1433
+ else:
1434
+ layer_outputs = decoder_layer(
1435
+ inputs_embeds,
1436
+ hidden_states,
1437
+ attention_mask=attention_mask,
1438
+ position_ids=position_ids,
1439
+ past_key_value=past_key_values,
1440
+ output_attentions=output_attentions,
1441
+ output_router_logits=output_router_logits,
1442
+ use_cache=use_cache,
1443
+ position_embeddings=position_embeddings,
1444
+ )
1445
+ if mtp_hidden_states is None:
1446
+ mtp_hidden_states = []
1447
+ hidden_states = layer_outputs[0]
1448
+ mtp_hidden_states.append(hidden_states)
1449
+
1450
+ if output_hidden_states:
1451
+ all_hidden_states += (hidden_states,)
1452
+
1453
+ if use_cache:
1454
+ next_decoder_cache = layer_outputs[2 if output_attentions else 1]
1455
+
1456
+ if output_attentions:
1457
+ all_self_attns += (layer_outputs[1],)
1458
+
1459
+ if output_router_logits and layer_outputs[-1] is not None:
1460
+ all_router_logits += (layer_outputs[-1],)
1461
+
1462
+ next_cache = None
1463
+ if use_cache:
1464
+ next_cache = next_decoder_cache
1465
+ if not return_dict:
1466
+ return tuple(
1467
+ v
1468
+ for v in [main_hidden_states, next_cache, all_hidden_states, all_self_attns, all_router_logits]
1469
+ if v is not None
1470
+ )
1471
+ return MoeV3ModelOutputWithPast(
1472
+ last_hidden_state=main_hidden_states,
1473
+ past_key_values=next_cache,
1474
+ hidden_states=all_hidden_states,
1475
+ mtp_hidden_states=mtp_hidden_states,
1476
+ attentions=all_self_attns,
1477
+ router_logits=all_router_logits,
1478
+ )
1479
+
1480
+
1481
+ class BailingMoeV3ForCausalLM(BailingMoeV3PreTrainedModel, GenerationMixin):
1482
+ _tied_weights_keys = ["lm_head.weight"]
1483
+
1484
+ def __init__(self, config: BailingMoeV3Config):
1485
+ super().__init__(config)
1486
+ self.model = BailingMoeV3Model(config)
1487
+ self.vocab_size = config.vocab_size
1488
+ self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
1489
+ self.num_nextn_predict_layers = config.num_nextn_predict_layers
1490
+ self.mtp_loss_scaling_factor = config.mtp_loss_scaling_factor
1491
+
1492
+ # Initialize weights and apply final processing
1493
+ self.post_init()
1494
+
1495
+ def get_input_embeddings(self):
1496
+ return self.model.word_embeddings
1497
+
1498
+ def set_input_embeddings(self, value):
1499
+ self.model.word_embeddings = value
1500
+
1501
+ def get_output_embeddings(self):
1502
+ return self.lm_head
1503
+
1504
+ def set_output_embeddings(self, new_embeddings):
1505
+ self.lm_head = new_embeddings
1506
+
1507
+ def set_decoder(self, decoder):
1508
+ self.model = decoder
1509
+
1510
+ def get_decoder(self):
1511
+ return self.model
1512
+
1513
+ @add_start_docstrings_to_model_forward(BAILINGMOEV3_INPUTS_DOCSTRING)
1514
+ @replace_return_docstrings(output_type=MoEV3CausalLMOutputWithPast, config_class=_CONFIG_FOR_DOC)
1515
+ def forward(
1516
+ self,
1517
+ input_ids: torch.LongTensor = None,
1518
+ attention_mask: Optional[torch.Tensor] = None,
1519
+ position_ids: Optional[torch.LongTensor] = None,
1520
+ past_key_values: Optional[List[torch.FloatTensor]] = None,
1521
+ inputs_embeds: Optional[torch.FloatTensor] = None,
1522
+ labels: Optional[torch.LongTensor] = None,
1523
+ use_cache: Optional[bool] = None,
1524
+ output_attentions: Optional[bool] = None,
1525
+ output_hidden_states: Optional[bool] = None,
1526
+ output_router_logits: Optional[bool] = None,
1527
+ return_dict: Optional[bool] = None,
1528
+ **kwargs,
1529
+ ) -> Union[Tuple, MoEV3CausalLMOutputWithPast]:
1530
+ r"""
1531
+ Args:
1532
+ labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
1533
+ Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,
1534
+ config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored
1535
+ (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.
1536
+ Returns:
1537
+ Example:
1538
+ ```python
1539
+ >>> from transformers import AutoTokenizer
1540
+ >>> model = BailingMoeV3ForCausalLM.from_pretrained(PATH_TO_CONVERTED_WEIGHTS)
1541
+ >>> tokenizer = AutoTokenizer.from_pretrained(PATH_TO_CONVERTED_TOKENIZER)
1542
+ >>> prompt = "Hey, are you conscious? Can you talk to me?"
1543
+ >>> inputs = tokenizer(prompt, return_tensors="pt")
1544
+ >>> # Generate
1545
+ >>> generate_ids = model.generate(inputs.input_ids, max_length=30)
1546
+ >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
1547
+ "Hey, are you conscious? Can you talk to me?\nI'm not conscious, but I can talk to you."
1548
+ ```"""
1549
+ output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
1550
+ output_hidden_states = (
1551
+ output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
1552
+ )
1553
+ output_router_logits = (
1554
+ output_router_logits if output_router_logits is not None else self.config.output_router_logits
1555
+ )
1556
+ return_dict = return_dict if return_dict is not None else self.config.use_return_dict
1557
+ # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
1558
+ outputs = self.model(
1559
+ input_ids=input_ids,
1560
+ attention_mask=attention_mask,
1561
+ position_ids=position_ids,
1562
+ past_key_values=past_key_values,
1563
+ inputs_embeds=inputs_embeds,
1564
+ use_cache=use_cache,
1565
+ output_attentions=output_attentions,
1566
+ output_hidden_states=output_hidden_states,
1567
+ output_router_logits=output_router_logits,
1568
+ return_dict=return_dict,
1569
+ **kwargs,
1570
+ )
1571
+
1572
+ loss = None
1573
+ all_mtp_loss = None
1574
+ aux_loss = None
1575
+ hidden_states = outputs[0]
1576
+ logits = self.lm_head(hidden_states)
1577
+ logits = logits.float()
1578
+
1579
+ if labels is not None:
1580
+ loss = self.loss_function(logits, labels, self.config.vocab_size, **kwargs)
1581
+
1582
+ all_mtp_logits = None
1583
+ if self.num_nextn_predict_layers > 0:
1584
+ mtp_hidden_states = outputs.mtp_hidden_states
1585
+ shift_labels_mtp = None
1586
+ for i in range(self.num_nextn_predict_layers):
1587
+ mtp_hidden_states = mtp_hidden_states[i]
1588
+ mtp_logits = self.lm_head(mtp_hidden_states).float()
1589
+ if all_mtp_logits is None:
1590
+ all_mtp_logits = []
1591
+ all_mtp_logits.append(mtp_logits)
1592
+ if labels is not None:
1593
+ if shift_labels_mtp is None:
1594
+ shift_labels_mtp = labels.clone()
1595
+ shift_labels_mtp, _ = roll_tensor(shift_labels_mtp, shifts=-1, dims=-1, fill_value=-100)
1596
+ mtp_logits_ = mtp_logits.view(-1, self.config.vocab_size)
1597
+ mtp_loss = self.loss_function(
1598
+ mtp_logits_, shift_labels_mtp.to(mtp_logits_.device).view(-1), self.config.vocab_size, **kwargs
1599
+ )
1600
+ if loss is not None:
1601
+ loss += self.mtp_loss_scaling_factor * mtp_loss
1602
+ else:
1603
+ loss = self.mtp_loss_scaling_factor * mtp_loss
1604
+
1605
+ if all_mtp_loss is None:
1606
+ all_mtp_loss = []
1607
+ all_mtp_loss.append(mtp_loss)
1608
+
1609
+ if not return_dict:
1610
+ output = (logits,) + outputs[1:]
1611
+ if output_router_logits:
1612
+ output = (aux_loss,) + output
1613
+ return (loss,) + output if loss is not None else output
1614
+
1615
+ return MoEV3CausalLMOutputWithPast(
1616
+ loss=loss,
1617
+ mtp_loss=all_mtp_loss,
1618
+ aux_loss=aux_loss,
1619
+ logits=logits,
1620
+ mtp_logits=all_mtp_logits,
1621
+ past_key_values=outputs.past_key_values,
1622
+ hidden_states=outputs.hidden_states,
1623
+ attentions=outputs.attentions,
1624
+ router_logits=outputs.router_logits,
1625
+ )
tokenizer.json ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:40fb9d7d7795b8bd305aeff39ce9963f3f450915b9553f2938e009be9a1fed60
3
+ size 12205732
tokenizer_config.json ADDED
@@ -0,0 +1,19 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "backend": "tokenizers",
3
+ "bos_token": "<|startoftext|>",
4
+ "clean_up_tokenization_spaces": false,
5
+ "cls_token": "[CLS]",
6
+ "eos_token": "<|role_end|>",
7
+ "fast_tokenizer": true,
8
+ "gmask_token": "[gMASK]",
9
+ "is_local": true,
10
+ "local_files_only": false,
11
+ "merges_file": null,
12
+ "model_max_length": 1000000000000000019884624838656,
13
+ "model_specific_special_tokens": {
14
+ "gmask_token": "[gMASK]"
15
+ },
16
+ "pad_token": "<|endoftext|>",
17
+ "tokenizer_class": "TokenizersBackend",
18
+ "tool_parser_type": "glm47"
19
+ }