hotheadhacker commited on
Commit
00a0cd9
·
verified ·
1 Parent(s): eef7c7b

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
chat_template.jinja ADDED
@@ -0,0 +1,280 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {{- bos_token -}}
2
+ {%- set preserve_thinking = preserve_thinking | default(false) -%}
3
+ {%- set enable_thinking = enable_thinking | default(true) -%}
4
+
5
+ {%- macro render_value(v) -%}
6
+ {%- if v is string -%}
7
+ {{- v -}}
8
+ {%- elif v is boolean -%}
9
+ {{- "true" if v else "false" -}}
10
+ {%- elif v is none -%}
11
+ {{- "null" -}}
12
+ {%- elif v is mapping or (v is sequence and v is not string) -%}
13
+ {{- v | tojson -}}
14
+ {%- else -%}
15
+ {{- v -}}
16
+ {%- endif -%}
17
+ {%- endmacro -%}
18
+
19
+ {%- macro parse_content(content) -%}
20
+ {%- if content is string -%}
21
+ {{- content -}}
22
+ {%- elif content is mapping and content.get("type") == "image" -%}
23
+ {{- "<|vision_start|><|image|><|vision_end|>" -}}
24
+ {%- elif content is mapping and content.get("type") == "video" -%}
25
+ {{- "<|vision_start|><|video|><|vision_end|>" -}}
26
+ {%- elif content is mapping and content.get("type") == "text" -%}
27
+ {{- (content.get("text") or "") | string -}}
28
+ {%- elif content is mapping -%}
29
+ {{- content | tojson -}}
30
+ {%- elif content is iterable -%}
31
+ {%- set _ns = namespace(result="") -%}
32
+ {%- for item in content -%}
33
+ {%- if item is string -%}
34
+ {%- set _ns.result = _ns.result + item -%}
35
+ {%- elif item is mapping and item.get("type") == "image" -%}
36
+ {%- set _ns.result = _ns.result + "<|vision_start|><|image|><|vision_end|>" -%}
37
+ {%- elif item is mapping and item.get("type") == "video" -%}
38
+ {%- set _ns.result = _ns.result + "<|vision_start|><|video|><|vision_end|>" -%}
39
+ {%- elif item is mapping and item.get("type") == "text" -%}
40
+ {%- set _ns.result = _ns.result + ((item.get("text") or "") | string) -%}
41
+ {%- else -%}
42
+ {%- set _ns.result = _ns.result + (item | tojson) -%}
43
+ {%- endif -%}
44
+ {%- endfor -%}
45
+ {{- _ns.result -}}
46
+ {%- elif content is not none -%}
47
+ {{- content | tojson -}}
48
+ {%- endif -%}
49
+ {%- endmacro -%}
50
+
51
+ {%- macro render_tool_calls(tool_calls) -%}
52
+ {%- for tool_call in tool_calls -%}
53
+ {%- if not loop.first -%}{{- "\n\n" -}}{%- endif -%}
54
+ {%- set func = tool_call["function"] if "function" in tool_call else tool_call -%}
55
+ {%- set func_args = func.get("arguments") -%}
56
+ {%- if func_args is string and (func_args | trim) not in ["", "{}", "null"] -%}
57
+ {{- raise_exception("Tool call arguments must be a mapping, got a JSON-encoded string: parse arguments with json.loads() before applying the chat template") -}}
58
+ {%- endif -%}
59
+ {{- "<tool_call>\n<function=" + func["name"] + ">\n" -}}
60
+ {%- if func_args is mapping -%}
61
+ {%- for arg_name, arg_value in func_args.items() -%}
62
+ {{- "<parameter=" + arg_name + ">\n" + render_value(arg_value) + "\n</parameter>\n" -}}
63
+ {%- endfor -%}
64
+ {%- endif -%}
65
+ {{- "</function>\n</tool_call>" -}}
66
+ {%- endfor -%}
67
+ {%- endmacro -%}
68
+
69
+ {%- macro tool_to_json(tool) -%}
70
+ {%- set fn = tool.function if tool.function is defined else tool -%}
71
+ {%- set ns_t = namespace(first=true) -%}
72
+ {{- "{" -}}
73
+ {%- for k, v in fn.items() -%}
74
+ {%- if k != "defer_loading" -%}
75
+ {%- if not ns_t.first -%}{{- ", " -}}{%- endif -%}
76
+ {%- set ns_t.first = false -%}
77
+ {{- k | tojson -}}{{- ": " -}}{{- v | tojson -}}
78
+ {%- endif -%}
79
+ {%- endfor -%}
80
+ {{- "}" -}}
81
+ {%- endmacro -%}
82
+
83
+ {%- macro is_tool_reference(item) -%}
84
+ {{- "1" if item is mapping and item.get("type") == "tool_reference" else "" -}}
85
+ {%- endmacro -%}
86
+
87
+ {%- macro referenced_tool_to_json(name) -%}
88
+ {%- set _ns = namespace(found=none) -%}
89
+ {%- for tool in tools or [] -%}
90
+ {%- set fn = tool.function if tool.function is defined else tool -%}
91
+ {%- if fn.name == name -%}
92
+ {%- set _ns.found = fn -%}
93
+ {%- endif -%}
94
+ {%- endfor -%}
95
+ {%- if _ns.found is none -%}
96
+ {{- raise_exception("Tool reference '" + name | string + "' not found in available tools. Every tool referenced by a tool_reference must be in `tools` on every request.") -}}
97
+ {%- endif -%}
98
+ {{- tool_to_json(_ns.found) -}}
99
+ {%- endmacro -%}
100
+
101
+ {%- macro render_tool_response(content) -%}
102
+ {%- if content is string and "<tool_response>" in content -%}
103
+ {{- content -}}
104
+ {%- else -%}
105
+ {{- "<tool_response>\n" -}}
106
+ {%- if content is string -%}
107
+ {{- content -}}
108
+ {%- elif content is mapping -%}
109
+ {{- parse_content(content) -}}
110
+ {%- elif content is iterable -%}
111
+ {%- for item in content -%}
112
+ {%- if not loop.first -%}{{- "\n" -}}{%- endif -%}
113
+ {%- if is_tool_reference(item) -%}
114
+ {%- if loop.first or not is_tool_reference(loop.previtem) -%}
115
+ {{- "<additional_tools>\n" -}}
116
+ {%- endif -%}
117
+ {{- referenced_tool_to_json(item.name) -}}
118
+ {%- if loop.last or not is_tool_reference(loop.nextitem) -%}
119
+ {{- "\n</additional_tools>" -}}
120
+ {%- endif -%}
121
+ {%- elif item is mapping and item.result is defined and item.result is mapping and item.result.tools is defined -%}
122
+ {{- "<additional_tools>\n" -}}
123
+ {%- set _tl = item.result.tools -%}
124
+ {%- if _tl is string -%}
125
+ {{- _tl -}}
126
+ {%- elif _tl is mapping -%}
127
+ {{- tool_to_json(_tl) -}}
128
+ {%- else -%}
129
+ {%- for t in _tl -%}
130
+ {%- if not loop.first -%}{{- "\n" -}}{%- endif -%}
131
+ {{- tool_to_json(t) -}}
132
+ {%- endfor -%}
133
+ {%- endif -%}
134
+ {{- "\n</additional_tools>" -}}
135
+ {%- else -%}
136
+ {{- parse_content(item) -}}
137
+ {%- endif -%}
138
+ {%- endfor -%}
139
+ {%- elif content is not none -%}
140
+ {{- content | tojson -}}
141
+ {%- endif -%}
142
+ {{- "\n</tool_response>" -}}
143
+ {%- endif -%}
144
+ {%- endmacro -%}
145
+
146
+ {%- set ns = namespace(system_prompt="", last_user_index=-1) -%}
147
+ {%- if messages and messages[0]["role"] == "system" -%}
148
+ {%- if messages[0].get("content") -%}
149
+ {%- set ns.system_prompt = parse_content(messages[0]["content"]) -%}
150
+ {%- endif -%}
151
+ {%- set messages = messages[1:] -%}
152
+ {%- endif -%}
153
+ {%- if tools -%}
154
+ {%- set tools_ns = namespace(lines=[]) -%}
155
+ {%- for tool in tools -%}
156
+ {%- set fn = tool.function if tool.function is defined else tool -%}
157
+ {%- set _dyn = (fn.defer_loading is defined and fn.defer_loading) or (tool.defer_loading is defined and tool.defer_loading) -%}
158
+ {%- if not _dyn -%}
159
+ {%- set tools_ns.lines = tools_ns.lines + [tool_to_json(fn)] -%}
160
+ {%- endif -%}
161
+ {%- endfor -%}
162
+ {%- if tools_ns.lines -%}
163
+ {%- set ns.system_prompt = ns.system_prompt + ("\n\n" if ns.system_prompt else "") + "<tools>\n" + (tools_ns.lines | join("\n")) + "\n</tools>" -%}
164
+ {%- endif -%}
165
+ {%- endif -%}
166
+ {%- if ns.system_prompt -%}
167
+ {{- "<|im_start|>system\n" + ns.system_prompt + "<|im_end|>\n" -}}
168
+ {%- endif -%}
169
+ {%- for message in messages -%}
170
+ {%- if message["role"] == "user" -%}
171
+ {%- set ns.last_user_index = loop.index0 -%}
172
+ {%- endif -%}
173
+ {%- endfor -%}
174
+ {%- set carry = namespace(pending=[]) -%}
175
+ {%- for message in messages -%}
176
+ {%- set _next = messages[loop.index0 + 1] if not loop.last else none -%}
177
+ {%- set _merge = message.role == "tool" and not message.get("additional_tools")
178
+ and _next is not none and _next.role == "tool" and not _next.get("additional_tools")
179
+ and message.get("tool_call_id") and _next.get("tool_call_id") == message.get("tool_call_id") -%}
180
+ {%- set _items = [] -%}
181
+ {%- if message.role == "tool" and not message.get("additional_tools") -%}
182
+ {%- set _c = message.get("content") -%}
183
+ {%- if _c is string -%}
184
+ {%- set _items = [{"type": "text", "text": _c}] if _c else [] -%}
185
+ {%- elif _c is mapping -%}
186
+ {%- set _items = [_c] -%}
187
+ {%- elif _c is iterable -%}
188
+ {%- set _items = _c | list -%}
189
+ {%- elif _c is not none -%}
190
+ {%- set _items = [_c] -%}
191
+ {%- endif -%}
192
+ {%- endif -%}
193
+ {%- if _merge -%}
194
+ {%- set carry.pending = carry.pending + _items -%}
195
+ {%- else -%}
196
+ {{- "<|im_start|>" + message.role + "\n" -}}
197
+ {%- if message.role == "assistant" -%}
198
+ {%- generation -%}
199
+ {%- set thinking = message.thinking if message.thinking is defined else message.reasoning if message.reasoning is defined else message.reasoning_content -%}
200
+ {%- set keep_thinking = preserve_thinking or loop.index0 > ns.last_user_index -%}
201
+ {%- if thinking is defined and keep_thinking -%}
202
+ {{- "<think>" + (thinking or "") + "</think>" -}}
203
+ {%- elif keep_thinking and "</think>" not in (message.get("content") if message.get("content") is string else "") -%}
204
+ {{- "<think></think>" -}}
205
+ {%- endif -%}
206
+ {%- set _cfm_tag = "CONTINUE_FINAL_MESSAGE_TAG " -%}
207
+ {%- set _has_cfm = false -%}
208
+ {%- set content = "" -%}
209
+ {%- if message.get("content") -%}
210
+ {%- set content = parse_content(message.content) -%}
211
+ {%- endif -%}
212
+ {%- if not keep_thinking -%}
213
+ {%- if "</think>" in content -%}
214
+ {%- set content = content.split("</think>")[-1] | trim -%}
215
+ {%- endif -%}
216
+ {%- set content = "<think></think>" + content -%}
217
+ {%- endif -%}
218
+ {%- if content.endswith(_cfm_tag) -%}
219
+ {%- set _has_cfm = true -%}
220
+ {%- set _trunc_len = (content | length) - (_cfm_tag | length) -%}
221
+ {%- set content = content[:_trunc_len] -%}
222
+ {%- endif -%}
223
+ {{- content -}}
224
+ {%- if message.tool_calls -%}
225
+ {{- render_tool_calls(message.tool_calls) -}}
226
+ {%- endif -%}
227
+ {%- if _has_cfm -%}
228
+ {{- _cfm_tag -}}
229
+ {%- endif -%}
230
+ {{- "<|im_end|>\n" -}}
231
+ {%- endgeneration -%}
232
+ {%- elif message.role == "tool" -%}
233
+ {%- if message.get("additional_tools") -%}
234
+ {%- set _at = message["additional_tools"] -%}
235
+ {%- if _at is mapping or (_at is sequence and _at is not string) -%}
236
+ {%- set _src = _at -%}
237
+ {%- else -%}
238
+ {%- set _src = message.get("content") -%}
239
+ {%- endif -%}
240
+ {%- if _src is string and "<tool_response>" in _src -%}
241
+ {{- _src -}}
242
+ {%- elif _src is string and "<additional_tools>" in _src -%}
243
+ {{- "<tool_response>\n" + _src + "\n</tool_response>" -}}
244
+ {%- elif _src -%}
245
+ {{- "<tool_response>\n<additional_tools>\n" -}}
246
+ {%- if _src is string -%}
247
+ {{- _src -}}
248
+ {%- elif _src is mapping -%}
249
+ {{- tool_to_json(_src) -}}
250
+ {%- else -%}
251
+ {%- for t in _src -%}
252
+ {%- if not loop.first -%}{{- "\n" -}}{%- endif -%}
253
+ {{- tool_to_json(t) -}}
254
+ {%- endfor -%}
255
+ {%- endif -%}
256
+ {{- "\n</additional_tools>\n</tool_response>" -}}
257
+ {%- endif -%}
258
+ {%- elif carry.pending -%}
259
+ {{- render_tool_response(carry.pending + _items) -}}
260
+ {%- set carry.pending = [] -%}
261
+ {%- else -%}
262
+ {{- render_tool_response("" if message.get("content") is none else message.get("content")) -}}
263
+ {%- endif -%}
264
+ {{- "<|im_end|>\n" -}}
265
+ {%- else %}
266
+ {%- if message.get("content") -%}
267
+ {{- parse_content(message["content"]) -}}
268
+ {%- endif -%}
269
+ {{- "<|im_end|>\n" -}}
270
+ {%- endif %}
271
+ {%- endif -%}
272
+ {%- endfor -%}
273
+ {%- if add_generation_prompt -%}
274
+ {{- "<|im_start|>assistant\n" -}}
275
+ {%- if enable_thinking -%}
276
+ {{- "<think>" -}}
277
+ {%- else -%}
278
+ {{- "<think></think>" -}}
279
+ {%- endif -%}
280
+ {%- endif -%}
config.json ADDED
@@ -0,0 +1,69 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "architectures": [
3
+ "Lfm3ForCausalLM"
4
+ ],
5
+ "auto_map": {
6
+ "AutoConfig": "configuration_lfm3.Lfm3Config",
7
+ "AutoModel": "modeling_lfm3.Lfm3Model",
8
+ "AutoModelForCausalLM": "modeling_lfm3.Lfm3ForCausalLM"
9
+ },
10
+ "bos_token_id": 124894,
11
+ "conv_L_cache": 7,
12
+ "conv_bias": false,
13
+ "dtype": "bfloat16",
14
+ "eos_token_id": 124900,
15
+ "head_dim": 128,
16
+ "hidden_size": 2048,
17
+ "initializer_range": 0.02,
18
+ "intermediate_size": 7680,
19
+ "layer_types": [
20
+ "conv",
21
+ "conv",
22
+ "full_attention",
23
+ "conv",
24
+ "sliding_attention",
25
+ "conv",
26
+ "full_attention",
27
+ "conv",
28
+ "sliding_attention",
29
+ "conv",
30
+ "full_attention",
31
+ "conv",
32
+ "sliding_attention",
33
+ "conv",
34
+ "full_attention",
35
+ "conv",
36
+ "sliding_attention",
37
+ "conv",
38
+ "full_attention",
39
+ "conv",
40
+ "sliding_attention",
41
+ "conv",
42
+ "full_attention",
43
+ "conv"
44
+ ],
45
+ "max_position_embeddings": 262144,
46
+ "model_type": "lfm3",
47
+ "norm_eps": 1e-05,
48
+ "num_attention_heads": 16,
49
+ "num_hidden_layers": 24,
50
+ "num_key_value_heads": 4,
51
+ "pad_token_id": 124893,
52
+ "rope_parameters": {
53
+ "full_attention": {
54
+ "partial_rotary_factor": 0.5,
55
+ "rope_theta": 1000000.0,
56
+ "rope_type": "default"
57
+ },
58
+ "sliding_attention": {
59
+ "rope_theta": 10000,
60
+ "rope_type": "default"
61
+ }
62
+ },
63
+ "sliding_window": 256,
64
+ "tie_word_embeddings": true,
65
+ "transformers_version": "5.20.2",
66
+ "use_cache": true,
67
+ "use_sandwich_norm": true,
68
+ "vocab_size": 128000
69
+ }
configuration_lfm3.py ADDED
@@ -0,0 +1,132 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Copyright 2026 The HuggingFace Team. All rights reserved.
2
+ #
3
+ # Licensed under the Apache License, Version 2.0 (the "License");
4
+ # you may not use this file except in compliance with the License.
5
+ # You may obtain a copy of the License at
6
+ #
7
+ # http://www.apache.org/licenses/LICENSE-2.0
8
+ #
9
+ # Unless required by applicable law or agreed to in writing, software
10
+ # distributed under the License is distributed on an "AS IS" BASIS,
11
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12
+ # See the License for the specific language governing permissions and
13
+ # limitations under the License.
14
+ from huggingface_hub.dataclasses import strict
15
+
16
+ from transformers.configuration_utils import PreTrainedConfig
17
+ from transformers.utils import auto_docstring
18
+
19
+
20
+ @auto_docstring(checkpoint="LiquidAI/LFM3-1.7B-Preview")
21
+ @strict
22
+ class Lfm3Config(PreTrainedConfig):
23
+ r"""
24
+ rope_parameters (`dict[str, RopeParameters]`, *optional*):
25
+ RoPE parameters keyed by attention layer type, for example `{"full_attention": {"rope_type": "default",
26
+ "rope_theta": 1000000.0, "partial_rotary_factor": 0.5}, "sliding_attention": {"rope_type": "default",
27
+ "rope_theta": 10000.0}}`. `conv` layers do not use RoPE and have no entry. Full-attention layers rotate the
28
+ first `partial_rotary_factor * head_dim` dimensions and leave the rest unrotated.
29
+ conv_bias (`bool`, *optional*, defaults to `False`):
30
+ Whether to use bias in the conv layers.
31
+ conv_L_cache (`int`, *optional*, defaults to 7):
32
+ L_cache dim in the conv layers.
33
+ layer_types (`list[str]`, *optional*):
34
+ Per-layer operator type, each one of `"full_attention"`, `"sliding_attention"` or `"conv"`. Defaults to the
35
+ LFM3 pattern of a full-attention layer every fourth layer starting from the third, a sliding-attention layer
36
+ in between, and gated short convolutions elsewhere.
37
+ sliding_window (`int`, *optional*, defaults to 256):
38
+ Sliding-window size of the `sliding_attention` layers.
39
+ use_sandwich_norm (`bool`, *optional*, defaults to `True`):
40
+ Whether to apply an RMSNorm to the operator and feed-forward outputs before they enter the residual stream.
41
+
42
+ ```python
43
+ >>> from transformers import Lfm3Model, Lfm3Config
44
+
45
+ >>> # Initializing a LFM3 model
46
+ >>> configuration = Lfm3Config()
47
+
48
+ >>> # Initializing a model from the LFM3-1.7B style configuration
49
+ >>> model = Lfm3Model(configuration)
50
+
51
+ >>> # Accessing the model configuration
52
+ >>> configuration = model.config
53
+ ```
54
+ """
55
+
56
+ model_type = "lfm3"
57
+ keys_to_ignore_at_inference = ["past_key_values"]
58
+ default_theta = 1000000.0
59
+
60
+ vocab_size: int = 128_000
61
+ hidden_size: int = 2048
62
+ intermediate_size: int = 7680
63
+ num_hidden_layers: int = 24
64
+ num_attention_heads: int = 16
65
+ num_key_value_heads: int = 4
66
+ max_position_embeddings: int = 128_000
67
+ initializer_range: float = 0.02
68
+ norm_eps: float = 0.00001
69
+ use_cache: bool = True
70
+ pad_token_id: int | None = 124_893
71
+ bos_token_id: int | None = 124_894
72
+ eos_token_id: int | list[int] | None = 124_900
73
+ tie_word_embeddings: bool = True
74
+ rope_parameters: dict | None = None
75
+ conv_bias: bool = False
76
+ conv_L_cache: int = 7
77
+ layer_types: list[str] | None = None
78
+ _attention_layer_types = ("full_attention", "sliding_attention")
79
+ head_dim: int | None = None
80
+ sliding_window: int | None = 256
81
+ use_sandwich_norm: bool = True
82
+
83
+ def __post_init__(self, **kwargs):
84
+ if self.head_dim is None:
85
+ self.head_dim = self.hidden_size // self.num_attention_heads
86
+ if self.rope_parameters is None:
87
+ self.rope_parameters = {
88
+ "full_attention": {"rope_type": "default", "rope_theta": 1_000_000.0, "partial_rotary_factor": 0.5},
89
+ "sliding_attention": {"rope_type": "default", "rope_theta": 10_000.0},
90
+ }
91
+ if self.layer_types is None:
92
+ self.layer_types = [
93
+ "full_attention" if i % 4 == 2 else "sliding_attention" if i % 4 == 0 and i > 0 else "conv"
94
+ for i in range(self.num_hidden_layers)
95
+ ]
96
+
97
+ self.tie_word_embeddings = kwargs.pop("tie_embedding", self.tie_word_embeddings)
98
+ super().__post_init__(**kwargs)
99
+
100
+ def validate_architecture(self):
101
+ """Part of ``@strict``-powered validation."""
102
+ if len(self.layer_types) != self.num_hidden_layers:
103
+ raise ValueError(
104
+ f"`layer_types` has {len(self.layer_types)} entries but `num_hidden_layers` is {self.num_hidden_layers}."
105
+ )
106
+ allowed_layer_types = (*self._attention_layer_types, "conv")
107
+ invalid_layer_types = [layer_type for layer_type in self.layer_types if layer_type not in allowed_layer_types]
108
+ if invalid_layer_types:
109
+ raise ValueError(f"`layer_types` entries must be one of {allowed_layer_types}; got {invalid_layer_types}.")
110
+ if "sliding_attention" in self.layer_types and self.sliding_window is None:
111
+ raise ValueError("`layer_types` has `sliding_attention` layers but `sliding_window` is None.")
112
+ missing_rope = [
113
+ layer_type
114
+ for layer_type in self._attention_layer_types
115
+ if layer_type in self.layer_types and layer_type not in self.rope_parameters
116
+ ]
117
+ if missing_rope:
118
+ raise ValueError(f"`rope_parameters` has no entry for the {missing_rope} layers in `layer_types`.")
119
+ for layer_type, rope_params in self.rope_parameters.items():
120
+ if not isinstance(rope_params, dict):
121
+ continue
122
+ if (
123
+ rope_params.get("rope_type", "default") in ("dynamic", "longrope")
124
+ and rope_params.get("partial_rotary_factor", 1.0) < 1
125
+ ):
126
+ raise ValueError(
127
+ f"`{rope_params['rope_type']}` RoPE recomputes its frequencies during the forward pass over the "
128
+ f"rotated dimension only and cannot keep LFM3's partial rotation on `{layer_type}` layers."
129
+ )
130
+
131
+
132
+ __all__ = ["Lfm3Config"]
generation_config.json ADDED
@@ -0,0 +1,16 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "_from_model_config": true,
3
+ "bos_token_id": 124894,
4
+ "do_sample": true,
5
+ "eos_token_id": [
6
+ 124900
7
+ ],
8
+ "output_attentions": false,
9
+ "output_hidden_states": false,
10
+ "pad_token_id": 124893,
11
+ "repetition_penalty": 1.1,
12
+ "temperature": 0.1,
13
+ "top_k": 50,
14
+ "transformers_version": "5.20.2",
15
+ "use_cache": true
16
+ }
lfm3_conv.py ADDED
@@ -0,0 +1,202 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Copyright 2026 The HuggingFace Team. All rights reserved.
2
+ #
3
+ # Licensed under the Apache License, Version 2.0 (the "License");
4
+ # you may not use this file except in compliance with the License.
5
+ # You may obtain a copy of the License at
6
+ #
7
+ # http://www.apache.org/licenses/LICENSE-2.0
8
+ #
9
+ # Unless required by applicable law or agreed to in writing, software
10
+ # distributed under the License is distributed on an "AS IS" BASIS,
11
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12
+ # See the License for the specific language governing permissions and
13
+ # limitations under the License.
14
+ """Inference kernels for LFM3's width-7 depthwise convolution."""
15
+
16
+ import torch
17
+
18
+ from transformers.utils import is_triton_available
19
+
20
+
21
+ if is_triton_available():
22
+ import triton
23
+ import triton.language as tl
24
+
25
+ # the length and strides are runtime arguments, or every prompt length would compile its own kernel
26
+ @triton.jit(do_not_specialize=["T", "X0", "X1", "X2", "S0", "S1", "ST"])
27
+ def _prefill(
28
+ X,
29
+ W,
30
+ BIAS,
31
+ SEQ,
32
+ Y,
33
+ T,
34
+ D: tl.constexpr,
35
+ X0,
36
+ X1,
37
+ X2,
38
+ S0,
39
+ S1,
40
+ ST,
41
+ HAS_BIAS: tl.constexpr,
42
+ PACKED: tl.constexpr,
43
+ BLOCK: tl.constexpr,
44
+ ):
45
+ # int64 offsets: batch * channels * length can exceed 2**31 elements
46
+ b, d = tl.program_id(0).to(tl.int64), tl.program_id(1).to(tl.int64)
47
+ t = tl.program_id(2).to(tl.int64) * BLOCK + tl.arange(0, BLOCK)
48
+ acc = tl.full((BLOCK,), 0, tl.float32)
49
+ if PACKED:
50
+ doc = tl.load(SEQ + b * S0 + (t - T + ST) * S1, (t < T) & (t >= T - ST), other=-1)
51
+ for k in tl.static_range(7):
52
+ pos = t + k - 6
53
+ valid = (t < T) & (pos >= 0)
54
+ if PACKED:
55
+ prev = tl.load(SEQ + b * S0 + (pos - T + ST) * S1, valid & (pos >= T - ST), other=-1)
56
+ valid = valid & (doc == prev)
57
+ x = tl.load(X + b * X0 + d * X1 + pos * X2, valid, other=0).to(tl.float32)
58
+ w = tl.load(W + d * 7 + k).to(tl.float32)
59
+ acc += x * w
60
+ if HAS_BIAS:
61
+ acc += tl.load(BIAS + d).to(tl.float32)
62
+ tl.store(Y + (b * D + d) * T + t, acc, t < T)
63
+
64
+ @triton.jit
65
+ def _decode(
66
+ X,
67
+ W,
68
+ BIAS,
69
+ STATE,
70
+ Y,
71
+ D: tl.constexpr,
72
+ K: tl.constexpr,
73
+ X0: tl.constexpr,
74
+ X1: tl.constexpr,
75
+ S0: tl.constexpr,
76
+ S1: tl.constexpr,
77
+ S2: tl.constexpr,
78
+ HAS_BIAS: tl.constexpr,
79
+ BLOCK: tl.constexpr,
80
+ ):
81
+ b = tl.program_id(0).to(tl.int64)
82
+ d = tl.program_id(1).to(tl.int64) * BLOCK + tl.arange(0, BLOCK)
83
+ k = tl.arange(0, 8)
84
+ # Load the complete state before writing the shifted window.
85
+ old = tl.load(STATE + b * S0 + d[:, None] * S1 + k[None, :] * S2, (d[:, None] < D) & (k[None, :] < K), other=0)
86
+ x = tl.load(X + b * X0 + d * X1, d < D, other=0)
87
+ shifted = tl.gather(old, tl.broadcast_to(tl.minimum(k + 1, 7)[None, :], (BLOCK, 8)), 1)
88
+ shifted = tl.where(k[None, :] == K - 1, x[:, None], shifted)
89
+ # A cache may store either K=6 previous inputs or K=7 including the oldest discarded input.
90
+ positions = k + K - 6
91
+ values = tl.gather(old, tl.broadcast_to(tl.minimum(positions, 7)[None, :], (BLOCK, 8)), 1)
92
+ values = tl.where(k[None, :] == 6, x[:, None], values)
93
+ weights = tl.load(W + d[:, None] * 7 + k[None, :], (d[:, None] < D) & (k[None, :] < 7), other=0)
94
+ acc = tl.sum(values.to(tl.float32) * weights.to(tl.float32), 1)
95
+ if HAS_BIAS:
96
+ acc += tl.load(BIAS + d, d < D, other=0).to(tl.float32)
97
+ tl.store(STATE + b * S0 + d[:, None] * S1 + k[None, :] * S2, shifted, (d[:, None] < D) & (k[None, :] < K))
98
+ tl.store(Y + b * D + d, acc, d < D)
99
+
100
+
101
+ def can_use_lfm3_conv(hidden_states, weight):
102
+ return (
103
+ is_triton_available()
104
+ and hasattr(tl, "gather") # the decode kernel needs Triton >= 3.3
105
+ and hidden_states.is_cuda
106
+ and not torch.is_grad_enabled()
107
+ and hidden_states.dtype == weight.dtype
108
+ and hidden_states.dtype in (torch.float16, torch.bfloat16, torch.float32)
109
+ and weight.shape[-1] == 7
110
+ )
111
+
112
+
113
+ def lfm3_causal_conv1d(hidden_states, weight, bias=None, seq_idx=None):
114
+ batch, channels, length = hidden_states.shape
115
+ # the kernels index every tensor by the batch and channel of `hidden_states` and cannot bound the others
116
+ if weight.shape != (channels, 7) or (bias is not None and bias.shape != (channels,)):
117
+ raise ValueError(
118
+ f"`weight` must have shape ({channels}, 7) and `bias` ({channels},), got {tuple(weight.shape)}"
119
+ )
120
+ if seq_idx is not None and (
121
+ seq_idx.ndim != 2 or seq_idx.shape[0] not in (1, batch) or seq_idx.dtype not in (torch.int32, torch.int64)
122
+ ):
123
+ raise ValueError(
124
+ f"`seq_idx` must be an int32 or int64 tensor of shape (1, length) or ({batch}, length), got "
125
+ f"{seq_idx.dtype} {tuple(seq_idx.shape)}"
126
+ )
127
+ if any(t is not None and t.device != hidden_states.device for t in (weight, bias, seq_idx)):
128
+ raise ValueError(
129
+ f"`weight`, `bias` and `seq_idx` must be on the device of `hidden_states` ({hidden_states.device})"
130
+ )
131
+ output = torch.empty_like(hidden_states, memory_format=torch.contiguous_format)
132
+ # Triton launches on the current device; the model may live on another one (`device_map`, `.to("cuda:1")`).
133
+ with torch.cuda.device(hidden_states.device):
134
+ _prefill[(batch, channels, triton.cdiv(length, 128))](
135
+ hidden_states,
136
+ # the kernels index `bias` and `weight` densely
137
+ weight.contiguous(),
138
+ bias.contiguous() if bias is not None else None,
139
+ seq_idx,
140
+ output,
141
+ length,
142
+ channels,
143
+ *hidden_states.stride(),
144
+ *(
145
+ (0 if seq_idx.shape[0] == 1 else seq_idx.stride(0), seq_idx.stride(1))
146
+ if seq_idx is not None
147
+ else (0, 0)
148
+ ),
149
+ seq_idx.shape[-1] if seq_idx is not None else 0,
150
+ bias is not None,
151
+ seq_idx is not None,
152
+ 128,
153
+ enable_fp_fusion=False,
154
+ )
155
+ return output
156
+
157
+
158
+ def lfm3_causal_conv1d_update(hidden_states, conv_state, weight, bias=None):
159
+ batch, channels, length = hidden_states.shape
160
+ # the kernel addresses one position per row and bounds nothing by the other tensors' shapes
161
+ if length != 1:
162
+ raise ValueError(f"`hidden_states` must hold one position per row, got {tuple(hidden_states.shape)}")
163
+ # the 8-wide tile holds a 6- or 7-wide state (see `positions` in `_decode`)
164
+ if (
165
+ conv_state.ndim != 3
166
+ or conv_state.shape[:2] != (batch, channels)
167
+ or conv_state.shape[-1] not in (6, 7)
168
+ or weight.shape != (channels, 7)
169
+ ):
170
+ raise ValueError(
171
+ f"`conv_state` must have shape ({batch}, {channels}, 6 or 7) and `weight` ({channels}, 7), got "
172
+ f"{tuple(conv_state.shape)} and {tuple(weight.shape)}"
173
+ )
174
+ if bias is not None and bias.shape != (channels,):
175
+ raise ValueError(f"`bias` must have shape ({channels},), got {tuple(bias.shape)}")
176
+ # the kernel writes the state in place: rows or channels sharing memory (an `expand`ed cache) would race
177
+ extent = 0
178
+ for stride, size in sorted(zip(conv_state.stride(), conv_state.shape)):
179
+ if size > 1 and stride <= extent:
180
+ raise ValueError(f"`conv_state` must not overlap itself, got strides {conv_state.stride()}")
181
+ extent += stride * (size - 1)
182
+ if any(t is not None and t.device != hidden_states.device for t in (conv_state, weight, bias)):
183
+ raise ValueError(
184
+ f"`conv_state`, `weight` and `bias` must be on the device of `hidden_states` ({hidden_states.device})"
185
+ )
186
+ output = torch.empty_like(hidden_states, memory_format=torch.contiguous_format)
187
+ with torch.cuda.device(hidden_states.device):
188
+ _decode[(batch, triton.cdiv(channels, 32))](
189
+ hidden_states,
190
+ weight.contiguous(),
191
+ bias.contiguous() if bias is not None else None,
192
+ conv_state,
193
+ output,
194
+ channels,
195
+ conv_state.shape[-1],
196
+ *hidden_states.stride()[:2],
197
+ *conv_state.stride(),
198
+ bias is not None,
199
+ 32,
200
+ enable_fp_fusion=False,
201
+ )
202
+ return output
model.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:135922caef64e81e52af88ede8b7c3ddaa9d6da220a9737fb3065275dbfb4330
3
+ size 3457635448
modeling_lfm3.py ADDED
@@ -0,0 +1,881 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Copyright 2026 The HuggingFace Team. All rights reserved.
2
+ #
3
+ # Licensed under the Apache License, Version 2.0 (the "License");
4
+ # you may not use this file except in compliance with the License.
5
+ # You may obtain a copy of the License at
6
+ #
7
+ # http://www.apache.org/licenses/LICENSE-2.0
8
+ #
9
+ # Unless required by applicable law or agreed to in writing, software
10
+ # distributed under the License is distributed on an "AS IS" BASIS,
11
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12
+ # See the License for the specific language governing permissions and
13
+ # limitations under the License.
14
+ import copy
15
+ from collections.abc import Callable
16
+
17
+ import torch
18
+ import torch.nn.functional as F
19
+ from torch import nn
20
+
21
+ from transformers import initialization as init
22
+ from transformers.activations import ACT2FN
23
+ from transformers.cache_utils import Cache, DynamicCache
24
+ from transformers.configuration_utils import PreTrainedConfig
25
+ from transformers.generation import GenerationMixin
26
+ from transformers.integrations import (
27
+ use_kernel_forward_from_hub,
28
+ use_kernel_func_from_hub_with_fallback,
29
+ use_kernelized_func,
30
+ )
31
+ from transformers.integrations.accelerate import force_accelerate_hooks
32
+ from transformers.masking_utils import (
33
+ create_causal_mask,
34
+ create_recurrent_attention_mask,
35
+ create_sliding_window_causal_mask,
36
+ )
37
+ from transformers.modeling_layers import GradientCheckpointingLayer
38
+ from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast
39
+ from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update
40
+ from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel
41
+ from transformers.processing_utils import Unpack
42
+ from transformers.utils import TransformersKwargs, auto_docstring, can_return_tuple, torch_compilable_check
43
+ from transformers.utils.generic import merge_with_config_defaults
44
+ from transformers.utils.import_utils import is_tracing
45
+ from transformers.utils.output_capturing import capture_outputs
46
+
47
+ from .configuration_lfm3 import Lfm3Config
48
+ from .lfm3_conv import can_use_lfm3_conv, lfm3_causal_conv1d, lfm3_causal_conv1d_update
49
+
50
+
51
+ @use_kernel_forward_from_hub("RMSNorm")
52
+ class Lfm3RMSNorm(nn.Module):
53
+ def __init__(self, hidden_size, eps: float = 1e-6) -> None:
54
+ """
55
+ Lfm3RMSNorm is equivalent to T5LayerNorm
56
+ """
57
+ super().__init__()
58
+ self.weight = nn.Parameter(torch.ones(hidden_size))
59
+ self.variance_epsilon = eps
60
+
61
+ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
62
+ input_dtype = hidden_states.dtype
63
+ hidden_states = hidden_states.to(torch.float32)
64
+ variance = hidden_states.pow(2).mean(-1, keepdim=True)
65
+ hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
66
+ return self.weight * hidden_states.to(input_dtype)
67
+
68
+ def extra_repr(self):
69
+ return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}"
70
+
71
+
72
+ class Lfm3RotaryEmbedding(nn.Module):
73
+ """
74
+ Rotary embedding per attention layer type, like Gemma 3; `conv` layers register no buffers. Partial RoPE keeps the
75
+ frequency progression of the full head dimension and zeroes the tail instead of shrinking the rotated dimension.
76
+ """
77
+
78
+ def __init__(self, config: Lfm3Config, device=None):
79
+ super().__init__()
80
+ self.max_seq_len_cached = config.max_position_embeddings
81
+ self.original_max_seq_len = config.max_position_embeddings
82
+ self.config = config
83
+ self.layer_types = sorted(set(config.layer_types) & set(config.rope_parameters))
84
+ self.rope_type = {}
85
+ for layer_type in self.layer_types:
86
+ rope_params = self.config.rope_parameters[layer_type]
87
+ self.rope_type[layer_type] = rope_params["rope_type"]
88
+ curr_inv_freq, curr_attention_scaling = self.compute_rope_parameters(config, device, layer_type=layer_type)
89
+ setattr(self, f"{layer_type}_inv_freq", nn.Buffer(curr_inv_freq, persistent=False))
90
+ setattr(self, f"{layer_type}_original_inv_freq", nn.Buffer(curr_inv_freq.clone(), persistent=False))
91
+ setattr(self, f"{layer_type}_attention_scaling", curr_attention_scaling)
92
+
93
+ @staticmethod
94
+ def compute_default_rope_parameters(
95
+ config: Lfm3Config, device=None, layer_type: str | None = None, **kwargs
96
+ ) -> tuple[torch.Tensor, float]:
97
+ """
98
+ Inverse frequencies of the default RoPE with LFM3's zero-tail partial rotation: the first
99
+ `partial_rotary_factor * head_dim // 2` frequencies follow the progression of the full head dimension and the
100
+ rest are zero, leaving those dimensions unrotated.
101
+
102
+ Returns:
103
+ Tuple of (`torch.Tensor`, `float`): the inverse frequencies and the attention scaling factor (always 1).
104
+ """
105
+ base = config.rope_parameters[layer_type]["rope_theta"]
106
+ partial_rotary_factor = config.rope_parameters[layer_type].get("partial_rotary_factor", 1.0)
107
+ dim = getattr(config, "head_dim", None) or config.hidden_size // config.num_attention_heads
108
+ rope_angles = int(partial_rotary_factor * dim) // 2
109
+ inv_freq = 1.0 / (base ** (torch.arange(0, 2 * rope_angles, 2, dtype=torch.float) / dim))
110
+ inv_freq = torch.cat((inv_freq, torch.zeros(dim // 2 - rope_angles, dtype=torch.float)))
111
+ return inv_freq.to(device), 1.0
112
+
113
+ @torch.no_grad()
114
+ @dynamic_rope_update # power user: used with advanced RoPE types (e.g. dynamic rope)
115
+ def forward(self, x, position_ids, layer_type):
116
+ inv_freq = getattr(self, f"{layer_type}_inv_freq")
117
+ attention_scaling = getattr(self, f"{layer_type}_attention_scaling")
118
+
119
+ freqs = position_ids[..., None].float() * inv_freq.to(device=x.device, dtype=torch.float)
120
+ emb = torch.cat((freqs, freqs), dim=-1)
121
+ cos = emb.cos() * attention_scaling
122
+ sin = emb.sin() * attention_scaling
123
+
124
+ return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)
125
+
126
+ @classmethod
127
+ def compute_rope_parameters(
128
+ cls, config: Lfm3Config, device=None, layer_type: str | None = None
129
+ ) -> tuple[torch.Tensor, float]:
130
+ """Inverse frequencies of any RoPE type in the zero-tail layout: scaled types are evaluated over the full head
131
+ dimension and truncated to the rotated prefix."""
132
+ rope_params = config.rope_parameters[layer_type]
133
+ if rope_params["rope_type"] == "default":
134
+ return cls.compute_default_rope_parameters(config, device, layer_type=layer_type)
135
+ full_config = copy.copy(config)
136
+ full_config.rope_parameters = {
137
+ **config.rope_parameters,
138
+ layer_type: {**rope_params, "partial_rotary_factor": 1.0},
139
+ }
140
+ inv_freq, attention_factor = ROPE_INIT_FUNCTIONS[rope_params["rope_type"]](
141
+ full_config, device, layer_type=layer_type
142
+ )
143
+ dim = getattr(config, "head_dim", None) or config.hidden_size // config.num_attention_heads
144
+ rope_angles = int(rope_params.get("partial_rotary_factor", 1.0) * dim) // 2
145
+ inv_freq = torch.cat((inv_freq[:rope_angles], inv_freq.new_zeros(dim // 2 - rope_angles)))
146
+ return inv_freq, attention_factor
147
+
148
+
149
+ class Lfm3MLP(nn.Module):
150
+ def __init__(self, config: Lfm3Config):
151
+ super().__init__()
152
+ self.w1 = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
153
+ self.w3 = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
154
+ self.w2 = nn.Linear(config.intermediate_size, config.hidden_size, bias=False)
155
+
156
+ def forward(self, x):
157
+ return self.w2(F.silu(self.w1(x)) * self.w3(x))
158
+
159
+
160
+ def rotate_half(x):
161
+ """Rotates half the hidden dims of the input."""
162
+ x1 = x[..., : x.shape[-1] // 2]
163
+ x2 = x[..., x.shape[-1] // 2 :]
164
+ return torch.cat((-x2, x1), dim=-1)
165
+
166
+
167
+ def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
168
+ """
169
+ This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch,
170
+ num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim)
171
+ """
172
+ batch, num_key_value_heads, slen, head_dim = hidden_states.shape
173
+ if n_rep == 1:
174
+ return hidden_states
175
+ hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim)
176
+ return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)
177
+
178
+
179
+ def eager_attention_forward(
180
+ module: nn.Module,
181
+ query: torch.Tensor,
182
+ key: torch.Tensor,
183
+ value: torch.Tensor,
184
+ attention_mask: torch.Tensor | None,
185
+ scaling: float,
186
+ dropout: float = 0.0,
187
+ **kwargs: Unpack[TransformersKwargs],
188
+ ):
189
+ key_states = repeat_kv(key, module.num_key_value_groups)
190
+ value_states = repeat_kv(value, module.num_key_value_groups)
191
+
192
+ attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling
193
+ if attention_mask is not None:
194
+ attn_weights = attn_weights + attention_mask
195
+
196
+ attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype)
197
+ attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training)
198
+ attn_output = torch.matmul(attn_weights, value_states)
199
+ attn_output = attn_output.transpose(1, 2).contiguous()
200
+
201
+ return attn_output, attn_weights
202
+
203
+
204
+ def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1):
205
+ """Applies RoPE to the query and key tensors in float32, like the reference implementation."""
206
+ cos = cos.unsqueeze(unsqueeze_dim).float()
207
+ sin = sin.unsqueeze(unsqueeze_dim).float()
208
+ q_embed = (q.float() * cos) + (rotate_half(q.float()) * sin)
209
+ k_embed = (k.float() * cos) + (rotate_half(k.float()) * sin)
210
+ return q_embed.to(q.dtype), k_embed.to(k.dtype)
211
+
212
+
213
+ class Lfm3Attention(nn.Module):
214
+ """Llama attention with QK RMSNorm, a sigmoid output gate per head, a sliding window on `sliding_attention`
215
+ layers and RoPE in float32."""
216
+
217
+ def __init__(self, config: Lfm3Config, layer_idx: int):
218
+ super().__init__()
219
+ self.config = config
220
+ self.layer_idx = layer_idx
221
+ self.head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads)
222
+ self.num_key_value_groups = config.num_attention_heads // config.num_key_value_heads
223
+ self.scaling = self.head_dim**-0.5
224
+ self.is_causal = True
225
+ self.q_proj = nn.Linear(config.hidden_size, config.num_attention_heads * self.head_dim, bias=False)
226
+ self.k_proj = nn.Linear(config.hidden_size, config.num_key_value_heads * self.head_dim, bias=False)
227
+ self.v_proj = nn.Linear(config.hidden_size, config.num_key_value_heads * self.head_dim, bias=False)
228
+ self.o_proj = nn.Linear(config.num_attention_heads * self.head_dim, config.hidden_size, bias=False)
229
+ self.q_norm = Lfm3RMSNorm(self.head_dim, eps=config.norm_eps)
230
+ self.k_norm = Lfm3RMSNorm(self.head_dim, eps=config.norm_eps)
231
+ self.gate_proj = nn.Linear(config.hidden_size, config.num_attention_heads, bias=False)
232
+ self.sliding_window = config.sliding_window if config.layer_types[layer_idx] == "sliding_attention" else None
233
+
234
+ def forward(
235
+ self,
236
+ hidden_states: torch.Tensor,
237
+ position_embeddings: tuple[torch.Tensor, torch.Tensor],
238
+ attention_mask: torch.Tensor | None,
239
+ past_key_values: Cache | None = None,
240
+ **kwargs: Unpack[TransformersKwargs],
241
+ ) -> tuple[torch.Tensor, torch.Tensor | None]:
242
+ input_shape = hidden_states.shape[:-1]
243
+ hidden_shape = (*input_shape, -1, self.head_dim)
244
+
245
+ query_states = self.q_norm(self.q_proj(hidden_states).view(*hidden_shape)).transpose(1, 2)
246
+ key_states = self.k_norm(self.k_proj(hidden_states).view(*hidden_shape)).transpose(1, 2)
247
+ value_states = self.v_proj(hidden_states).view(*hidden_shape).transpose(1, 2)
248
+
249
+ cos, sin = position_embeddings
250
+ query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)
251
+
252
+ if past_key_values is not None:
253
+ key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx)
254
+
255
+ attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface(
256
+ self.config._attn_implementation, eager_attention_forward
257
+ )
258
+
259
+ attn_output, attn_weights = attention_interface(
260
+ self,
261
+ query_states,
262
+ key_states,
263
+ value_states,
264
+ attention_mask,
265
+ dropout=0.0,
266
+ scaling=self.scaling,
267
+ sliding_window=self.sliding_window,
268
+ **kwargs,
269
+ )
270
+ gate = torch.sigmoid(self.gate_proj(hidden_states)).unsqueeze(-1)
271
+ attn_output = (attn_output * gate).reshape(*input_shape, -1).contiguous()
272
+ attn_output = self.o_proj(attn_output)
273
+ return attn_output, attn_weights
274
+
275
+
276
+ def apply_mask_to_padding_states(hidden_states, attention_mask):
277
+ """
278
+ Tunes out the hidden states for padding tokens, see https://github.com/state-spaces/mamba/issues/66
279
+ """
280
+ # NOTE: attention mask is a 2D boolean tensor
281
+ if attention_mask is not None:
282
+ dtype = hidden_states.dtype
283
+ hidden_states = (hidden_states * attention_mask[:, :, None]).to(dtype)
284
+
285
+ return hidden_states
286
+
287
+
288
+ @use_kernel_func_from_hub_with_fallback("causal_conv1d_update", "causal_conv1d")
289
+ def causal_conv1d_update(
290
+ hidden_states: torch.Tensor,
291
+ conv_state: torch.Tensor,
292
+ weight: nn.Parameter,
293
+ bias: nn.Parameter | None = None,
294
+ activation: str | None = None,
295
+ ):
296
+ _, hidden_size, seq_len = hidden_states.shape
297
+ state_len = conv_state.shape[-1]
298
+
299
+ hidden_states_new = torch.cat([conv_state, hidden_states], dim=-1).to(weight.dtype)
300
+ conv_state.copy_(hidden_states_new[:, :, -state_len:])
301
+ out = F.conv1d(hidden_states_new, weight.unsqueeze(1), bias, padding=0, groups=hidden_size)
302
+ out = out[:, :, -seq_len:]
303
+ if activation is not None:
304
+ out = ACT2FN[activation](out)
305
+ return out.to(hidden_states.dtype)
306
+
307
+
308
+ @use_kernel_func_from_hub_with_fallback("causal_conv1d_fn", "causal_conv1d")
309
+ def causal_conv1d_fn(
310
+ hidden_states: torch.Tensor,
311
+ weight: nn.Parameter,
312
+ bias: nn.Parameter | None = None,
313
+ activation: str | None = None,
314
+ **kwargs,
315
+ ):
316
+ _, hidden_size, seq_len = hidden_states.shape
317
+ padding = weight.shape[-1] - 1
318
+
319
+ out = F.conv1d(
320
+ hidden_states.to(weight.dtype),
321
+ weight=weight.unsqueeze(1),
322
+ bias=bias,
323
+ padding=padding,
324
+ groups=hidden_size,
325
+ )[:, :, :seq_len]
326
+ if activation is not None:
327
+ out = ACT2FN[activation](out)
328
+ return out.to(hidden_states.dtype)
329
+
330
+
331
+ @use_kernelized_func([causal_conv1d_update, causal_conv1d_fn])
332
+ class Lfm3ShortConv(nn.Module):
333
+ """LFM2 gated short convolution with packed-document boundaries (`seq_idx`), the padding mask re-applied after a
334
+ biased `in_proj`, and width-7 Triton inference kernels with a PyTorch fallback."""
335
+
336
+ def __init__(
337
+ self,
338
+ config: Lfm3Config,
339
+ layer_idx: int,
340
+ ):
341
+ super().__init__()
342
+ self.config = config
343
+ self.layer_idx = layer_idx
344
+ self.conv_kernel_size = config.conv_L_cache
345
+ self.bias = config.conv_bias
346
+
347
+ self.conv = nn.Conv1d(
348
+ in_channels=config.hidden_size,
349
+ out_channels=config.hidden_size,
350
+ kernel_size=self.conv_kernel_size,
351
+ groups=config.hidden_size,
352
+ bias=self.bias,
353
+ padding=self.conv_kernel_size - 1,
354
+ )
355
+ self.in_proj = nn.Linear(config.hidden_size, 3 * config.hidden_size, bias=self.bias)
356
+ self.out_proj = nn.Linear(config.hidden_size, config.hidden_size, bias=self.bias)
357
+
358
+ self.layer_type = config.layer_types[layer_idx]
359
+
360
+ @force_accelerate_hooks("conv")
361
+ def forward(
362
+ self,
363
+ hidden_states: torch.Tensor,
364
+ past_key_values: Cache | None = None,
365
+ attention_mask: torch.Tensor | None = None,
366
+ seq_idx: torch.IntTensor | None = None,
367
+ ):
368
+ seq_len = hidden_states.shape[1]
369
+
370
+ hidden_states = apply_mask_to_padding_states(hidden_states, attention_mask)
371
+ BCx = self.in_proj(hidden_states).transpose(-1, -2)
372
+ B, C, x = BCx.chunk(3, dim=-2)
373
+ hidden_states = B * x
374
+ if attention_mask is not None:
375
+ # with `conv_bias` the projection makes the padded positions non-zero again
376
+ hidden_states = hidden_states * attention_mask[:, None, :].to(hidden_states.dtype)
377
+
378
+ use_precomputed_states = past_key_values is not None and past_key_values.has_previous_state(self.layer_idx)
379
+ use_kernels = 2 <= self.conv_kernel_size <= 4
380
+ use_wide_kernel = can_use_lfm3_conv(hidden_states, self.conv.weight)
381
+
382
+ if use_precomputed_states and seq_len == 1 and not past_key_values.layers[self.layer_idx].record_past:
383
+ conv_state = past_key_values.layers[self.layer_idx].conv_states[0]
384
+ if use_wide_kernel and conv_state.shape[-1] in (6, 7):
385
+ hidden_states = lfm3_causal_conv1d_update(
386
+ hidden_states, conv_state, self.conv.weight.squeeze(1), self.conv.bias
387
+ )
388
+ elif use_kernels:
389
+ hidden_states = causal_conv1d_update(
390
+ hidden_states, conv_state, self.conv.weight.squeeze(1), self.conv.bias
391
+ )
392
+ else:
393
+ hidden_states = torch.cat([conv_state, hidden_states], dim=-1)
394
+ conv_state.copy_(hidden_states[:, :, -conv_state.shape[-1] :])
395
+ hidden_states = self._conv(hidden_states)[:, :, -seq_len:]
396
+ else:
397
+ if past_key_values is not None and seq_idx is not None:
398
+ # `seq_idx` is only set on an empty cache, and the decode steps read the cached window without it:
399
+ # only the last document (after any positional left padding) may enter the window
400
+ last_document = (seq_idx == seq_idx[:, -1:])[:, None].to(hidden_states.dtype)
401
+ past_key_values.update_conv_state(
402
+ hidden_states * last_document, self.layer_idx, conv_kernel_size=self.conv_kernel_size
403
+ )
404
+ elif past_key_values is not None:
405
+ hidden_states = past_key_values.update_conv_state(
406
+ hidden_states, self.layer_idx, conv_kernel_size=self.conv_kernel_size
407
+ )
408
+
409
+ if use_wide_kernel:
410
+ hidden_states = lfm3_causal_conv1d(hidden_states, self.conv.weight.squeeze(1), self.conv.bias, seq_idx)
411
+ elif seq_idx is not None:
412
+ hidden_states = self._segmented_conv(hidden_states, seq_idx)
413
+ elif use_kernels:
414
+ hidden_states = causal_conv1d_fn(hidden_states, self.conv.weight.squeeze(1), self.conv.bias)
415
+ else:
416
+ hidden_states = self._conv(hidden_states)
417
+
418
+ # Drop the additional previous states
419
+ if past_key_values is not None:
420
+ hidden_states = hidden_states[:, :, -seq_len:]
421
+
422
+ y = C * hidden_states
423
+ y = y.transpose(-1, -2).contiguous()
424
+ y = self.out_proj(y)
425
+ return y
426
+
427
+ def _conv(self, hidden_states: torch.Tensor) -> torch.Tensor:
428
+ """Causal depthwise conv in torch: the `causal_conv1d` kernels only accept widths 2 to 4."""
429
+ conv_out = self.conv(hidden_states.to(self.conv.weight.dtype))[..., : hidden_states.shape[-1]]
430
+ return conv_out.to(hidden_states.dtype)
431
+
432
+ def _segmented_conv(self, hidden_states: torch.Tensor, seq_idx: torch.Tensor) -> torch.Tensor:
433
+ """Causal depthwise conv that only reads positions sharing the current token's document index."""
434
+ kernel_size = self.conv_kernel_size
435
+ seq_idx = F.pad(seq_idx, (hidden_states.shape[-1] - seq_idx.shape[-1] + kernel_size - 1, 0), value=-1)
436
+ windows = F.pad(hidden_states, (kernel_size - 1, 0)).unfold(-1, kernel_size, 1)
437
+ seq_idx_windows = seq_idx.unfold(-1, kernel_size, 1)
438
+ same_document = (seq_idx_windows == seq_idx_windows[..., -1:]).unsqueeze(1)
439
+ conv_out = torch.einsum("bdtk,dk->bdt", windows * same_document.to(windows.dtype), self.conv.weight.squeeze(1))
440
+ if self.bias:
441
+ conv_out = conv_out + self.conv.bias[None, :, None]
442
+ return conv_out.to(hidden_states.dtype)
443
+
444
+
445
+ class Lfm3DecoderLayer(GradientCheckpointingLayer):
446
+ def __init__(self, config: Lfm3Config, layer_idx: int):
447
+ super().__init__()
448
+ self.is_attention_layer = config.layer_types[layer_idx] in config._attention_layer_types
449
+ self.use_sandwich_norm = config.use_sandwich_norm
450
+
451
+ if self.is_attention_layer:
452
+ self.self_attn = Lfm3Attention(config, layer_idx)
453
+ else:
454
+ self.conv = Lfm3ShortConv(config, layer_idx)
455
+ self.feed_forward = Lfm3MLP(config)
456
+ self.operator_norm = Lfm3RMSNorm(config.hidden_size, eps=config.norm_eps)
457
+ self.ffn_norm = Lfm3RMSNorm(config.hidden_size, eps=config.norm_eps)
458
+ if self.use_sandwich_norm:
459
+ self.post_operator_norm = Lfm3RMSNorm(config.hidden_size, eps=config.norm_eps)
460
+ self.post_ffn_norm = Lfm3RMSNorm(config.hidden_size, eps=config.norm_eps)
461
+
462
+ def forward(
463
+ self,
464
+ hidden_states: torch.Tensor,
465
+ position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None,
466
+ attention_mask: torch.Tensor | None = None,
467
+ position_ids: torch.LongTensor | None = None,
468
+ past_key_values: Cache | None = None,
469
+ seq_idx: torch.IntTensor | None = None,
470
+ **kwargs,
471
+ ) -> torch.Tensor:
472
+ residual = hidden_states
473
+ if self.is_attention_layer:
474
+ hidden_states, _ = self.self_attn(
475
+ hidden_states=self.operator_norm(hidden_states),
476
+ position_embeddings=position_embeddings,
477
+ attention_mask=attention_mask,
478
+ position_ids=position_ids,
479
+ past_key_values=past_key_values,
480
+ **kwargs,
481
+ )
482
+ else:
483
+ hidden_states = self.conv(
484
+ hidden_states=self.operator_norm(hidden_states),
485
+ past_key_values=past_key_values,
486
+ attention_mask=attention_mask,
487
+ seq_idx=seq_idx,
488
+ )
489
+ if self.use_sandwich_norm:
490
+ hidden_states = self.post_operator_norm(hidden_states)
491
+ hidden_states = hidden_states + residual
492
+
493
+ ffn_out = self.feed_forward(self.ffn_norm(hidden_states))
494
+ if self.use_sandwich_norm:
495
+ ffn_out = self.post_ffn_norm(ffn_out)
496
+ hidden_states = hidden_states + ffn_out
497
+
498
+ return hidden_states
499
+
500
+
501
+ def create_masks_by_layer_type(
502
+ config: PreTrainedConfig,
503
+ inputs_embeds: torch.Tensor,
504
+ attention_mask: torch.Tensor | None,
505
+ past_key_values: Cache | None,
506
+ position_ids: torch.Tensor | None = None,
507
+ ) -> dict:
508
+ """The cache sizes each mask from the layer it is asked about, and the generic mask utilities ask about the first
509
+ non-sliding layer, which for a leading `conv` layer redirects to the first attention layer of any type: name the
510
+ first layer of each mask's own type instead."""
511
+ config = config.get_text_config()
512
+ mask_kwargs = {
513
+ "config": config,
514
+ "inputs_embeds": inputs_embeds,
515
+ "attention_mask": attention_mask,
516
+ "past_key_values": past_key_values,
517
+ "position_ids": position_ids,
518
+ }
519
+ layer_types, causal_mask_mapping = config.layer_types, {}
520
+ if "full_attention" in layer_types:
521
+ layer_idx = layer_types.index("full_attention")
522
+ causal_mask_mapping["full_attention"] = create_causal_mask(**mask_kwargs, layer_idx=layer_idx)
523
+ if "sliding_attention" in layer_types:
524
+ layer_idx = layer_types.index("sliding_attention")
525
+ causal_mask_mapping["sliding_attention"] = create_sliding_window_causal_mask(
526
+ **mask_kwargs, layer_idx=layer_idx
527
+ )
528
+ if "conv" in layer_types:
529
+ causal_mask_mapping["conv"] = create_recurrent_attention_mask(**mask_kwargs)
530
+ return causal_mask_mapping
531
+
532
+
533
+ @auto_docstring
534
+ class Lfm3PreTrainedModel(PreTrainedModel):
535
+ config: Lfm3Config
536
+ base_model_prefix = "model"
537
+ supports_gradient_checkpointing = True
538
+ _no_split_modules = ["Lfm3DecoderLayer"]
539
+ _skip_keys_device_placement = ["past_key_values"]
540
+ _supports_flash_attn = True
541
+ _supports_sdpa = True
542
+ _supports_flex_attn = True
543
+
544
+ _can_compile_fullgraph = True
545
+ _supports_attention_backend = True
546
+ _can_record_outputs = {
547
+ "hidden_states": Lfm3DecoderLayer,
548
+ "attentions": Lfm3Attention,
549
+ }
550
+
551
+ @torch.no_grad()
552
+ def _init_weights(self, module):
553
+ super()._init_weights(module)
554
+ if isinstance(module, Lfm3RotaryEmbedding):
555
+ for layer_type in module.layer_types:
556
+ curr_inv_freq, _ = module.compute_rope_parameters(module.config, layer_type=layer_type)
557
+ init.copy_(getattr(module, f"{layer_type}_inv_freq"), curr_inv_freq)
558
+ init.copy_(getattr(module, f"{layer_type}_original_inv_freq"), curr_inv_freq)
559
+
560
+ @staticmethod
561
+ def create_masks_for_generate(
562
+ config: PreTrainedConfig,
563
+ inputs_embeds: torch.Tensor,
564
+ attention_mask: torch.Tensor | None,
565
+ past_key_values: Cache | None,
566
+ position_ids: torch.Tensor | None = None,
567
+ **kwargs,
568
+ ) -> dict:
569
+ return create_masks_by_layer_type(config, inputs_embeds, attention_mask, past_key_values, position_ids)
570
+
571
+
572
+ def packed_under_padding(position_ids: torch.Tensor, padding_mask: torch.Tensor, filled) -> torch.Tensor:
573
+ """Padding repeats position 0, so packing is a restart between two kept tokens, or a kept position 0 on a filled
574
+ cache whose row holds a kept token; attention under a padding mask cannot isolate it."""
575
+ restarts = torch.diff(position_ids, prepend=position_ids[:, :1] - 1, dim=-1) != 1
576
+ valid = padding_mask[:, -restarts.shape[-1] :].bool()
577
+ packed = restarts[:, 1:] & valid[:, 1:] & valid[:, :-1]
578
+ if filled:
579
+ # a row whose cached tokens are all padding (chunked prefill of a left-padded batch) has no document to pack
580
+ # against; a mask covering only the chunk keeps the strict rule
581
+ cached = padding_mask[:, : padding_mask.shape[-1] - restarts.shape[-1]]
582
+ cached_valid = cached.bool().any(-1, keepdim=True) if cached.shape[-1] else torch.ones_like(valid[:, :1])
583
+ packed = packed | (valid[:, :1] & (position_ids[:, :1] == 0) & cached_valid)
584
+ return packed
585
+
586
+
587
+ @auto_docstring
588
+ class Lfm3Model(Lfm3PreTrainedModel):
589
+ """LFM2 model with sliding-attention layers as a third layer type: the mask and the rotary embedding are built per
590
+ layer type, and packed documents (`position_ids` restarting) reset the short-conv window through `seq_idx`."""
591
+
592
+ def __init__(self, config: Lfm3Config):
593
+ super().__init__(config)
594
+ self.padding_idx = config.pad_token_id
595
+ self.vocab_size = config.vocab_size
596
+
597
+ self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)
598
+ self.layers = nn.ModuleList(
599
+ [Lfm3DecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]
600
+ )
601
+ self.rotary_emb = Lfm3RotaryEmbedding(config=config)
602
+ self.gradient_checkpointing = False
603
+ self.embedding_norm = Lfm3RMSNorm(config.hidden_size, eps=config.norm_eps)
604
+
605
+ # Initialize weights and apply final processing
606
+ self.post_init()
607
+
608
+ @merge_with_config_defaults
609
+ @capture_outputs
610
+ @auto_docstring
611
+ def forward(
612
+ self,
613
+ input_ids: torch.LongTensor | None = None,
614
+ attention_mask: torch.Tensor | None = None,
615
+ position_ids: torch.LongTensor | None = None,
616
+ past_key_values: Cache | None = None,
617
+ inputs_embeds: torch.FloatTensor | None = None,
618
+ use_cache: bool | None = None,
619
+ **kwargs: Unpack[TransformersKwargs],
620
+ ) -> BaseModelOutputWithPast:
621
+ r"""
622
+ attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)` or `dict`, *optional*):
623
+ Mask to avoid performing attention on padding token indices (1 keeps a token, 0 masks it), or the mask
624
+ dict of `create_masks_for_generate`. A mask dict is used as given: `generate` refuses packed
625
+ `position_ids` before its prefill.
626
+ """
627
+ if (input_ids is None) ^ (inputs_embeds is not None):
628
+ raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
629
+ # `seq_idx` is a shared kwarg (`DataCollatorWithFlattening(return_seq_idx=True)`, read by Lfm2); LFM3 derives
630
+ # the documents from the restarting `position_ids` that come with it
631
+ if kwargs.pop("seq_idx", None) is not None and position_ids is None:
632
+ raise ValueError("Pass packed documents as `position_ids` restarting at 0; LFM3 does not read `seq_idx`")
633
+
634
+ if inputs_embeds is None:
635
+ inputs_embeds = self.embed_tokens(input_ids)
636
+
637
+ if use_cache and past_key_values is None:
638
+ past_key_values = DynamicCache(config=self.config)
639
+ filled = past_key_values is not None and (
640
+ past_key_values.get_seq_length() > 0
641
+ if any(t != "conv" for t in self.config.layer_types)
642
+ else past_key_values.has_previous_state()
643
+ )
644
+ packed_cache = getattr(past_key_values, "lfm3_packed", None)
645
+ if isinstance(packed_cache, torch.Tensor) or packed_cache:
646
+ # Reset layers have no history to mix. Ordinary caches carry a Python `False` and skip this, so the decode
647
+ # steps do no tensor work here; only a compiled packed prefill leaves a tensor flag, and `filled` stays
648
+ # tensor-valued for static caches.
649
+ packed_message = (
650
+ "A cache filled by a packed-document forward cannot be continued: its attention spans the documents"
651
+ )
652
+ packed_cache = packed_cache & filled if isinstance(packed_cache, torch.Tensor) else filled
653
+ if isinstance(packed_cache, torch.Tensor):
654
+ torch_compilable_check(~packed_cache, packed_message)
655
+ if not is_tracing(packed_cache):
656
+ past_key_values.lfm3_packed = bool(packed_cache)
657
+ elif packed_cache:
658
+ raise ValueError(packed_message)
659
+ else:
660
+ past_key_values.lfm3_packed = False
661
+
662
+ # The mask utilities detect packing only with `attention_mask=None`, so a multi-token forward with explicit
663
+ # positions drops the all-ones mask tokenizers emit (a data-dependent check, hence not when traced).
664
+ explicit_positions = position_ids is not None and inputs_embeds.shape[1] > 1
665
+ if (
666
+ explicit_positions
667
+ and isinstance(attention_mask, torch.Tensor)
668
+ and attention_mask.ndim == 2
669
+ and not is_tracing(position_ids)
670
+ and bool(attention_mask.all())
671
+ ):
672
+ attention_mask = None
673
+ seq_idx = None
674
+ if position_ids is None:
675
+ past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
676
+ position_ids = torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device) + past_seen_tokens
677
+ position_ids = position_ids.unsqueeze(0)
678
+ elif explicit_positions and not isinstance(attention_mask, dict):
679
+ # a mask dict comes from `generate`, which refused packing before its prefill
680
+ if isinstance(attention_mask, torch.Tensor) and attention_mask.ndim == 2:
681
+ # data-dependent: a traced forward skips it (`generate` refuses packing eagerly before its prefill)
682
+ if not is_tracing(position_ids):
683
+ torch_compilable_check(
684
+ ~packed_under_padding(position_ids, attention_mask, filled),
685
+ "Packed documents (`position_ids` restarting) cannot be combined with an attention mask: "
686
+ "attention would not isolate the documents. Pass `attention_mask=None` without a static cache.",
687
+ )
688
+ else:
689
+ restarts = torch.diff(position_ids, prepend=position_ids[:, :1] - 1, dim=-1) != 1
690
+ document_ids = restarts.cumsum(-1)
691
+ if is_tracing(position_ids) or bool((document_ids[..., -1] > 0).any()):
692
+ seq_idx = document_ids.to(torch.int32)
693
+ # Packed documents are prefilled in one forward on an empty cache (cached tokens would join the first
694
+ # document), and their cache is never continued
695
+ if filled:
696
+ torch_compilable_check(
697
+ (~restarts[:, 1:]).all() & (position_ids[:, 0] != 0).all(),
698
+ "Packed documents must be prefilled in one forward; the cache already holds tokens",
699
+ )
700
+ # Ordinary continuation must retain the convolution window from the cached document.
701
+ seq_idx = None
702
+ if attention_mask is None:
703
+ # without a mask only the cacheless packed-sequence mask isolates repeated zero positions (left
704
+ # padding): a cache-sized mask and every later decode step would attend them, so they count too
705
+ packed = (document_ids[..., -1] > 0).any()
706
+ else:
707
+ # a prepared 4D mask is the caller's and carries the padding (repeated zero positions), but may
708
+ # or may not isolate documents, so the cache records the packing and is never continued
709
+ steps = torch.diff(position_ids, dim=-1)
710
+ packed = ((steps != 1) & ~((steps == 0) & (position_ids[:, :-1] == 0))).any()
711
+ if past_key_values is not None:
712
+ # resolved here, where the prefill already syncs, so the decode steps carry a Python flag: no
713
+ # per-step host sync, data-dependent op or CUDA-graph output on the cache
714
+ past_key_values.lfm3_packed = packed if is_tracing(position_ids) else bool(packed)
715
+
716
+ if not isinstance(causal_mask_mapping := attention_mask, dict):
717
+ mask_kwargs = {
718
+ "config": self.config,
719
+ "inputs_embeds": inputs_embeds,
720
+ "attention_mask": attention_mask,
721
+ "past_key_values": past_key_values,
722
+ "position_ids": position_ids,
723
+ }
724
+ # The mask utilities only isolate packed documents without a cache, and an empty cache adds nothing.
725
+ layer_types = self.config.layer_types
726
+ if seq_idx is not None and past_key_values is not None:
727
+ attention_layer = next((i for i, t in enumerate(layer_types) if t != "conv"), None)
728
+ if attention_layer is not None and past_key_values.get_seq_length(attention_layer) == 0:
729
+ if past_key_values.is_compileable:
730
+ # static layers always return their full length, so the masks are sized by the cache, which
731
+ # isolates neither documents nor positional-only left padding
732
+ torch_compilable_check(
733
+ ~packed,
734
+ "Packed documents, or left padding without an attention mask, cannot be prefilled into a "
735
+ "static cache: its masks are sized by the cache. Pass the padding as `attention_mask`, or "
736
+ "use a `DynamicCache` for packed documents",
737
+ )
738
+ else:
739
+ mask_kwargs["past_key_values"] = None
740
+ causal_mask_mapping = create_masks_by_layer_type(**mask_kwargs)
741
+
742
+ hidden_states = inputs_embeds
743
+ position_embeddings = {
744
+ layer_type: self.rotary_emb(hidden_states, position_ids=position_ids, layer_type=layer_type)
745
+ for layer_type in self.rotary_emb.layer_types
746
+ }
747
+
748
+ for i, decoder_layer in enumerate(self.layers[: self.config.num_hidden_layers]):
749
+ layer_type = self.config.layer_types[i]
750
+ hidden_states = decoder_layer(
751
+ hidden_states,
752
+ attention_mask=causal_mask_mapping[layer_type],
753
+ position_embeddings=position_embeddings.get(layer_type),
754
+ position_ids=position_ids,
755
+ past_key_values=past_key_values,
756
+ seq_idx=seq_idx,
757
+ **kwargs,
758
+ )
759
+
760
+ hidden_states = self.embedding_norm(hidden_states)
761
+
762
+ return BaseModelOutputWithPast(
763
+ last_hidden_state=hidden_states,
764
+ past_key_values=past_key_values,
765
+ )
766
+
767
+
768
+ def reject_packed_generation(config: PreTrainedConfig, input_ids, generation_config, model_kwargs, is_first_iteration):
769
+ """`generate` decodes from the cache its prefill fills, and no mask it builds isolates packed documents
770
+ (`position_ids` restarting): the prompt is refused before the prefill, one forward or every chunk, writes to it."""
771
+ position_ids, attention_mask = model_kwargs.get("position_ids"), model_kwargs.get("attention_mask")
772
+ cache = model_kwargs.get("past_key_values")
773
+ if cache is None or position_ids is None or position_ids.ndim != 2:
774
+ return
775
+ if attention_mask is not None and attention_mask.ndim != 2:
776
+ return
777
+ has_attention = any(t != "conv" for t in config.get_text_config().layer_types)
778
+ filled = bool(cache.get_seq_length() > 0 if has_attention else cache.has_previous_state())
779
+ inputs_embeds = model_kwargs.get("inputs_embeds") if is_first_iteration else None
780
+ length = input_ids.shape[1] if inputs_embeds is None else inputs_embeds.shape[1]
781
+ if generation_config.prefill_chunk_size is not None:
782
+ # the chunks run from the start of the prompt
783
+ position_ids = position_ids[:, :length]
784
+ else:
785
+ if filled and has_attention:
786
+ # the slice `_prefill` feeds: a prompt that spans the cache is cut to its uncached part
787
+ mask_length = getattr(generation_config, "_mask_length", -1)
788
+ if attention_mask is not None:
789
+ mask_length = attention_mask.shape[1]
790
+ if inputs_embeds is not None or mask_length == length:
791
+ length -= int(cache.get_seq_length())
792
+ position_ids = position_ids[:, position_ids.shape[1] - length :]
793
+ if attention_mask is None:
794
+ attention_mask = torch.ones_like(position_ids)
795
+ if bool(packed_under_padding(position_ids, attention_mask, filled).any()):
796
+ raise ValueError(
797
+ "`generate` cannot decode packed documents (`position_ids` restarting, or starting at 0 on a cache that "
798
+ "holds tokens): attention would span them. Generate each document in its own row of a left-padded batch."
799
+ )
800
+
801
+
802
+ @auto_docstring
803
+ class Lfm3ForCausalLM(Lfm3PreTrainedModel, GenerationMixin):
804
+ _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"}
805
+ _tp_plan = {"lm_head": "colwise_gather_output"}
806
+ _pp_plan = {"lm_head": (["hidden_states"], ["logits"])}
807
+ _fsdp_plan = {"lm_head": "keep_full_weight"}
808
+
809
+ def __init__(self, config):
810
+ super().__init__(config)
811
+ self.model = Lfm3Model(config)
812
+ self.vocab_size = config.vocab_size
813
+ self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
814
+
815
+ # Initialize weights and apply final processing
816
+ self.post_init()
817
+
818
+ @can_return_tuple
819
+ @auto_docstring
820
+ def forward(
821
+ self,
822
+ input_ids: torch.LongTensor | None = None,
823
+ attention_mask: torch.Tensor | None = None,
824
+ position_ids: torch.LongTensor | None = None,
825
+ past_key_values: Cache | None = None,
826
+ inputs_embeds: torch.FloatTensor | None = None,
827
+ labels: torch.LongTensor | None = None,
828
+ use_cache: bool | None = None,
829
+ logits_to_keep: int | torch.Tensor = 0,
830
+ **kwargs: Unpack[TransformersKwargs],
831
+ ) -> CausalLMOutputWithPast:
832
+ r"""
833
+ Example:
834
+
835
+ ```python
836
+ >>> from transformers import AutoTokenizer, Lfm3ForCausalLM
837
+
838
+ >>> model = Lfm3ForCausalLM.from_pretrained("LiquidAI/LFM3-1.7B-Preview")
839
+ >>> tokenizer = AutoTokenizer.from_pretrained("LiquidAI/LFM3-1.7B-Preview")
840
+
841
+ >>> inputs = tokenizer.apply_chat_template(
842
+ ... [{"role": "user", "content": "What is C. elegans?"}], add_generation_prompt=True, return_tensors="pt"
843
+ ... )
844
+ >>> generate_ids = model.generate(**inputs, max_new_tokens=64)
845
+ >>> tokenizer.decode(generate_ids[0, inputs["input_ids"].shape[1] :], skip_special_tokens=True)
846
+ ```"""
847
+ outputs: BaseModelOutputWithPast = self.model(
848
+ input_ids=input_ids,
849
+ attention_mask=attention_mask,
850
+ position_ids=position_ids,
851
+ past_key_values=past_key_values,
852
+ inputs_embeds=inputs_embeds,
853
+ use_cache=use_cache,
854
+ **kwargs,
855
+ )
856
+
857
+ hidden_states = outputs.last_hidden_state
858
+ # Only compute necessary logits, and do not upcast them to float if we are not computing the loss
859
+ slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
860
+ logits = self.lm_head(hidden_states[:, slice_indices, :])
861
+
862
+ loss = None
863
+ if labels is not None:
864
+ loss = self.loss_function(logits=logits, labels=labels, vocab_size=self.config.vocab_size, **kwargs)
865
+
866
+ return CausalLMOutputWithPast(
867
+ loss=loss,
868
+ logits=logits,
869
+ past_key_values=outputs.past_key_values,
870
+ hidden_states=outputs.hidden_states,
871
+ attentions=outputs.attentions,
872
+ )
873
+
874
+ def _prefill(self, input_ids, generation_config, model_kwargs, is_first_iteration=True):
875
+ # every prefill passes here once before its first forward: a check per chunk would leave the earlier chunks
876
+ # in the cache, and a dynamic cache would only refuse packing at the first decode step
877
+ reject_packed_generation(self.config, input_ids, generation_config, model_kwargs, is_first_iteration)
878
+ return super()._prefill(input_ids, generation_config, model_kwargs, is_first_iteration=is_first_iteration)
879
+
880
+
881
+ __all__ = ["Lfm3ForCausalLM", "Lfm3Model", "Lfm3PreTrainedModel"]
tokenizer.json ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:41927bd314d17d7d7fadb44e10cd0919f7b71528d7741596d8cbd15c6945ec44
3
+ size 17908634
tokenizer_config.json ADDED
@@ -0,0 +1,64 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "audio_out_token": "<|audio_out|>",
3
+ "audio_token": "<|audio|>",
4
+ "backend": "tokenizers",
5
+ "bos_token": "<|startoftext|>",
6
+ "eos_token": "<|im_end|>",
7
+ "extra_special_tokens": [],
8
+ "image_token": "<|image|>",
9
+ "is_local": false,
10
+ "local_files_only": false,
11
+ "model_max_length": 1000000000000000019884624838656,
12
+ "model_specific_special_tokens": {
13
+ "audio_out_token": "<|audio_out|>",
14
+ "audio_token": "<|audio|>",
15
+ "image_token": "<|image|>",
16
+ "video_token": "<|video|>",
17
+ "vision_end_token": "<|vision_end|>",
18
+ "vision_start_token": "<|vision_start|>"
19
+ },
20
+ "pad_token": "<|pad|>",
21
+ "response_template": {
22
+ "defaults": {
23
+ "role": "assistant"
24
+ },
25
+ "fields": {
26
+ "content": {
27
+ "close_pattern": "<\\|im_end\\|>\\s*",
28
+ "content": "text"
29
+ },
30
+ "reasoning_content": {
31
+ "close_pattern": "</think>\\s*",
32
+ "content": "text",
33
+ "open_pattern": "<think>\\n?"
34
+ },
35
+ "tool_calls": {
36
+ "close_pattern": "</tool_call>\\s*",
37
+ "content": "xml-inline",
38
+ "content_args": {
39
+ "tag_pattern": "<parameter=(?P<key>[^>\\n]+)>\\n?(?P<value>.*?)\\n?</parameter>(?=\\s*(?:<parameter=|</function>\\s*\\Z|\\Z))",
40
+ "value_parser": {
41
+ "args": {
42
+ "strip": false
43
+ },
44
+ "name": "text"
45
+ }
46
+ },
47
+ "open_pattern": "<tool_call>\\s*<function=(?P<name>[^\\n>]+)>",
48
+ "repeats": true,
49
+ "transform": {
50
+ "function": {
51
+ "arguments": "{content}",
52
+ "name": "{name}"
53
+ },
54
+ "type": "function"
55
+ }
56
+ }
57
+ },
58
+ "start_anchor_pattern": "<\\|im_start\\|>assistant\\n"
59
+ },
60
+ "tokenizer_class": "TokenizersBackend",
61
+ "video_token": "<|video|>",
62
+ "vision_end_token": "<|vision_end|>",
63
+ "vision_start_token": "<|vision_start|>"
64
+ }