chrisboulton's picture
Add tiny Kimi-K2.7-Code (0.7B total / 0.4B activated)
b8b1172 verified
|
Raw
History Blame Contribute Delete
6.54 kB
metadata
license: mit
base_model:
  - moonshotai/Kimi-K2.7-Code
library_name: transformers

Kimi-K2.7-Code-0.7B-A0.4B

This is a tiny version of moonshotai/Kimi-K2.7-Code created for testing and development.

Model Details

  • Base Model: moonshotai/Kimi-K2.7-Code
  • Architecture: kimi_k25
  • Total Parameters: 0.678B
  • Activated Parameters: 0.413B (8 of 64 routed experts per token, plus the shared expert)

The architecture of the base model is preserved: a MoonViT3d vision tower, a patchmerger multimodal projector, and a DeepseekV3-style text decoder using multi-head latent attention (MLA) with YaRN rope and a noaux_tc sigmoid router. The text stack keeps first_k_dense_replace=1, so layer 0 is a dense MLP and the remaining layers are MoE — one of each layer type in the original 61-layer stack.

Unlike the base checkpoint, which is int4 compressed-tensors (pack-quantized), this model is dense bfloat16. In the base model only the routed experts were quantized — self_attn, shared_experts, the dense-layer MLP, lm_head, vision_tower and mm_projector are all in its ignore list and were already bfloat16 — so the routed experts are the only tensors whose format differs.

Configuration Changes

The following parameters were reduced from the original model:

text_config

Parameter Original Tiny
num_hidden_layers 61 4
hidden_size 7168 1024
intermediate_size 18432 2048
moe_intermediate_size 2048 512
n_routed_experts 384 64
num_attention_heads 64 8
num_key_value_heads 64 8
q_lora_rank 1536 512

vision_config

Parameter Original Tiny
vt_num_hidden_layers 27 2
vt_hidden_size 1152 512
vt_intermediate_size 4304 1024
vt_num_attention_heads 16 8
mm_hidden_size 1152 512
text_hidden_size 7168 1024

Left unchanged so routing and attention behave as in the original: kv_lora_rank (512), num_experts_per_tok (8), n_shared_experts (1), first_k_dense_replace (1), qk_nope_head_dim (128), qk_rope_head_dim (64), v_head_dim (128), topk_method (noaux_tc), vocab_size (163840), patch_size (14), merge_kernel_size (2x2) and the YaRN rope_scaling block.

kv_lora_rank is held at 512 deliberately. Inference runtimes built on FlashInfer compile the MLA kernel with HEAD_DIM_CKV as a compile-time constant, and the ahead-of-time kernel cache is built for 512; shrinking it would push the model onto a JIT-compiled or missing kernel. Keeping it costs about 3M parameters and keeps qk_head_dim at the standard 576 = 512 + 64.

Checkpoint Structure

Single-file model.safetensors (1.4 GB, bfloat16, 668 tensors); the base model is sharded across 64 files with an index. The tensor naming matches the base checkpoint exactly — language_model.model.layers.N.*, language_model.lm_head.weight, vision_tower.* and mm_projector.* — with one expected difference: the base model's routed experts are stored as int4 weight_packed / weight_scale / weight_shape triples, whereas this model stores a plain weight per expert because it is not quantized.

Usage

from transformers import AutoModelForCausalLM, AutoTokenizer

model = AutoModelForCausalLM.from_pretrained("Kimi-K2.7-Code-0.7B-A0.4B", device_map="auto", trust_remote_code=True)
tokenizer = AutoTokenizer.from_pretrained("Kimi-K2.7-Code-0.7B-A0.4B", trust_remote_code=True)

input_ids = tokenizer("According to all known laws", return_tensors="pt").input_ids.to(model.device)
output = model.generate(input_ids, max_new_tokens=20)
print(tokenizer.decode(output[0]))

Creation Process

This model was created using the llm-compressor create-tiny-model claude skill.

  1. The config was shrunk as above and the model randomly initialized with skip_weights_download, so the ~1 TB base checkpoint is never downloaded. MoEGate.weight and MoEGate.e_score_correction_bias are bare torch.empty parameters that _init_weights does not cover, so every parameter is re-initialized explicitly.
  2. Fine-tuned on the skill's toy copypasta dataset until the perplexity stopping criterion fired (190 steps, training loss 12.2 -> 0.0016).
  3. Cast to bfloat16 and re-saved.

Validation (validate_tiny_model.py):

Success: 1.001359462738037 <= 10.0

==================================================
Generating sample text:
According to all known laws of aviation, there is no way a bee should be able to fly.
==================================================

Notes

The modeling code vendored from the base repo needed four fixes to work under transformers 5.14, all applied to modeling_deepseek.py in this repo:

  • _tied_weights_keys was a list; transformers 5 expects a {tied_key: source_key} dict and raises AttributeError on save.
  • DynamicCache.from_legacy_cache() and Cache.to_legacy_cache() were removed in transformers 5; these now use DynamicCache(...) and the cache object directly.
  • The MoE was inference-only: MoEGate.forward asserted not self.training and DeepseekV3MoE.moe_infer is wrapped in torch.no_grad(), so the routed experts could never receive a gradient. A differentiable moe_train path was added that recombines expert outputs identically to moe_infer (verified to agree to 6e-8) and is used only when the module is in training mode. Inference numerics are unchanged.

Two further points affect anyone re-saving this model:

  • _attn_implementation is set to eager in config.json because the base repo defaults the vision tower to flash_attention_2. transformers strips this key in to_dict(), so it must be re-added after every save_pretrained or the model will fail to load without flash-attn installed.
  • transformers 5.14 ships a native kimi_k25 weight-conversion mapping. When a config object is passed explicitly to from_pretrained, that mapping is attached to this remote-code model and save_pretrained reverses it, renaming layers to blocks, mangling lm_head, and silently collapsing the vision norm1/fc1 weights onto norm0/fc0. Setting model._weight_conversions = None before saving avoids this.

The vision tower is randomly initialized and was not trained — fine-tuning used the text-only toy dataset, per the skill. Only the text path has been validated.