lablup / lablup/mlxcel

feat(models): add Lille 130M (lille-130m) fused-QKV GQA text model with in-sublayer norms

Open
#1,363 0 comments 0 reactions 0 assignees View on GitHub
arch:dense area:docs area:inference area:models modelsize:small modeltype:text priority:low status:ready type:enhancement
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

Open the contributing 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.