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
---
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](https://huggingface.co/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
```python
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.