diff --git a/.gitattributes b/.gitattributes new file mode 100644 index 0000000000000000000000000000000000000000..62375aaf6f0f97fad2b39e33dbc99fee014c7d8d --- /dev/null +++ b/.gitattributes @@ -0,0 +1,37 @@ +*.7z filter=lfs diff=lfs merge=lfs -text +*.arrow filter=lfs diff=lfs merge=lfs -text +*.bin filter=lfs diff=lfs merge=lfs -text +*.bz2 filter=lfs diff=lfs merge=lfs -text +*.ckpt filter=lfs diff=lfs merge=lfs -text +*.ftz filter=lfs diff=lfs merge=lfs -text +*.gz filter=lfs diff=lfs merge=lfs -text +*.h5 filter=lfs diff=lfs merge=lfs -text +*.joblib filter=lfs diff=lfs merge=lfs -text +*.lfs.* filter=lfs diff=lfs merge=lfs -text +*.mlmodel filter=lfs diff=lfs merge=lfs -text +*.model filter=lfs diff=lfs merge=lfs -text +*.msgpack filter=lfs diff=lfs merge=lfs -text +*.npy filter=lfs diff=lfs merge=lfs -text +*.npz filter=lfs diff=lfs merge=lfs -text +*.onnx filter=lfs diff=lfs merge=lfs -text +*.ot filter=lfs diff=lfs merge=lfs -text +*.parquet filter=lfs diff=lfs merge=lfs -text +*.pb filter=lfs diff=lfs merge=lfs -text +*.pickle filter=lfs diff=lfs merge=lfs -text +*.pkl filter=lfs diff=lfs merge=lfs -text +*.pt filter=lfs diff=lfs merge=lfs -text +*.pth filter=lfs diff=lfs merge=lfs -text +*.rar filter=lfs diff=lfs merge=lfs -text +*.safetensors filter=lfs diff=lfs merge=lfs -text +saved_model/**/* filter=lfs diff=lfs merge=lfs -text +*.tar.* filter=lfs diff=lfs merge=lfs -text +*.tar filter=lfs diff=lfs merge=lfs -text +*.tflite filter=lfs diff=lfs merge=lfs -text +*.tgz filter=lfs diff=lfs merge=lfs -text +*.wasm filter=lfs diff=lfs merge=lfs -text +*.xz filter=lfs diff=lfs merge=lfs -text +*.zip filter=lfs diff=lfs merge=lfs -text +*.zst filter=lfs diff=lfs merge=lfs -text +*tfevents* filter=lfs diff=lfs merge=lfs -text +model.safetensors.index.json filter=lfs diff=lfs merge=lfs -text +assets/k3-hero.jpg filter=lfs diff=lfs merge=lfs -text diff --git a/LICENSE b/LICENSE new file mode 100644 index 0000000000000000000000000000000000000000..97c0111bba229b0981b1d74e34bacef723ca8f7b --- /dev/null +++ b/LICENSE @@ -0,0 +1,52 @@ +Kimi K3 License + +Copyright (c) 2026 Moonshot AI + +Permission is hereby granted, free of charge, to any person (the "Licensee") +obtaining a copy of this software β€” including the model weights, parameters, +configuration files, inference and training code, and associated documentation +(collectively, the "Software") β€” to deal in the Software without restriction. +This includes, without limitation, the rights to use, copy, modify, merge, +publish, distribute, sublicense, and/or sell copies of the Software; to run, +deploy, fine-tune, or otherwise modify the Software and create derivative works +from it; and to permit persons to whom the Software is furnished to do so, in +each case subject to the following conditions: + +1. The above copyright notice and this permission notice shall be included in +all copies or substantial portions of the Software. Licensee's use of the +Software must comply with applicable laws and regulations. + +2. "Model as a Service" means giving a third party access to language model +inference or fine-tuning (e.g., via API) in a manner that allows such third +party to exercise meaningful control over the inputs, parameters, or training +data. This does not include (a) end-user products with model capabilities solely +embedded within specific features or harnesses, or (b) mere relaying of requests +to models hosted by others. + +If the Licensee or any of its affiliates operates a Model as a Service business, +and the aggregate revenue of the Licensee and its affiliates exceeds 20 million +US dollars (or the equivalent in other currencies) in total over any consecutive +12 months, the Licensee must enter into a separate agreement with Moonshot AI +before using the Software or its derivative works for any commercial purpose. + +3. If the Software (or any derivative works thereof) is used for any of the +Licensee's commercial products or services that have more than 100 million +monthly active users, or more than 20 million US dollars (or equivalent in other +currencies) in monthly revenue, "Kimi K3" must be prominently displayed on the +user interface of such product or service. + +4. The requirements set forth in Sections 2 and 3 do not apply to: (a) internal +use of the Software, defined as any use that does not make the Software, its +outputs, or its underlying capabilities available to third parties; or (b) any +use of the Software accessed through Moonshot AI's official products or +certified inference partners. + +5. THE SOFTWARE AND ANY OUTPUT AND RESULTS THEREFROM ARE PROVIDED ON AN β€œAS IS” +BASIS, WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT +LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE +AND NONINFRINGEMENT. IN NO EVENT SHALL MOONSHOT AI OR ITS AFFILIATES OR +COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER +IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN +CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. + +For any questions regarding this license, please contact . \ No newline at end of file diff --git a/README.md b/README.md new file mode 100644 index 0000000000000000000000000000000000000000..ff9d49d0a735330cfbff26ae1c9d5e5c9dd5b2eb --- /dev/null +++ b/README.md @@ -0,0 +1,679 @@ +--- +tags: +- compressed-tensors +license: other +license_name: "kimi-k3" +library_name: transformers +pipeline_tag: image-text-to-text +--- +
+ + Kimi K3 + +
+
+
+ Chat + Homepage +
+ +
+ Hugging Face + Twitter Follow + Discord + ModelScope +
+
+ License +
+ + +

+πŸ“°  Tech Blog |     πŸ“„  Full Report +

