feat(models): add Lille 130M (lille-130m) fused-QKV GQA text model with in-sublayer norms
- Dominant language
- Rust
- Stars
- 467
- Forks
- 54
- Avg merge
- 4h 25m
- Merged PRs (30d)
- 310
Description
## Summary
`model_type: "`lille-130m`"` (Nikity's Lille 130M base/instruct models) is not detected by mlxcel. It is a small Llama-style decoder with three layout quirks: a fused `qkv_proj` with GQA (10 query heads, 2 KV heads), RMSNorms that live inside the attention and MLP sub-modules (`attention.norm`, `feed_forward.norm`) rather than on the block, traditional (interleaved-pair) RoPE, and an MLP width derived from `n_embd` by a rounding formula rather than stored in the config. Embeddings are always tied. This issue adds `src/models/lille_130m.rs`, detection, registration and docs.
## Current behavior
`"`lille-130m`"` hits the `_ => Err("Unsupported model type")` arm at `src/models/detection.rs:379`; nothing under `src/models/` references it. The nearest existing code is `src/models/helium.rs` (Llama-shaped dense decoder with traditional RoPE) and `src/models/phi3.rs` (pre-fused `qkv_proj`, loaded through `FusedQKVLinear::from_weights_fused` at `src/lib/mlxcel-core/src/layers.rs:2769`).
## Expected behavior
`mlxcel generate -m models/lille-130m-instruct-8bit -p "..."` loads and generates; greedy ids match an oracle run of the same checkpoint.
### Config (`config.json`)
| key | type | default | published value |
|---|---|---|---|
| `model_type` | str | required | `"`lille-130m`"` (note the hyphen) |
| `n_embd` | int | 768 | 640 |
| `n_head` | int | 12 | 10 |
| `n_kv_heads` | int | 12 | 2 |
| `n_layer` | int | 12 | 24 |
| `vocab_size` | int | 50304 | 32768 |
| `block_size` | int | 2048 | 512 (max context) |
| `layer_norm_eps` | float | 1e-5 | 1e-5 (used as the RMSNorm eps) |
| `rope_theta` | float | 10000.0 | 10000.0 |
| `tie_word_embeddings` | bool | true | absent (tied) |
| `dropout` | float | ignored | 0.0 |
Derived: `head_dim = n_embd / n_head` (64); `mlp_hidden = 256 * round((int(8 * n_embd / 3)) / 256)` where `round` is banker's rounding on a float (`640 -> int(1706.67) = 1706 -> 1706 / 256 = 6.664 -> round = 7 -> 1792`; `768 -> 2048`). Validate the derived width against the `down_proj` weight shape at load and fail with a clear error on mismatch.
### Forward pass
```
h = tok_embeddings(tokens)
for each layer:
# attention sub-module, norm INSIDE it
a = RMSNorm_attn(h; layer_norm_eps)
qkv = qkv_proj(a) # [B, L, (n_head + 2 n_kv) * head_dim], no bias
q, k, v = split(qkv, [n_head * hd, n_kv * hd, n_kv * hd])
q = rope(q, offset), k = rope(k, offset) # traditional = TRUE (interleaved pairs), base rope_theta, full head_dim
h = h + out_proj(sdpa(q, k, v, scale = hd ** -0.5, causal))
# MLP sub-module, norm INSIDE it
m = RMSNorm_ffn(h; layer_norm_eps)
h = h + down_proj(silu(gate_proj(m)) * up_proj(m))
h = RMSNorm_final(h)
logits = tok_embeddings.as_linear(h) # always tied
```
No QK norm, no bias anywhere, no softcap.
### Weight keys
```
transformer.tok_embeddings.weight [vocab, n_embd] (+ .scales/.biases when quantized)
transformer.layers.{i}.attention.norm.weight [n_embd]
transformer.layers.{i}.attention.qkv_proj.weight [(n_head + 2 n_kv) * hd, n_embd] = [896, 640]
transformer.layers.{i}.attention.out_proj.weight [n_embd, n_head * hd]
transformer.layers.{i}.feed_forward.norm.weight [n_embd]
transformer.layers.{i}.feed_forward.gate_proj.weight [mlp_hidden, n_embd]
transformer.layers.{i}.feed_forward.up_proj.weight [mlp_hidden, n_embd]
transformer.layers.{i}.feed_forward.down_proj.weight [n_embd, mlp_hidden]
transformer.norm.weight [n_embd]
```
Sanitize rules (idempotent): drop every key containing `rotary_emb` (the bf16 original stores RoPE buffers); drop `lm_head.*` if present (tied). Nothing else.
### Tokenizer and template
Byte-level BPE (`vocab.json` + `merges.txt` + `tokenizer.json`), 32768 entries. The instruct checkpoint ships `chat_template.jinja`; follow it as-is through the existing template path.
## Implementation plan
1. `src/models/lille_130m.rs` (new):
- `ModelArgs` serde struct with the table's defaults; `fn mlp_hidden(&self) -> usize` implementing the rounding formula with `f32::round` semantics matched to Python's `round` (ties-to-even; the only published size, 640, is not a tie, but add a unit test for the formula at 640 and 768).
- `Attention { norm: RMSNorm, qkv: FusedQKVLinear (from_weights_fused with prefix "...attention", n_heads, n_kv_heads, head_dim), out_proj: UnifiedLinear, rope_base, scale }`. Note the fused weight key is `qkv_proj` under the `attention` prefix, which is exactly what `FusedQKVLinear::from_weights_fused(weights, "transformer.layers.{i}.attention", ...)` expects. RoPE: `mlxcel_core::fast_rope(x, head_dim, /*traditional*/ true, rope_theta, 1.0, offset)`. Build the causal mask locally when the caller passes none (maskless-prefill contract).
- `FeedForward { norm: RMSNorm, gate_proj, up_proj, down_proj: UnifiedLinear }`.
- `Block { attention, feed_forward }` with the residual adds shown above (no block-level norms).
- `Lille130mModel { tok_embeddings: UnifiedEmbedding, layers, norm: RMSNorm }` implementing `LanguageModel`; logits via `tok_embeddings.as_linear`; `make_caches`: one `KVCache` per layer.
- `sanitize_weights`, `load(&Path)`.
2. Detection: `"`lille-130m`" => Ok(ModelType::Lille130m)` in `src/models/detection.rs` (the string has a hyphen; also accept `"lille_130m"`), test `lille_130m_model_type_is_detected` in `src/models/detection_tests.rs`.
3. Registration: `ModelType::Lille130m` in `src/models/mod.rs` (enum, text list, description `("Lille 130M (fused-QKV GQA, in-sublayer RMSNorm, traditional RoPE)", "Specialized")`), `LoadedModel::Lille130m` in `src/loaded_model.rs`, `for_each_model_registration!` row in `src/model_metadata.rs` (`kind: Text, directory: ConfigBacked`), arch-string arm in `src/distributed/tensor_parallel/inference.rs` (not TP-enabled).
4. `docs/supported-models.md` entry.
## Validation
(a) `src/models/lille_130m_tests.rs`:
- `mlp_hidden_formula`: `n_embd = 640 -> 1792`, `768 -> 2048`.
- `config_defaults`: omitting optional keys yields the table defaults; `tie_word_embeddings` defaults true.
- `sanitize_drops_rotary_and_lm_head_and_is_idempotent`.
- `qkv_split_respects_gqa`: random fused weight `[896, 640]`, assert q has 10 heads and k/v have 2.
- `tiny_model_prefill_is_causal`: 2 layers, random weights, 96-token prefill with `mask == None`; first-48 logits equal those of a 48-token prefill within 1e-4.
(b) Real checkpoint: `mlx-community/lille-130m-instruct-8bit` (about 150 MB; `mlx-community/lille-130m-instruct-bf16` and the original `Nikity/lille-130m-instruct` are the unquantized references).
```
./target/release/mlxcel download mlx-community/lille-130m-instruct-8bit
./target/release/mlxcel generate -m models/lille-130m-instruct-8bit -p "Write a haiku about the sea." -n 64 --temp 0
```
Acceptance: finite, fluent output; greedy token ids match an oracle run of the same checkpoint through at least the first 32 tokens on three prompts. Repeat once with `Nikity/lille-130m-instruct` (f32 weights, converted to f16 at the load boundary) to check the unquantized path and the `rotary_emb` drop.
## Acceptance criteria
- [ ] `src/models/lille_130m.rs` (+ `_tests.rs`) implements fused-QKV GQA, in-sublayer norms, traditional RoPE, derived MLP width and tied embeddings as specified
- [ ] `mlx-community/lille-130m-instruct-8bit` and `Nikity/lille-130m-instruct` both load and match oracle greedy ids over the deterministic prefix
- [ ] `mlxcel list` reports the family
- [ ] docs/supported-models.md updated
- [ ] detection table in src/models/detection.rs updated with a test
- [ ] cargo test --workspace --profile test-fast --features metal,accelerate passes
- [ ] cargo clippy --workspace --all-targets -- -D warnings and cargo fmt --all -- --check pass
Contributor guide
Research direction
Start with src/models/helium.rs, src/models/phi3.rs, src/models/detection.rs, and the registration files listed in the plan to understand model loading and fused QKV conventions. Add src/models/lille_130m.rs and its tests, then run the specified unit tests, checkpoint generation commands, cargo test, clippy, and fmt checks. Done means both checkpoints load, detection and listing work, docs are updated, and deterministic outputs match the oracle.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- rust
- Domain
- machine-learning
- Issue type
- Feature
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Active
- Clarity
- Clearly specified
- Newbie friendliness
- 48/100