lablup / lablup/mlxcel

feat(models): add nanochat (nanochat) weightless-RMSNorm, relu-squared, softcapped text model

Open
#1,368 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: "`nanochat`"` (the `nanochat` d20/d32 speedrun models, 0.56B to 1.9B parameters) is not detected by mlxcel. The architecture is a GPT-2-shaped decoder with four non-standard choices: RMSNorm without learnable weight everywhere (including on q and k after RoPE, and on the embedding output), a `relu(x)^2` MLP, a mirrored (negative-angle) RoPE, and a `15 * tanh(logits / 15)` softcap on the output logits. Every one of those yields fluent but wrong text if skipped. This issue adds `src/models/`nanochat`.rs`, detection, registration and docs.

## Current behavior

`"`nanochat`"` hits the `_ => Err("Unsupported model type")` arm at `src/models/detection.rs:379`; no file under `src/models/` references it. The nearest existing code is `src/models/gpt2.rs` (c_fc / c_proj naming, `transformer.h.{i}` prefixes) and `src/models/bitnet.rs:281` (a `relu2` MLP). The primitives are already in `mlxcel-core`: `fast_rms_norm_no_weight` (`src/lib/mlxcel-core/src/lib.rs:1245`), `relu_squared` (`src/lib/mlxcel-core/src/utils.rs:846`), `fast_rope_with_freqs` (`src/lib/mlxcel-core/src/lib.rs:1174`), and the compiled softcap used by `src/models/gemma2.rs:27`.

## Expected behavior

`mlxcel generate -m models/nanochat-d20-q8-mlx -p "..."` loads and generates; greedy ids match an oracle run of the same checkpoint.

### Config (`config.json`)

| key | type | default | d20 value |
|---|---|---|---|
| `model_type` | str | required | `"`nanochat`"` |
| `hidden_size` | int | 1280 | 1280 (d32: 2048) |
| `num_hidden_layers` | int | 20 | 20 (d32: 32) |
| `num_attention_heads` | int | 10 | 10 |
| `num_key_value_heads` | int | 10 | 10 |
| `vocab_size` | int | 65536 | 65536 |
| `max_position_embeddings` | int | 2048 | 2048 |
| `intermediate_size` | int | 5120 (4 * hidden) | 5120 |
| `rope_theta` | float | 10000.0 | 10000.0 |
| `rms_norm_eps` | float | 1e-5 | absent in the MLX conversions; 1e-6 in transformers-format configs |
| `logits_soft_cap` (alias `logits_softcap`) | float or null | 15.0 | 15.0 |
| `tie_word_embeddings` | bool | false | false (separate `lm_head.weight`) |
| `bos_token_id` / `eos_token_id` | int | from tokenizer config | `<|bos|>` = 65527, `<|assistant_end|>` = 65531 |

`head_dim = hidden_size / num_attention_heads` (128). `attention_bias` is always false; no bias anywhere.

### Forward pass

```
norm(x) = x * rsqrt(mean(x^2, -1) + eps) # no weight, eps = rms_norm_eps
h = norm(wte(tokens)) # embedding output IS normed
for each block:
a = norm(h)
q = c_q(a) -> [B, L, H, 128]; k = c_k(a) -> [B, L, Hkv, 128]; v = c_v(a)
q = rope(q, offset); k = rope(k, offset) # rotate FIRST
q = norm(q); k = norm(k) # then weightless RMSNorm per head (QK-norm AFTER RoPE)
h = h + c_proj(sdpa(q, k, v, scale = 128 ** -0.5, causal))
h = h + mlp.c_proj(relu(mlp.c_fc(norm(h))) ** 2)
h = norm(h) # final weightless norm
logits = lm_head(h)
logits = cap * tanh(logits / cap), cap = logits_soft_cap (15.0); skip when null
```

RoPE: the rotation is the mirror image of the usual one. With `half = head_dim / 2`, per-pair frequency `f_i = base ** (i / half)` for `i in 0..half`, and the angle for position `p` and pair `i` is `-p / f_i`. In `mlxcel_core` terms this is exactly

