Text Generation
Transformers
Safetensors
English
lfm3
llm
language-model
open-source
instruction-following
conversational
custom_code
Instructions to use cosmicstack/cosmic-lightning with libraries, inference providers, notebooks, and local apps. Follow these links to get started.
- Libraries
- Transformers
How to use cosmicstack/cosmic-lightning with Transformers:
# Use a pipeline as a high-level helper from transformers import pipeline pipe = pipeline("text-generation", model="cosmicstack/cosmic-lightning", trust_remote_code=True) messages = [ {"role": "user", "content": "Who are you?"}, ] pipe(messages)# pip install -U transformers accelerate # Load model directly from transformers import AutoModelForCausalLM model = AutoModelForCausalLM.from_pretrained("cosmicstack/cosmic-lightning", trust_remote_code=True, device_map="auto") - Notebooks
- Google Colab
- Kaggle
- Local Apps Settings
- vLLM
How to use cosmicstack/cosmic-lightning with vLLM:
Install from pip and serve model
# Install vLLM from pip: pip install vllm # Start the vLLM server: vllm serve "cosmicstack/cosmic-lightning" # Call the server using curl (OpenAI-compatible API): curl -X POST "http://localhost:8000/v1/chat/completions" \ -H "Content-Type: application/json" \ --data '{ "model": "cosmicstack/cosmic-lightning", "messages": [ { "role": "user", "content": "What is the capital of France?" } ] }'Use Docker
docker model run hf.co/cosmicstack/cosmic-lightning
- SGLang
How to use cosmicstack/cosmic-lightning with SGLang:
Install from pip and serve model
# Install SGLang from pip: pip install sglang # Start the SGLang server: python3 -m sglang.launch_server \ --model-path "cosmicstack/cosmic-lightning" \ --host 0.0.0.0 \ --port 30000 # Call the server using curl (OpenAI-compatible API): curl -X POST "http://localhost:30000/v1/chat/completions" \ -H "Content-Type: application/json" \ --data '{ "model": "cosmicstack/cosmic-lightning", "messages": [ { "role": "user", "content": "What is the capital of France?" } ] }'Use Docker images
docker run --gpus all \ --shm-size 32g \ -p 30000:30000 \ -v ~/.cache/huggingface:/root/.cache/huggingface \ --env "HF_TOKEN=<secret>" \ --ipc=host \ lmsysorg/sglang:latest \ python3 -m sglang.launch_server \ --model-path "cosmicstack/cosmic-lightning" \ --host 0.0.0.0 \ --port 30000 # Call the server using curl (OpenAI-compatible API): curl -X POST "http://localhost:30000/v1/chat/completions" \ -H "Content-Type: application/json" \ --data '{ "model": "cosmicstack/cosmic-lightning", "messages": [ { "role": "user", "content": "What is the capital of France?" } ] }' - Docker Model Runner
How to use cosmicstack/cosmic-lightning with Docker Model Runner:
docker model run hf.co/cosmicstack/cosmic-lightning
Upload folder using huggingface_hub
Browse files- .gitattributes +1 -0
- chat_template.jinja +280 -0
- config.json +69 -0
- configuration_lfm3.py +132 -0
- generation_config.json +16 -0
- lfm3_conv.py +202 -0
- model.safetensors +3 -0
- modeling_lfm3.py +881 -0
- tokenizer.json +3 -0
- tokenizer_config.json +64 -0
.gitattributes
CHANGED
|
@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
|
|
| 33 |
*.zip filter=lfs diff=lfs merge=lfs -text
|
| 34 |
*.zst filter=lfs diff=lfs merge=lfs -text
|
| 35 |
*tfevents* filter=lfs diff=lfs merge=lfs -text
|
|
|
|
|
|
| 33 |
*.zip filter=lfs diff=lfs merge=lfs -text
|
| 34 |
*.zst filter=lfs diff=lfs merge=lfs -text
|
| 35 |
*tfevents* filter=lfs diff=lfs merge=lfs -text
|
| 36 |
+
tokenizer.json filter=lfs diff=lfs merge=lfs -text
|
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 |
+
}
|