+ + + +## 1. Model Introduction + +Kimi K3 is an open-weight, native multimodal agentic model and our most capable model to date. It is a 2.8T-parameter model built on Kimi Delta Attention (KDA) and Attention Residuals (AttnRes), with native vision capabilities and a 1-million-token context window. It is the world's first open 3T-class model, designed for frontier intelligence across long-horizon coding, knowledge work, and reasoning. + +### Key Features +- **New Architecture**: Kimi K3 is built on Kimi Delta Attention (KDA) and Attention Residuals (AttnRes), and scales up MoE sparsity with a Stable LatentMoE framework that activates 16 out of 896 experts β€” yielding an approximate 2.5Γ— improvement in overall scaling efficiency over Kimi K2. +- **Long-Horizon Coding**: Operating with minimal human oversight, Kimi K3 sustains long engineering sessions, navigates massive repositories, and orchestrates terminal tools β€” from GPU kernel optimization and compiler development to vision-in-the-loop game dev, CAD, and even chip design. +- **Agentic Knowledge Work**: Kimi K3 advances end-to-end knowledge work, producing deep research with interactive visualizations, widgets and dashboards, and motion design and video editing, powered by its native multimodal architecture. +- **Native Multimodality & Long Context**: Kimi K3 understands text, images, and video within the same model, and supports a 1-million-token context window. +- **Open Frontier Weights**: We release the full Kimi K3 model weights under the Kimi K3 License, making frontier intelligence openly available for research, deployment, and further innovation. +## 2. Model Summary + +
+ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + +
ArchitectureMixture-of-Experts (MoE)
Total Parameters2.8T
Activated Parameters104B
Number of Layers93
Number of Dense Layers1
Attention-Layer Composition69 KDA + 24 Gated MLA
Attention Hidden Dimension7168
Number of Attention Heads96
Latent MoE Dimension3584
MoE Hidden Dimension (per Expert)3072
Number of Experts896
Selected Experts per Token16
Number of Shared Experts2
Vocabulary Size160K
Context Length1048576
Attention MechanismKDA & Gated MLA
Activation FunctionSiTU-GLU
Vision EncoderMoonViT-V2
Parameters of Vision Encoder401M
QuantizationMXFP4 weights / MXFP8 activations
(quantization-aware training)
ModalityText, Image
+
+ + +## 3. Evaluation Results + +
+ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + +
BenchmarkKimi K3
(max)		Claude Fable 5
(max, w/ fallback)		GPT-5.6 Sol
(max)		Claude Opus 4.8
(max)		GPT-5.5
(xhigh)		GLM-5.2
(max)
Reasoning & Knowledge
GPQA Diamond93.592.694.191.093.591.2
CritPt23.428.632.320.927.120.9
AA-LCR74.770.073.767.774.371.3
HLE-Full43.5 / 56.053.3 / 63.044.5 / 58.049.8 / 57.941.4 / 52.2β€”
Coding
DeepSWE67.570.073.059.067.046.2
ProgramBench77.876.877.671.970.863.7
Terminal-Bench 2.188.388.088.884.683.482.7
FrontierSWE81.286.671.366.764.967.3
SWE-Marathon42.035.039.040.014.013.0
PostTrainBench36.641.434.634.128.434.3
MLS-Bench-Lite48.349.946.242.835.540.4
SciCode58.760.256.153.556.150.5
Kimi Code Bench 2.072.976.964.871.769.064.2
Agentic
BrowseComp91.288.090.484.384.4β€”
DeepSearchQA (F1)95.094.2β€”93.1β€”β€”
ResearchRubrics76.2β€”73.873.564.071.1
GDPval-AA v2 (Elo)168617471736159314911510
Toolathlon-Verified76.577.974.976.273.559.9
MCPMark-Verified94.587.492.976.492.9β€”
MCP-Atlas84.284.783.683.682.882.6
AutomationBench30.829.129.727.222.712.9
JobBench54.357.445.448.438.343.4
AA-Briefcase (Elo)154815831495135411581260
Agents' Last Exam28.325.7†29.627.026.620.4
APEX-Agents41.043.339.939.438.535.6
OfficeQA Pro63.369.963.263.960.941.4
SpreadsheetBench 234.834.732.431.629.128.1
OSWorld-Verified84.885.083.083.479.0β€”
OSWorld 2.058.366.162.655.749.5β€”
SaaS-Bench60.1β€”61.456.143.8β€”
τ³-Banking33.426.833.027.631.326.8
Harvey Lab-AA94.693.687.291.186.391.0
CorpFin v271.671.864.466.768.466.1
Finance Agent v254.456.353.853.951.849.7
Legal Research Bench44.249.548.143.840.431.3
Vision
WorldVQA ForceAnswer51.056.741.839.138.5β€”
OmniDocBench91.189.885.887.989.4β€”
PerceptionBench58.557.259.747.255.8β€”
Video-MME (w. sub)90.0β€”89.586.089.3β€”
MMVU82.1β€”81.279.281.7β€”
BabyVision w/ python85.790.588.981.283.6β€”
MMMU-Pro81.6 / 83.481.2 / 86.583.0 / 84.678.9 / 82.781.2 / 83.2β€”
CharXiv (RQ)84.8 / 91.388.9 / 93.584.6 / 89.180.5 / 89.984.1 / 89.0β€”
MathVision94.3 / 97.894.8 / 98.695.8 / 97.886.7 / 97.192.2 / 96.8β€”
ZeroBench (pass@5)23.0 / 41.023.0 / 46.017.0 / 35.017.0 / 34.022.0 / 41.0β€”
+
+ +
+Footnotes + +All Kimi K3 results are obtained with reasoning effort set to 'max' and temperature = 1.0. For single-step tasks, such as GPQA Diamond, HLE-Full, and vision benchmarks without tools, we set top-p = 0.95; for agentic tasks, we set top-p = 1.0. For HLE-Full, MMMU-Pro, CharXiv (RQ), MathVision, and ZeroBench, each cell reports the scores without and with tool augmentation (general tools for HLE-Full, Python for the vision benchmarks), in that order. + +1. **Reasoning & knowledge benchmarks** + - **CritPt and AA-LCR.** Scores are cited from [Artificial Analysis](https://artificialanalysis.ai/) as of July 23, 2026. +2. **Coding benchmarks** + - **DeepSWE.** Kimi K3 is evaluated with the Kimi Code harness. The GLM-5.2 score is taken from the [GLM-5.2 release blog](https://z.ai/blog/glm-5.2); all remaining scores are from the official [DeepSWE leaderboard](https://deepswe.datacurve.ai/), under which Kimi K3 attains 67.3 with the mini-SWE-agent harness. We report the DeepSWE v1.1 tasks. + - **Terminal-Bench 2.1.** Kimi K3 is evaluated with the Kimi Code harness. For all other models, we report the best score across harnesses: GLM-5.2 with Claude Code ([GLM-5.2 release blog](https://z.ai/blog/glm-5.2)); Claude Opus 4.8 and Claude Fable 5 with Terminus 2 ([Artificial Analysis](https://artificialanalysis.ai/evaluations/terminalbench-v2-1)); GPT-5.5 and GPT-5.6 Sol with Codex ([OpenAI](https://openai.com/index/previewing-gpt-5-6-sol/)). + - **ProgramBench.** Kimi K3 is evaluated with the Kimi Code harness. The GLM-5.2 score is from the [GLM-5.2 release blog](https://z.ai/blog/glm-5.2); all other scores are from [Vals AI](https://www.vals.ai/benchmarks/programbench). + - **SWE-Marathon.** Kimi K3, Claude Opus 4.8, and Claude Fable 5 are evaluated with the Claude Code harness; GPT-5.6 Sol is evaluated with the Codex harness. The GLM-5.2 score is from the [GLM-5.2 release blog](https://z.ai/blog/glm-5.2). Our evaluation is based on an H20-calibrated branch of the [official tasks](https://www.swe-marathon.org/) as of July 9, 2026, prior to the final v1.1 release: the Docker images, performance gates, and reference oracles for the GPU tasks have been recalibrated for H20, while the correctness and anti-cheat validators remain unchanged. Additionally, Claude Fable 5 hit fallbacks on 35% of the tasks in our evaluation, which may have negatively impacted its measured performance. + - **FrontierSWE.** Kimi K3 is evaluated with the Kimi Code harness and GPT-5.6 Sol with the Codex harness; all other results are from [FrontierSWE](https://www.frontierswe.com/). Dominance scores are recomputed from the raw scores using the official evaluation script and are current as of July 16, 2026. + - **PostTrainBench.** Scores for GLM-5.2, GPT-5.5, and Claude Opus 4.8 are adopted from the official [PostTrainBench](https://posttrainbench.com/) results. Kimi K3, Claude Fable 5, and GPT-5.6 Sol are evaluated with the official Harbor implementation at maximum reasoning effort, averaged over three runs on H20 GPUs (instead of H100 in the official setting) β€” Kimi K3 and Claude Fable 5 with the Claude Code harness, and GPT-5.6 Sol with the Codex harness. + - **MLS-Bench-Lite.** Kimi K3 is evaluated with the Kimi Code harness; GLM-5.2 and the Claude models with the Claude Code harness; GPT-5.5 and GPT-5.6 Sol with the Codex harness. + - **SciCode.** Scores are cited from [Artificial Analysis](https://artificialanalysis.ai/) as of July 23, 2026. + - **Kimi Code Bench 2.0 (in-house).** Kimi K3 is evaluated with the Kimi Code harness (it attains 73.7 with the Claude Code harness); GLM-5.2, Claude Opus 4.8, and Claude Fable 5 with the Claude Code harness; GPT-5.5 and GPT-5.6 Sol with the Codex harness. All models are evaluated at maximum reasoning effort, except GPT-5.5, which uses the "xhigh" setting. As the benchmark includes cybersecurity and safety-related tasks, we also disclose the fraction of refused or fallback tasks: Claude Fable 5 hit 13 fallbacks and 1 refusal out of 80 tasks; 10 refusals out of 80 tasks entered GPT-5.6 Sol's cyber guard; GPT-5.5 had 3 refusals out of 80 tasks. +3. **Agentic benchmarks** + - **OfficeQA Pro.** Each test case provides the agent with the entire PDF corpus, with all PDFs rendered as images and no machine-readable text available. + - **OfficeQA Pro and SpreadsheetBench 2.** Kimi K3, GLM-5.2, Claude Opus 4.8, and Claude Fable 5 are evaluated with the Claude Code harness; GPT-5.5 and GPT-5.6 Sol are evaluated with the Codex harness. + - **MCP-Atlas.** All models are evaluated on the 500-task public subset with a 100-turn limit, using Gemini 3.1 Pro as the judge. + - **AutomationBench.** All models are evaluated on the 600-task public subset, following the official GitHub setup in all other respects. + - **BrowseComp.** We adopt a context-compaction strategy triggered at 300K tokens. When evaluated with the full 1M-token context window and no context management, Kimi K3 achieves a score of 90.4. The results of Claude Fable 5, Claude Opus 4.8, GPT-5.6 Sol, and GPT-5.5 are cited from [Anthropic](https://www.anthropic.com/news/claude-fable-5-mythos-5) and [OpenAI](https://openai.com/index/gpt-5-6/). + - **GDPval-AA v2, AA-Briefcase, τ³-Banking, Harvey Lab-AA, and APEX-Agents.** Scores are cited from [Artificial Analysis](https://artificialanalysis.ai/) and the [APEX-Agents leaderboard](https://www.mercor.com/apex/apex-agents-leaderboard/) as of July 23, 2026. For Harvey Lab-AA, we report the criterion pass rate. + - **CorpFin v2, Finance Agent v2, and Legal Research Bench.** Scores are cited from [Vals AI](https://www.vals.ai/). + - **Agents' Last Exam.** Scores are cited from the [official leaderboard](https://agents-last-exam.org/leaderboard) as of July 23, 2026; we report the leaderboard's primary pass-rate metric. On the leaderboard, each model is paired with a specific harness: Kimi K3 with Kimi Code; GPT-5.6 Sol and GPT-5.5 with Codex; Claude Fable 5, Claude Opus 4.8, and GLM-5.2 with Claude Code. † The Claude Fable 5 entry runs at xhigh effort with 40% of tasks annotated as downgraded. +4. **Multimodal benchmarks** + - Except for ZeroBench, which follows the official setting and is run five times, all multimodal scores are averaged over three runs. MMMU-Pro is evaluated following the official protocol, preserving the original input order and prepending images to the text input. + - **PerceptionBench** is an in-house benchmark that focuses on atomic visual perception capabilities. + +
+ +## 4. Native MXFP4 Quantization + +Kimi K3 applies quantization-aware training from the SFT stage onward, using MXFP4 weights with MXFP8 activations for broad hardware compatibility. + +## 5. Deployment + +> [!Note] +> You can access Kimi K3's API on https://platform.kimi.ai by selecting `kimi-k3`, and we provide OpenAI/Anthropic-compatible API for you. Currently, Kimi K3 is recommended to run on the following inference engines: + +- [vLLM](https://github.com/vllm-project/vllm) β€” see [recipes](https://recipes.vllm.ai/moonshotai/Kimi-K3) +- [SGLang](https://github.com/sgl-project/sglang) β€” see [cookbook](https://docs.sglang.io/cookbook/autoregressive/Moonshotai/Kimi-K3) +- [TokenSpeed](https://lightseek.org/tokenspeed) β€” see [recipes](https://lightseek.org/tokenspeed/recipes/models#kimi-k3) + +--- +## 6. Model Usage + +Kimi K3 always has thinking enabled, and will return `reasoning_content`. Thinking effort is configured with the top-level `reasoning_effort` request field, which supports `"low"`, `"high"`, and `"max"` (default `"max"`). + +Kimi K3 was trained in the preserved thinking history mode. For multi-turn conversations and tool calls, Kimi K3 requires the complete assistant message returned by the API to be passed back to `messages` as-is β€” including `reasoning_content` and `tool_calls`, not just `content`: + +```python +import openai + +def chat_with_preserved_thinking(client: openai.OpenAI, model_name: str): + messages = [ + { + "role": "user", + "content": "Tell me three random numbers." + }, + { + "role": "assistant", + "reasoning_content": "I'll start by listing five numbers: 473, 921, 235, 215, 222, and I'll tell you the first three.", + "content": "473, 921, 235" + }, + { + "role": "user", + "content": "What are the other two numbers you have in mind?" + } + ] + + response = client.chat.completions.create( + model=model_name, + messages=messages, + stream=False, + max_tokens=4096, + reasoning_effort="max", + ) + # the assistant should mention 215 and 222 that appear in the prior reasoning content + print(f"response: {response.choices[0].message.reasoning}") + return response.choices[0].message.content +``` + +For full guides and examples (vision input, structured output, partial mode, tool choice, dynamic tool loading, context caching), see the [Kimi K3 Quickstart](https://platform.kimi.ai/docs/guide/kimi-k3-quickstart) and [Thinking Effort](https://platform.kimi.ai/docs/guide/use-thinking-effort). + +### Coding Agent Framework + +Kimi K3 works best with [Kimi Code CLI](https://www.kimi.com/code) as its agent framework. We warmly invite you to give it a try β€” run Kimi Code in your terminal and select Kimi K3 using the `/model` command. We hope you enjoy building with Kimi K3, and we would love to hear your feedback! + + +--- + +## 7. License + +Both the code repository and the model weights are released under the [Kimi K3 License](https://huggingface.co/moonshotai/Kimi-K3/blob/main/LICENSE). + +--- + +## 8. Contact Us + +If you have any questions, please reach out at [support@moonshot.ai](mailto:support@moonshot.ai). diff --git a/assets/kimi-logo.png b/assets/kimi-logo.png new file mode 100644 index 0000000000000000000000000000000000000000..870b8be6e07cc2c46f7173e800fbaff8af0af5d1 Binary files /dev/null and b/assets/kimi-logo.png differ diff --git a/config.json b/config.json new file mode 100644 index 0000000000000000000000000000000000000000..d7f26ead420b1d967f2759679dbebc65edfcff93 --- /dev/null +++ b/config.json @@ -0,0 +1,300 @@ +{ + "architectures": [ + "KimiK3ForConditionalGeneration" + ], + "auto_map": { + "AutoConfig": "configuration_kimi_k3.KimiK3Config", + "AutoModel": "modeling_kimi_k3.KimiK3ForConditionalGeneration", + "AutoModelForCausalLM": "modeling_kimi_k3.KimiK3ForConditionalGeneration" + }, + "bos_token_id": 163584, + "dtype": "bfloat16", + "eos_token_id": 163586, + "ignore_index": -100, + "image_placeholder": "<|kimi_image_placeholder|>", + "media_placeholder_token_id": 163605, + "model_type": "kimi_k3", + "pad_token_id": 163839, + "text_config": { + "_name_or_path": "", + "activation_situ_beta": 4.0, + "activation_situ_linear_beta": 25.0, + "add_cross_attention": false, + "architectures": [ + "KimiLinearForCausalLM" + ], + "attn_res_block_size": 12, + "auto_map": { + "AutoConfig": "configuration_kimi_k3.KimiLinearConfig", + "AutoModel": "modeling_kimi_linear.KimiLinearModel", + "AutoModelForCausalLM": "modeling_kimi_linear.KimiLinearForCausalLM" + }, + "bad_words_ids": null, + "begin_suppress_tokens": null, + "bos_token_id": 163584, + "chunk_size_feed_forward": 0, + "cross_attention_hidden_size": null, + "decoder_start_token_id": null, + "diversity_penalty": 0.0, + "do_sample": false, + "dtype": "bfloat16", + "early_stopping": false, + "encoder_no_repeat_ngram_size": 0, + "eos_token_id": 163586, + "exponential_decay_length_penalty": null, + "finetuning_task": null, + "first_k_dense_replace": 1, + "forced_bos_token_id": null, + "forced_eos_token_id": null, + "hidden_act": "situ", + "hidden_size": 7168, + "id2label": { + "0": "LABEL_0", + "1": "LABEL_1" + }, + "initializer_range": 0.02, + "intermediate_size": 33792, + "is_decoder": false, + "is_encoder_decoder": false, + "kv_lora_rank": 512, + "label2id": { + "LABEL_0": 0, + "LABEL_1": 1 + }, + "latent_moe_use_norm": true, + "length_penalty": 1.0, + "linear_attn_config": { + "full_attn_layers": [ + 4, + 8, + 12, + 16, + 20, + 24, + 28, + 32, + 36, + 40, + 44, + 48, + 52, + 56, + 60, + 64, + 68, + 72, + 76, + 80, + 84, + 88, + 92, + 93 + ], + "gate_lower_bound": -5.0, + "head_dim": 128, + "kda_layers": [ + 1, + 2, + 3, + 5, + 6, + 7, + 9, + 10, + 11, + 13, + 14, + 15, + 17, + 18, + 19, + 21, + 22, + 23, + 25, + 26, + 27, + 29, + 30, + 31, + 33, + 34, + 35, + 37, + 38, + 39, + 41, + 42, + 43, + 45, + 46, + 47, + 49, + 50, + 51, + 53, + 54, + 55, + 57, + 58, + 59, + 61, + 62, + 63, + 65, + 66, + 67, + 69, + 70, + 71, + 73, + 74, + 75, + 77, + 78, + 79, + 81, + 82, + 83, + 85, + 86, + 87, + 89, + 90, + 91 + ], + "num_heads": 96, + "short_conv_kernel_size": 4, + "use_full_rank_gate": true + }, + "max_length": 20, + "max_position_embeddings": 1048576, + "min_length": 0, + "mla_use_nope": true, + "mla_use_output_gate": true, + "model_type": "kimi_linear", + "moe_intermediate_size": 3072, + "moe_layer_freq": 1, + "moe_renormalize": true, + "moe_router_activation_func": "sigmoid", + "no_repeat_ngram_size": 0, + "num_attention_heads": 96, + "num_beam_groups": 1, + "num_beams": 1, + "num_expert_group": 1, + "num_experts": 896, + "num_experts_per_token": 16, + "num_hidden_layers": 93, + "num_key_value_heads": 96, + "num_nextn_predict_layers": 0, + "num_return_sequences": 1, + "num_shared_experts": 2, + "output_attentions": false, + "output_hidden_states": false, + "output_scores": false, + "pad_token_id": 163839, + "prefix": null, + "problem_type": null, + "pruned_heads": {}, + "q_lora_rank": 1536, + "qk_nope_head_dim": 128, + "qk_rope_head_dim": 64, + "quantization_config": { + "config_groups": { + "group_0": { + "format": "mxfp4-pack-quantized", + "input_activations": null, + "output_activations": null, + "targets": [ + "Linear" + ], + "weights": { + "actorder": null, + "block_structure": null, + "dynamic": false, + "group_size": 32, + "num_bits": 4, + "observer": "minmax", + "observer_kwargs": {}, + "scale_dtype": "torch.uint8", + "strategy": "group", + "symmetric": true, + "type": "float", + "zp_dtype": null + } + } + }, + "format": "mxfp4-pack-quantized", + "global_compression_ratio": null, + "ignore": [ + "re:.*self_attn.*", + "re:.*shared_experts.*", + "re:.*mlp\\.(gate|up|gate_up|down)_proj.*", + "re:.*lm_head.*", + "re:.*vision_tower.*", + "re:.*mm_projector.*" + ], + "kv_cache_scheme": null, + "quant_method": "compressed-tensors", + "quantization_status": "compressed" + }, + "remove_invalid_values": false, + "repetition_penalty": 1.0, + "return_dict": true, + "return_dict_in_generate": false, + "rms_norm_eps": 1e-05, + "routed_expert_hidden_size": 3584, + "routed_scaling_factor": 1.0, + "sep_token_id": null, + "suppress_tokens": null, + "task_specific_params": null, + "temperature": 1.0, + "tf_legacy_loss": false, + "tie_encoder_decoder": false, + "tie_word_embeddings": false, + "tokenizer_class": null, + "top_k": 50, + "top_p": 1.0, + "topk_group": 1, + "topk_method": "noaux_tc", + "torchscript": false, + "transformers_version": "4.56.2", + "typical_p": 1.0, + "use_bfloat16": false, + "use_cache": true, + "use_grouped_topk": true, + "v_head_dim": 128, + "vocab_size": 163840 + }, + "tie_word_embeddings": false, + "vision_config": { + "_attn_implementation": "flash_attention_2", + "activation_func": "gelu_pytorch_tanh", + "attn_bias": false, + "init_pos_emb_height": 64, + "init_pos_emb_time": 4, + "init_pos_emb_width": 64, + "linear_bias": false, + "merge_kernel_size": [ + 2, + 2 + ], + "merge_type": "sd2_tpool", + "mlp_type": "mlp2", + "mm_hidden_size": 1024, + "mm_projector_type": "patchmergerv2", + "norm_type": "rmsnorm", + "patch_embed_proj_bias": false, + "patch_size": 14, + "pos_emb_interpolation_mode": "bilinear", + "pos_emb_type": "divided_fixed", + "projector_hidden_act": "gelu", + "projector_ln_eps": 1e-05, + "qkv_hidden_size": 1536, + "text_hidden_size": 7168, + "vt_hidden_size": 1024, + "vt_intermediate_size": 4096, + "vt_num_attention_heads": 12, + "vt_num_hidden_layers": 27 + } +} \ No newline at end of file diff --git a/configuration_kimi_k3.py b/configuration_kimi_k3.py new file mode 100644 index 0000000000000000000000000000000000000000..b38c4f8fcab1ce295aea236772ea1cba57f978c9 --- /dev/null +++ b/configuration_kimi_k3.py @@ -0,0 +1,285 @@ +from typing import Optional + +from transformers.configuration_utils import PretrainedConfig + + +class KimiLinearConfig(PretrainedConfig): + model_type = "kimi_linear" + keys_to_ignore_at_inference = ["past_key_values"] + + def __init__( + self, + model_type="kimi_linear", + vocab_size=163840, + hidden_size=4096, + head_dim=None, + intermediate_size=11008, + num_hidden_layers=32, + num_attention_heads=32, + num_key_value_heads=None, + hidden_act="silu", + initializer_range=0.02, + rms_norm_eps=1e-6, + use_cache=True, + pad_token_id=0, + bos_token_id=1, + eos_token_id=2, + rope_theta=10000.0, + rope_scaling=None, + tie_word_embeddings=False, + moe_intermediate_size: Optional[int] = None, + moe_renormalize: bool = True, + moe_router_activation_func: str = "sigmoid", + num_experts: Optional[int] = None, + num_experts_per_token: Optional[int] = None, + num_shared_experts: int = 0, + routed_scaling_factor: float = 1.0, + first_k_dense_replace: int = 0, + moe_layer_freq: int = 1, + use_grouped_topk: bool = True, + num_expert_group: int = 1, + topk_group: int = 1, + q_lora_rank: Optional[int] = None, + kv_lora_rank: Optional[int] = None, + qk_nope_head_dim: Optional[int] = None, + qk_rope_head_dim: Optional[int] = None, + v_head_dim: Optional[int] = None, + mla_use_nope: Optional[bool] = False, + mla_use_output_gate: Optional[bool] = False, + num_nextn_predict_layers: int = 0, + linear_attn_config: Optional[dict] = None, + attn_res_block_size: Optional[int] = None, + latent_moe_use_norm: bool = False, + activation_situ_beta: Optional[float] = None, + activation_situ_linear_beta: Optional[float] = None, + max_position_embeddings: int = 4096, + routed_expert_hidden_size: Optional[int] = None, + topk_method: str = "noaux_tc", + **kwargs, + ): + self.model_type = model_type + self.vocab_size = vocab_size + self.hidden_size = hidden_size + self.head_dim = ( + head_dim if head_dim is not None else hidden_size // num_attention_heads + ) + self.intermediate_size = intermediate_size + self.num_hidden_layers = num_hidden_layers + self.num_attention_heads = num_attention_heads + + # for backward compatibility + if num_key_value_heads is None: + num_key_value_heads = num_attention_heads + + self.num_key_value_heads = num_key_value_heads + self.hidden_act = hidden_act + self.initializer_range = initializer_range + self.rms_norm_eps = rms_norm_eps + self.use_cache = use_cache + self.rope_theta = rope_theta + self.rope_scaling = rope_scaling + + self.q_lora_rank = q_lora_rank + self.kv_lora_rank = kv_lora_rank + self.qk_nope_head_dim = qk_nope_head_dim + self.qk_rope_head_dim = qk_rope_head_dim + self.v_head_dim = v_head_dim + self.mla_use_nope = mla_use_nope + self.mla_use_output_gate = mla_use_output_gate + # moe config + self.num_experts = num_experts + self.num_experts_per_token = num_experts_per_token + self.moe_renormalize = moe_renormalize + self.num_shared_experts = num_shared_experts + self.routed_scaling_factor = routed_scaling_factor + self.moe_router_activation_func = moe_router_activation_func + assert self.moe_router_activation_func in ("softmax", "sigmoid") + self.moe_intermediate_size = moe_intermediate_size + self.first_k_dense_replace = first_k_dense_replace + self.moe_layer_freq = moe_layer_freq + self.use_grouped_topk = use_grouped_topk + self.num_expert_group = num_expert_group + self.topk_group = topk_group + self.num_nextn_predict_layers = num_nextn_predict_layers + + self.attn_res_block_size = attn_res_block_size + self.latent_moe_use_norm = latent_moe_use_norm + self.activation_situ_beta = activation_situ_beta + self.activation_situ_linear_beta = activation_situ_linear_beta + self.max_position_embeddings = max_position_embeddings + self.routed_expert_hidden_size = routed_expert_hidden_size + self.topk_method = topk_method + + if linear_attn_config is not None: + assert linear_attn_config["kda_layers"] is not None + assert linear_attn_config["full_attn_layers"] is not None + self.linear_attn_config = linear_attn_config + + super().__init__( + pad_token_id=pad_token_id, + bos_token_id=bos_token_id, + eos_token_id=eos_token_id, + tie_word_embeddings=tie_word_embeddings, + **kwargs, + ) + + @property + def is_mla(self): + return ( + self.q_lora_rank is not None + or self.kv_lora_rank is not None + or self.qk_nope_head_dim is not None + or self.qk_rope_head_dim is not None + or self.v_head_dim is not None + or self.mla_use_nope is True + ) + + @property + def is_moe(self): + return self.num_experts is not None + + @property + def is_linear_attn(self) -> bool: + return not ( + self.linear_attn_config is None + or ( + isinstance(self.linear_attn_config, dict) + and self.linear_attn_config["kda_layers"] is not None + and len(self.linear_attn_config["kda_layers"]) == 0 + ) + ) + + def is_kda_layer(self, layer_idx: int): + return ( + self.linear_attn_config is not None + and (layer_idx + 1) in self.linear_attn_config["kda_layers"] + ) + + +class KimiK3VisionConfig(PretrainedConfig): + + def __init__( + self, + patch_size: int = 14, + init_pos_emb_height: int = 64, + init_pos_emb_width: int = 64, + init_pos_emb_time: int = 4, + pos_emb_type: str = 'divided_fixed', + vt_num_attention_heads: int = 12, + vt_num_hidden_layers: int = 27, + vt_hidden_size: int = 1024, + vt_intermediate_size: int = 4096, + merge_kernel_size: tuple = (2, 2), + merge_type: str = 'sd2_tpool', + _attn_implementation: str = 'flash_attention_2', + # MM Projector parameters + mm_projector_type: str = 'patchmergerv2', + mm_hidden_size: int | None = None, + projector_hidden_act: str = "gelu", + projector_ln_eps: float = 1e-5, + # vision tower parameters + qkv_hidden_size: int = 1536, + norm_type: str = 'rmsnorm', + attn_bias: bool = False, + patch_embed_proj_bias: bool = False, + mlp_type: str = 'mlp2', + linear_bias: bool = False, + activation_func: str = 'gelu_pytorch_tanh', + pos_emb_interpolation_mode: str = 'bilinear', + # Other parameters + ignore_index: int = -100, + media_placeholder_token_id: int = 163605, + pad_token_id: int = 0, + text_hidden_size=7168, + **kwargs): + + self.patch_size = patch_size + self.init_pos_emb_height = init_pos_emb_height + self.init_pos_emb_width = init_pos_emb_width + self.init_pos_emb_time = init_pos_emb_time + self.pos_emb_type = pos_emb_type + self.vt_num_attention_heads = vt_num_attention_heads + self.vt_num_hidden_layers = vt_num_hidden_layers + self.vt_hidden_size = vt_hidden_size + self.vt_intermediate_size = vt_intermediate_size + self.merge_kernel_size = merge_kernel_size + self.merge_type = merge_type + self._attn_implementation = _attn_implementation + + # MM Projector config + self.mm_projector_type = mm_projector_type + self.mm_hidden_size = mm_hidden_size if mm_hidden_size is not None else vt_hidden_size + self.projector_hidden_act = projector_hidden_act + self.projector_ln_eps = projector_ln_eps + self.text_hidden_size = text_hidden_size + + # vision tower parameters + self.qkv_hidden_size = qkv_hidden_size + self.norm_type = norm_type + self.attn_bias = attn_bias + self.patch_embed_proj_bias = patch_embed_proj_bias + self.mlp_type = mlp_type + self.linear_bias = linear_bias + self.activation_func = activation_func + self.pos_emb_interpolation_mode = pos_emb_interpolation_mode + + super().__init__(**kwargs) + + +class KimiK3Config(PretrainedConfig): + """Kimi-K3 model configuration. + + Args: + text_config (dict | KimiLinearConfig): Configuration for the text model. + + Vision Tower Parameters (from MoonViT3dConfig): + patch_size (int): Patch size for vision tower. + init_pos_emb_height (int): Initial position embedding height. + init_pos_emb_width (int): Initial position embedding width. + init_pos_emb_time (int): Initial position embedding time dimension. + pos_emb_type (str): Type of position embedding. + vt_num_attention_heads (int): Number of attention heads in vision tower. + vt_num_hidden_layers (int): Number of hidden layers in vision tower. + vt_hidden_size (int): Hidden size of vision tower. + vt_intermediate_size (int): Intermediate size in vision tower FFN. + merge_kernel_size (tuple): Kernel size for patch merging. + merge_type (str): Type of merge operation. + _attn_implementation (str): Attention implementation type. + + MM Projector Parameters (from MultiModalProjectorConfig): + mm_projector_type (str): Type of multimodal projector. + mm_hidden_size (int): Hidden size from vision tower (should match vt_hidden_size). + projector_hidden_act (str): Activation function for projector. + projector_ln_eps (float): Layer norm epsilon for projector. + + Other Parameters: + ignore_index (int): The ignore index for the loss function. + media_placeholder_token_id (int): The token ID to use for media placeholders. + pad_token_id (int): The token ID to use for padding. + """ + + model_type = "kimi_k3" + + def __init__( + self, + text_config: dict | KimiLinearConfig = None, + vision_config: dict | KimiK3VisionConfig = None, + # Other parameters + ignore_index: int = -100, + media_placeholder_token_id: int = 163605, + pad_token_id: int = 0, + **kwargs, + ): + if isinstance(text_config, dict): + text_config = KimiLinearConfig(**text_config) + if isinstance(vision_config, dict): + vision_config = KimiK3VisionConfig(**vision_config) + self.text_config = text_config + self.vision_config = vision_config + # Other config + self.ignore_index = ignore_index + self.media_placeholder_token_id = media_placeholder_token_id + if getattr(self.text_config, "quantization_config", None) is not None: + self.quantization_config = self.text_config.quantization_config + + super().__init__(pad_token_id=pad_token_id, **kwargs) diff --git a/encoding_k3.py b/encoding_k3.py new file mode 100644 index 0000000000000000000000000000000000000000..44a56b7371fab32513da0eaed6b9a8a7c56499a2 --- /dev/null +++ b/encoding_k3.py @@ -0,0 +1,647 @@ +"""Kimi K3 XTML encoding helpers. + +This module keeps chat rendering in Python. +Callers that need token IDs should consume ``EncodeSegment`` objects directly: +structural markers may be encoded as tiktoken special tokens, while user/tool +text and attribute values are encoded as ordinary text. +""" + +from __future__ import annotations + +import json +from dataclasses import dataclass +from typing import Any, Iterable, Optional + +OPEN_TOKEN = "<|open|>" +CLOSE_TOKEN = "<|close|>" +SEP_TOKEN = "<|sep|>" +END_OF_MSG_TOKEN = "<|end_of_msg|>" +IMAGE_PLACEHOLDER = "<|kimi_image_placeholder|>" + +_VALID_THINKING_EFFORTS = {"low", "high", "max"} + + +@dataclass(frozen=True) +class EncodeSegment: + text: str + allow_special: bool = False + + +class _ImagePromptState: + def __init__(self, image_prompts: Optional[list[str]] = None): + self.image_prompts = image_prompts + self.index = 0 + + def next_prompt(self) -> str: + if self.image_prompts is None: + return IMAGE_PLACEHOLDER + if self.index >= len(self.image_prompts): + raise ValueError("More image placeholders than image prompts.") + prompt = self.image_prompts[self.index] + self.index += 1 + return prompt + + def assert_consumed(self) -> None: + if self.image_prompts is None: + return + if self.index != len(self.image_prompts): + raise ValueError( + f"image prompt count {len(self.image_prompts)} != " + f"consumed placeholder count {self.index}" + ) + + +def _segment(text: Any, *, allow_special: bool = False) -> list[EncodeSegment]: + text = str(text) + if not text: + return [] + return [EncodeSegment(text, allow_special=allow_special)] + + +def _control(text: str) -> list[EncodeSegment]: + return _segment(text, allow_special=True) + + +def _text(text: Any) -> list[EncodeSegment]: + return _segment(text, allow_special=False) + + +def _append_text( + segments: list[EncodeSegment], + text: Any, + image_state: _ImagePromptState, +) -> None: + text = str(text) + if text == "": + return + if image_state.image_prompts is None or IMAGE_PLACEHOLDER not in text: + segments.extend(_text(text)) + return + + parts = text.split(IMAGE_PLACEHOLDER) + for i, part in enumerate(parts): + segments.extend(_text(part)) + if i < len(parts) - 1: + segments.extend(_segment(image_state.next_prompt(), + allow_special=True)) + + +def _escape_attr_value(value: Any) -> str: + return str(value).replace("&", "&").replace('"', """) + + +def _attr(key: str, value: Any) -> list[EncodeSegment]: + return ( + _text(f" {key}") + + _text('="') + + _text(_escape_attr_value(value)) + + _text('"') + ) + + +def _open_tag(tag: str, attrs: Iterable[tuple[str, Any]] = ()) -> list[EncodeSegment]: + segments: list[EncodeSegment] = [] + segments.extend(_control(OPEN_TOKEN)) + segments.extend(_text(tag)) + for key, value in attrs: + segments.extend(_attr(key, value)) + segments.extend(_control(SEP_TOKEN)) + return segments + +def _close_tag(tag: str) -> list[EncodeSegment]: + segments: list[EncodeSegment] = [] + segments.extend(_control(CLOSE_TOKEN)) + segments.extend(_text(tag)) + segments.extend(_control(SEP_TOKEN)) + return segments + + +def _end_of_msg() -> list[EncodeSegment]: + return _control(END_OF_MSG_TOKEN) + + +def _json_compact(value: Any) -> str: + return json.dumps(value, ensure_ascii=False, separators=(",", ":")) + + +def _is_mapping(value: Any) -> bool: + return isinstance(value, dict) + + +def _xtml_type(value: Any) -> str: + if isinstance(value, bool): + return "boolean" + if value is None: + return "null" + if isinstance(value, (int, float)) and not isinstance(value, bool): + return "number" + if isinstance(value, str): + return "string" + if _is_mapping(value): + return "object" + return "array" + + +def _xtml_value(value: Any) -> str: + if isinstance(value, str): + return value + return json.dumps(value, ensure_ascii=False) + + +def _get_value(obj: Any, key: str, default: Any = None) -> Any: + if isinstance(obj, dict): + return obj.get(key, default) + return getattr(obj, key, default) + + +def extract_response_schema(response_format: Any) -> Any: + if response_format is None: + return None + + json_schema = _get_value(response_format, "json_schema") + if json_schema is None: + return None + + if isinstance(json_schema, dict): + return json_schema.get( + "schema", + json_schema.get("json_schema", json_schema), + ) + + schema = _get_value(json_schema, "schema") + if schema is not None: + return schema + + schema = _get_value(json_schema, "json_schema") + if schema is not None: + return schema + + return json_schema + + +def deep_sort_dict(obj: Any) -> Any: + if isinstance(obj, dict): + return {k: deep_sort_dict(v) for k, v in sorted(obj.items())} + if isinstance(obj, list): + return [deep_sort_dict(item) for item in obj] + return obj + + +def normalize_tool_arguments(arguments: Any) -> tuple[dict[str, Any], Optional[str]]: + if arguments is None: + return {}, None + if isinstance(arguments, dict): + return arguments, None + if isinstance(arguments, str): + if not arguments.strip(): + return {}, None + try: + parsed = json.loads(arguments) + except json.JSONDecodeError: + return {}, arguments + if not isinstance(parsed, dict): + raise ValueError("Kimi K3 tool call arguments must be a JSON object.") + return parsed, None + raise TypeError( + "Kimi K3 tool call arguments must be a dict or a JSON object string." + ) + + +def normalize_message(message: Any) -> Any: + if not isinstance(message, dict): + return message + + normalized = dict(message) + + tools = normalized.get("tools") + if tools is not None: + normalized["tools"] = deep_sort_dict(tools) + + tool_calls = normalized.get("tool_calls") + if not tool_calls: + return normalized + + normalized_calls = [] + for tool_call in tool_calls: + if not isinstance(tool_call, dict): + normalized_calls.append(tool_call) + continue + + tc = dict(tool_call) + function = tc.get("function") + if isinstance(function, dict): + fn = dict(function) + arguments, json_block = normalize_tool_arguments(fn.get("arguments")) + fn["arguments"] = arguments + if json_block is None: + fn.pop("_xtml_json_block", None) + else: + fn["_xtml_json_block"] = json_block + tc["function"] = fn + else: + arguments, json_block = normalize_tool_arguments(tc.get("arguments")) + tc["arguments"] = arguments + if json_block is None: + tc.pop("_xtml_json_block", None) + else: + tc["_xtml_json_block"] = json_block + normalized_calls.append(tc) + + normalized["tool_calls"] = normalized_calls + return normalized + + +def normalize_conversation(conversation: Any) -> Any: + if not isinstance(conversation, list): + return conversation + + def normalize_messages(messages: list[Any]) -> list[Any]: + return [normalize_message(message) for message in messages] + + if conversation and isinstance(conversation[0], list): + return [normalize_messages(messages) for messages in conversation] + return normalize_messages(conversation) + + +def _tool_call_id_index(tool_calls: Any) -> dict: + """Map assistant ``tool_calls[].id`` to ``(1-based position, function name)``. + + The position mirrors the chat template's enumeration over ``tool_calls`` + (every entry advances the position, even an id-less one). Duplicate ids keep + their first occurrence. + """ + index: dict = {} + if not isinstance(tool_calls, list): + return index + for position, tool_call in enumerate(tool_calls, start=1): + if not isinstance(tool_call, dict): + continue + call_id = tool_call.get("id") + if call_id is None: + continue + key = str(call_id) + if key in index: + continue + function = tool_call.get("function") + name = ( + function.get("name") if isinstance(function, dict) else tool_call.get("name") + ) + index[key] = (position, name) + return index + + +def normalize_xtml_tool_result_messages(messages: list[Any]) -> list[Any]: + """Re-sort K3 XTML tool results into assistant ``tool_calls`` order. + + Serving frameworks generally deliver tool results already in call order. A + direct Transformers caller, however, may pass OpenAI-style tool messages in any + order, so each run of consecutive tool messages is matched against the most + recent preceding assistant ``tool_calls`` by opaque ``tool_call_id`` == + ``tool_calls[].id`` (K3 drops the ``func:index`` format requirement) and + sorted by the matched 1-based position. The matched call is authoritative, + so each matched message's ``tool`` is set to that call's function name -- + this keeps an explicit (and possibly stale) ``tool``/``name`` from drifting + out of sync with the reordered position. ``index`` is still derived from the + rendered position by the chat template. A run that cannot be fully matched is + left untouched. Re-running is idempotent. + + This function is side-effect free: matched tool messages are shallow-copied + before their ``tool``/``name`` is rewritten, and every other message is + appended to the output as-is. The input list and its message objects are + never mutated. + """ + if not isinstance(messages, list): + return messages + + output: list[Any] = [] + current_index: dict = {} + i = 0 + n = len(messages) + + while i < n: + message = messages[i] + + if isinstance(message, dict) and message.get("role") == "assistant": + tool_calls = message.get("tool_calls") + current_index = _tool_call_id_index(tool_calls) if tool_calls else {} + output.append(message) + i += 1 + continue + + if not isinstance(message, dict) or message.get("role") != "tool": + output.append(message) + i += 1 + continue + + run: list[tuple] = [] # (position, original_offset, message, name) + unresolved = False + offset = 0 + while ( + i < n and isinstance(messages[i], dict) and messages[i].get("role") == "tool" + ): + tool_message = messages[i] + call_id = tool_message.get("tool_call_id", tool_message.get("id")) + matched = current_index.get(str(call_id)) if call_id is not None else None + if matched is None: + unresolved = True + run.append((None, offset, tool_message, None)) + else: + position, name = matched + run.append((position, offset, tool_message, name)) + offset += 1 + i += 1 + + if unresolved: + output.extend(item[2] for item in run) + else: + run.sort(key=lambda item: (item[0], item[1])) + for _, _, tool_message, name in run: + if name is None: + output.append(tool_message) + continue + # The id-matched call is authoritative: align tool (and any + # explicit name) so the rendered XTML tool attribute cannot + # disagree with the reordered position. Copy first so the + # caller's message object is never mutated. + resolved = dict(tool_message) + resolved["tool"] = name + if "name" in resolved: + resolved["name"] = name + output.append(resolved) + + return output + + +def is_batched_conversation(conversation: Any) -> bool: + return ( + isinstance(conversation, list) + and bool(conversation) + and isinstance(conversation[0], list) + ) + + +def _render_content_segments( + content: Any, + image_state: _ImagePromptState, +) -> list[EncodeSegment]: + segments: list[EncodeSegment] = [] + if isinstance(content, str): + _append_text(segments, content, image_state) + elif content is not None: + for part in content: + if part["type"] in ["image", "image_url"]: + segments.extend( + _segment(image_state.next_prompt(), allow_special=True)) + else: + _append_text(segments, part["text"], image_state) + return segments + + +def _internal_system_message(message_type: str, body: str) -> list[EncodeSegment]: + segments: list[EncodeSegment] = [] + segments.extend(_open_tag("message", [("role", "system"), ("type", message_type)])) + segments.extend(_text(body.strip())) + segments.extend(_close_tag("message")) + segments.extend(_end_of_msg()) + return segments + + +def _render_assistant_segments( + message: dict[str, Any], + image_state: _ImagePromptState, + thinking: bool = True, +) -> list[EncodeSegment]: + segments: list[EncodeSegment] = [] + # The channel is structural: in thinking mode every assistant + # message carries the open/close tags even when there is no reasoning + # content to fill in. In non-thinking mode the channel is dropped + # entirely. + if thinking: + reasoning_content = message.get("reasoning_content") or message.get( + "reasoning" + ) + segments.extend(_open_tag("think")) + if reasoning_content is not None and str(reasoning_content).strip(): + _append_text(segments, reasoning_content, image_state) + segments.extend(_close_tag("think")) + + segments.extend(_open_tag("response")) + segments.extend(_render_content_segments(message.get("content"), image_state)) + segments.extend(_close_tag("response")) + + tool_calls = message.get("tool_calls") + if tool_calls: + segments.extend(_open_tag("tools")) + for index, tool_call in enumerate(tool_calls, start=1): + fn = tool_call.get("function", tool_call) + segments.extend( + _open_tag("call", [("tool", fn["name"]), ("index", index)]) + ) + args = fn.get("arguments", {}) + json_block = fn.get("_xtml_json_block") + if json_block is not None: + segments.extend(_open_tag("json", [("type", "object")])) + _append_text(segments, json_block, image_state) + segments.extend(_close_tag("json")) + elif _is_mapping(args): + for key, value in args.items(): + segments.extend( + _open_tag( + "argument", + [("key", key), ("type", _xtml_type(value))], + ) + ) + _append_text(segments, _xtml_value(value), image_state) + segments.extend(_close_tag("argument")) + segments.extend(_close_tag("call")) + segments.extend(_close_tag("tools")) + + return segments + + +def _render_tool_declare(tools: Any, *, dynamic: bool = False) -> list[EncodeSegment]: + if dynamic: + body = ( + "## New Tools Available\n" + "The system dynamically extends the toolset via lazy-loading.\n" + "You have access to all existing and extended tools.\n" + "Here are the specs for the extended tools.\n\n" + "```json\n" + f"{_json_compact(tools)}\n" + "```" + ) + else: + body = ( + "# Tools\n" + "Here are the available tools, described in JSONSchema.\n\n" + "```json\n" + f"{_json_compact(tools)}\n" + "```" + ) + segments: list[EncodeSegment] = [] + segments.extend(_open_tag("message", [("role", "system"), ("type", "tool-declare")])) + segments.extend(_text(body)) + segments.extend(_close_tag("message")) + segments.extend(_end_of_msg()) + return segments + + +def build_chat_segments( + messages: list[Any], + tools: Optional[list[dict]] = None, + *, + add_generation_prompt: bool = True, + thinking: bool = True, + image_prompts: Optional[list[str]] = None, + **kwargs: Any, +) -> list[EncodeSegment]: + # Re-sort tool results by tool_call_id at the lowest layer so every caller + # (processor or direct tokenizer) gets correctly ordered XTML. The helper is + # side-effect free, so the caller's message objects are left untouched. + messages = normalize_xtml_tool_result_messages(messages) + messages = normalize_conversation(messages) + tools = deep_sort_dict(tools) + + kwargs = dict(kwargs) + response_format = kwargs.get("response_format") + if "response_schema" not in kwargs: + response_schema = extract_response_schema(response_format) + if response_schema is not None: + kwargs["response_schema"] = response_schema + if kwargs.get("response_schema") is not None: + kwargs["response_schema"] = deep_sort_dict(kwargs["response_schema"]) + + image_state = _ImagePromptState(image_prompts) + segments: list[EncodeSegment] = [] + + tool_calls = None + tool_index = 0 + + if tools: + segments.extend(_render_tool_declare(tools)) + + thinking_effort = kwargs.get("thinking_effort") + if thinking and thinking_effort is not None: + assert thinking_effort in _VALID_THINKING_EFFORTS, ( + f"Unsupported thinking_effort={thinking_effort!r}; " + f"supported values are {sorted(_VALID_THINKING_EFFORTS)}." + ) + if thinking and thinking_effort in _VALID_THINKING_EFFORTS: + segments.extend( + _internal_system_message( + "thinking-effort", + "`thinking_effort` guides on how much to think in your " + "thinking channel (not including the response channel), " + "supported values include `low`, `medium`, `high`, and `max`.\n" + f"Now the system is invoked with `thinking_effort={thinking_effort}`.", + ) + ) + + for message_index, message in enumerate(messages): + if not isinstance(message, dict): + continue + + role = message["role"] + if role == "user": + attrs = [("role", "user")] + if message.get("name"): + attrs.append(("name", message["name"])) + segments.extend(_open_tag("message", attrs)) + segments.extend(_render_content_segments(message.get("content"), image_state)) + segments.extend(_close_tag("message")) + segments.extend(_end_of_msg()) + elif role == "system" and message.get("tools"): + segments.extend(_render_tool_declare(message["tools"], dynamic=True)) + elif role == "system": + attrs = [("role", "system")] + if message.get("name"): + attrs.append(("name", message["name"])) + segments.extend(_open_tag("message", attrs)) + segments.extend(_render_content_segments(message.get("content"), image_state)) + segments.extend(_close_tag("message")) + segments.extend(_end_of_msg()) + elif role == "tool": + tool_index += 1 + tool_name = message.get("tool", message.get("name")) + if ( + tool_name is None + and tool_calls is not None + and tool_index <= len(tool_calls) + ): + tc = tool_calls[tool_index - 1] + fn = tc.get("function", tc) + tool_name = fn["name"] + if tool_name is None: + raise ValueError( + "Kimi K3 tool messages need a resolvable tool name: " + "carry `tool`/`name`, or match a preceding assistant " + "tool_call by order." + ) + segments.extend( + _open_tag( + "message", + [("role", "tool"), ("tool", tool_name), ("index", tool_index)], + ) + ) + segments.extend(_render_content_segments(message.get("content"), image_state)) + segments.extend(_close_tag("message")) + segments.extend(_end_of_msg()) + elif role == "assistant": + tool_calls = message.get("tool_calls") + tool_index = 0 + attrs = [("role", "assistant")] + if message.get("name"): + attrs.append(("name", message["name"])) + segments.extend(_open_tag("message", attrs)) + segments.extend(_render_assistant_segments(message, image_state, thinking)) + segments.extend(_close_tag("message")) + segments.extend(_end_of_msg()) + + tool_choice = kwargs.get("tool_choice") + if tool_choice == "required": + segments.extend( + _internal_system_message( + "tool-choice", + "The system is invoked with `tool_choice=required`.\n" + "You MUST call tools in the next message.", + ) + ) + elif tool_choice == "none": + segments.extend( + _internal_system_message( + "tool-choice", + "The system is invoked with `tool_choice=none`.\n" + "You MUST NOT call any tools in the next message.", + ) + ) + + rf = kwargs.get("response_format") + rf_type = _get_value(rf, "type", rf) if isinstance(rf, dict) else rf + if rf_type == "json_object": + segments.extend( + _internal_system_message( + "response-format", + "The system is invoked with `response_format=json_object`.\n" + "Your response must be raw JSON data without markdown code " + "blocks (```json) or any additional formatting.", + ) + ) + elif rf_type == "json_schema": + schema = _json_compact(kwargs.get("response_schema")) + segments.extend( + _internal_system_message( + "response-format", + "The system is invoked with `response_format=json_schema`.\n" + "Your response must be raw JSON data without markdown code " + "blocks (```json) or any additional formatting.\n" + "The JSON data must match the following schema:\n" + f"```json\n{schema}\n```", + ) + ) + + if add_generation_prompt: + segments.extend(_open_tag("message", [("role", "assistant")])) + segments.extend(_open_tag("think" if thinking else "response")) + + image_state.assert_consumed() + return segments diff --git a/generation_config.json b/generation_config.json new file mode 100644 index 0000000000000000000000000000000000000000..3908a4a5477122aa4b09386acab760fb7887c0fc --- /dev/null +++ b/generation_config.json @@ -0,0 +1,4 @@ +{ + "max_length": 1048576, + "eos_token_id": 163586 +} \ No newline at end of file diff --git a/kimi_k3_processor.py b/kimi_k3_processor.py new file mode 100644 index 0000000000000000000000000000000000000000..033a5773885c144509788e89a445217b76d6c046 --- /dev/null +++ b/kimi_k3_processor.py @@ -0,0 +1,187 @@ +"""Kimi-K3 processor: wraps vision processor + tokenizer into a single interface. + +Chat rendering (including XTML tool-result ordering) is handled by the +tokenizer's Python encoder; this processor adds multimodal media preprocessing. +""" + +from transformers.feature_extraction_utils import BatchFeature +from transformers.processing_utils import ProcessorMixin +from transformers.utils import logging + +from .media_utils import ensure_media_type + +logger = logging.get_logger(__name__) + +# ── KimiK3Processor ─────────────────────────────────────────────────── + + +class KimiK3Processor(ProcessorMixin): + r""" + Constructs a KimiK3 processor which wraps a KimiK3 image processor + and a tokenizer into a single processor. + + [`KimiK3Processor`] offers all the functionalities of + [`KimiK3VisionProcessor`] and [`TikTokenTokenizer`]. + + Args: + image_processor ([`KimiK3VisionProcessor`], *optional*): + The image processor is a required input. + tokenizer ([`TikTokenTokenizer`], *optional*): + The tokenizer is a required input. + chat_template (`str`, *optional*): Kept for ProcessorMixin + compatibility. Kimi K3 chat encoding is implemented in Python by + the tokenizer. + """ + + attributes = ["image_processor", "tokenizer"] + valid_kwargs = ["chat_template"] + image_processor_class = "AutoImageProcessor" + tokenizer_class = "AutoTokenizer" + + def __init__( + self, + image_processor=None, + tokenizer=None, + chat_template=None, + **kwargs, + ): + super().__init__(image_processor, + tokenizer, + chat_template=chat_template) + self.media_processor = image_processor + self.image_placeholder = "<|kimi_image_placeholder|>" + + # ── Media preprocessing ──────────────────────────────────────────── + + def update_raw_text(self, text: str, image_prompts: list[str]) -> str: + # Replace image placeholders + image_count = text.count(self.image_placeholder) + if image_count > 0: + assert image_count == len(image_prompts), ( + f"image placeholder count {image_count} != " + f"image_prompts count {len(image_prompts)}") + text_parts = text.split(self.image_placeholder) + assert len(text_parts) == len(image_prompts) + 1 + text = "".join([ + text_parts[i] + image_prompts[i] + for i in range(len(image_prompts)) + ]) + text += text_parts[-1] + + return text + + def preprocess_medias(self, + medias: list[dict]) -> tuple[list[dict], list[str]]: + """Process media items and generate corresponding prompts. + + Returns: + A tuple of (updated_medias, image_prompts). + """ + updated_medias = [] + image_prompts = [] + for media in medias: + if media['type'] == 'image': + updated_medias.append(media) + img = ensure_media_type( + media, + transparent_bg_config=self.media_processor. + _transparent_bg_config, + transparent_bg_fill_stage=self.media_processor. + _transparent_bg_fill_stage, + )['image'] + w, h = img.size + image_prompts.append( + self.media_processor.make_image_prompt(w, h)) + else: + raise ValueError(f"unsupported media type: {media['type']}") + return updated_medias, image_prompts + + # ── Main entry points ────────────────────────────────────────────── + + def __call__(self, + messages: list[dict] = None, + medias: list[dict] = None, + text: str = None, + return_tensors: str = "pt", + **kwargs) -> BatchFeature: + """ + Process multimodal inputs for Kimi-K3 model. + + Args: + messages: List of message dicts with 'role' and 'content' fields. + If provided, medias and text will be extracted automatically. + medias: Pre-extracted list of media dicts. + text: Pre-formatted text string. + return_tensors: Format of returned tensors. Default: 'pt'. + **kwargs: Additional arguments passed to apply_chat_template. + + Returns: + BatchFeature with fields: input_ids, attention_mask, + pixel_values, grid_thws. + """ + if messages is None and (medias is None or text is None): + raise ValueError( + "Provide either 'messages' or both 'medias' and 'text'") + + if medias is not None and text is not None: + updated_medias, image_prompts = (self.preprocess_medias(medias)) + preprocessed = self.media_processor.preprocess( + updated_medias, return_tensors=return_tensors) + text = self.update_raw_text(text, image_prompts) + text_inputs = self.tokenizer(text, return_tensors=return_tensors) + return BatchFeature(data={**text_inputs, **preprocessed.data}) + + if medias is None: + medias = self._extract_medias_from_messages(messages) + updated_medias, image_prompts = (self.preprocess_medias(medias)) + preprocessed = self.media_processor.preprocess( + updated_medias, return_tensors=return_tensors) + + if text is None: + text_inputs = self.tokenizer.apply_chat_template( + messages, + tokenize=True, + return_tensors=return_tensors, + return_dict=True, + image_prompts=image_prompts, + **kwargs) + return BatchFeature(data={**text_inputs, **preprocessed.data}) + + text = self.update_raw_text(text, image_prompts) + text_inputs = self.tokenizer(text, return_tensors=return_tensors) + return BatchFeature(data={**text_inputs, **preprocessed.data}) + + @staticmethod + def _extract_medias_from_messages(messages: list[dict]) -> list[dict]: + """Extract media items from messages in a single pass.""" + medias = [] + for msg in messages: + if msg['role'] != 'user' or not msg.get('content'): + continue + + for content_part in msg['content']: + if not isinstance(content_part, dict): + continue + + content_type = content_part.get('type') + if content_type in ['image_url', 'image']: + image_data = content_part.get(content_type) + assert image_data is not None, f"image data is missing for content part: {content_part}" + medias.append({ + 'type': 'image', + 'image': image_data, + }) + return medias + + def apply_chat_template(self, messages, **kwargs): + return self.tokenizer.apply_chat_template(messages, **kwargs) + + def batch_decode(self, *args, **kwargs): + return self.tokenizer.batch_decode(*args, **kwargs) + + def decode(self, *args, **kwargs): + return self.tokenizer.decode(*args, **kwargs) + + @property + def model_input_names(self): + return ['input_ids', 'attention_mask', 'pixel_values', 'grid_thws'] diff --git a/kimi_k3_vision_processing.py b/kimi_k3_vision_processing.py new file mode 100644 index 0000000000000000000000000000000000000000..58154424ff951289570c8fb5578e7372bfbfaa8b --- /dev/null +++ b/kimi_k3_vision_processing.py @@ -0,0 +1,179 @@ +"""Image processor class for Kimi-K3. +""" + +import json +from typing import Any, Dict, Optional, Union + +import numpy as np +import torch +from PIL import Image +from transformers.image_processing_utils import (BaseImageProcessor, + BatchFeature) +from transformers.utils import TensorType + +from .media_utils import (MediaInput, TransparentBgConfig, _to_tensor, + ensure_media_type, image_to_np, navit_patchify, + navit_resize_image, normalize) + + +class KimiK3VisionProcessor(BaseImageProcessor): + model_type = "kimi_k3" + + def __init__( + self, + media_proc_cfg: dict, + **kwargs, + ): + super().__init__(**kwargs) + self.media_proc_cfg = media_proc_cfg + + @property + def _transparent_bg_config(self) -> Optional[TransparentBgConfig]: + cfg = self.media_proc_cfg.get("transparent_bg_config") + if cfg is None: + return None + if isinstance(cfg, TransparentBgConfig): + return cfg + return TransparentBgConfig(**cfg) + + @property + def _transparent_bg_fill_stage(self) -> str: + return self.media_proc_cfg.get("transparent_bg_fill_stage", + "before_resize") + + def media_tokens_calculator(self, media: MediaInput): + media = ensure_media_type( + media, + transparent_bg_config=self._transparent_bg_config, + transparent_bg_fill_stage=self._transparent_bg_fill_stage, + ) + ret = self.get_resize_config(media) + return ret['num_tokens'] + + @classmethod + def make_image_prompt(cls, width: int, height: int) -> str: + """Build the K3 image placeholder with resolution info.""" + return (f"<|media_begin|>image {width}x{height}" + f"<|media_content|><|media_pad|><|media_end|>") + + def get_resize_config(self, media_input: MediaInput) -> dict: + if media_input['type'] == 'image': + w, h = media_input['image'].size + ret = navit_resize_image( + w, h, self.media_proc_cfg['patch_size'], + self.media_proc_cfg['merge_kernel_size'], + self.media_proc_cfg['in_patch_limit'], + self.media_proc_cfg['patch_limit_on_one_side'], + self.media_proc_cfg['fixed_output_tokens']) + return ret + else: + raise ValueError("Unsupported type: {}".format( + media_input['type'])) + + def resize_image(self, image: Image.Image, new_width: int, new_height: int, + pad_width: int, pad_height: int) -> np.ndarray: + image_np = image_to_np( + image, + (new_width, new_height), + "resize", + transparent_bg_config=self._transparent_bg_config, + transparent_bg_fill_stage=self._transparent_bg_fill_stage, + ) + image_np = np.pad( + image_np, + ((0, pad_height), (0, pad_width), (0, 0)), + mode="constant", + constant_values=0, + ) + return image_np + + def preprocess( + self, + medias: list[MediaInput], + return_tensors: Optional[Union[str, TensorType]] = None, + ) -> BatchFeature: + """ + Preprocess a atom vision input (images) into model-ready tensors. + + Args: + medias: List of MediaInput. + return_tensors: Desired output format ('pt', 'np', 'tf', or None). + + Returns: + BatchFeature containing 'pixel_values' and 'grid_thws' tensors. + """ + if not isinstance(medias, list): + medias = [medias] + if medias: + pixel_values = [] + for item in medias: + item = ensure_media_type( + item, + transparent_bg_config=self._transparent_bg_config, + transparent_bg_fill_stage=self._transparent_bg_fill_stage, + ) + resize_config = self.get_resize_config(item) + new_width, new_height, pad_width, pad_height = resize_config[ + 'new_width'], resize_config['new_height'], resize_config[ + 'pad_width'], resize_config['pad_height'] + if item['type'] == 'image': + image = item['image'] + image_np = self.resize_image(image, new_width, new_height, + pad_width, pad_height) + pixel_values.append(np.expand_dims(image_np, axis=0)) + else: + raise ValueError("Unsupported type: {}".format( + item['type'])) + normalized_pixel_values = [] + image_std_inv = 1.0 / np.array(self.media_proc_cfg['image_std']) + image_mean = np.array(self.media_proc_cfg['image_mean']) + for pixels in pixel_values: + pixels = normalize(pixels, image_mean, image_std_inv) + pixels_and_thw = navit_patchify( + pixels, + self.media_proc_cfg['patch_size'], + ) + normalized_pixel_values.append(pixels_and_thw) + + pixel_values = torch.cat([ + _to_tensor(pixel_value['pixel_values']) + for pixel_value in normalized_pixel_values + ]) + grid_thws = torch.cat([ + _to_tensor(pixel_value['grid_thw'], + dtype=torch.int64).unsqueeze(0) + for pixel_value in normalized_pixel_values + ]) + + data = { + 'pixel_values': pixel_values, + 'grid_thws': grid_thws, + } + + else: + data = {} + + return BatchFeature(data=data, tensor_type=return_tensors) + + def __repr__(self): + return f"KimiK3VisionProcessor(media_proc_cfg={self.media_proc_cfg})" + + def to_dict(self) -> Dict[str, Any]: + output = super().to_dict() + output["media_proc_cfg"] = self.media_proc_cfg + if "media_processor" in output: + del output["media_processor"] + return output + + @classmethod + def from_dict(cls, config_dict: Dict[str, Any], **kwargs): + config = config_dict.copy() + media_proc_cfg = config.pop("media_proc_cfg", {}) + return cls(media_proc_cfg=media_proc_cfg, **config, **kwargs) + + def to_json_string(self): + dictionary = self.to_dict() + for key, value in dictionary.items(): + if hasattr(value, 'tolist'): + dictionary[key] = value.tolist() + return json.dumps(dictionary, indent=2, sort_keys=True) + "\n" diff --git a/media_utils.py b/media_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..e8b2ae127e254e364e6b568b1093a7e12ac576e0 --- /dev/null +++ b/media_utils.py @@ -0,0 +1,376 @@ +import base64 +import functools +import io +import math +from dataclasses import dataclass +from typing import Literal, TypedDict + +import numpy as np +from PIL import Image + + +class ImageInput(TypedDict): + type: Literal['image'] + image: Image.Image + + +MediaInput = ImageInput + + +@dataclass +class TransparentBgConfig: + """The config of the transparent background.""" + + pattern: Literal["white", "black", "gray", "chessboard"] = "black" + """The pattern of the transparent background.""" + + chessboard_square_size: int = 16 + """The size of the squares in the chessboard background.""" + + chessboard_square_on_top_left: bool = True + """Whether to start the chessboard with a white square on the top left.""" + + chessboard_white_value: int = 255 + """The value of the white pixels in the background.""" + + chessboard_gray_value: int = 200 + """The value of the gray pixels in the background.""" + + +@functools.lru_cache(maxsize=256) +def _create_chessboard_background( + height: int, + width: int, + square_size: int, + square_on_top_left: bool, + white_value: int, + gray_value: int, +) -> np.ndarray: + """Create a chessboard background.""" + bg = np.ones((height, width, 3), dtype=np.uint8) * white_value + for y in range(0, height, square_size): + for x in range(0, width, square_size): + if (y // square_size + x // square_size) % 2 == ( + 1 if square_on_top_left else 0): + bg[y:y + square_size, x:x + square_size] = gray_value + return bg + + +def fill_transparent_bg_with( + image: Image.Image, + transparent_bg_config: TransparentBgConfig | None = None, +) -> Image.Image: + """Composite a (possibly) transparent image onto a configured background. + + When ``transparent_bg_config`` is ``None``, the image is simply converted + to RGB (preserving the historical behavior). Otherwise the alpha channel + is alpha-composited over a background generated according to the config. + """ + if transparent_bg_config is None: + return image.convert("RGB") + + if image.mode == "RGB": + return image + + has_alpha = "A" in image.getbands() or "transparency" in image.info + if not has_alpha: + return image.convert("RGB") + + img = np.array(image.convert("RGBA")) + height, width = img.shape[:2] + bg_pattern = transparent_bg_config.pattern + if bg_pattern == "white": + bg = np.full((height, width, 3), 255, dtype=np.uint8) + elif bg_pattern == "black": + bg = np.zeros((height, width, 3), dtype=np.uint8) + elif bg_pattern == "gray": + bg = np.full((height, width, 3), 128, dtype=np.uint8) + elif bg_pattern == "chessboard": + bg = _create_chessboard_background( + height, + width, + transparent_bg_config.chessboard_square_size, + transparent_bg_config.chessboard_square_on_top_left, + transparent_bg_config.chessboard_white_value, + transparent_bg_config.chessboard_gray_value, + ) + else: + raise ValueError(f"Invalid background pattern: {bg_pattern}") + + alpha = img[:, :, 3] + img_rgb = img[:, :, :3] + alpha_normalized = alpha.astype(np.float32) / 255.0 + alpha_3d = np.stack([alpha_normalized] * 3, axis=2) + result = alpha_3d * img_rgb + (1 - alpha_3d) * bg + result = result.astype(np.uint8) + return Image.fromarray(result) + + +def navit_resize_image( + width: int, + height: int, + patch_size: int, + merge_kernel_size: int, + in_patch_limit: int, + patch_limit_on_one_side: int, + fixed_output_tokens: int | None, +): + # Apply the patch limits. + s1 = math.sqrt( + in_patch_limit / + (max(1.0, width // patch_size) * max(1.0, height // patch_size))) + s2 = patch_limit_on_one_side * patch_size / width + s3 = patch_limit_on_one_side * patch_size / height + scale = min(1.0, s1, s2, s3) + new_w, new_h = max(1, int(width * scale)), max(1, int(height * scale)) + new_w = min(new_w, patch_limit_on_one_side * patch_size) + new_h = min(new_h, patch_limit_on_one_side * patch_size) + + # Calculate the padding to make the height and width divisible by the merge kernel size and patch size. + factor = merge_kernel_size * patch_size + + pad_height = (factor - new_h % factor) % factor + pad_width = (factor - new_w % factor) % factor + + if fixed_output_tokens is not None: + num_tokens = fixed_output_tokens + else: + # Calculate new dimensions after padding and patching + token_height = (new_h + pad_height) // factor + token_width = (new_w + pad_width) // factor + + assert token_height * merge_kernel_size <= patch_limit_on_one_side, ( + f"token_height {token_height} * merge_kernel_size {merge_kernel_size} > patch_limit_on_one_side {patch_limit_on_one_side}" + ) + assert token_width * merge_kernel_size <= patch_limit_on_one_side, ( + f"token_width {token_width} * merge_kernel_size {merge_kernel_size} > patch_limit_on_one_side {patch_limit_on_one_side}" + ) + + num_tokens = token_height * token_width + return { + "num_tokens": num_tokens, + "new_width": new_w, + "new_height": new_h, + "pad_width": pad_width, + "pad_height": pad_height, + "sampled_nframes": 1, + } + + +def _to_pil( + data: str | bytes | Image.Image, + transparent_bg_config: TransparentBgConfig | None = None, + to_rgb: bool = True, +) -> Image.Image: + """Load an image and (optionally) composite its transparent background. + + Args: + data: A PIL Image, a base64 ``data:`` URL, a file path, or raw bytes. + transparent_bg_config: The config used to fill the transparent + background. ``None`` keeps the historical behavior of converting + to RGB without compositing. + to_rgb: If ``False`` the image is returned as-is (the + ``transparent_bg_config`` is ignored). The caller is then + expected to call :func:`fill_transparent_bg_with` later β€” e.g. + after a resize. + """ + if isinstance(data, Image.Image): + image = data + elif isinstance(data, str): + if data.startswith("data:"): + raw_base64 = data.split(",")[1] + image = Image.open(io.BytesIO(base64.b64decode(raw_base64))) + else: + image = Image.open(data) + elif isinstance(data, bytes): + image = Image.open(io.BytesIO(data)) + else: + raise ValueError(f"Unsupported data type: {type(data)}") + + if not to_rgb: + return image + + return fill_transparent_bg_with(image, transparent_bg_config) + + +def ensure_media_type( + media: MediaInput, + transparent_bg_config: TransparentBgConfig | None = None, + transparent_bg_fill_stage: Literal["before_resize", + "after_resize"] = "before_resize", +) -> MediaInput: + if media['type'] == 'image': + media['image'] = _to_pil( + media['image'], + transparent_bg_config=transparent_bg_config, + to_rgb=transparent_bg_fill_stage == "before_resize", + ) + return media + else: + raise ValueError(f"Unsupported media type: {media['type']}") + + +def image_to_np( + image: Image.Image, + resize_to: tuple[int, int] | None = None, + mode: str = "resize", + raise_error_for_ill_resize: bool = True, + transparent_bg_config: TransparentBgConfig | None = None, + transparent_bg_fill_stage: Literal["before_resize", + "after_resize"] = "before_resize", +) -> np.ndarray: + """Convert an image to a numpy array. + + Args: + content: The image to convert. + resize_to: The size to resize the image to. + mode: The mode to resize the image to. + raise_error_for_ill_resize: Whether to raise an error for ill-sized resize. + transparent_bg_config: The config of the transparent background. Only + used when ``transparent_bg_fill_stage == "after_resize"`` (the + caller is responsible for filling before resize otherwise). + transparent_bg_fill_stage: When to composite the transparent + background β€” before or after the resize step. + + Returns: + A numpy array. + """ + assert isinstance(image, Image.Image), "image must be a PIL Image" + if resize_to is not None: + if mode == "resize": + image = image.resize(resize_to, resample=Image.Resampling.BICUBIC) + if transparent_bg_fill_stage == "after_resize": + image = fill_transparent_bg_with(image, transparent_bg_config) + + elif mode == "rescale_and_pad_to_center": + scale = min(resize_to[0] / image.width, + resize_to[1] / image.height, 1.0) + new_width = round(image.width * scale) + new_height = round(image.height * scale) + if new_width == 0 or new_height == 0: + if raise_error_for_ill_resize: + raise ValueError( + f"Invalid resize to: {resize_to}, from image size: {image.size}" + ) + else: + return np.zeros((resize_to[1], resize_to[0], 3), + dtype=np.uint8) + + image = image.resize((new_width, new_height), + resample=Image.Resampling.BICUBIC) + if transparent_bg_fill_stage == "after_resize": + image = fill_transparent_bg_with(image, transparent_bg_config) + padding_left = (resize_to[0] - new_width) // 2 + padding_right = resize_to[0] - new_width - padding_left + padding_top = (resize_to[1] - new_height) // 2 + padding_bottom = resize_to[1] - new_height - padding_top + image = np.asarray(image) + image = np.pad( + image, + ((padding_top, padding_bottom), (padding_left, padding_right), + (0, 0)), + mode="constant", + constant_values=0, + ) + assert image.shape == (resize_to[1], resize_to[0], 3) + + elif mode == "rescale_and_pad_to_rightbottom": + scale = min(resize_to[0] / image.width, + resize_to[1] / image.height, 1.0) + new_width = round(image.width * scale) + new_height = round(image.height * scale) + if new_width == 0 or new_height == 0: + if raise_error_for_ill_resize: + raise ValueError( + f"Invalid resize to: {resize_to}, from image size: {image.size}" + ) + else: + return np.zeros((resize_to[1], resize_to[0], 3), + dtype=np.uint8) + + image = image.resize((new_width, new_height), + resample=Image.Resampling.BICUBIC) + if transparent_bg_fill_stage == "after_resize": + image = fill_transparent_bg_with(image, transparent_bg_config) + padding_right = resize_to[0] - new_width + padding_bottom = resize_to[1] - new_height + image = np.asarray(image) + image = np.pad( + image, + ((0, padding_bottom), (0, padding_right), (0, 0)), + mode="constant", + constant_values=0, + ) + assert image.shape == (resize_to[1], resize_to[0], 3) + + else: + raise ValueError(f"Invalid mode: {mode}") + + if isinstance(image, Image.Image): + return np.asarray(image) + else: + return image + + +def navit_patchify(pixel_values: np.ndarray, + patch_size: int) -> dict[str, np.ndarray]: + """Reshape the pixel values to a navit shape. + + Args: + pixel_values: np.ndarray, shape (t, h, w, c) + patch_size: int + + Returns: + dict[str, np.ndarray] + - patches: np.ndarray, shape (t * h//patch_size * w//patch_size, c, patch_size, patch_size) + - grid_thw: np.ndarray, (t, h//patch_size, w//patch_size) + """ + T, H, W, C = pixel_values.shape + assert C == 3, "pixel_values must have 3 channels" + + patches = pixel_values.reshape(T, H // patch_size, patch_size, + W // patch_size, patch_size, C) + # (T, H//patch_size, W//patch_size, C, patch_size, patch_size) + patches = patches.transpose(0, 1, 3, 5, 2, 4) + patches = patches.reshape(-1, C, patch_size, patch_size) + grid_thw = np.array([T, H // patch_size, W // patch_size]) + return {"pixel_values": patches, "grid_thw": grid_thw} + + +def normalize(x: np.ndarray, + mean, + std_inv, + pixels_dtype: np.dtype = np.float32) -> np.ndarray: + """Normalize the image. + + Args: + x: The image to normalize. The shape is (..., 3). The dtype is uint8. The range is [0, 255]. + mean: The mean of the image. + std_inv: The inverse of the std of the image. + pixels_dtype: The dtype of the image. + Returns: + The normalized image. The shape is (..., 3). The dtype is determined by the pixels_dtype. + """ + x = (x / 255.0).astype(pixels_dtype) + x -= mean + x *= std_inv + return x + + +def _to_tensor(data, **kwargs): + import torch + + if isinstance(data, np.ndarray): + return torch.from_numpy(data).to(**kwargs) + elif isinstance(data, torch.Tensor): + return data.to(**kwargs) + elif isinstance(data, list): + return [_to_tensor(item, **kwargs) for item in data] + elif isinstance(data, tuple): + return tuple(_to_tensor(item, **kwargs) for item in data) + elif isinstance(data, dict): + return {k: _to_tensor(v, **kwargs) for k, v in data.items()} + elif data is None: + return None + else: + raise ValueError(f"Unsupported data type: {type(data)}") diff --git a/model-00001-of-000096.safetensors b/model-00001-of-000096.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..fb73b5d245e724c978fbef56a653c3b44a9e2b83 --- /dev/null +++ b/model-00001-of-000096.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:975584c00f85a95fce8ae0f840af8cef69c2ef4db00d34cab3e2cbdfc60f6e51 +size 2341216112 diff --git a/model-00002-of-000096.safetensors b/model-00002-of-000096.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..62066be9801f922833a5e3b37c1085a2b5602b59 --- /dev/null +++ b/model-00002-of-000096.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:26a3284e1d2cb567934ebef002e6a1813551d646739e8bcb1e9e3fe7f878e0f5 +size 16990911504 diff --git a/model-00003-of-000096.safetensors b/model-00003-of-000096.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..11dbdb67d83d32e90d8dfe29ce313a1e387d6d10 --- /dev/null +++ b/model-00003-of-000096.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:e54af9de4c554956082364010f732443bcd5097390f0121a33fb35e37280b5a9 +size 16990911504 diff --git a/model-00004-of-000096.safetensors b/model-00004-of-000096.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..943054919b3c6f1ed3a4fd046bec91c8a80d1d3e --- /dev/null +++ b/model-00004-of-000096.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:5955fd8feda89b1af8400c25e885e7177d47edff155f54b318beb8dd1cec5c05 +size 16567501776 diff --git a/model-00005-of-000096.safetensors b/model-00005-of-000096.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..fcf12d84b8089c5275705af47eb7b8b150f92275 --- /dev/null +++ b/model-00005-of-000096.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:d60d68ad0381ffd2d8716d5991a65993cc81837455bf05d1c76ea3a75662e268 +size 16990911504 diff --git a/model-00006-of-000096.safetensors b/model-00006-of-000096.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..543ee8bd58b737c39ef71ecdbdf3fceaa18159c6 --- /dev/null +++ b/model-00006-of-000096.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:b1d4805767471a9721cd087d2047843ab9262d4f9bbe0d4a306e72c07179f939 +size 16990911504 diff --git a/model-00007-of-000096.safetensors b/model-00007-of-000096.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..b2210fadbce11a8e93591e5a6528a9f076bd6e72 --- /dev/null +++ b/model-00007-of-000096.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:fb1120fef34c0416e73a3259ae6114b9a82b6081871b368a199e0b265cac7df7 +size 16990911504 diff --git a/model-00008-of-000096.safetensors b/model-00008-of-000096.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..6878ea442e2b46f6e561841623f3e88a76f14e87 --- /dev/null +++ b/model-00008-of-000096.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:2318dda54fc1985b63d6068de54fbe312ac87317c902666c9075aa5e10e8a1c9 +size 16567501776 diff --git a/model-00009-of-000096.safetensors b/model-00009-of-000096.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..32c661c5b32beaecb97a3e8ab9f14f001dd9fd53 --- /dev/null +++ b/model-00009-of-000096.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:8b66cdde34f5130cf4f99ec35b16f1f6d33bf0ec946cf2a5612004f5d399f998 +size 16990911504 diff --git a/model-00010-of-000096.safetensors b/model-00010-of-000096.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..38188c9c6177bd85b897d7eafaac4706cdb327ab --- /dev/null +++ b/model-00010-of-000096.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:d34f55f7b734a1eaecc986fc7f92bcaa1b7d18fad90a33befd3fae4f1b2771a9 +size 16990911504 diff --git a/model-00011-of-000096.safetensors b/model-00011-of-000096.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..6c394f8462cfd4bf745d7fc9b6467a1947a59138 --- /dev/null +++ b/model-00011-of-000096.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:1738856a4cca2e356510892de4accf3c3f94b9f598b0394a87e7951740ba4a21 +size 16990916912 diff --git a/model-00012-of-000096.safetensors b/model-00012-of-000096.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..74c5b37aa917f406405aff5286d75f5b86abd2a2 --- /dev/null +++ b/model-00012-of-000096.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:c6b9bef38415f509898fa7b5c04d20c6fee84b0e02c05d74259d201fe6f3868b +size 16567507176 diff --git a/model-00013-of-000096.safetensors b/model-00013-of-000096.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..18297583bc00ef577f01b47cf83b43d2fe8a0c92 --- /dev/null +++ b/model-00013-of-000096.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:3cbf43d56d9c80be0f848a4866e290741c697e09d885dc2351b0a11de514d9c8 +size 16990916912 diff --git a/model-00014-of-000096.safetensors b/model-00014-of-000096.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..bea8ae805f6a18dac1041f0249438577686667af --- /dev/null +++ b/model-00014-of-000096.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:ce5f343f07d408c64fb247bbe8691edb84d54e3976d3ff586ebcf5929fd143ad +size 16990916912 diff --git a/model-00015-of-000096.safetensors b/model-00015-of-000096.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..f713a6e7d41c9b53ddafd9218789b9c84a70835d --- /dev/null +++ b/model-00015-of-000096.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:b55caa8013348498525eef2b1416a657f0c28264ada9b245df19d2eea168b4c1 +size 16990916912 diff --git a/model-00016-of-000096.safetensors b/model-00016-of-000096.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..255a25ad439a5019cb62d052191a35d3d668e92b --- /dev/null +++ b/model-00016-of-000096.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:5a63e63ced659c1996b02ccabdea977546a39ecce025af1e54552de3bfc0489e +size 16567507176 diff --git a/model-00017-of-000096.safetensors b/model-00017-of-000096.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..39496e6cbd0d52a3b04f047c48201bea637ca794 --- /dev/null +++ b/model-00017-of-000096.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:622bfa605205f4ef65e44dac07246bea79e18123ae8b84ab33eec89b6364c3a9 +size 16990916912 diff --git a/model-00018-of-000096.safetensors b/model-00018-of-000096.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..8f57732e2c81476d79b091f0f6d9c6c90f7ae374 --- /dev/null +++ b/model-00018-of-000096.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:838b265ebb86a1320e589685779125d224663160823873b9e5930cf63fafe796 +size 16990916912 diff --git a/model-00019-of-000096.safetensors b/model-00019-of-000096.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..9bbb4ff2ad0caeffbfdcd6868a395b08661d9de9 --- /dev/null +++ b/model-00019-of-000096.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:599a8ecf88f21d4c91295b4160e7564e501807535e522195ee45050a04f17ded +size 16990916912 diff --git a/model-00020-of-000096.safetensors b/model-00020-of-000096.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..3c903a84b0d76fdc5f97e02d10d80f812f68e6e2 --- /dev/null +++ b/model-00020-of-000096.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:8e01b61ab76655ca0c5bd6645303a00c05fb58da94e00f978a9ef1e8114a6f0f +size 16567507176 diff --git a/model-00021-of-000096.safetensors b/model-00021-of-000096.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..c69fcadca7e713d99556c4806fb06daf53deac81 --- /dev/null +++ b/model-00021-of-000096.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:1944265bd024ea6caa23796326e420260a801e5c52e6253155f22ce97e9578d0 +size 16990916912 diff --git a/model-00022-of-000096.safetensors b/model-00022-of-000096.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..b5ea22963fad44b0c5cbe613749340db39cc8c74 --- /dev/null +++ b/model-00022-of-000096.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:2d32d3e3c8da0764f73830289c6a22224f755c0cf058df0262fb0f83b3777029 +size 16990916912 diff --git a/model-00023-of-000096.safetensors b/model-00023-of-000096.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..6ddaa6ffc4d8ab253d6f0c3fe85678caeead19ea --- /dev/null +++ b/model-00023-of-000096.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:c291f2ec15788cb24cdf791c6d68dbf8803b26d01415cff3fab367cba581a169 +size 16990916912 diff --git a/model-00024-of-000096.safetensors b/model-00024-of-000096.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..f91eba1cc7bdb1496664904dc5f9a937d3d68527 --- /dev/null +++ b/model-00024-of-000096.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:278855ab81d42337e2a1857f136d740911fe3f5b1813dc21ce2cb03dc5c3cde3 +size 16567507176 diff --git a/model-00025-of-000096.safetensors b/model-00025-of-000096.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..aed3b0067083b86911455049c636f4dac09b41e8 --- /dev/null +++ b/model-00025-of-000096.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:6e3ae6ad868f5b38fa7ddface520796ea251568b07447deee3d164a0a32d862b +size 16990916912 diff --git a/model-00026-of-000096.safetensors b/model-00026-of-000096.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..31096ed20d34361ea900c56939a339c78dfa2d60 --- /dev/null +++ b/model-00026-of-000096.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:cdd79fb52c7a3350926c4b947ae82cccd0af17fa3fb537c5694a4cedf2b59811 +size 16990916912 diff --git a/model-00027-of-000096.safetensors b/model-00027-of-000096.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..f6409df643698e5663faa7bfadb5f09f5d7cc8f7 --- /dev/null +++ b/model-00027-of-000096.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:1d974b40c4aeb5cc0adeba6f9fbad25fc99f8cd39f9dd3513aeb2fb30a9c4188 +size 16990916912 diff --git a/model-00028-of-000096.safetensors b/model-00028-of-000096.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..825f5a263e10cee9a59c886ad2cef9208501c522 --- /dev/null +++ b/model-00028-of-000096.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:1bec58e89ab7835d3ca40e581cc6f047cb93de47fda804eaea4e2407e01c3ed3 +size 16567507176 diff --git a/model-00029-of-000096.safetensors b/model-00029-of-000096.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..7624d11f6b45f8e7a33bdbc20083ed5a7a4b610f --- /dev/null +++ b/model-00029-of-000096.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:c3bf0e738aa5fbd6ca3ecad982aba8d09efffd58b40cd8fb62d02fe0492f5ceb +size 16990916912 diff --git a/model-00030-of-000096.safetensors b/model-00030-of-000096.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..bd47ee12954bb638046a959273c4893b7dfe97d5 --- /dev/null +++ b/model-00030-of-000096.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:7b996410482ad3c21700ed35d3f2b36282540db92b5631a31aeda1c5e90d6343 +size 16990916912 diff --git a/model-00031-of-000096.safetensors b/model-00031-of-000096.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..063ef5d60b4b6fffde45af096a6a3f9d6bc37b55 --- /dev/null +++ b/model-00031-of-000096.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:c26689540a2471fe690c441a6cf9bec8f00471930e3131962206fd992e0d377c +size 16990916912 diff --git a/model-00032-of-000096.safetensors b/model-00032-of-000096.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..30304e6f56e2fb1c74ea23e343806ae49d9c107e --- /dev/null +++ b/model-00032-of-000096.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:f7dc9e726d46d5c7c2c35af2eea4860cc89dd38a52b60591358a5ca7899c5159 +size 16567507176 diff --git a/model-00033-of-000096.safetensors b/model-00033-of-000096.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..3414c231e7af746d160acb447d78732e2f8f7b67 --- /dev/null +++ b/model-00033-of-000096.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:615afa33b69cbaa42122bb3bcce92e40d6de86713a301c43c5d1bc2ac4342c1b +size 16990916912 diff --git a/model-00034-of-000096.safetensors b/model-00034-of-000096.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..ea19991579b3ab509ba54259a82190e5b35b9c5b --- /dev/null +++ b/model-00034-of-000096.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:a53b27fe92df11e77aec7c5de93ac148b9e2a7bac95020535f68854a7977bde7 +size 16990916912 diff --git a/model-00035-of-000096.safetensors b/model-00035-of-000096.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..d4372968528e113907b2fab21e157bbb0380c836 --- /dev/null +++ b/model-00035-of-000096.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:9f4b44c89e4965a973f4fa71eaa0f8ada110f588065541e83bf52692feaa886d +size 16990916912 diff --git a/model-00036-of-000096.safetensors b/model-00036-of-000096.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..5c90e662ea8e92f9e6823bb8a55fe195f8d5ce4f --- /dev/null +++ b/model-00036-of-000096.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:55aef33fab36df93731289ba38d0ce516d9880c100be5f69f933d0af80bf3c6e +size 16567507176 diff --git a/model-00037-of-000096.safetensors b/model-00037-of-000096.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..778228fe1bcca4709b799e27dfc6be51553a7dd5 --- /dev/null +++ b/model-00037-of-000096.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:e95dd3599d6918f5a22bc7b1256aaee9147f40aa05f43135d083d511eb5cbd80 +size 16990916912 diff --git a/model-00038-of-000096.safetensors b/model-00038-of-000096.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..1b02ce648941b23ccd13e7aff007e08a05a94cc2 --- /dev/null +++ b/model-00038-of-000096.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:1ce472771309248fd1155313dc3f72daa6fcc4793fba86d6c5d6a758ff6254ae +size 16990916912 diff --git a/model-00039-of-000096.safetensors b/model-00039-of-000096.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..da15741818cc137e6d4c887afe6699abc9a42094 --- /dev/null +++ b/model-00039-of-000096.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:9c75b18c0d3a5088e13973dde519e74ead5c7b1df216dcd623b7ee561f378057 +size 16990916912 diff --git a/model-00040-of-000096.safetensors b/model-00040-of-000096.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..33f008e4267f660476dcd8d8de2d994a8c691c04 --- /dev/null +++ b/model-00040-of-000096.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:6f664598c00aee095174db1f860182652ff9b808f53c7764f70e071c6f8a4893 +size 16567507176 diff --git a/model-00041-of-000096.safetensors b/model-00041-of-000096.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..aeb30f19d64fa990c2effd911eac7505f1f73b27 --- /dev/null +++ b/model-00041-of-000096.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:375f05f94a59d1509bef0976269fceb8f28d40e161d4903e64f735b1ede405c2 +size 16990916912 diff --git a/model-00042-of-000096.safetensors b/model-00042-of-000096.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..f8f9b47828f40c67c1723c656b6fa34399e5f1d2 --- /dev/null +++ b/model-00042-of-000096.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:b65947611d6c150394955f4637d4b05007f20a9b1c0f0f02bcf45d0ae2a428ba +size 16990916912 diff --git a/model-00043-of-000096.safetensors b/model-00043-of-000096.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..0df6593bd018483d9bc676530c6df91347cc26a5 --- /dev/null +++ b/model-00043-of-000096.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:b5a425f100bbfaf804ac2dcca197ccf8fb3a910b6aa79d93e3a103b0146a8309 +size 16990916912 diff --git a/model-00044-of-000096.safetensors b/model-00044-of-000096.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..eb294c411393878c601bb1882762df76352e121e --- /dev/null +++ b/model-00044-of-000096.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:113ee012044206e9156dd683cb97ee5f27b07d0a4f936cbda457cd0bcf8ca336 +size 16567507176 diff --git a/model-00045-of-000096.safetensors b/model-00045-of-000096.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..44629b0dac0678fd6d6a7aae765cb363e9af02f0 --- /dev/null +++ b/model-00045-of-000096.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:eb0698659da59c469020a03b3aa8d9ef847fc6359b2513b6720746a44da6f4de +size 16990916912 diff --git a/model-00046-of-000096.safetensors b/model-00046-of-000096.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..20535fcac2b69ea1b2207173a51f722833cf60dd --- /dev/null +++ b/model-00046-of-000096.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:0270727a399c553a602360b04bc1b5feaf09b22e12e3bb7eea5261801d4a1024 +size 16990916912 diff --git a/model-00047-of-000096.safetensors b/model-00047-of-000096.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..9b43b4eb67af7fd9ed3f82dc2ec6675dcc024fab --- /dev/null +++ b/model-00047-of-000096.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:b38f63eb08036cdf766a078c93ee037630d6481ded26065150915b8deaf03bb8 +size 16990916912 diff --git a/model-00048-of-000096.safetensors b/model-00048-of-000096.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..b55b1e9c0321c6198c6e266528350406a73089c2 --- /dev/null +++ b/model-00048-of-000096.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:131e243c02cf2dc713de108edb8bfdb01c7ee28ce59d0a6615a630fb7cf0e4ae +size 16567507176 diff --git a/model-00049-of-000096.safetensors b/model-00049-of-000096.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..21ee9d95120c6dbb6f9364d92d8eb4c5f7a7e2b6 --- /dev/null +++ b/model-00049-of-000096.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:72c91dcf2909ead55b6b5234e725227effbcd59fa5ce9f552e91c48d423cef0f +size 16990916912 diff --git a/model-00050-of-000096.safetensors b/model-00050-of-000096.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..02535cb833d2b1c918c622d6b68cc83b9c9c8ae3 --- /dev/null +++ b/model-00050-of-000096.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:0a627a082cd39ee7d55f7440a5012c1fe084313f970e7d4351bb08d8ac7f5f2a +size 16990916912 diff --git a/model-00051-of-000096.safetensors b/model-00051-of-000096.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..be8d4fe1badfb318c32557e69a9b75956c26b383 --- /dev/null +++ b/model-00051-of-000096.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:38a37bcec20abe1ddfaf68a636cd61cd291b71d554419a57736fd97d08a8cbff +size 16990916912 diff --git a/model-00052-of-000096.safetensors b/model-00052-of-000096.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..aa12e39511200a56c57aa8814d983884f7673534 --- /dev/null +++ b/model-00052-of-000096.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:9703b632174112ea051aaa2ca8b67558ffbf41a4699d042580ff11f26ec28e0a +size 16567507176 diff --git a/model-00053-of-000096.safetensors b/model-00053-of-000096.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..56ffdc55049fb256b9a6b420ed60e0563ce2d227 --- /dev/null +++ b/model-00053-of-000096.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:413ec9ea0b693161c5e8a393a7a6be40a9f3c27114a919e1302fe94daeeee621 +size 16990916912 diff --git a/model-00054-of-000096.safetensors b/model-00054-of-000096.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..60b29945b6ffe530f0b5225aad4a4df26eda4475 --- /dev/null +++ b/model-00054-of-000096.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:0e3da201b765929787e40c2fb8cd0cfb3476bfb9b95d7eb59b7d1066960606a2 +size 16990916912 diff --git a/model-00055-of-000096.safetensors b/model-00055-of-000096.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..0f9178af860d57de6fd0e96d2fe3171f6a1dcfc7 --- /dev/null +++ b/model-00055-of-000096.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:e7c9f4e44f8ac95f73048078c406781fb0418831089fd2f7c55459887606204b +size 16990916912 diff --git a/model-00056-of-000096.safetensors b/model-00056-of-000096.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..53d93c226b24e873ba9e26213ac917881e10d92a --- /dev/null +++ b/model-00056-of-000096.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:efd6176016b9fab27cca7710dce7ed599f06bbba1a8b126cfc76781b42c0e880 +size 16567507176 diff --git a/model-00057-of-000096.safetensors b/model-00057-of-000096.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..0c956ab94b50b40d88a89c17257a48fbdd370780 --- /dev/null +++ b/model-00057-of-000096.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:e6982c96ac9101b75ab32c56568f5740ddae771acad29d92750b95aeeb3a7a30 +size 16990916912 diff --git a/model-00058-of-000096.safetensors b/model-00058-of-000096.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..7ffcfee6535254e6cbc83e46f73ca02f96705f93 --- /dev/null +++ b/model-00058-of-000096.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:c93a23cbd6530391dcd9ce5617d24fb759e380f5ca676f83d00c8d6a9cda896f +size 16990916912 diff --git a/model-00059-of-000096.safetensors b/model-00059-of-000096.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..02e15cecc6b10aa6c0688f9eee4b133bbe6d609d --- /dev/null +++ b/model-00059-of-000096.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:a62eb822071036421fa4aab662dafe3a3873b1b926378ebf489cddfd48a4d5de +size 16990916912 diff --git a/model-00060-of-000096.safetensors b/model-00060-of-000096.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..be38f21c5e00549725bb670dd08132d441a10b8f --- /dev/null +++ b/model-00060-of-000096.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:9bccbaa71b98f8526ca4c6cb091a489063ec4734c84bb5f6d11eeb0512360b5f +size 16567507176 diff --git a/model-00061-of-000096.safetensors b/model-00061-of-000096.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..11593e56df31db7e584b8ab27b88c2ca29141448 --- /dev/null +++ b/model-00061-of-000096.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:1ae3969540fc676480804c37adc54df77e5ebb923952abb46d398b3a2d612e9d +size 16990916912 diff --git a/model-00062-of-000096.safetensors b/model-00062-of-000096.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..410768cdd1578fcaf7a922a38ff63dff01accbc3 --- /dev/null +++ b/model-00062-of-000096.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:96babdc24f2203eed68bd88f8534546696ab35bb1d0db7a5f15aaee295deee68 +size 16990916912 diff --git a/model-00063-of-000096.safetensors b/model-00063-of-000096.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..1c8e823c753fb2e2d12332393932864c006a5615 --- /dev/null +++ b/model-00063-of-000096.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:8a81caa697a7baf491393d83e3b85e14a0d025f5ee2ca91561528e8aec5c63ea +size 16990916912 diff --git a/model-00064-of-000096.safetensors b/model-00064-of-000096.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..82ceb4851db3848b168002d09542a2041cc902f3 --- /dev/null +++ b/model-00064-of-000096.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:325c72d6ca5aaefd020da647ac317621903c8576a97d1de664b93638b7172f39 +size 16567507176 diff --git a/model-00065-of-000096.safetensors b/model-00065-of-000096.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..d3308bd76ca17cab0c4d0dc7cac20a43fe09b0e3 --- /dev/null +++ b/model-00065-of-000096.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:276d1cce1d8d494526938d878d4afb713ca379f98064dcd3c505325c954639db +size 16990916912 diff --git a/model-00066-of-000096.safetensors b/model-00066-of-000096.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..2e3440ee5d7eb87f755ceda1ce91a278114e37e2 --- /dev/null +++ b/model-00066-of-000096.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:2ebd83fea628d8cf7496a936f66a8a7f1b430bc8438644ae6ac9761ac056194e +size 16990916912 diff --git a/model-00067-of-000096.safetensors b/model-00067-of-000096.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..8f98317bb409dae4f979b7f1f09a0130a109997f --- /dev/null +++ b/model-00067-of-000096.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:f0228892f8199c29c5053e7e5388a60dacf1458b53eb6f471e1f548f02478dd1 +size 16990916912 diff --git a/model-00068-of-000096.safetensors b/model-00068-of-000096.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..1b7a6481c9acc432f76dc521f1b7598be81f0efb --- /dev/null +++ b/model-00068-of-000096.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:fa75764056d16855c42baffd474bcc01045deda7ed5f773d88ff8887ea8458b6 +size 16567507176 diff --git a/model-00069-of-000096.safetensors b/model-00069-of-000096.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..8e638ff2ef159075945ef0d2b0922b08dd561d94 --- /dev/null +++ b/model-00069-of-000096.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:9375584663bd3073cd95df18c3dc99fd4346c0c71dfdeec686a4891b7768a738 +size 16990916912 diff --git a/model-00070-of-000096.safetensors b/model-00070-of-000096.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..151dfcd787ce2488c41b0524b4b780c0104f3aff --- /dev/null +++ b/model-00070-of-000096.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:ab53464148721a209d4941dc3ab6f698574f913179e729726d00ed6232144893 +size 16990916912 diff --git a/model-00071-of-000096.safetensors b/model-00071-of-000096.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..4a45f5a377f46cd05498c915aa0b303bc41617b1 --- /dev/null +++ b/model-00071-of-000096.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:28ac0d3286ffc5c8567cc3b265500777ba2b75a0a5c6a273c87a5c9cfd8a8a5b +size 16990916912 diff --git a/model-00072-of-000096.safetensors b/model-00072-of-000096.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..12d5f2788c028e62986b5a6aaad1e8a1ea24700e --- /dev/null +++ b/model-00072-of-000096.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:0a269faaf8ea2daccef499e7c1501089ea05386ca7c6a5a9d4150fd38e37673b +size 16567507176 diff --git a/model-00073-of-000096.safetensors b/model-00073-of-000096.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..100f8d616285312d99476c9425d70198336fa3e9 --- /dev/null +++ b/model-00073-of-000096.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:a1b2e79e1bb7c9e0ade3f0825b748024a5207317589fc57f91dc2540443d716a +size 16990916912 diff --git a/model-00074-of-000096.safetensors b/model-00074-of-000096.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..bea94ea6a44cb84ec1396ec42017fe7aafb63b13 --- /dev/null +++ b/model-00074-of-000096.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:df945022b493373bdcacc3903cc406978342dff2c1b3ace5278b1909fb1d2479 +size 16990916912 diff --git a/model-00075-of-000096.safetensors b/model-00075-of-000096.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..afcd9fde1aa23140c2ad6ef7d218adf881aa139e --- /dev/null +++ b/model-00075-of-000096.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:8c6d7cb12f7c51aa88207c5b9410fccfa44e4a9022efc75d73c06054c1296c05 +size 16990916912 diff --git a/model-00076-of-000096.safetensors b/model-00076-of-000096.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..863e97918eb7d10b0ab6c17e7551368f67d84324 --- /dev/null +++ b/model-00076-of-000096.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:d5174eb5de19274b0fe62d466bfb21bac3b03d6d9676b801c6fea1ed57657998 +size 16567507176 diff --git a/model-00077-of-000096.safetensors b/model-00077-of-000096.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..4697b7ce01b08ca7496e034e920cc75b5f88eea8 --- /dev/null +++ b/model-00077-of-000096.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:8707eacfd69e8be65c7052551834e7ac26945b9ddf915939e08f5eb813c6a262 +size 16990916912 diff --git a/model-00078-of-000096.safetensors b/model-00078-of-000096.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..40435d907bef3087c5f1797e2a5944bbfcdfa259 --- /dev/null +++ b/model-00078-of-000096.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:2773d41de168c7373bf2f7f9994cbbb2712c379503578d8462aa783db7cba128 +size 16990916912 diff --git a/model-00079-of-000096.safetensors b/model-00079-of-000096.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..383f052daef55f2b2beb3a988bf907792b1ef21b --- /dev/null +++ b/model-00079-of-000096.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:da02c8a46b81f42a64330e0e31478dcd7b020633a65251b535817431bba1d96b +size 16990916912 diff --git a/model-00080-of-000096.safetensors b/model-00080-of-000096.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..1f34e8f6c8d06c7d17eb1da347ff72f7d498bdc5 --- /dev/null +++ b/model-00080-of-000096.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:96b5accdf3bc2b4651979361fa7f6191197b8573c36ab6a58effce487402b177 +size 16567507176 diff --git a/model-00081-of-000096.safetensors b/model-00081-of-000096.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..5e291172acefbb537cfff06214dac30ba354e317 --- /dev/null +++ b/model-00081-of-000096.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:01562aa616ea3dc398a7c40b4d9c3c4700885541efeaed6dbd1ba8076cae9c4b +size 16990916912 diff --git a/model-00082-of-000096.safetensors b/model-00082-of-000096.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..40861d8cb1d8636f454c7009ed2856f38bdee030 --- /dev/null +++ b/model-00082-of-000096.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:8cba090734d9a3900276cca442bd6fa60f5206ce0c70884f7b2df3df877485bd +size 16990916912 diff --git a/model-00083-of-000096.safetensors b/model-00083-of-000096.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..d2af798c15472194af193c0778be3bd523b35368 --- /dev/null +++ b/model-00083-of-000096.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:6ceebf8ce621712dd3185943224ca83a45f5bc310c272df7dee4b69e3dfc87fe +size 16990916912 diff --git a/model-00084-of-000096.safetensors b/model-00084-of-000096.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..0ad001a656f757ec8a5816e5689eb39b7282de41 --- /dev/null +++ b/model-00084-of-000096.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:c3f1318e7e1cc6893230a14c229a0c9ec0037cb1d0912a395f9222a3e8a8bc2d +size 16567507176 diff --git a/model-00085-of-000096.safetensors b/model-00085-of-000096.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..05516da484fbbe2b50eb1750b3289d16a1c3d174 --- /dev/null +++ b/model-00085-of-000096.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:633b2e3b86ec35f87fa4216b14bba49fc4058a1296433bc22b0576925118b6ad +size 16990916912 diff --git a/model-00086-of-000096.safetensors b/model-00086-of-000096.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..5377e379a7fba8d894d0bb4128c796a85f1cd648 --- /dev/null +++ b/model-00086-of-000096.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:add4056ab3ecaa7befdac2e00280d294ce470d541917ac90df8fc50497df7fb0 +size 16990916912 diff --git a/model-00087-of-000096.safetensors b/model-00087-of-000096.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..9b177cea465dc91cb1ccb741d6b67dfb4727cb1c --- /dev/null +++ b/model-00087-of-000096.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:f6d608f2c40b584b775db184013e37740d4710abad414c584408e0ad09fb486b +size 16990916912 diff --git a/model-00088-of-000096.safetensors b/model-00088-of-000096.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..3f1ba88842e255e9ac2645af0dc640eaeb7c135e --- /dev/null +++ b/model-00088-of-000096.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:87afe43b8a71d2cc56787a976f52f8c35bcfbe85df7d0c445194bc1f4d7e0c4a +size 16567507176 diff --git a/model-00089-of-000096.safetensors b/model-00089-of-000096.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..4a43e788426be09b73fe8321dfb65440da24a624 --- /dev/null +++ b/model-00089-of-000096.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:24016b28cfdf9420ae2f411145ef919cf2e1ec667830aded439bfb761c61755f +size 16990916912 diff --git a/model-00090-of-000096.safetensors b/model-00090-of-000096.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..d3b13bfdf970016b589603672e90e1c568b3bee3 --- /dev/null +++ b/model-00090-of-000096.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:1ecd85dbd77ce081f6238b894f8a7da2ff40346da7042c427a52d83562dbb5ea +size 16990916912 diff --git a/model-00091-of-000096.safetensors b/model-00091-of-000096.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..6ff05f19ebbdfde7b21185599b7261e0397f1a93 --- /dev/null +++ b/model-00091-of-000096.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:a4e666132aed052b4bf174f56853b70cbb2ada898fe568f60e0cdf3ad46a0a4a +size 16990916912 diff --git a/model-00092-of-000096.safetensors b/model-00092-of-000096.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..a218498552333a8cc603b80b15b712725836c42f --- /dev/null +++ b/model-00092-of-000096.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:359848294be55e1c8949f0a4c92098bc73276895289fd1123f10769ba685e248 +size 16567507176 diff --git a/model-00093-of-000096.safetensors b/model-00093-of-000096.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..be7e5b5f279deadea7f3d9d08803f9cc9e91ca4f --- /dev/null +++ b/model-00093-of-000096.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:d31d58d1bd3fcb8c350ef497d2acd23e0fe768ab5c0621e60f20b03d611f64e6 +size 16567507176 diff --git a/model-00094-of-000096.safetensors b/model-00094-of-000096.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..16a9eb06af0c4ab4bb83ffec0bd78d4def45ac3c --- /dev/null +++ b/model-00094-of-000096.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:ad66e1cb96b86963e63d6a0a466b6a407b13c9815cb480fe612480cc6bb3b6e1 +size 4697664072 diff --git a/model-00095-of-000096.safetensors b/model-00095-of-000096.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..edb58d17e62992bd1817de893897e9083d82547b --- /dev/null +++ b/model-00095-of-000096.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:01d41139abb8cf3b5288a97318cc4ab92676671b4eb141031cd80d7c3ced6122 +size 92289328 diff --git a/model-00096-of-000096.safetensors b/model-00096-of-000096.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..da296080d7800b10d6f98c8b5e093cb35923bd23 --- /dev/null +++ b/model-00096-of-000096.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:9d10c74fc10161bef9463a8541a634a97f521f43c99368ea7243ce0c79cdbf7c +size 802448352 diff --git a/model.safetensors.index.json b/model.safetensors.index.json new file mode 100644 index 0000000000000000000000000000000000000000..2e7014d727cf518d32bf21674e3b16687b2ba21b --- /dev/null +++ b/model.safetensors.index.json @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:a1c5210650ce71d2d3ae9ec5a101ac4afd3cf4b10091be589853437eb967febd +size 59764096 diff --git a/modeling_kimi_k3.py b/modeling_kimi_k3.py new file mode 100644 index 0000000000000000000000000000000000000000..be85612710534c2e41065b67933a9639e7817d7f --- /dev/null +++ b/modeling_kimi_k3.py @@ -0,0 +1,1317 @@ +# coding=utf-8 +# Copyright 2025-2026 The Moonshot AI Team and HuggingFace Inc. team. All rights reserved. +# +# The code is based on llava (llava/modeling_llava.py), but modified for Kimi-K3. +# +# Licensing Information: +# - Code derived from llava (llava/modeling_llava.py) is licensed under the Apache License, Version 2.0. +# - Other parts of the code are licensed under the Kimi K3 License (see the LICENSE file in this repository). +# +# Apache License, Version 2.0: +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +# NOTE: Reference implementation for model architecture; see the model card for production deployment. +import math +from collections.abc import Sequence +from copy import deepcopy +from typing import Optional + +import numpy as np +import torch +import torch.nn as nn +import torch.nn.functional as F +from transformers import activations + +try: + from transformers.activations import PytorchGELUTanh +except ImportError: + from transformers.activations import GELUTanh + activations.PytorchGELUTanh = GELUTanh + PytorchGELUTanh = GELUTanh +from transformers.activations import PytorchGELUTanh +from transformers.configuration_utils import PretrainedConfig +from transformers.modeling_utils import PreTrainedModel +from transformers.models.llava.modeling_llava import \ + LlavaCausalLMOutputWithPast +from transformers.utils import is_flash_attn_2_available + +from .configuration_kimi_k3 import KimiK3Config +from .modeling_kimi_linear import KimiLinearForCausalLM + +# Flash attention imports +if is_flash_attn_2_available(): + from flash_attn import flash_attn_varlen_func +else: + flash_attn_varlen_func = None + + +def multihead_attention( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + q_cu_seqlens: torch.Tensor | None = None, + k_cu_seqlens: torch.Tensor | None = None, + max_seqlen_q: int | None = None, + max_seqlen_k: int | None = None, + deterministic: bool = False, +): + """Multi-head attention using flash attention 2. + + Args: + q, k, v: tensor of shape (batch_size, seqlen, num_heads, head_dim), + or (tot_seqlens, num_heads, head_dim) if packing. + q_cu_seqlens (torch.Tensor): cumulative sequence lengths of q. + The first element should be 0 and the last element should be q.shape[0]. + k_cu_seqlens (torch.Tensor): cumulative sequence lengths of k. + The first element should be 0 and the last element should be k.shape[0]. + + Returns: + output: shape (batch_size, seqlen, dim) or (tot_seqlens, dim) if packing, + where dim = num_heads * head_dim + """ + attn_out = flash_attn_varlen_func( + q, + k, + v, + q_cu_seqlens, + k_cu_seqlens, + max_seqlen_q, + max_seqlen_k, + causal=False, + deterministic=deterministic, + ) + if isinstance(attn_out, tuple): + attn_out = attn_out[0] + + attn_out = attn_out.flatten(start_dim=-2) + + return attn_out + + +def eager_attention( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + q_cu_seqlens: Optional[torch.Tensor] = None, + k_cu_seqlens: Optional[torch.Tensor] = None, + **kwargs, +) -> torch.Tensor: + seq_length = q.shape[0] + attention_mask = torch.zeros([1, seq_length, seq_length], + device=q.device, + dtype=torch.bool) + for i in range(1, len(q_cu_seqlens)): + attention_mask[ + ..., + q_cu_seqlens[i - 1]:q_cu_seqlens[i], + q_cu_seqlens[i - 1]:q_cu_seqlens[i], + ] = True + q = q.transpose(0, 1) + k = k.transpose(0, 1) + v = v.transpose(0, 1) + + attn_weight = q @ k.transpose(-2, -1) / math.sqrt(q.shape[-1]) + attn_weight = attn_weight.masked_fill( + ~attention_mask, torch.finfo(attn_weight.dtype).min) + attn_weight = torch.softmax(attn_weight, dim=-1, + dtype=torch.float32).to(q.dtype) + + attn_output = attn_weight @ v + attn_output = attn_output.transpose(0, 1) + attn_output = attn_output.reshape(seq_length, -1) + return attn_output + + +VL_VISION_ATTENTION_FUNCTIONS = { + "flash_attention_2": multihead_attention, + "eager": eager_attention, +} + + +def _apply_rope_input_validation(x, freqs_cis): + assert x.ndim == freqs_cis.ndim + 1, (x.shape, freqs_cis.shape) + assert x.shape[:-2] == freqs_cis.shape[:-1], (x.shape, freqs_cis.shape) + assert x.shape[-1] == 2 * freqs_cis.shape[-1], (x.shape, freqs_cis.shape) + assert freqs_cis.dtype == torch.complex64, freqs_cis.dtype + + +def get_rope_shape_decorate(func): + _get_rope_shape_first_call_flag = set() + + def wrapper(org, interpolation_mode, shape): + key = (org.requires_grad, torch.is_grad_enabled(), interpolation_mode) + if key not in _get_rope_shape_first_call_flag: + _get_rope_shape_first_call_flag.add(key) + _ = func(org, interpolation_mode, shape=(64, 64)) + return func(org, interpolation_mode, shape) + + return wrapper + + +@get_rope_shape_decorate +@torch.compile(dynamic=True) +def get_rope_shape(org, interpolation_mode, shape): + return (F.interpolate( + org.permute((2, 0, 1)).unsqueeze(0), + size=shape, + mode=interpolation_mode, + ).squeeze(0).permute((1, 2, 0)).flatten(end_dim=1)) + + +def apply_rope(xq: torch.Tensor, xk: torch.Tensor, + freqs_cis: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """ + Args: (The leading dimensions of all inputs should be the same) + xq: query, tensor of shape (..., num_heads, head_dim) + xk: key, tensor of shape (..., num_heads, head_dim) + freqs_cis: tensor of shape (..., head_dim/2), dtype=torch.complex64. It contains the precomputed cis(freqs) for each position in the 2D grid. + Returns: + xq_out, xk_out: tensors of shape (..., num_heads, head_dim) + """ + _apply_rope_input_validation(xq, freqs_cis) + _apply_rope_input_validation(xk, freqs_cis) + + freqs_cis = freqs_cis.unsqueeze(-2) # ..., 1, head_dim/2 + # ..., num_heads, head_dim/2 + xq_ = torch.view_as_complex(xq.float().view(*xq.shape[:-1], -1, 2)) + xk_ = torch.view_as_complex(xk.float().view(*xq.shape[:-1], -1, 2)) + xq_out = torch.view_as_real(xq_ * freqs_cis).flatten( + -2) # ..., num_heads, head_dim + xk_out = torch.view_as_real(xk_ * freqs_cis).flatten( + -2) # ..., num_heads, head_dim + return xq_out.type_as(xq), xk_out.type_as(xk) + + +def get_1d_sincos_pos_embed_from_grid(embed_dim, pos): + """ + From: + https://github.com/OpenGVLab/InternVideo/blob/421f6d2361fc8f61a3394244571f2601a4e99e29/InternVideo2/multi_modality/models/backbones/internvideo2/pos_embed.py#L86 + embed_dim: output dimension for each position + pos: a list of positions to be encoded: size (M,) + out: (M, D) + """ + assert embed_dim % 2 == 0 + omega = np.arange(embed_dim // 2, dtype=np.float32) + omega /= embed_dim / 2.0 + omega = 1.0 / 10000**omega # (D/2,) + + pos = pos.reshape(-1) # (M,) + out = np.einsum('m,d->md', pos, omega) # (M, D/2), outer product + + emb_sin = np.sin(out) # (M, D/2) + emb_cos = np.cos(out) # (M, D/2) + + emb = np.concatenate([emb_sin, emb_cos], axis=1) # (M, D) + return emb + + +def get_1d_sincos_pos_embed(embed_dim, t_size, cls_token=False): + """ + t_size: int of the temporal size + return: + pos_embed: [t_size, embed_dim] or [1+t_size, embed_dim] (w/ or w/o cls_token) + """ + grid_t = np.arange(t_size, dtype=np.float32) + pos_embed = get_1d_sincos_pos_embed_from_grid(embed_dim, grid_t) + if cls_token: + pos_embed = np.concatenate([np.zeros([1, embed_dim]), pos_embed], + axis=0) + return pos_embed + + +class Learnable2DInterpPosEmbDivided_fixed(nn.Module): + + def __init__(self, + height: int, + width: int, + num_frames: int, + dim: int, + interpolation_mode: str = 'bicubic') -> None: + super().__init__() + self.height = height + self.width = width + self.num_frames = num_frames + self.dim = dim + self.interpolation_mode = interpolation_mode + self.weight = nn.Parameter(torch.empty(height, width, dim)) + self.register_buffer('time_weight', + torch.from_numpy( + get_1d_sincos_pos_embed( + self.dim, + self.num_frames)).float().unsqueeze(1), + persistent=False) + + self.reset_parameters() + + def reset_parameters(self): + nn.init.normal_(self.weight) + + def forward(self, x: torch.Tensor, + grid_thws: torch.Tensor) -> torch.Tensor: + pos_embs = [] + for t, h, w in grid_thws.tolist(): + assert t <= self.num_frames, f't:{t} > self.num_frames:{self.num_frames}' + if (h, w) == self.weight.shape[:-1]: + pos_emb_2d = self.weight.flatten(end_dim=1) + else: + pos_emb_2d = get_rope_shape( + self.weight, + interpolation_mode=self.interpolation_mode, + shape=(h, w), + ) + + if t == 1: + pos_emb_3d = pos_emb_2d + else: + pos_emb_3d = pos_emb_2d.unsqueeze(0).repeat( + t, 1, 1) + self.time_weight[0:t] + + pos_embs.append(pos_emb_3d.reshape(-1, pos_emb_3d.shape[-1])) + + out = x + torch.cat(pos_embs) + return out + + +class MoonVision3dPatchEmbed(nn.Module): + + def __init__(self, + out_dim: int, + in_dim: int = 3, + patch_size: int | tuple[int, int] = (14, 14), + pos_emb_height: int = 14, + pos_emb_width: int = 14, + pos_emb_time: int = 4, + pos_emb_type: str = 'divided_fixed', + patch_embed_proj_bias: bool = True, + pos_emb_interpolation_mode: str = 'bicubic'): + super().__init__() + assert isinstance( + patch_size, + int | Sequence), f'Invalid patch_size type: {type(patch_size)}' + if isinstance(patch_size, int): + patch_size = (patch_size, patch_size) + assert (len(patch_size) == 2 + ), f'Expected patch_size to be a tuple of 2, got {patch_size}' + self.patch_size = patch_size + + self.proj = nn.Conv2d(in_dim, + out_dim, + kernel_size=patch_size, + stride=patch_size, + bias=patch_embed_proj_bias) + + if pos_emb_type == 'divided_fixed': + self.pos_emb = Learnable2DInterpPosEmbDivided_fixed( + height=pos_emb_height, + width=pos_emb_width, + num_frames=pos_emb_time, + dim=out_dim, + interpolation_mode=pos_emb_interpolation_mode) + else: + raise NotImplementedError( + f'Not support pos_emb_type: {pos_emb_type}') + + def forward(self, x: torch.Tensor, + grid_thws: torch.Tensor) -> torch.Tensor: + """ + Args: + x (L, Channels): input tensor + grid_hws (N, 3): temporal, height and width + + Returns: + (L, Cout) tensor + """ + x = self.proj(x).view(x.size(0), -1) + # apply positional embedding + x = self.pos_emb(x, grid_thws) + return x + + +class Rope2DPosEmbRepeated(nn.Module): + """2D rotary position embedding with multi-resolution support. + + This class is intended to be used in the following way: + 1. Before training, create an instance of Rope2DPosEmb. This instance will hold the precomputed cis. + 2. Before each forward pass, call `get_freqs_cis_by_*` to get the `freqs_cis` tensor for this iteration. + 3. During the forward pass, pass the `freqs_cis` tensor to each attention layer, and call `apply` just before each attention operation. + The rope is shared across all attention layers and all heads. + + Refs: + - RoFormer: https://arxiv.org/abs/2104.09864 + - VisionLLaMA: https://arxiv.org/abs/2403.00522 + - https://github.com/Meituan-AutoML/VisionLLaMA/blob/main/dit/models.py + + Args: + dim (int): usually the multi-head attention dimension, should be divisible by 4 (TODO: relax this constraint if needed) + max_height (int): the maximum height of the 2D grid + max_width (int): the maximum width of the 2D grid + theta_base (float): the base of the theta + device (str): the device to store the precomputed cis + """ + + def __init__(self, + dim: int, + max_height: int, + max_width: int, + theta_base=10000): + super().__init__() + self.dim = dim + assert self.dim % 4 == 0, 'dim must be divisible by 4' + self.max_height = max_height + self.max_width = max_width + self.theta_base = theta_base + + def extra_repr(self): + return f'dim={self.dim}, max_height={self.max_height}, max_width={self.max_width}, theta_base={self.theta_base}' + + def _precompute_freqs_cis(self, device: torch.device) -> torch.Tensor: + """Calculate the cis(freqs) for each position in the 2D grid. + + Return: complex tensor of shape (max_height, max_width, dim//2) and value: + height axis: ret[h, w, 2*i] = cis(h * theta_base**(-4*i/dim)) + weight axis: ret[h, w, 2*i+1] = cis(w * theta_base**(-4*i/dim)) with (i in [0, dim//4)) + note: `cis` is a mathematical notation defined by cis x = cos x + i sin x, + """ + N = self.max_height * self.max_width + flat_pos = torch.arange(0, N).float().to(device) + x_pos = flat_pos % self.max_width + y_pos = flat_pos // self.max_width + dim_range = (torch.arange(0, self.dim, + 4)[:(self.dim // 4)].float().to(device) + ) # C/4 + freqs = 1.0 / (self.theta_base**(dim_range / self.dim)) + x_freqs = torch.outer(x_pos, freqs).float() # N, C/4 + y_freqs = torch.outer(y_pos, freqs).float() # N, C/4 + x_cis = torch.polar(torch.ones_like(x_freqs), x_freqs) # N, C/4 + y_cis = torch.polar(torch.ones_like(y_freqs), y_freqs) # N, C/4 + # N, C/4, 2 + freqs_cis = torch.cat( + [x_cis.unsqueeze(dim=-1), + y_cis.unsqueeze(dim=-1)], dim=-1) + # max_height, max_width, C/2 + freqs_cis = freqs_cis.reshape(self.max_height, self.max_width, -1) + return freqs_cis + + def get_freqs_cis(self, grid_thws: torch.Tensor, + device: torch.device) -> torch.Tensor: + """ + Args: + grid_thws (torch.Tensor): grid time, height and width + + Returns: + freqs_cis: tensor of shape (sum(t * height * width), dim//2) + """ + if not hasattr(self, 'freqs_cis'): + self.register_buffer('freqs_cis', + self._precompute_freqs_cis(device), + persistent=False) + + shapes = grid_thws.tolist() + assert all(1 <= h <= self.max_height and 1 <= w <= self.max_width + for t, h, w in shapes), ( + shapes, + self.max_height, + self.max_width, + ) + freqs_cis = torch.cat( + [ + self.freqs_cis[:h, :w].reshape(-1, self.dim // 2).repeat(t, 1) + for t, h, w in shapes + ], + dim=0, + ) + return freqs_cis + + +class MLP2(nn.Module): + """ + Args: + dims: [in_dim, hidden_dim, out_dim] + bias: whether to use bias in linear layer. + """ + + def __init__(self, dims: list[int], activation, bias=True): + super().__init__() + assert len(dims) == 3 + self.fc0 = nn.Linear(dims[0], dims[1], bias=bias) + self.fc1 = nn.Linear(dims[1], dims[2], bias=bias) + self.activation = activation + for m in [self.fc0, self.fc1]: + nn.init.trunc_normal_(m.weight, std=math.sqrt(2 / m.in_features)) + if m.bias is not None: + nn.init.zeros_(m.bias) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x = self.fc0(x) + x = self.activation(x) + return self.fc1(x) + + +class MoonViTEncoderLayer(nn.Module): + + def __init__( + self, + num_heads: int, + hidden_dim: int, + mlp_dim: int, + qkv_hidden_size: int | None = None, + norm_type: str = 'layernorm', + mlp_type: str = 'mlp2', + *, + attn_implementation: str = 'flash_attention_2', + activation=F.gelu, + attn_bias: bool = False, + linear_bias: bool = True, + use_deterministic_attn: bool = False, + ): + super().__init__() + self.num_heads = num_heads + self.hidden_dim = hidden_dim + self.qkv_hidden_size = hidden_dim if qkv_hidden_size is None else qkv_hidden_size + self.hidden_size_per_attention_head = self.qkv_hidden_size // self.num_heads + self.attn_implementation = attn_implementation + self.use_deterministic_attn = use_deterministic_attn + + if norm_type == "layernorm": + self.norm0 = nn.LayerNorm(hidden_dim) + self.norm1 = nn.LayerNorm(hidden_dim) + elif norm_type == "rmsnorm": + self.norm0 = nn.RMSNorm(hidden_dim) + self.norm1 = nn.RMSNorm(hidden_dim) + else: + raise NotImplementedError(f"Not support norm_type: {norm_type}") + + if mlp_type == "mlp2": + self.mlp = MLP2([hidden_dim, mlp_dim, hidden_dim], + activation, + bias=linear_bias) + else: + raise NotImplementedError(f"Not support mlp_type: {mlp_type}") + + self.wqkv = nn.Linear(hidden_dim, + self.qkv_hidden_size * 3, + bias=attn_bias) + self.wo = nn.Linear(self.qkv_hidden_size, hidden_dim, bias=attn_bias) + + def attention_qkvpacked( + self, + x: torch.Tensor, + cu_seqlens: torch.Tensor, + max_seqlen: torch.Tensor, + rope_freqs_cis: torch.Tensor | None = None, + ): + """ + Args: + x (torch.Tensor): (batch_size, seqlen, hidden_dim) + cu_seqlens (torch.Tensor): + """ + xqkv = self.wqkv(x) + + qkv_shape = xqkv.size()[:-1] + ( + 3, + self.num_heads, + self.hidden_size_per_attention_head, + ) + # xqkv: (batch_size, seqlen, 3, nheads, headdim) + xqkv = xqkv.view(*qkv_shape) + xq, xk, xv = torch.unbind(xqkv, dim=-3) + + xq, xk = apply_rope(xq, xk, rope_freqs_cis) + + attn_func = VL_VISION_ATTENTION_FUNCTIONS[self.attn_implementation] + attn_out = attn_func(xq, + xk, + xv, + q_cu_seqlens=cu_seqlens, + k_cu_seqlens=cu_seqlens, + max_seqlen_k=max_seqlen, + max_seqlen_q=max_seqlen, + deterministic=self.use_deterministic_attn) + + attn_out = self.wo(attn_out) + return attn_out + + def forward( + self, + hidden_states: torch.Tensor, + cu_seqlens: torch.Tensor, + max_seqlen: int, + rope_freqs_cis: torch.Tensor | None = None, + ): + residual = hidden_states + hidden_states = self.norm0(hidden_states) + + hidden_states = self.attention_qkvpacked(hidden_states, cu_seqlens, + max_seqlen, rope_freqs_cis) + hidden_states = residual + hidden_states + + residual = hidden_states + hidden_states = self.norm1(hidden_states) + hidden_states = self.mlp(hidden_states) + hidden_states = residual + hidden_states + + return hidden_states + + +class MoonViT3dEncoder(nn.Module): + + def __init__(self, + hidden_dim: int, + num_layers: int, + block_cfg: dict, + use_deterministic_attn: bool = False) -> None: + super().__init__() + self.use_deterministic_attn = use_deterministic_attn + + qkv_hidden_size = block_cfg['hidden_dim'] if block_cfg.get( + 'qkv_hidden_size') is None else block_cfg['qkv_hidden_size'] + self.rope_2d = Rope2DPosEmbRepeated( + qkv_hidden_size // block_cfg['num_heads'], 512, 512) + self.blocks = nn.ModuleList([ + MoonViTEncoderLayer( + **block_cfg, + use_deterministic_attn=self.use_deterministic_attn) + for _ in range(num_layers) + ]) + norm_type = block_cfg.get('norm_type', 'layernorm') + if norm_type == "layernorm": + self.final_layernorm = nn.LayerNorm(hidden_dim) + elif norm_type == "rmsnorm": + self.final_layernorm = nn.RMSNorm(hidden_dim) + else: + raise NotImplementedError(f"Not support norm_type: {norm_type}") + + def forward( + self, + hidden_states: torch.Tensor, + grid_thws: torch.Tensor, + ) -> torch.Tensor: + rope_freqs_cis = self.rope_2d.get_freqs_cis( + grid_thws=grid_thws, device=hidden_states.device) + + lengths = torch.cat(( + torch.zeros(1, dtype=grid_thws.dtype, device=grid_thws.device), + grid_thws[:, 0] * grid_thws[:, 1] * grid_thws[:, 2], + )) + + max_seqlen = lengths.max() + cu_seqlens = lengths.to(hidden_states.device).cumsum(dim=0, + dtype=torch.int32) + for block in self.blocks: + hidden_states = block(hidden_states, + cu_seqlens, + max_seqlen, + rope_freqs_cis=rope_freqs_cis) + + hidden_states = self.final_layernorm(hidden_states) + return hidden_states + + +def tpool_patch_merger( + x: torch.Tensor, + grid_thws: torch.Tensor, + merge_kernel_size: tuple[int, int] = (2, 2), +) -> list[torch.Tensor]: + d_model = x.size(-1) + + outputs = [] + pre_sum = 0 + for t, h, w in grid_thws.tolist(): + # Get the current sequence + seq = x[pre_sum:pre_sum + t * h * w] + # Reshape along self.merge_kernel_size and concat to the last dimension + kernel_height, kernel_width = merge_kernel_size + new_height, new_width = h // kernel_height, w // kernel_width + reshaped_seq = seq.view(t, new_height, kernel_height, new_width, + kernel_width, d_model) + reshaped_seq = reshaped_seq.permute(0, 1, + 3, 2, 4, 5).contiguous().mean( + dim=0) # temporal pooling + padded_seq = reshaped_seq.view(new_height * new_width, + kernel_height * kernel_width, -1) + outputs.append(padded_seq) + pre_sum += t * h * w + + return outputs + + +class MoonViT3dPretrainedModel(PreTrainedModel): + config_class = None + model_type = 'moonvit3d' + _no_split_modules = ['MoonViTEncoderLayer'] + _supports_flash_attn_2 = True + _supports_sdpa = True + + def __init__(self, config, *inputs, **kwargs): + super().__init__(config, *inputs, **kwargs) + config = deepcopy(config) + self.merge_kernel_size = config.merge_kernel_size + self.patch_size = config.patch_size + self.merge_type = config.merge_type + + self.patch_embed = MoonVision3dPatchEmbed( + out_dim=config.hidden_size, + patch_size=config.patch_size, + pos_emb_height=config.init_pos_emb_height, + pos_emb_width=config.init_pos_emb_width, + pos_emb_time=config.init_pos_emb_time, + pos_emb_type=config.pos_emb_type, + patch_embed_proj_bias=getattr(config, 'patch_embed_proj_bias', + True), + pos_emb_interpolation_mode=getattr( + config, 'pos_emb_interpolation_mode', 'bicubic'), + ) + + self.encoder = MoonViT3dEncoder( + hidden_dim=config.hidden_size, + num_layers=config.num_hidden_layers, + block_cfg={ + 'num_heads': config.num_attention_heads, + 'hidden_dim': config.hidden_size, + 'qkv_hidden_size': getattr(config, 'qkv_hidden_size', None), + 'mlp_dim': config.intermediate_size, + 'norm_type': getattr(config, 'norm_type', 'layernorm'), + 'mlp_type': getattr(config, 'mlp_type', 'mlp2'), + 'activation': PytorchGELUTanh(), + 'attn_bias': getattr(config, 'attn_bias', True), + 'linear_bias': getattr(config, 'linear_bias', True), + 'attn_implementation': config._attn_implementation, + }, + use_deterministic_attn=getattr(self, 'use_deterministic_attn', + False)) + + def forward(self, pixel_values: torch.Tensor, + grid_thws: torch.Tensor) -> torch.Tensor: + """ + Args: + pixel_values (torch.Tensor): The input pixel values. + grid_thws (torch.Tensor): Temporal, height and width. + + Returns: + torch.Tensor: The output tokens. + """ + # grid_thws = grid_thws.to('cpu') + assert grid_thws.ndim == 2, f'grid_thws should be 2D, got {grid_thws.ndim}' + assert grid_thws.size(1) == 3, f'No support for thw: {grid_thws}' + hidden_states = self.patch_embed(pixel_values, grid_thws) + hidden_states = self.encoder(hidden_states, grid_thws) + if self.merge_type == 'sd2_tpool': # spatial downsampling 2x with temporal pooling all + hidden_states = tpool_patch_merger( + hidden_states, + grid_thws, + merge_kernel_size=self.merge_kernel_size) + else: + raise NotImplementedError(f'Not support {self.merge_type}') + + return hidden_states + + +# ============================================================================ +# MM Projector Helper Classes (from mm_projector/modeling_mm_projectors.py) +# ============================================================================ + + +class IdentityMap(nn.Module): + + def __init__(self): + super().__init__() + + def forward(self, x, *args, **kwargs): + return x + + +class MLP(nn.Module): + + def __init__(self, config): + super().__init__() + # TODO, use faster LayerNorm + self.pre_norm = nn.LayerNorm(config.mm_hidden_size) + self.proj = nn.Sequential( + nn.Linear(config.mm_hidden_size, config.hidden_size), nn.GELU(), + nn.Linear(config.hidden_size, config.hidden_size)) + + def forward(self, x, *args, **kwargs): + assert isinstance(x, + list | tuple), f'x is not a list or tuple: {type(x)}' + lengths = [item.shape[0] for item in x] + x = torch.cat(x, dim=0) + x = self.pre_norm(x) + x = self.proj(x) + x = torch.split(x, lengths, dim=0) + + return x + + +class PatchMergerMLP(nn.Module): + + def __init__(self, config): + super().__init__() + eps = config.projector_ln_eps + self.hidden_size = config.mm_hidden_size * ( + config.merge_kernel_size[0] * config.merge_kernel_size[1]) + self.pre_norm = nn.LayerNorm(config.mm_hidden_size, eps=eps) + self.proj = nn.Sequential( + nn.Linear(self.hidden_size, self.hidden_size), + nn.GELU(), + nn.Linear(self.hidden_size, config.hidden_size), + ) + + def forward(self, x, *args, **kwargs): + if isinstance(x, list) or isinstance(x, tuple): + x = [ + self.proj(self.pre_norm(item).view(item.shape[0], -1)) + for item in x + ] + else: + # B, N, N_k, C = x.shape + B = x.shape[0] + x = self.proj(self.pre_norm(x).view(B, -1, self.hidden_size)) + return x + + +class PatchMergerMLPV2(nn.Module): + + def __init__(self, config): + super().__init__() + eps = config.projector_ln_eps + self.hidden_size = config.mm_hidden_size * ( + config.merge_kernel_size[0] * config.merge_kernel_size[1]) + self.proj = nn.Sequential( + nn.Linear(self.hidden_size, self.hidden_size, bias=False), + nn.GELU(), + nn.Linear(self.hidden_size, config.hidden_size, bias=False), + ) + self.post_norm = nn.RMSNorm(config.hidden_size, eps=eps) + for m in self.proj.modules(): + if isinstance(m, nn.Linear): + nn.init.trunc_normal_(m.weight, + std=math.sqrt(2 / m.in_features)) + if m.bias is not None: + nn.init.zeros_(m.bias) + + def forward(self, x, *args, **kwargs): + if isinstance(x, list) or isinstance(x, tuple): + lengths = [item.shape[0] for item in x] + x = torch.concat([item.view(item.shape[0], -1) for item in x], + dim=0) + x = self.post_norm(self.proj(x)) + x = torch.split(x, lengths, dim=0) + else: + # B, N, N_k, C = x.shape + B = x.shape[0] + x = self.proj(x.view(B, -1, self.hidden_size)) + x = self.post_norm(x) + return x + + +class KimiK3PreTrainedModel(PreTrainedModel): + config_class = KimiK3Config + base_model_prefix = "model" + _no_split_modules = [ + "MoonViT3dPretrainedModel", + "MoonViTEncoderLayer", + "KimiDecoderLayer", + "PatchMergerMLP", + "PatchMergerMLPV2", + ] + _skip_keys_device_placement = "past_key_values" + _supports_flash_attn_2 = True + _supports_sdpa = False + + def _init_weights(self, module): + # important: this ported version of Llava isn't meant for training from scratch - only + # inference and fine-tuning - so the proper init weights code has been removed - the original codebase + # https://github.com/haotian-liu/LLaVA/tree/main/llava should serve for that purpose + std = (self.config.initializer_range if hasattr( + self.config, "initializer_range") else + self.config.text_config.initializer_range) + + if hasattr(module, "class_embedding"): + module.class_embedding.data.normal_(mean=0.0, std=std) + + if isinstance(module, (nn.Linear, nn.Conv2d)): + module.weight.data.normal_(mean=0.0, std=std) + if module.bias is not None: + module.bias.data.zero_() + elif isinstance(module, nn.Embedding): + module.weight.data.normal_(mean=0.0, std=std) + if module.padding_idx is not None: + module.weight.data[module.padding_idx].zero_() + + +class VisionTowerConfig(PretrainedConfig): + model_type = 'moonvit3d' + + def __init__(self, config: KimiK3Config, **kwargs): + super().__init__(**kwargs) + self.patch_size = config.patch_size + self.init_pos_emb_height = config.init_pos_emb_height + self.init_pos_emb_width = config.init_pos_emb_width + self.init_pos_emb_time = config.init_pos_emb_time + self.pos_emb_type = config.pos_emb_type + self.num_attention_heads = config.vt_num_attention_heads + self.num_hidden_layers = config.vt_num_hidden_layers + self.hidden_size = config.vt_hidden_size + self.intermediate_size = config.vt_intermediate_size + self.merge_kernel_size = config.merge_kernel_size + self.merge_type = config.merge_type + self._attn_implementation = config._attn_implementation + self.qkv_hidden_size = getattr(config, 'qkv_hidden_size', None) + self.norm_type = getattr(config, 'norm_type', 'layernorm') + self.attn_bias = getattr(config, 'attn_bias', True) + self.patch_embed_proj_bias = getattr(config, 'patch_embed_proj_bias', + True) + self.mlp_type = getattr(config, 'mlp_type', 'mlp2') + self.linear_bias = getattr(config, 'linear_bias', True) + self.pos_emb_interpolation_mode = getattr( + config, 'pos_emb_interpolation_mode', 'bilinear') + + +class ProjectorConfig: + + def __init__(self, config: KimiK3Config): + self.mm_projector_type = config.mm_projector_type + self.mm_hidden_size = config.mm_hidden_size + self.hidden_size = config.text_hidden_size + self.merge_kernel_size = config.merge_kernel_size + self.projector_hidden_act = config.projector_hidden_act + self.projector_ln_eps = config.projector_ln_eps + + +# ref https://github.com/huggingface/transformers/blob/78b2929c0554b79e0489b451ce4ece14d265ead2/src/transformers/models/llava/modeling_llava.py#L240 +class KimiK3ForConditionalGeneration(KimiK3PreTrainedModel): + + @classmethod + def _supports_default_dynamic_cache(cls) -> bool: + return False + + def __init__(self, config: KimiK3Config): + super().__init__(config) + + vt_config = VisionTowerConfig(config.vision_config) + self.vision_tower = MoonViT3dPretrainedModel(vt_config) + + proj_config = ProjectorConfig(config.vision_config) + if proj_config.mm_projector_type == 'identity': + self.mm_projector = IdentityMap() + elif proj_config.mm_projector_type == 'mlp': + self.mm_projector = MLP(proj_config) + elif proj_config.mm_projector_type == 'patchmerger': + self.mm_projector = PatchMergerMLP(proj_config) + elif proj_config.mm_projector_type == 'patchmergerv2': + self.mm_projector = PatchMergerMLPV2(proj_config) + else: + raise ValueError( + f"Unsupported mm_projector_type: {proj_config.mm_projector_type}" + ) + + self.language_model = KimiLinearForCausalLM(config.text_config) + self.post_init() + + if hasattr(self.language_model, 'dtype'): + target_dtype = self.language_model.dtype + self.vision_tower = self.vision_tower.to(dtype=target_dtype) + self.mm_projector = self.mm_projector.to(dtype=target_dtype) + + def get_input_embeddings(self): + return self.language_model.get_input_embeddings() + + def set_input_embeddings(self, value): + self.language_model.set_input_embeddings(value) + + def get_output_embeddings(self): + return self.language_model.get_output_embeddings() + + def set_output_embeddings(self, new_embeddings): + self.language_model.set_output_embeddings(new_embeddings) + + def set_decoder(self, decoder): + self.language_model.set_decoder(decoder) + + def get_decoder(self): + return self.language_model.get_decoder() + + def tie_weights(self): + return self.language_model.tie_weights() + + def resize_token_embeddings(self, + new_num_tokens: int | None = None, + pad_to_multiple_of=None) -> nn.Embedding: + model_embeds = self.language_model.resize_token_embeddings( + new_num_tokens, pad_to_multiple_of) + # update vocab size + self.config.text_config.vocab_size = model_embeds.num_embeddings + self.vocab_size = model_embeds.num_embeddings + return model_embeds + + def _merge_input_ids_with_image_features( + self, + image_features: list[torch.Tensor], + inputs_embeds: torch.Tensor, + input_ids: torch.Tensor, + attention_mask: torch.Tensor, + labels: torch.Tensor | None = None, + ): + """ + Args: + image_features (:obj:`torch.Tensor` of shape :obj:`(num_image_tokens, embed_dim)`): + The image features to merge with the input embeddings. + inputs_embeds (:obj:`torch.Tensor` of shape :obj:`(batch_size, sequence_length, embed_dim)`): + The input embeddings. + input_ids (:obj:`torch.Tensor` of shape :obj:`(batch_size, sequence_length)`): + The input ids. + attention_mask (:obj:`torch.Tensor` of shape :obj:`(batch_size, sequence_length)`): + The attention mask. + labels (:obj:`torch.Tensor` of shape :obj:`(batch_size, sequence_length)`, *optional*): + The labels. + """ + _, embed_dim = image_features[0].shape + feature_lengths = [x.shape[0] for x in image_features] + image_features = torch.cat(image_features, dim=0) + + image_token_index: int = self.config.media_placeholder_token_id + pad_token_id: int = self.config.pad_token_id + ignore_index: int = self.config.ignore_index + + batch_size, sequence_length = input_ids.shape + left_padding = not torch.sum( + input_ids[:, -1] == torch.tensor(pad_token_id)) + + # 1. Create a mask to know where special image tokens are + _token_occupation_table = torch.ones_like(input_ids.flatten()) + _token_occupation_table[input_ids.flatten() == + image_token_index] = torch.tensor( + feature_lengths, + dtype=torch.long, + device=input_ids.device) + _token_occupation_table = _token_occupation_table.reshape( + input_ids.shape) + + max_embed_dim = _token_occupation_table.sum(-1).max().item() + assert ( + max_embed_dim >= sequence_length + ), f"The maximum embedding dimension ({max_embed_dim}) is less than the sequence length ({sequence_length})" + batch_indices, non_image_indices = torch.where( + input_ids != image_token_index) + + # 2. Compute the positions where text should be written + # Calculate new positions for text tokens in merged image-text sequence. + new_token_positions = torch.cumsum(_token_occupation_table, -1) - 1 + nb_image_pad = max_embed_dim - 1 - new_token_positions[:, -1] + if left_padding: + new_token_positions += nb_image_pad[:, + None] # offset for left padding + text_to_overwrite = new_token_positions[batch_indices, + non_image_indices] + + # 3. Create the full embedding, already padded to the maximum position + final_embedding = torch.zeros( + batch_size, + max_embed_dim, + embed_dim, + dtype=inputs_embeds.dtype, + device=inputs_embeds.device, + ) + final_attention_mask = torch.zeros(batch_size, + max_embed_dim, + dtype=attention_mask.dtype, + device=inputs_embeds.device) + if labels is not None: + final_labels = torch.full( + (batch_size, max_embed_dim), + ignore_index, + dtype=input_ids.dtype, + device=input_ids.device, + ) + # In case the Vision model or the Language model has been offloaded to CPU, we need to manually + # set the corresponding tensors into their correct target device. + target_device = inputs_embeds.device + batch_indices, non_image_indices, text_to_overwrite = ( + batch_indices.to(target_device), + non_image_indices.to(target_device), + text_to_overwrite.to(target_device), + ) + attention_mask = attention_mask.to(target_device) + + # 4. Fill the embeddings based on the mask. + final_embedding[batch_indices, + text_to_overwrite] = inputs_embeds[batch_indices, + non_image_indices] + final_attention_mask[batch_indices, + text_to_overwrite] = attention_mask[ + batch_indices, non_image_indices] + if labels is not None: + final_labels[batch_indices, + text_to_overwrite] = labels[batch_indices, + non_image_indices] + + # 5. Fill the embeddings corresponding to the images. Anything that is not `text_positions` needs filling (#29835) + image_to_overwrite = torch.full((batch_size, max_embed_dim), + True, + dtype=torch.bool, + device=inputs_embeds.device) + image_to_overwrite[batch_indices, text_to_overwrite] = False + image_to_overwrite &= image_to_overwrite.cumsum( + -1) - 1 >= nb_image_pad[:, None].to(target_device) + + if image_to_overwrite.sum() != image_features.shape[:-1].numel(): + raise ValueError( + f"The input provided to the model are wrong. The number of image tokens is {image_to_overwrite.sum()} while" + f" the number of image features given to the model is {image_features.shape[:-1].numel()}. " + "This prevents correct indexing and breaks batch generation.") + + final_embedding[image_to_overwrite] = ( + image_features.contiguous().reshape(-1, + embed_dim).to(target_device)) + final_attention_mask |= image_to_overwrite + position_ids = (final_attention_mask.cumsum(-1) - 1).masked_fill_( + (final_attention_mask == 0), 1) + + # 6. Mask out the embedding at padding positions, as we later use the past_key_value value to determine the non-attended tokens. + batch_indices, pad_indices = torch.where(input_ids == pad_token_id) + indices_to_mask = new_token_positions[batch_indices, pad_indices] + + final_embedding[batch_indices, indices_to_mask] = 0 + + if labels is None: + final_labels = None + + return final_embedding, final_attention_mask, final_labels, position_ids + + def _extract_image_features(self, pixel_values: torch.Tensor, + grid_thws: torch.Tensor) -> list[torch.Tensor]: + """ + Args: + pixel_values (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, num_channels, height, width)`): + The pixel values of the images processed by image processor. + grid_thws (:obj:`torch.Tensor` of shape :obj:`(batch_size, 3)`): + The grid, height, width of the images. + + Returns: + selected_image_feature (:obj:`torch.FloatTensor` of shape :obj:`(num_image_tokens, embed_dim)`): + The selected image features to use as input to the projector head. + + """ + + target_dtype = self.vision_tower.patch_embed.proj.weight.dtype + pixel_values = pixel_values.to(target_dtype) + + image_features = self.vision_tower(pixel_values, grid_thws) + return image_features + + def forward( + self, + input_ids: torch.LongTensor | None = None, + pixel_values: torch.FloatTensor | list[torch.FloatTensor] + | None = None, + grid_thws: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + position_ids: torch.LongTensor | None = None, + past_key_values: list[torch.FloatTensor] | None = None, + inputs_embeds: torch.FloatTensor | None = None, + labels: torch.LongTensor | None = None, + use_cache: bool | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + return_dict: bool | None = None, + ) -> tuple | LlavaCausalLMOutputWithPast: + r""" + Args: + labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): + Labels for computing the masked language modeling loss. Indices should either be in `[0, ..., + config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored + (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`. + + ```""" + assert self.vision_tower is not None, "vision_tower is not loaded" + output_attentions = (output_attentions if output_attentions is not None + else self.config.output_attentions) + output_hidden_states = (output_hidden_states + if output_hidden_states is not None else + self.config.output_hidden_states) + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + if inputs_embeds is None: + # 1. Extra the input embeddings + inputs_embeds = self.get_input_embeddings()(input_ids) + + # 2. Merge text and images + if pixel_values is not None and len( + pixel_values) > 0 and input_ids.shape[1] != 1: + image_features = self._extract_image_features( + pixel_values, grid_thws) + if self.mm_projector: + image_features = self.mm_projector(image_features) + + inputs_embeds = inputs_embeds.to( + image_features[0].dtype) # num_tokens, embed_dim + inputs_embeds, attention_mask, labels, position_ids = ( + self._merge_input_ids_with_image_features( + image_features, + inputs_embeds, + input_ids, + attention_mask, + labels, + )) + + # In case input_ids.shape[1] == 1 & pixel_values==None & past_key_values != None, we are in the case of + # generation with cache + elif (past_key_values is not None and pixel_values is not None + and input_ids.shape[1] == 1): + # Retrieve the first layer to inspect the logits and mask out the hidden states + # that are set to 0 + first_layer_past_key_value = past_key_values[0][0][:, :, :, 0] + + # Sum all dimensions of head_dim (-2) to avoid random errors such as: https://github.com/huggingface/transformers/pull/28032#issuecomment-1863691941 + batch_index, non_attended_tokens = torch.where( + first_layer_past_key_value.float().sum(-2) == 0) + + # Get the target length + target_length = input_ids.shape[1] + past_length = first_layer_past_key_value.shape[-1] + + extended_attention_mask = torch.ones( + (attention_mask.shape[0], past_length), + dtype=attention_mask.dtype, + device=attention_mask.device, + ) + + # Filter out only the tokens that can be un-attended, this can happen + # if one uses Llava + Fused modules where the cache on the + # first iteration is already big enough, or if one passes custom cache + valid_indices = non_attended_tokens < extended_attention_mask.size( + -1) + new_batch_index = batch_index[valid_indices] + new_non_attended_tokens = non_attended_tokens[valid_indices] + + # Zero-out the places where we don't need to attend + extended_attention_mask[new_batch_index, + new_non_attended_tokens] = 0 + + attention_mask = torch.cat( + (extended_attention_mask, attention_mask[:, + -target_length:]), + dim=1) + position_ids = torch.sum(attention_mask, + dim=1).unsqueeze(-1) - 1 + + outputs = self.language_model( + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + + logits = outputs[0] + + loss = None + if labels is not None: + # Shift so that tokens < n predict n + if attention_mask is not None: + shift_attention_mask = attention_mask[..., 1:] + shift_logits = logits[..., :-1, :][shift_attention_mask.to( + logits.device) != 0].contiguous() + shift_labels = labels[..., 1:][shift_attention_mask.to( + labels.device) != 0].contiguous() + else: + shift_logits = logits[..., :-1, :].contiguous() + shift_labels = labels[..., 1:].contiguous() + # Flatten the tokens + loss_fct = nn.CrossEntropyLoss() + loss = loss_fct( + shift_logits.view(-1, shift_logits.size(-1)), + shift_labels.view(-1).to(shift_logits.device), + ) + + if not return_dict: + output = (logits, ) + outputs[1:] + return (loss, ) + output if loss is not None else output + + return LlavaCausalLMOutputWithPast( + loss=loss, + logits=logits, + past_key_values=outputs.past_key_values, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + ) + + def prepare_inputs_for_generation( + self, + input_ids, + past_key_values=None, + inputs_embeds=None, + pixel_values=None, + grid_thws=None, + attention_mask=None, + **kwargs, + ): + if past_key_values is not None: + if hasattr(past_key_values, "get_seq_length"): + cache_length = past_key_values.get_seq_length() + past_length = getattr(past_key_values, 'seen_tokens', + cache_length) + else: + cache_length = past_length = past_key_values[0][0].shape[2] + + # Keep only the unprocessed tokens: + # 1 - If the length of the attention_mask exceeds the length of input_ids, then we are in a setting where + # some of the inputs are exclusively passed as part of the cache (e.g. when passing input_embeds as + # input) + if attention_mask is not None and attention_mask.shape[ + 1] > input_ids.shape[1]: + input_ids = input_ids[:, -(attention_mask.shape[1] - + past_length):] + # 2 - If the past_length is smaller than input_ids', then input_ids holds all input tokens. We can discard + # input_ids based on the past_length. + elif past_length < input_ids.shape[1]: + input_ids = input_ids[:, past_length:] + # 3 - Otherwise (past_length >= input_ids.shape[1]), let's assume input_ids only has unprocessed tokens. + elif self.config.media_placeholder_token_id in input_ids: + input_ids = input_ids[:, input_ids.shape[1] - 1:] + # If the cache has seen more tokens than it can hold, then the cache has a size limit. Let's discard the + # older attention values, as their corresponding values are not part of the input. + if cache_length < past_length and attention_mask is not None: + attention_mask = attention_mask[:, -(cache_length + + input_ids.shape[1]):] + + position_ids = kwargs.get("position_ids", None) + if attention_mask is not None and position_ids is None: + # create position_ids on the fly for batch generation + position_ids = attention_mask.long().cumsum(-1) - 1 + position_ids.masked_fill_(attention_mask == 0, 1) + if past_key_values: + position_ids = position_ids[:, -input_ids.shape[1]:] + + # if `inputs_embeds` are passed, we only want to use them in the 1st generation step + if inputs_embeds is not None and past_key_values is None: + model_inputs = {"inputs_embeds": inputs_embeds} + else: + model_inputs = {"input_ids": input_ids} + + model_inputs.update({ + "position_ids": position_ids, + "past_key_values": past_key_values, + "use_cache": kwargs.get("use_cache"), + "attention_mask": attention_mask, + "pixel_values": pixel_values, + "grid_thws": grid_thws, + }) + return model_inputs + + def _reorder_cache(self, *args, **kwargs): + return self.language_model._reorder_cache(*args, **kwargs) diff --git a/modeling_kimi_linear.py b/modeling_kimi_linear.py new file mode 100644 index 0000000000000000000000000000000000000000..b8c41e8bfce768d74d8da3a37e693f5ee43876a0 --- /dev/null +++ b/modeling_kimi_linear.py @@ -0,0 +1,1314 @@ +# coding=utf-8 +# Copyright 2025-2026 The Moonshot AI Team, DeepSeek-AI, and HuggingFace Inc. team. All rights reserved. +# +# The multi-head latent attention, MoE gating and sparse MoE block in this file are +# adapted from DeepSeek-V3 (DeepSeek-V3/modeling_deepseek.py). They have been +# extensively modified and extended for the Kimi-Linear architecture. +# +# Licensing Information: +# - Code adapted from DeepSeek-V3 (DeepSeek-V3/modeling_deepseek.py) is licensed under the Apache License, Version 2.0. +# - Other parts of the code are licensed under the Kimi K3 License (see the LICENSE file in this repository). +# +# Apache License, Version 2.0: +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +import math +from collections.abc import Callable +from typing import Any + +import torch +import torch.nn.functional as F +import transformers +from einops import rearrange +from packaging import version +from torch import nn +from transformers.activations import ACT2FN +from transformers.cache_utils import Cache +from transformers.generation import GenerationMixin +from transformers.masking_utils import create_causal_mask +from transformers.modeling_flash_attention_utils import FlashAttentionKwargs +from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast +from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel +from transformers.processing_utils import Unpack +from transformers.pytorch_utils import ALL_LAYERNORM_LAYERS +from transformers.utils import TransformersKwargs, auto_docstring, can_return_tuple, logging +from transformers.utils.generic import OutputRecorder, check_model_inputs + +try: + from fla.modules import FusedRMSNormGated, ShortConvolution + from fla.ops.kda import chunk_kda, fused_recurrent_kda + # from fla.ops.kda.gate import fused_kda_gate # deprecated, gate is now computed inside chunk_kda/fused_recurrent_kda + from fla.ops.utils.index import prepare_cu_seqlens_from_mask, prepare_lens_from_mask + from fla.utils import tensor_cache +except ImportError: + raise ImportError("Plese run `pip install -U fla-core`") + +from .configuration_kimi_k3 import KimiLinearConfig + +assert version.parse(transformers.__version__) >= version.parse("4.56.0"), \ + "Please upgrade transformers to >= 4.56.0" + +logger = logging.get_logger(__name__) + + +# Register Moonshot-specific activation functions +class SituAndMul(nn.Module): + """ + SituAndMul activation: beta * tanh(gate / beta) * sigmoid(gate) * up + When linear_beta is set, up is also transformed by linear_beta * tanh(up / linear_beta). + """ + + def __init__(self, beta: float = 1.0, linear_beta: float | None = None): + super().__init__() + self.beta = beta + self.linear_beta = linear_beta + + def forward(self, x: torch.Tensor) -> torch.Tensor: + d = x.shape[-1] // 2 + gate = x[..., :d].to(torch.float32) + up = x[..., d:].to(torch.float32) + situ_a = self.beta * torch.tanh(gate / self.beta) * torch.sigmoid(gate) + if self.linear_beta is not None: + up = self.linear_beta * torch.tanh(up / self.linear_beta) + return (situ_a * up).to(x.dtype) + + +ACT2FN["situ"] = SituAndMul + + +def _get_situ_activation_params(config: KimiLinearConfig): + beta = getattr(config, "activation_situ_beta", None) + linear_beta = getattr(config, "activation_situ_linear_beta", None) + return beta or 1.0, linear_beta + + +def index_first_axis(x, indices): + return x[indices] + + +@tensor_cache +def get_unpad_data( + attention_mask: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor, int]: + lens = prepare_lens_from_mask(attention_mask) + indices = torch.nonzero(attention_mask.flatten(), as_tuple=False).flatten() + max_seqlen_in_batch = lens.max().item() + cu_seqlens = prepare_cu_seqlens_from_mask(attention_mask) + return indices, cu_seqlens, max_seqlen_in_batch + + +def pad_input( + hidden_states: torch.Tensor, + indices: torch.LongTensor, + batch_size: int, + seq_len: int, +) -> torch.Tensor: + out = hidden_states.new_zeros((batch_size * seq_len, *hidden_states.shape[1:])) + out[indices] = hidden_states + return out.view(batch_size, seq_len, *hidden_states.shape[1:]) + + +class KimiDynamicCache: + """ + Dynamic cache for Kimi model. + Inspired by Qwen3-Next + """ + is_compileable = False + + def __init__(self, config: KimiLinearConfig): + super().__init__() + self.config = config + + if config.linear_attn_config is not None: + self.layer_types = [] + for i in range(config.num_hidden_layers): + if config.is_kda_layer(i): + self.layer_types.append("linear_attention") + else: + self.layer_types.append("full_attention") + else: + self.layer_types = ["full_attention"] * config.num_hidden_layers + + self.transformer_layers = [ + i for i in range(config.num_hidden_layers) if self.layer_types[i] == "full_attention" + ] + + linear_layers = [i for i in range( + config.num_hidden_layers) if self.layer_types[i] == "linear_attention"] + self.last_linear_layer = linear_layers[-1] if linear_layers else -1 + + self.conv_states = [None for _ in range(config.num_hidden_layers)] + self.recurrent_states = [None for _ in range(config.num_hidden_layers)] + self.key_cache = [None for _ in range(config.num_hidden_layers)] + self.value_cache = [None for _ in range(config.num_hidden_layers)] + + def __len__(self): + return len(self.layer_types) + + def update( + self, + key_states: torch.Tensor, + value_states: torch.Tensor, + layer_idx: int, + cache_kwargs: dict[str, Any] | None = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + if self.key_cache[layer_idx] is None: + self.key_cache[layer_idx] = key_states + self.value_cache[layer_idx] = value_states + else: + self.key_cache[layer_idx] = torch.cat( + [self.key_cache[layer_idx], key_states], dim=2) + self.value_cache[layer_idx] = torch.cat( + [self.value_cache[layer_idx], value_states], dim=2) + + return self.key_cache[layer_idx], self.value_cache[layer_idx] + + def reorder_cache(self, beam_idx: torch.LongTensor): + """Reorders the cache for beam search, given the selected beam indices.""" + for layer_idx in range(len(self.key_cache)): + if self.key_cache[layer_idx] is not None: + device = self.key_cache[layer_idx].device + beam_idx = beam_idx.to(device) + self.key_cache[layer_idx] = self.key_cache[layer_idx].index_select( + 0, beam_idx) + self.value_cache[layer_idx] = self.value_cache[layer_idx].index_select( + 0, beam_idx) + + if self.conv_states[layer_idx] is not None: + device = self.conv_states[layer_idx][0].device + beam_idx = beam_idx.to(device) + q_conv, k_conv, v_conv = self.conv_states[layer_idx] + self.conv_states[layer_idx] = ( + q_conv.index_select(0, beam_idx), + k_conv.index_select(0, beam_idx), + v_conv.index_select(0, beam_idx), + ) + self.recurrent_states[layer_idx] = self.recurrent_states[layer_idx].index_select( + 0, beam_idx) + + def get_seq_length(self, layer_idx: int | None = 0) -> int: + """Returns the sequence length of the cached states. A layer index can be optionally passed.""" + # take any layer that contains cache and not empty tensor + layer_idx = self.transformer_layers[0] if layer_idx not in self.transformer_layers else layer_idx + if len(self.key_cache) <= layer_idx or self.key_cache[layer_idx] is None: + return 0 + return self.key_cache[layer_idx].shape[-2] + + def get_mask_sizes(self, cache_position: torch.Tensor, layer_idx: int) -> tuple[int, int]: + """ + Return a tuple (kv_length, kv_offset) corresponding to the length and offset that will be returned for + the given layer at `layer_idx`. + The masks are then prepared according to the given lengths (kv_length, kv_offset) and patterns for each layer. + """ + kv_offset = 0 + query_length = cache_position.shape[0] + past_seen_tokens = self.get_seq_length(layer_idx) + kv_length = query_length + past_seen_tokens + return kv_length, kv_offset + + @property + def has_previous_state(self): + """We have a previous state if the last linear (conv) layer was already updated.""" + if self.last_linear_layer == -1: + return False + return self.conv_states[self.last_linear_layer] is not None + + +class KimiRMSNorm(nn.Module): + def __init__(self, hidden_size, eps=1e-6): + super().__init__() + self.weight = nn.Parameter(torch.ones(hidden_size)) + self.variance_epsilon = eps + + def forward(self, hidden_states): + dtype = hidden_states.dtype + x = hidden_states.float() + x = x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.variance_epsilon) + return self.weight * x.to(dtype) + + +ALL_LAYERNORM_LAYERS.append(KimiRMSNorm) + + +class KimiBlockSparseMLP(nn.Module): + def __init__(self, config: KimiLinearConfig, hidden_size=None, intermediate_size=None): + super().__init__() + self.config = config + self.ffn_dim = config.intermediate_size if intermediate_size is None else intermediate_size + self.hidden_dim = config.hidden_size if hidden_size is None else hidden_size + + self.w1 = nn.Linear(self.hidden_dim, self.ffn_dim, bias=False) # gate + self.w2 = nn.Linear(self.ffn_dim, self.hidden_dim, bias=False) # down + self.w3 = nn.Linear(self.hidden_dim, self.ffn_dim, bias=False) # up + + if config.hidden_act == "situ": + beta, linear_beta = _get_situ_activation_params(config) + self.act_fn = SituAndMul( + beta=beta, + linear_beta=linear_beta, + ) + else: + self.act_fn = ACT2FN[config.hidden_act] + + def forward(self, hidden_states): + if self.config.hidden_act == "situ": + gate_up = torch.cat([self.w1(hidden_states), self.w3(hidden_states)], dim=-1) + current_hidden_states = self.act_fn(gate_up) + else: + current_hidden_states = self.act_fn( + self.w1(hidden_states)) * self.w3(hidden_states) + current_hidden_states = self.w2(current_hidden_states) + return current_hidden_states + + +class KimiMLP(nn.Module): + def __init__(self, config: KimiLinearConfig, hidden_size=None, intermediate_size=None): + super().__init__() + self.config = config + self.hidden_size = config.hidden_size if hidden_size is None else hidden_size + self.intermediate_size = config.intermediate_size if intermediate_size is None else intermediate_size + self.gate_proj = nn.Linear( + self.hidden_size, self.intermediate_size, bias=False) + self.up_proj = nn.Linear( + self.hidden_size, self.intermediate_size, bias=False) + self.down_proj = nn.Linear( + self.intermediate_size, self.hidden_size, bias=False) + if config.hidden_act == "situ": + beta, linear_beta = _get_situ_activation_params(config) + self.act_fn = SituAndMul( + beta=beta, + linear_beta=linear_beta, + ) + else: + self.act_fn = ACT2FN[config.hidden_act] + + def forward(self, x): + if self.config.hidden_act == "situ": + gate_up = torch.cat([self.gate_proj(x), self.up_proj(x)], dim=-1) + down_proj = self.down_proj(self.act_fn(gate_up)) + else: + down_proj = self.down_proj(self.act_fn( + self.gate_proj(x)) * self.up_proj(x)) + return down_proj + + +def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor: + """Expand the key/value heads from `num_key_value_heads` to `num_attention_heads`.""" + if n_rep == 1: + return hidden_states + return torch.repeat_interleave(hidden_states, dim=1, repeats=n_rep) + + +def eager_attention_forward( + module: nn.Module, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attention_mask: torch.Tensor | None, + scaling: float, + dropout: float = 0.0, + **kwargs: Unpack[TransformersKwargs], +): + key = repeat_kv(key, module.num_key_value_groups) + value = repeat_kv(value, module.num_key_value_groups) + + scores = torch.einsum("bhqd,bhkd->bhqk", query, key) * scaling + if attention_mask is not None: + scores = scores + attention_mask[:, :, :, : key.shape[-2]] + + probs = F.softmax(scores, dim=-1, dtype=torch.float32).to(query.dtype) + probs = F.dropout(probs, p=dropout, training=module.training) + out = torch.einsum("bhqk,bhkd->bhqd", probs, value).transpose(1, 2).contiguous() + + return out, probs + + +class KimiMLAAttention(nn.Module): + """ + Multi-Latent Attention adapted from deepseek-v3 + """ + + def __init__(self, config: KimiLinearConfig, layer_idx: int): + nn.Module.__init__(self) + self.config = config + self.layer_idx = layer_idx + self.hidden_size = config.hidden_size + self.num_heads = config.num_attention_heads + self.num_key_value_heads = config.num_key_value_heads + self.num_key_value_groups = self.num_heads // self.num_key_value_heads + + self.attention_dropout = getattr(config, "attention_dropout", 0.0) + + try: + self.q_lora_rank = config.q_lora_rank + self.qk_rope_head_dim = config.qk_rope_head_dim + self.kv_lora_rank = config.kv_lora_rank + self.v_head_dim = config.v_head_dim + self.qk_nope_head_dim = config.qk_nope_head_dim + self.q_head_dim = self.qk_nope_head_dim + self.qk_rope_head_dim + self.use_nope = config.mla_use_nope + self.scaling = self.q_head_dim ** (-0.5) + except Exception as e: + raise ValueError( + f"Kimi MLA config is not found or not properly formatted: {e}") + + if self.q_lora_rank is not None: + self.q_a_proj = nn.Linear( + self.hidden_size, self.q_lora_rank, bias=False, + ) + self.q_a_layernorm = KimiRMSNorm(self.q_lora_rank) + self.q_b_proj = nn.Linear( + self.q_lora_rank, + self.num_heads * self.q_head_dim, + bias=False, + ) + else: + self.q_proj = nn.Linear( + self.hidden_size, self.num_heads * self.q_head_dim, bias=False, + ) + self.kv_a_proj_with_mqa = nn.Linear( + self.hidden_size, + self.kv_lora_rank + self.qk_rope_head_dim, + bias=False, + ) + self.kv_a_layernorm = KimiRMSNorm(self.kv_lora_rank) + self.kv_b_proj = nn.Linear( + self.kv_lora_rank, + self.num_heads + * (self.q_head_dim - self.qk_rope_head_dim + self.v_head_dim), + bias=False, + ) + self.o_proj = nn.Linear( + self.num_heads * self.v_head_dim, + self.hidden_size, + bias=False, + ) + self.is_causal = True + assert self.use_nope + + self.use_output_gate = getattr(config, "mla_use_output_gate", False) + if self.use_output_gate: + projection_size = self.num_heads * self.v_head_dim + self.g_proj = nn.Linear(self.hidden_size, projection_size, bias=False) + + self.rotary_emb = None + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None = None, + position_ids: torch.LongTensor | None = None, + past_key_values: Cache | None = None, + **kwargs, + ) -> tuple[torch.Tensor, torch.Tensor | None, tuple[torch.Tensor] | None]: + batch_size, seq_length = hidden_states.shape[:-1] + query_shape = (batch_size, seq_length, -1, self.q_head_dim) + key_shape = (batch_size, seq_length, -1, + self.qk_nope_head_dim + self.v_head_dim) + + if self.q_lora_rank is not None: + q_states = self.q_b_proj(self.q_a_layernorm(self.q_a_proj(hidden_states))) + else: + q_states = self.q_proj(hidden_states) + q_states = q_states.view(query_shape).transpose(1, 2) + q_pass, q_rot = torch.split( + q_states, [self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1) + + compressed_kv = self.kv_a_proj_with_mqa(hidden_states) + k_pass, k_rot = torch.split( + compressed_kv, [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1) + + k_pass = self.kv_b_proj(self.kv_a_layernorm( + k_pass)).view(key_shape).transpose(1, 2) + k_pass, value_states = torch.split( + k_pass, [self.qk_nope_head_dim, self.v_head_dim], dim=-1) + + k_rot = k_rot.view(batch_size, 1, seq_length, self.qk_rope_head_dim) + + k_rot = k_rot.expand(*k_pass.shape[:-1], -1) + + query_states = torch.cat((q_pass, q_rot), dim=-1) + key_states = torch.cat((k_pass, k_rot), dim=-1) + + if past_key_values is not None: + key_states, value_states = past_key_values.update( + key_states, value_states, self.layer_idx) + + if self.config._attn_implementation == "flash_attention_2" and self.q_head_dim != self.v_head_dim: + value_states = F.pad( + value_states, [0, self.q_head_dim - self.v_head_dim]) + + attention_interface: Callable = eager_attention_forward + if self.config._attn_implementation != "eager": + attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation] + + attn_output, _ = attention_interface( + self, + query_states, + key_states, + value_states, + attention_mask, + dropout=0.0 if not self.training else self.attention_dropout, + scaling=self.scaling, + **kwargs, + ) + + if self.config._attn_implementation == "flash_attention_2" and self.q_head_dim != self.v_head_dim: + attn_output = attn_output[:, :, :, : self.v_head_dim] + + attn_output = attn_output.reshape( + batch_size, seq_length, -1).contiguous() + if self.use_output_gate: + g = self.g_proj(hidden_states).sigmoid() + attn_output = attn_output * g + attn_output = self.o_proj(attn_output) + return attn_output + + +class KimiDeltaAttention(nn.Module): + def __init__(self, config: KimiLinearConfig, layer_idx: int): + super().__init__() + self.config = config + self.mode = "chunk" + + self.hidden_size = config.hidden_size + self.conv_size = config.linear_attn_config["short_conv_kernel_size"] + self.head_dim = config.linear_attn_config["head_dim"] + self.num_heads = config.linear_attn_config["num_heads"] + self.head_k_dim = self.head_dim + self.num_k_heads = self.num_heads + + self.layer_idx = layer_idx + + assert self.mode in [ + 'chunk', 'fused_recurrent'], f"Not supported mode `{self.mode}`." + + projection_k_size = self.head_k_dim * self.num_k_heads + projection_size = self.head_dim * self.num_heads + + self.q_proj = nn.Linear( + self.hidden_size, projection_k_size, bias=False) + self.k_proj = nn.Linear( + self.hidden_size, projection_k_size, bias=False) + self.v_proj = nn.Linear(self.hidden_size, projection_size, bias=False) + + self.q_conv1d = ShortConvolution( + hidden_size=projection_k_size, + kernel_size=self.conv_size, + activation='silu', + ) + self.k_conv1d = ShortConvolution( + hidden_size=projection_k_size, + kernel_size=self.conv_size, + activation='silu', + ) + self.v_conv1d = ShortConvolution( + hidden_size=projection_size, + kernel_size=self.conv_size, + activation='silu', + ) + + self.A_log = torch.nn.Parameter(torch.log(torch.empty( + self.num_heads, dtype=torch.float32).uniform_(1, 16))) + + self.f_a_proj = nn.Linear(self.hidden_size, self.head_dim, bias=False) + self.f_b_proj = nn.Linear(self.head_dim, projection_size, bias=False) + + self.dt_bias = nn.Parameter( + torch.empty(projection_size, dtype=torch.float32)) + + self.b_proj = nn.Linear(self.hidden_size, self.num_heads, bias=False) + + self.use_full_rank_gate = config.linear_attn_config.get("use_full_rank_gate", False) + self.gate_lower_bound = config.linear_attn_config.get("gate_lower_bound", None) + if self.use_full_rank_gate: + self.g_proj = nn.Linear(self.hidden_size, projection_size, bias=False) + else: + self.g_a_proj = nn.Linear(self.hidden_size, self.head_dim, bias=False) + self.g_b_proj = nn.Linear(self.head_dim, projection_size, bias=False) + + self.o_norm = FusedRMSNormGated( + self.head_dim, eps=config.rms_norm_eps, activation='sigmoid') + self.o_proj = nn.Linear(projection_size, self.hidden_size, bias=False) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None = None, + cache_params: KimiDynamicCache | None = None, + **kwargs: Unpack[dict], + ) -> tuple[torch.Tensor, torch.Tensor | None, Cache | None]: + if attention_mask is not None: + if attention_mask.dim() != 2: + attention_mask = kwargs.get("padding_mask") + + if attention_mask is not None and attention_mask.dim() != 2: + raise ValueError( + "attention_mask must be a 0-1 matrix of shape [batch_size, seq_len] " + "(0 = padding). 3D masks are not supported here.", + ) + use_cache = cache_params is not None + batch_size, q_len, _ = hidden_states.shape + mode = 'fused_recurrent' if use_cache and q_len == 1 else self.mode + if self.training: + assert mode == 'chunk', "Only chunk mode is supported in training." + + cu_seqlens = kwargs.get('cu_seqlens') + indices = None + if attention_mask is not None: + indices, cu_seqlens, _ = get_unpad_data(attention_mask[:, -q_len:]) + hidden_states = index_first_axis( + rearrange(hidden_states, "b s ... -> (b s) ..."), indices).unsqueeze(0) + + conv_state_q, conv_state_k, conv_state_v = None, None, None + recurrent_state = None + if cache_params is not None: + if cache_params.conv_states[self.layer_idx] is not None: + conv_state_q, conv_state_k, conv_state_v = cache_params.conv_states[ + self.layer_idx] + recurrent_state = cache_params.recurrent_states[self.layer_idx] + + q_proj_states = self.q_proj(hidden_states) + k_proj_states = self.k_proj(hidden_states) + v_proj_states = self.v_proj(hidden_states) + q, conv_state_q = self.q_conv1d( + x=q_proj_states, + cache=conv_state_q, + output_final_state=use_cache, + cu_seqlens=cu_seqlens, + ) + k, conv_state_k = self.k_conv1d( + x=k_proj_states, + cache=conv_state_k, + output_final_state=use_cache, + cu_seqlens=cu_seqlens, + ) + v, conv_state_v = self.v_conv1d( + x=v_proj_states, + cache=conv_state_v, + output_final_state=use_cache, + cu_seqlens=cu_seqlens, + ) + g = self.f_b_proj(self.f_a_proj(hidden_states)) + g = rearrange(g, '... (h d) -> ... h d', d=self.head_dim) + beta = self.b_proj(hidden_states).float() + + q, k = map(lambda x: rearrange( + x, '... (h d) -> ... h d', d=self.head_k_dim), (q, k)) + v = rearrange(v, '... (h d) -> ... h d', d=self.head_dim) + + if mode == 'chunk': + o, recurrent_state = chunk_kda( + q=q, + k=k, + v=v, + g=g, + beta=beta, + A_log=self.A_log, + dt_bias=self.dt_bias, + initial_state=recurrent_state, + output_final_state=True, + use_qk_l2norm_in_kernel=True, + use_gate_in_kernel=True, + use_beta_sigmoid_in_kernel=True, + safe_gate=self.gate_lower_bound is not None, + lower_bound=self.gate_lower_bound, + transpose_state_layout=True, + cu_seqlens=cu_seqlens, + ) + else: + o, recurrent_state = fused_recurrent_kda( + q=q, + k=k, + v=v, + g=g, + beta=beta, + A_log=self.A_log, + dt_bias=self.dt_bias, + initial_state=recurrent_state, + output_final_state=True, + use_qk_l2norm_in_kernel=True, + use_gate_in_kernel=True, + use_beta_sigmoid_in_kernel=True, + lower_bound=self.gate_lower_bound, + transpose_state_layout=True, + cu_seqlens=cu_seqlens, + ) + if cache_params is not None: + cache_params.recurrent_states[self.layer_idx] = recurrent_state + cache_params.conv_states[self.layer_idx] = ( + conv_state_q, conv_state_k, conv_state_v) + + if self.use_full_rank_gate: + g = self.g_proj(hidden_states) + else: + g = self.g_b_proj(self.g_a_proj(hidden_states)) + g = rearrange(g, '... (h d) -> ... h d', d=self.head_dim) + o = self.o_norm(o, g) + + o = rearrange(o, 'b t h d -> b t (h d)') + o = self.o_proj(o) + if attention_mask is not None: + o = pad_input(o.squeeze(0), indices, batch_size, q_len) + + return o + + +class KimiMoEGate(nn.Module): + """ + MoEGate adapted from Deepseek-V3. + Parameter correspondences: + num_experts -> n_routed_experts + num_experts_per_token -> num_experts_per_tok + num_expert_group -> n_group + moe_router_activation_func -> scoring_func + """ + + def __init__(self, config: KimiLinearConfig): + super().__init__() + self.config = config + self.top_k = config.num_experts_per_token + self.num_experts = config.num_experts + self.routed_scaling_factor = config.routed_scaling_factor + self.moe_router_activation_func = config.moe_router_activation_func + self.num_expert_group = getattr(config, "num_expert_group", 1) + self.topk_group = getattr(config, "topk_group", 1) + + # topk selection algorithm + self.moe_renormalize = config.moe_renormalize + self.gating_dim = config.hidden_size + self.weight = nn.Parameter( + torch.empty((self.num_experts, self.gating_dim)), + ) + + self.e_score_correction_bias = nn.Parameter( + torch.empty(self.num_experts), + ) + self.reset_parameters() + + def reset_parameters(self) -> None: + import torch.nn.init as init + + init.kaiming_uniform_(self.weight, a=math.sqrt(5)) + + def forward(self, hidden_states): + bsz, seq_len, h = hidden_states.shape + # compute gating score + hidden_states = hidden_states.view(-1, h) + logits = F.linear( + hidden_states.type(torch.float32), self.weight.type( + torch.float32), None, + ) + if self.moe_router_activation_func == "sigmoid": + scores = logits.sigmoid() + elif self.moe_router_activation_func == "softmax": + scores = logits.softmax(dim=1) + else: + raise NotImplementedError( + f"insupportable scoring function for MoE gating: {self.moe_router_activation_func}", + ) + + # select top-k experts + assert not self.training + scores = scores.view(bsz * seq_len, -1) + scores_for_choice = scores + self.e_score_correction_bias.unsqueeze(0) + if self.num_expert_group > 1 and self.num_expert_group > self.topk_group: + group_scores = ( + scores_for_choice.view( + bsz * seq_len, self.num_expert_group, -1).topk(2, dim=-1)[0].sum(dim=-1) + ) # [n, num_expert_group] + group_idx = torch.topk( + group_scores, k=self.topk_group, dim=-1, sorted=False, + )[ + 1 + ] # [n, top_k_group] + group_mask = torch.zeros_like(group_scores) # [n, num_expert_group] + group_mask.scatter_(1, group_idx, 1) # [n, num_expert_group] + score_mask = ( + group_mask.unsqueeze(-1) + .expand( + bsz * seq_len, self.num_expert_group, self.num_experts // self.num_expert_group, + ) + .reshape(bsz * seq_len, -1) + ) # [n, e] + tmp_scores = scores_for_choice.masked_fill( + ~score_mask.bool(), float("-inf")) # [n, e] + else: + tmp_scores = scores_for_choice + _, topk_idx = torch.topk( + tmp_scores, k=self.top_k, dim=-1, sorted=False, + ) + topk_weight = scores.gather(1, topk_idx) + + # norm gate to sum 1 + if self.top_k > 1 and self.moe_renormalize: + denominator = topk_weight.sum(dim=-1, keepdim=True) + 1e-20 + topk_weight = topk_weight / denominator + # must multiply the scaling factor + topk_weight = topk_weight * self.routed_scaling_factor + + return topk_idx, topk_weight + + +class KimiSparseMoeBlock(nn.Module): + """ + Adapted from Deepseek-V3's MOE implementation + The namings are consistent with Kimi's version. + """ + + def __init__(self, config: KimiLinearConfig): + super().__init__() + self.config = config + self.hidden_dim = config.hidden_size + self.num_experts = config.num_experts + self.top_k = config.num_experts_per_token + self.moe_renormalize = config.moe_renormalize + + self.use_latent_moe = getattr(config, "routed_expert_hidden_size", None) is not None + self.moe_hidden_size = ( + config.routed_expert_hidden_size + if self.use_latent_moe else config.hidden_size + ) + self.latent_moe_use_norm = getattr(config, "latent_moe_use_norm", False) + + self.ep_size = 1 + self.experts_per_rank = config.num_experts + self.ep_rank = 0 + self.experts = nn.ModuleList( + [ + KimiBlockSparseMLP( + config, + hidden_size=self.moe_hidden_size, + intermediate_size=config.moe_intermediate_size, + ) + for _ in range(config.num_experts) + ], + ) + self.gate = KimiMoEGate(config) + if config.num_shared_experts is not None: + intermediate_size = config.moe_intermediate_size * config.num_shared_experts + self.shared_experts = KimiMLP( + config=config, intermediate_size=intermediate_size, + ) + + if self.use_latent_moe: + self.routed_expert_down_proj = nn.Linear( + config.hidden_size, self.moe_hidden_size, bias=False, + ) + self.routed_expert_up_proj = nn.Linear( + self.moe_hidden_size, config.hidden_size, bias=False, + ) + if self.latent_moe_use_norm: + self.routed_expert_norm = KimiRMSNorm( + self.moe_hidden_size, eps=config.rms_norm_eps, + ) + + def forward(self, hidden_states): + identity = hidden_states + orig_shape = hidden_states.shape + topk_idx, topk_weight = self.gate(hidden_states) + hidden_states = hidden_states.view(-1, hidden_states.shape[-1]) + + if self.use_latent_moe: + hidden_states = self.routed_expert_down_proj(hidden_states) + + if not self.training: + y = self.moe_infer(hidden_states, topk_idx, topk_weight) + else: + raise NotImplementedError("Training mode is not supported in KimiSparseMoeBlock") + + if self.use_latent_moe: + if self.latent_moe_use_norm: + y = self.routed_expert_norm(y) + y = self.routed_expert_up_proj(y) + + y = y.view(*orig_shape) + + if self.config.num_shared_experts is not None: + y = y + self.shared_experts(identity) + return y + + @torch.no_grad() + def moe_infer(self, x, topk_ids, topk_weight): + cnts = topk_ids.new_zeros((topk_ids.shape[0], len(self.experts))) + cnts.scatter_(1, topk_ids, 1) + tokens_per_expert = cnts.sum(dim=0) + idxs = topk_ids.view(-1).argsort() + sorted_tokens = x[idxs // topk_ids.shape[1]] + + tokens_per_expert = tokens_per_expert.cpu().numpy() + + outputs = [] + start_idx = 0 + for i, num_tokens in enumerate(tokens_per_expert): + end_idx = start_idx + num_tokens + if num_tokens == 0: + continue + expert = self.experts[i + self.ep_rank * self.experts_per_rank] + tokens_for_this_expert = sorted_tokens[start_idx:end_idx] + expert_out = expert(tokens_for_this_expert) + outputs.append(expert_out) + start_idx = end_idx + + outs = torch.cat(outputs, dim=0) if len( + outputs) else sorted_tokens.new_empty(0) + + new_x = torch.empty_like(outs) + new_x[idxs] = outs + final_out = ( + new_x.view(*topk_ids.shape, -1) + .type(topk_weight.dtype) + .mul_(topk_weight.unsqueeze(dim=-1)) + .sum(dim=1) + .type(new_x.dtype) + ) + return final_out + + +class KimiDecoderLayer(nn.Module): + def __init__(self, config: KimiLinearConfig, layer_idx: int): + super().__init__() + self.hidden_size = config.hidden_size + self.config = config + self.layer_idx = layer_idx + if config.is_kda_layer(layer_idx): + self.is_linear_attn = True + self.self_attn = KimiDeltaAttention( + config=config, layer_idx=layer_idx) + elif config.is_mla: + self.is_linear_attn = False + self.self_attn = KimiMLAAttention( + config=config, layer_idx=layer_idx) + else: + raise NotImplementedError + if ( + config.num_experts is not None + and layer_idx >= config.first_k_dense_replace + and layer_idx % getattr(config, "moe_layer_freq", 1) == 0 + ): + self.block_sparse_moe = KimiSparseMoeBlock(config) + else: + self.mlp = KimiMLP(config) + self.input_layernorm = KimiRMSNorm( + config.hidden_size, eps=config.rms_norm_eps) + self.post_attention_layernorm = KimiRMSNorm( + config.hidden_size, eps=config.rms_norm_eps) + + # Attention residual + self.use_attn_residuals = getattr(config, "attn_res_block_size", None) is not None + if self.use_attn_residuals: + self.attn_res_block_size = config.attn_res_block_size + self.self_attention_res_norm = KimiRMSNorm( + config.hidden_size, eps=config.rms_norm_eps) + self.mlp_res_norm = KimiRMSNorm( + config.hidden_size, eps=config.rms_norm_eps) + self.self_attention_res_proj = nn.Linear( + config.hidden_size, 1, bias=False) + self.mlp_res_proj = nn.Linear( + config.hidden_size, 1, bias=False) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None = None, + position_ids: torch.LongTensor | None = None, + past_key_values: tuple[torch.Tensor] | None = None, + output_attentions: bool | None = False, + use_cache: bool | None = False, + block_residual: torch.Tensor | None = None, + **kwargs: Unpack[FlashAttentionKwargs], + ): + if self.use_attn_residuals: + return self._forward_attn_residual( + hidden_states, attention_mask, position_ids, + past_key_values, output_attentions, use_cache, + block_residual, **kwargs) + + residual = hidden_states + + hidden_states = self.input_layernorm(hidden_states) + + # Self Attention + if self.is_linear_attn is False: + hidden_states = self.self_attn( + hidden_states=hidden_states, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + output_attentions=output_attentions, + use_cache=use_cache, + **kwargs, + ) + else: + hidden_states = self.self_attn( + hidden_states=hidden_states, + attention_mask=attention_mask, + cache_params=past_key_values, + output_attentions=output_attentions, + use_cache=use_cache, + **kwargs, + ) + hidden_states = residual + hidden_states + + # Fully Connected + residual = hidden_states + hidden_states = self.post_attention_layernorm(hidden_states) + if hasattr(self, "block_sparse_moe"): + hidden_states = self.block_sparse_moe(hidden_states) + else: + hidden_states = self.mlp(hidden_states) + hidden_states = residual + hidden_states + + return hidden_states + + def _forward_attn_residual( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None = None, + position_ids: torch.LongTensor | None = None, + past_key_values: tuple[torch.Tensor] | None = None, + output_attentions: bool | None = False, + use_cache: bool | None = False, + block_residual: torch.Tensor | None = None, + **kwargs: Unpack[FlashAttentionKwargs], + ): + batch_size, seq_len, hidden_size = hidden_states.shape + prefix_sum = hidden_states + + if block_residual is not None and block_residual.shape[1] > 0: + hidden_states = _apply_attn_res( + prefix_sum.view(-1, hidden_size), + block_residual, + self.self_attention_res_proj, + self.self_attention_res_norm, + ).view(batch_size, seq_len, hidden_size) + + if self.layer_idx % self.attn_res_block_size == 0: + block_residual = torch.cat( + [block_residual, prefix_sum.view(-1, hidden_size).unsqueeze(1)], dim=1) + prefix_sum = None + + hidden_states = self.input_layernorm(hidden_states) + + # Self Attention + if self.is_linear_attn is False: + hidden_states = self.self_attn( + hidden_states=hidden_states, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + output_attentions=output_attentions, + use_cache=use_cache, + **kwargs, + ) + else: + hidden_states = self.self_attn( + hidden_states=hidden_states, + attention_mask=attention_mask, + cache_params=past_key_values, + output_attentions=output_attentions, + use_cache=use_cache, + **kwargs, + ) + + if prefix_sum is not None: + prefix_sum = prefix_sum + hidden_states + else: + prefix_sum = hidden_states + + hidden_states = _apply_attn_res( + prefix_sum.view(-1, hidden_size), + block_residual, + self.mlp_res_proj, + self.mlp_res_norm, + ).view(batch_size, seq_len, hidden_size) + + hidden_states = self.post_attention_layernorm(hidden_states) + if hasattr(self, "block_sparse_moe"): + hidden_states = self.block_sparse_moe(hidden_states) + else: + hidden_states = self.mlp(hidden_states) + + if prefix_sum is None: + prefix_sum = hidden_states + else: + prefix_sum = prefix_sum + hidden_states + + return prefix_sum, block_residual + + +class KimiPreTrainedModel(PreTrainedModel): + config_class = KimiLinearConfig + base_model_prefix = "model" + supports_gradient_checkpointing = True + _no_split_modules = ["KimiDecoderLayer"] + _skip_keys_device_placement = "past_key_values" + _supports_flash_attn_2 = True + _can_record_outputs = { + "router_logits": OutputRecorder(KimiBlockSparseMLP, index=1), + "hidden_states": KimiDecoderLayer, + "attentions": KimiMLAAttention, + } + _is_stateful = True + + def _init_weights(self, module): + std = self.config.initializer_range + if isinstance(module, nn.Linear): + module.weight.data.normal_(mean=0.0, std=std) + if module.bias is not None: + module.bias.data.zero_() + elif isinstance(module, nn.Embedding): + module.weight.data.normal_(mean=0.0, std=std) + if module.padding_idx is not None: + module.weight.data[module.padding_idx].zero_() + + +def _apply_attn_res(prefix_sum, block_residual, proj, norm): + """ + prefix_sum: (num_tokens, hidden_size) + block_residual: (num_tokens, num_blocks, hidden_size) + """ + v = torch.cat((block_residual, prefix_sum.unsqueeze(1)), dim=1) + v_float = v.float() + variance = v_float.pow(2).mean(-1, keepdim=True) + k = v_float * torch.rsqrt(variance + norm.variance_epsilon) + score_weight = norm.weight.float() * proj.weight.squeeze(0).float() + scores = (k * score_weight).sum(-1) + probs = scores.softmax(-1).unsqueeze(1) + hidden_states = torch.matmul(probs, v_float).squeeze(1) + return hidden_states.to(v.dtype) + +class KimiLinearModel(KimiPreTrainedModel): + def __init__(self, config: KimiLinearConfig): + super().__init__(config) + self.padding_idx = config.pad_token_id + self.vocab_size = config.vocab_size + + self.embed_tokens = nn.Embedding( + config.vocab_size, config.hidden_size, self.padding_idx) + self.layers = nn.ModuleList([KimiDecoderLayer( + config, layer_idx) for layer_idx in range(config.num_hidden_layers)]) + self.norm = KimiRMSNorm( + config.hidden_size, eps=config.rms_norm_eps) + + self.use_attn_residuals = getattr(config, "attn_res_block_size", None) is not None + if self.use_attn_residuals: + self.output_attn_res_norm = KimiRMSNorm( + config.hidden_size, eps=config.rms_norm_eps) + self.output_attn_res_proj = nn.Linear( + config.hidden_size, 1, bias=False) + + if getattr(config, "_attn_implementation", None) is not None: + if config._attn_implementation != "flash_attention_2": + logger.warning_once( + f"Ignoring the provided attention implementation {config._attn_implementation}") + logger.warning_once("Using flash_attention_2 backend instead.") + config._attn_implementation = "flash_attention_2" + else: + config._attn_implementation = "flash_attention_2" + + self._use_flash_attention_2 = config._attn_implementation == "flash_attention_2" + self.gradient_checkpointing = False + # Initialize weights and apply final processing + self.post_init() + + def _update_linear_attn_mask(self, attention_mask, cache_position): + """ + NOTE: Left-padding is used for linear attention mask. + No need for zeroing states when + 1. Cached forward + 2. Attending to all inputs + """ + linear_attn_mask = attention_mask + if cache_position[0] > 0 or (attention_mask is not None and torch.all(attention_mask == 1)): + linear_attn_mask = None + return linear_attn_mask + + @check_model_inputs + # @auto_docstring + def forward( + self, + input_ids: torch.LongTensor = None, + attention_mask: torch.Tensor | None = None, + position_ids: torch.LongTensor | None = None, + past_key_values: Cache | None = None, + inputs_embeds: torch.FloatTensor | None = None, + cache_position: torch.LongTensor | None = None, + use_cache: bool | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPast: + + use_cache = use_cache if use_cache is not None else self.config.use_cache + + if (input_ids is None) and (inputs_embeds is None): + raise ValueError( + "You must specify exactly one of input_ids or inputs_embeds") + + # Get inputs_embeds + if inputs_embeds is None: + inputs_embeds = self.embed_tokens(input_ids) + + if use_cache and past_key_values is None: + past_key_values = KimiDynamicCache(config=self.config) + + if cache_position is None: + past_seen_tokens = past_key_values.get_seq_length( + ) if past_key_values is not None else 0 + cache_position: torch.Tensor = torch.arange( + past_seen_tokens, past_seen_tokens + inputs_embeds.shape[1], device=inputs_embeds.device, + ) + + if position_ids is None: + position_ids = cache_position.unsqueeze(0) + + causal_mask = create_causal_mask( + config=self.config, + input_embeds=inputs_embeds, + attention_mask=attention_mask, + cache_position=cache_position, + past_key_values=past_key_values, + position_ids=position_ids, + ) + linear_attn_mask = self._update_linear_attn_mask( + attention_mask, cache_position) + + hidden_states = inputs_embeds + if past_key_values is not None: + assert isinstance(past_key_values, KimiDynamicCache) + + block_residual = None + if self.use_attn_residuals: + block_residual = hidden_states.new_zeros( + hidden_states.shape[0] * hidden_states.shape[1], 0, + hidden_states.shape[2]) + + for decoder_layer in self.layers: + layer_mask = linear_attn_mask if decoder_layer.is_linear_attn else causal_mask + + if self.use_attn_residuals: + hidden_states, block_residual = decoder_layer( + hidden_states, + attention_mask=layer_mask, + past_key_values=past_key_values, + cache_position=cache_position, + block_residual=block_residual, + **kwargs, + ) + else: + hidden_states = decoder_layer( + hidden_states, + attention_mask=layer_mask, + past_key_values=past_key_values, + cache_position=cache_position, + **kwargs, + ) + + if self.use_attn_residuals: + hidden_states = self._apply_output_attn_res( + hidden_states, block_residual) + + hidden_states = self.norm(hidden_states) + + return BaseModelOutputWithPast( + last_hidden_state=hidden_states, + past_key_values=past_key_values, + ) + + def _apply_output_attn_res(self, hidden_states, block_residual): + batch_size, seq_len, hidden_size = hidden_states.shape + return _apply_attn_res( + hidden_states.view(-1, hidden_size), + block_residual, + self.output_attn_res_proj, + self.output_attn_res_norm, + ).view(batch_size, seq_len, hidden_size) + + +class KimiLinearForCausalLM(KimiPreTrainedModel, GenerationMixin): + @classmethod + def _supports_default_dynamic_cache(cls) -> bool: + return False + + _tied_weights_keys = ["lm_head.weight"] + + def __init__(self, config): + super().__init__(config) + self.model = KimiLinearModel(config) + self.vocab_size = config.vocab_size + self.lm_head = nn.Linear( + config.hidden_size, config.vocab_size, bias=False) + + # Initialize weights and apply final processing + self.post_init() + + @can_return_tuple + # @auto_docstring + def forward( + self, + input_ids: torch.LongTensor = None, + attention_mask: torch.Tensor | None = None, + position_ids: torch.LongTensor | None = None, + past_key_values: list[torch.FloatTensor] | None = None, + inputs_embeds: torch.FloatTensor | None = None, + labels: torch.LongTensor | None = None, + use_cache: bool | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + generation_mode: bool | None = None, + return_dict: bool | None = None, + cache_position: torch.LongTensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | CausalLMOutputWithPast: + r""" + Args: + labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): + Labels for computing the masked language modeling loss. Indices should either be in `[0, ..., + config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored + (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`. + """ + + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_hidden_states = ( + output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + ) + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + outputs = self.model( + input_ids=input_ids, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + cache_position=cache_position, + ) + + logits = outputs[0] + if generation_mode: + logits = logits[:, -1:] + logits = self.lm_head(logits) + + loss = None + if labels is not None: + loss = self.loss_function( + logits, labels, self.vocab_size, **kwargs) + + return CausalLMOutputWithPast( + loss=loss, + logits=logits, + past_key_values=outputs.past_key_values, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + ) diff --git a/preprocessor_config.json b/preprocessor_config.json new file mode 100644 index 0000000000000000000000000000000000000000..8be28c13c93fb8982f81928d3f4fe80568102286 --- /dev/null +++ b/preprocessor_config.json @@ -0,0 +1,38 @@ +{ + "media_proc_cfg": { + "in_patch_limit": 65536, + "patch_size": 14, + "image_mean": [ + 0.5, + 0.5, + 0.5 + ], + "image_std": [ + 0.5, + 0.5, + 0.5 + ], + "merge_kernel_size": 2, + "fixed_output_tokens": null, + "patch_limit_on_one_side": 512, + "in_patch_limit_each_frame": 16384, + "in_patch_limit_video": 655360, + "sample_fps": 8.0, + "max_num_frames_each_video": null, + "temporal_merge_kernel_size": 4, + "timestamp_mode": "hh:mm:ss.fff", + "transparent_bg_config": { + "pattern": "chessboard", + "chessboard_square_size": 8, + "chessboard_square_on_top_left": true, + "chessboard_white_value": 255, + "chessboard_gray_value": 180 + }, + "transparent_bg_fill_stage": "after_resize", + "config_type": "media_proc.processors.moonvit.MoonViTMediaProcessorConfig" + }, + "auto_map": { + "AutoProcessor": "kimi_k3_processor.KimiK3Processor", + "AutoImageProcessor": "kimi_k3_vision_processing.KimiK3VisionProcessor" + } +} \ No newline at end of file diff --git a/tiktoken.model b/tiktoken.model new file mode 100644 index 0000000000000000000000000000000000000000..b4149a6e17a01b6442187f39890f89bc2fe8d309 --- /dev/null +++ b/tiktoken.model @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:b6c497a7469b33ced9c38afb1ad6e47f03f5e5dc05f15930799210ec050c5103 +size 2795286 diff --git a/tokenization_kimi.py b/tokenization_kimi.py new file mode 100644 index 0000000000000000000000000000000000000000..5df9d462b43295245084fbab46fd956020d4b84e --- /dev/null +++ b/tokenization_kimi.py @@ -0,0 +1,408 @@ +import os +from logging import getLogger +from pathlib import Path +from shutil import copyfile +from typing import Dict, Iterator, List, Optional, Tuple, Union, cast + +import tiktoken +from tiktoken.load import load_tiktoken_bpe +from tokenizers import AddedToken +from transformers.convert_slow_tokenizer import bytes_to_unicode +from transformers.tokenization_utils import PreTrainedTokenizer + +try: + from .encoding_k3 import build_chat_segments, is_batched_conversation +except ImportError: # pragma: no cover - supports direct file execution/import. + from encoding_k3 import build_chat_segments, is_batched_conversation + +logger = getLogger(__name__) +VOCAB_FILES_NAMES = {"vocab_file": "tiktoken.model"} + + +class TikTokenTokenizer(PreTrainedTokenizer): + """ + Tokenizing and encoding/decoding text using the Tiktoken tokenizer. See megatron/tokenizer/tiktoken_tokenizer.py. + + This tokenizer inherits from [`PreTrainedTokenizer`] which contains most of the main methods. Users should refer to + this superclass for more information regarding those methods. + + Args: + vocab_file (`str`): + The path to the Tiktoken model file. + bos_token (`str` or `tokenizers.AddedToken`, *optional*, defaults to `"<|begin_of_text|>",`): + The beginning of sequence token that was used during pretraining. Can be used a sequence classifier token. + eos_token (`str` or `tokenizers.AddedToken`, *optional*, defaults to `"<|end_of_text|>"`): + The end of sequence token. + unk_token (`str` or `tokenizers.AddedToken`, *optional*, defaults to `"<|reserved_special_token_249|>"`): + The unknown token. A token that is not in the vocabulary cannot be converted to an ID and is set to be this + token instead. The second to last item in special_tokens. + pad_token (`str` or `tokenizers.AddedToken`, *optional*, defaults to `"<|reserved_special_token_250|>"`): + The token used for padding, for example when batching sequences of different lengths. + additional_special_tokens (list of `str`, *optional*): + A tuple or a list of additional tokens, which will be marked as `special`, meaning that they will be + skipped when decoding if `skip_special_tokens` is set to `True`. + """ + + vocab_files_names = VOCAB_FILES_NAMES + + model_input_names = ["input_ids", "attention_mask"] + + special_tokens: Dict[str, int] + + num_reserved_special_tokens = 256 + + pat_str = "|".join([ + r"""[\p{Han}]+""", + r"""[^\r\n\p{L}\p{N}]?[\p{Lu}\p{Lt}\p{Lm}\p{Lo}\p{M}&&[^\p{Han}]]*[\p{Ll}\p{Lm}\p{Lo}\p{M}&&[^\p{Han}]]+(?i:'s|'t|'re|'ve|'m|'ll|'d)?""", + r"""[^\r\n\p{L}\p{N}]?[\p{Lu}\p{Lt}\p{Lm}\p{Lo}\p{M}&&[^\p{Han}]]+[\p{Ll}\p{Lm}\p{Lo}\p{M}&&[^\p{Han}]]*(?i:'s|'t|'re|'ve|'m|'ll|'d)?""", + r"""\p{N}{1,3}""", + r""" ?[^\s\p{L}\p{N}]+[\r\n]*""", + r"""\s*[\r\n]+""", + r"""\s+(?!\S)""", + r"""\s+""", + ]) + + def __init__( + self, + vocab_file, + bos_token: Union[str, AddedToken] = "[BOS]", + eos_token: Union[str, AddedToken] = "[EOS]", + unk_token: Union[str, AddedToken, None] = None, + pad_token: Union[str, AddedToken, None] = None, + additional_special_tokens: List[str] = None, + added_tokens_decoder: Optional[dict] = None, + **kwargs, + ): + assert os.path.isfile(vocab_file), vocab_file + + if additional_special_tokens is None: + additional_special_tokens = [ + "<|im_end|>", + "<|im_user|>", + "<|im_assistant|>", + "<|start_header_id|>", + "<|end_header_id|>", + "[EOT]", + "<|im_system|>", + "<|im_middle|>", + ] + + if added_tokens_decoder: + special_tokens_mapping = { + i: added_tokens_decoder[i].content + for i in added_tokens_decoder + } + else: + special_tokens_mapping = {} + + self.vocab_file = vocab_file + mergeable_ranks = load_tiktoken_bpe(vocab_file) + num_base_tokens = len(mergeable_ranks) + self.special_tokens = { + special_tokens_mapping.get(i, f"<|reserved_token_{i}|>"): i + for i in range(num_base_tokens, num_base_tokens + + self.num_reserved_special_tokens) + } + + self.model = tiktoken.Encoding( + name=Path(vocab_file).name, + pat_str=self.pat_str, + mergeable_ranks=mergeable_ranks, + special_tokens=self.special_tokens, + ) + logger.info(f"Reloaded tiktoken model from {vocab_file}") + + self.n_words: int = self.model.n_vocab + # BOS / EOS token IDs + self.bos_id: int = self.special_tokens[str(bos_token)] + self.eos_id: int = self.special_tokens[str(eos_token)] + logger.info( + f"#words: {self.n_words} - BOS ID: {self.bos_id} - EOS ID: {self.eos_id}" + ) + + self.pad_id: int = self.special_tokens[str(pad_token)] + self.unk_id: int = self.special_tokens[str(unk_token)] + + self.byte_encoder = bytes_to_unicode() + self.byte_decoder = {v: k for k, v in self.byte_encoder.items()} + + self.decoder = {} + for i in range(self.n_words): + # Taken from https://gist.github.com/xenova/a452a6474428de0182b17605a98631ee + decoding = ''.join([ + self.byte_encoder[ord(char)] for char in + self.model.decode_single_token_bytes(i).decode('latin-1') + ]) + self.decoder[i] = decoding + + self.encoder = {} + for i in range(self.n_words): + if i in self.decoder: + self.encoder[self.decoder[i]] = i + + super().__init__( + bos_token=bos_token, + eos_token=eos_token, + unk_token=unk_token, + pad_token=pad_token, + additional_special_tokens=additional_special_tokens, + added_tokens_decoder=added_tokens_decoder, + **kwargs, + ) + self.all_special_ids_set = set(self.all_special_ids) + + def _encode_text_piece(self, text: str, + allow_special_tokens: bool = True) -> List[int]: + # The tiktoken tokenizer can handle <=400k chars without + # pyo3_runtime.PanicException. + TIKTOKEN_MAX_ENCODE_CHARS = 400_000 + + # https://github.com/openai/tiktoken/issues/195 + # Here we iterate over subsequences and split if we exceed the limit + # of max consecutive non-whitespace or whitespace characters. + MAX_NO_WHITESPACES_CHARS = 25_000 + + t: List[int] = [] + for i in range(0, len(text), TIKTOKEN_MAX_ENCODE_CHARS): + for substr in self._split_whitespaces_or_nonwhitespaces( + text[i:i + TIKTOKEN_MAX_ENCODE_CHARS], + MAX_NO_WHITESPACES_CHARS, + ): + if allow_special_tokens: + t.extend( + # structural markers: encode <|...|> as their special token IDs + self.model.encode( + substr, + allowed_special="all", + )) + else: + t.extend( + # user/tool text: encode any <|...|> as ordinary BPE tokens (never as control tokens) + self.model.encode( + substr, + disallowed_special=(), + )) + + return t + + def encode(self, + text: str, + allow_special_tokens: bool = True, + **kwargs) -> List[int]: + """ + Encodes a string into a list of token IDs. + + Args: + text (str): The input string to be encoded. + + Returns: + list[int]: A list of token IDs. + """ + # If there are other args, we should call super().encode because there are a lot of code + # to handle those args. supper().encode finally will call _tokenize and _convert_token_to_id. + # NOTE: our encode method is not compatible with the super().encode method, + # e.g. split_special_tokens' default is True in our encode method. + if len(kwargs) > 0: + logger.warning(f"Calling super().encode with {kwargs}") + return super().encode(text, **kwargs) + + assert type(text) is str + return self._encode_text_piece(text, + allow_special_tokens=allow_special_tokens) + + def decode(self, token_ids: Union[int, List[int]], **kwargs) -> str: + """ + Decodes a list of token IDs into a string. + + Args: + token_ids (List[int]): The list of token IDs to be decoded. + + Returns: + str: The decoded string. + """ + # If there are other args, we should call super().decode because there are a lot of code + # to handle those args. supper().encode finally will call convert_tokens_to_string and _convert_id_to_token. + if len(kwargs) > 0: + return super().decode(token_ids, **kwargs) + + if type(token_ids) is int: + token_ids = [token_ids] + + return self.model.decode(cast(List[int], token_ids)) + + @staticmethod + def _split_whitespaces_or_nonwhitespaces( + s: str, max_consecutive_slice_len: int) -> Iterator[str]: + """ + Splits the string `s` so that each substring contains no more than `max_consecutive_slice_len` + consecutive whitespaces or consecutive non-whitespaces. + """ + current_slice_len = 0 + current_slice_is_space = s[0].isspace() if len(s) > 0 else False + slice_start = 0 + + for i in range(len(s)): + is_now_space = s[i].isspace() + + if current_slice_is_space ^ is_now_space: + current_slice_len = 1 + current_slice_is_space = is_now_space + else: + current_slice_len += 1 + if current_slice_len > max_consecutive_slice_len: + yield s[slice_start:i] + slice_start = i + current_slice_len = 1 + yield s[slice_start:] + + def _encode_chat_segments(self, segments) -> List[int]: + token_ids: List[int] = [] + for segment in segments: + token_ids.extend( + self._encode_text_piece( + segment.text, + allow_special_tokens=segment.allow_special, + )) + return token_ids + + @staticmethod + def _truncate(ids: List[int], + truncation: bool = False, + max_length: Optional[int] = None) -> List[int]: + if truncation and max_length is not None: + return ids[:max_length] + return ids + + def _format_chat_token_output(self, + encoded_inputs: List[List[int]], + *, + is_batched: bool, + padding=False, + truncation: bool = False, + max_length: Optional[int] = None, + return_tensors=None, + return_dict: bool = False): + encoded_inputs = [ + self._truncate(ids, truncation=truncation, max_length=max_length) + for ids in encoded_inputs + ] + + needs_batch_encoding = ( + is_batched or padding or return_tensors is not None or return_dict) + if not needs_batch_encoding: + return encoded_inputs[0] + + features = [{ + "input_ids": ids, + "attention_mask": [1] * len(ids) + } for ids in encoded_inputs] + batch = self.pad(features, + padding=padding, + max_length=max_length if padding else None, + return_attention_mask=True, + return_tensors=return_tensors) + + if return_dict: + return batch + if is_batched: + return batch["input_ids"] + return batch["input_ids"][0] if return_tensors is None else batch[ + "input_ids"] + + """ ----- Below are the abstract methods required by PreTrainedTokenizer ----- """ + + @property + def vocab_size(self) -> int: + return self.n_words + + def get_vocab(self) -> Dict[str, int]: + return self.encoder + + def _tokenize(self, text: str, **kwargs) -> List[str]: + return [self.decoder[t] for t in self.encode(text)] + + def _convert_token_to_id(self, token: str) -> int: + return self.encoder.get(token, self.unk_id) + + def _convert_id_to_token(self, index: int) -> str: + return self.decoder.get(index) + + @staticmethod + def clean_up_tokenization(out_string: str) -> str: + return out_string + + def convert_tokens_to_string(self, tokens: List[str]) -> str: + text = ''.join(tokens) + text = bytearray([self.byte_decoder[c] + for c in text]).decode('utf-8', 'replace') + return text + + def save_vocabulary(self, + save_directory: str, + filename_prefix: Optional[str] = None) -> Tuple[str]: + if not os.path.isdir(save_directory): + raise ValueError( + f"vocabulary path ({save_directory}) should be a directory") + out_vocab_file = os.path.join( + save_directory, + (filename_prefix + "-" if filename_prefix else "") + + VOCAB_FILES_NAMES["vocab_file"]) + + if os.path.abspath(self.vocab_file) != os.path.abspath( + out_vocab_file) and os.path.isfile(self.vocab_file): + copyfile(self.vocab_file, out_vocab_file) + + return (out_vocab_file, ) + + def apply_chat_template(self, + conversation, + tools: Optional[list[dict]] = None, + tokenize: bool = False, + add_generation_prompt: bool = True, + thinking: bool = True, + padding=False, + truncation: bool = False, + max_length: Optional[int] = None, + return_tensors=None, + return_dict: bool = False, + **kwargs): + # Tokenizer-level rendering reorders tool result messages to match + # assistant tool_calls, normalizes per-call arguments and response + # schema, then encodes the resulting XTML structure segment-by-segment. + is_batched = is_batched_conversation(conversation) + conversations = conversation if is_batched else [conversation] + image_prompts = kwargs.pop("image_prompts", None) + if is_batched and image_prompts is not None: + raise ValueError("image_prompts is only supported for one chat.") + + # by default set thinking effort to max + kwargs.setdefault("thinking_effort", "max") + + segment_batches = [ + build_chat_segments( + messages, + tools=tools, + add_generation_prompt=add_generation_prompt, + thinking=thinking, + image_prompts=image_prompts, + **kwargs, + ) for messages in conversations + ] + + if not tokenize: + rendered = ["".join(segment.text for segment in segments) + for segments in segment_batches] + return rendered if is_batched else rendered[0] + + encoded_inputs = [ + self._encode_chat_segments(segments) for segments in segment_batches + ] + return self._format_chat_token_output( + encoded_inputs, + is_batched=is_batched, + padding=padding, + truncation=truncation, + max_length=max_length, + return_tensors=return_tensors, + return_dict=return_dict, + ) diff --git a/tokenizer_config.json b/tokenizer_config.json new file mode 100644 index 0000000000000000000000000000000000000000..418aaacbd996a6674efa89d26be06dae11433e1c --- /dev/null +++ b/tokenizer_config.json @@ -0,0 +1,157 @@ +{ + "added_tokens_decoder": { + "163584": { + "content": "[BOS]", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": true + }, + "163585": { + "content": "[EOS]", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": true + }, + "163586": { + "content": "<|end_of_msg|>", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": true + }, + "163587": { + "content": "<|open|>", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "163588": { + "content": "<|close|>", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "163589": { + "content": "<|sep|>", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "163590": { + "content": "[start_header_id]", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": true + }, + "163591": { + "content": "[end_header_id]", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": true + }, + "163593": { + "content": "[EOT]", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": true + }, + "163602": { + "content": "<|media_begin|>", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": true + }, + "163603": { + "content": "<|media_content|>", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": true + }, + "163604": { + "content": "<|media_end|>", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": true + }, + "163605": { + "content": "<|media_pad|>", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": true + }, + "163649": { + "content": "", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": true + }, + "163838": { + "content": "[UNK]", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": true + }, + "163839": { + "content": "[PAD]", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": true + } + }, + "additional_special_tokens": [ + "<|end_of_msg|>", + "[start_header_id]", + "[end_header_id]", + "[EOT]", + "<|media_begin|>", + "<|media_content|>", + "<|media_end|>", + "<|media_pad|>", + "" + ], + "bos_token": "[BOS]", + "clean_up_tokenization_spaces": false, + "eos_token": "[EOS]", + "extra_special_tokens": {}, + "model_max_length": 1000000000000000019884624838656, + "pad_token": "[PAD]", + "tokenizer_class": "TikTokenTokenizer", + "unk_token": "[UNK]", + "auto_map": { + "AutoTokenizer": [ + "tokenization_kimi.TikTokenTokenizer", + null + ] + } +} \ No newline at end of file