```
freqs = -(base ** (arange(half) / half)) # f32, shape [half], NEGATIVE
q = fast_rope_with_freqs(q, head_dim, /*traditional*/ false, /*scale*/ 1.0, offset, &freqs)
```

(`fast_rope_with_freqs` divides the position by each entry of `freqs`, so negating the table negates the angle; `traditional = false` pairs feature `i` with `i + half`, which is the required half-split layout.) Using the ordinary `fast_rope(..., base)` gives the opposite rotation direction and fluent but wrong text.

### Weight keys (MLX conversions, the target layout)

```
transformer.wte.weight [vocab, hidden]
transformer.h.{i}.attn.c_q.weight [H * 128, hidden]
transformer.h.{i}.attn.c_k.weight [Hkv * 128, hidden]
transformer.h.{i}.attn.c_v.weight [Hkv * 128, hidden]
transformer.h.{i}.attn.c_proj.weight [hidden, H * 128]
transformer.h.{i}.mlp.c_fc.weight [intermediate, hidden]
transformer.h.{i}.mlp.c_proj.weight [hidden, intermediate]
lm_head.weight [vocab, hidden]
```

plus `.scales` / `.biases` beside every `.weight` in quantized exports (including `transformer.wte` and `lm_head`; use `UnifiedEmbedding` / `UnifiedLinear`). There are no norm weights at all.

Sanitize rules (idempotent):

1. transformers-format exports use a second layout; map it onto the first: `model.embed_tokens.weight -> transformer.wte.weight`, `model.layers.{i}.self_attn.{q,k,v,o}_proj -> transformer.h.{i}.attn.{c_q,c_k,c_v,c_proj}`, `model.layers.{i}.mlp.fc1 -> transformer.h.{i}.mlp.c_fc`, `model.layers.{i}.mlp.fc2 -> transformer.h.{i}.mlp.c_proj`, `lm_head.weight` unchanged.
2. Drop any `*rotary_emb.inv_freq*` key.
3. No fusion or transposition is needed.

### Tokenizer and template

BPE `tokenizer.json` with 65536 entries; special tokens `<|bos|>` (65527), `<|user_start|>` (65528), `<|user_end|>` (65529), `<|assistant_start|>` (65530), `<|assistant_end|>` (65531), `<|python_start|>`, `<|python_end|>`, `<|output_start|>`, `<|output_end|>`. The chat template (present in `madmag77/nanochat-d20-q8-mlx`, absent in `dnakov/nanochat-d20-mlx`) starts with `{{- bos_token }}` and wraps turns in the start/end tokens; a system message is folded into the first user turn. Stop on `<|assistant_end|>`. For a raw `-p` prompt the BOS token must be prepended (the model was never trained without it); follow whatever `tokenizer_config.json` declares via the existing tokenizer path, and document that `--no-chat-template` prompts should start with `<|bos|>`.

## Implementation plan

1. `src/models/`nanochat`.rs` (new):
- `ModelArgs` serde struct with the defaults in the table (`#[serde(alias = "logits_softcap")]` on `logits_soft_cap`, `Option`).
- `Attention { c_q, c_k, c_v: UnifiedLinear (or FusedQKVLinear::from_weights_separate over the three), c_proj: UnifiedLinear, freqs: UniquePtr, head_dim, scale }`. `forward`: project, reshape to `[B, H, L, D]`, `fast_rope_with_freqs` with the negative table, `fast_rms_norm_no_weight(q, eps)` / `(k, eps)`, `KVCache::update_and_fetch`, fast SDPA with the causal mask built locally when the caller passes none (maskless-prefill contract), `c_proj`.
- `MLP { c_fc, c_proj }`: `c_proj(relu_squared(c_fc(x)))`.
- `TransformerBlock` with no norm weights; call `fast_rms_norm_no_weight` inline.
- `NanoChatModel { wte: UnifiedEmbedding, h: Vec, lm_head: UnifiedLinear, softcap: Option }` implementing `LanguageModel` (`make_caches`: one `KVCache` per layer). Apply the softcap with the compiled softcap helper used by `gemma2.rs` when `softcap.is_some()`.
- `sanitize_weights` and `load(&Path)`.
2. Detection: `"`nanochat`" => Ok(ModelType::NanoChat)` in `src/models/detection.rs` plus `nanochat_model_type_is_detected` in `src/models/detection_tests.rs`.
3. Registration: `ModelType::NanoChat` in `src/models/mod.rs` (enum, text list, description `("`nanochat` (weightless RMSNorm, relu^2 MLP, softcapped logits)", "Specialized")`), `LoadedModel::NanoChat` in `src/loaded_model.rs`, the `for_each_model_registration!` row in `src/model_metadata.rs` (`kind: Text, directory: ConfigBacked`), and the arch-string arm in `src/distributed/tensor_parallel/inference.rs` (not TP-enabled).
4. `docs/supported-models.md` entry naming the validation checkpoints and the BOS requirement.

## Validation

(a) `src/models/nanochat_tests.rs`:

- `config_defaults_match_d20` (omit every optional key; assert 1280 / 20 / 10 / 5120 / 15.0 / 1e-5).
- `sanitize_maps_transformers_layout_and_is_idempotent`.
- `negative_freq_rope_rotates_backwards`: for a single pair at position 1 with base 10000, assert `rope(q)[0] == q0*cos(1) + q1*sin(1)` and `rope(q)[half] == -q0*sin(1) + q1*cos(1)` (the sign-flipped rotation), within 1e-5.
- `qk_norm_runs_after_rope`: build a 1-head attention with a known q; assert the normed-after-rope value differs from normed-before-rope for a non-zero offset (guards the ordering).
- `softcap_bounds_logits`: random logits of magnitude 1e3 come out in `(-15, 15)`; with `logits_soft_cap = null` they pass through.
- `tiny_model_prefill_is_causal`: 2 layers, random weights, 96-token prefill with `mask == None`; first-48-position logits equal those of a 48-token prefill within 1e-4.

(b) Real checkpoints: `madmag77/nanochat-d20-q8-mlx` (8-bit, about 0.6 GB, ships a chat template and `eos_token_id` 65531) and `dnakov/nanochat-d20-mlx` (bf16, 1.1 GB, no chat template).

```
./target/release/mlxcel download madmag77/nanochat-d20-q8-mlx
./target/release/mlxcel generate -m models/nanochat-d20-q8-mlx -p "Explain what a transformer is in two sentences." -n 64 --temp 0
```

Acceptance: output is finite and fluent English, stops at `<|assistant_end|>`, and the greedy token ids match an oracle run of the same checkpoint through at least the first 32 tokens on three prompts. Repeat with `dnakov/nanochat-d20-mlx` using `--no-chat-template` and a prompt that begins with `<|bos|>`.

## Acceptance criteria

- [ ] `src/models/`nanochat`.rs` (+ `_tests.rs`) implements weightless RMSNorm (embedding, pre-attn, pre-MLP, final, and post-RoPE q/k), negative-angle RoPE, relu-squared MLP and logit softcap as specified
- [ ] both weight layouts (`transformer.*` and `model.layers.*`) load through `sanitize_weights`
- [ ] `madmag77/nanochat-d20-q8-mlx` generates fluently and matches 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

## Out of scope

- Converting the original `.pt` training checkpoints (`karpathy/nanochat-d32` ships only `model_000650.pt` and a pickled tokenizer); only safetensors exports are targeted.
- The `<|python_start|>` tool-call loop.

Contributor guide

Open the contributing guide

Research direction

Start with src/models/gpt2.rs, src/models/bitnet.rs, the listed mlxcel-core primitives, and the detection and registration files named in the plan. Run the specified nanochat detection, sanitization, RoPE, causal-prefill, and softcap tests as they are added, then validate both checkpoints with the provided deterministic generation commands. Done means the model is listed, both weight layouts load, and greedy outputs match the oracle while workspace checks pass.

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
35/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.