tjdoomer/nanochat-mlx-d4-experiment

nanochat-mlx d4 — laptop pipeline smoke test

A 37M-parameter GPT trained end-to-end on a single M4 Max MacBook (128GB) using scasella/nanochat-mlx, the MLX port of Karpathy's nanochat.

This is a smoke-test artifact, not a useful model. Published to document a full pipeline run (data → tokenizer → pretrain → SFT → chat) on Apple Silicon and to expose the hardware utilisation envelope at small depths.

Sample chat output: "The capital of France is the the the the..." — mode-collapsed onto the most-frequent token. Working as expected at this scale.

What's inside

Path	Contents	Size
pretrain/model.safetensors	50-step pretrain on FineWeb-EDU	140 MB
pretrain/meta.json	Model + training metadata	<1 KB
sft/model.safetensors	256,946-step SFT on smoltalk	140 MB
sft/meta.json	SFT metadata	<1 KB
tokenizer/tokenizer.pkl	BPE tokenizer (vocab=32,768)	403 KB
tokenizer/token_bytes.npy	Token byte map	128 KB

Optimizer state (*_optim.safetensors) is not included — only inference weights.

Architecture

depth        = 4
model_dim    = 256        (depth * 64)
num_heads    = 2          (model_dim / 128)
num_kv_heads = 2          (no GQA at this scale)
vocab_size   = 32,768
seq_len      = 512   (pretrain)
seq_len      = 2048  (SFT)

Features inherited from nanochat: RoPE, QK-norm, ReLU², GQA-ready, sliding window, logit softcap, value embeddings, per-layer residual scaling.

Param breakdown at d4 (36.7M total):

• wte (token embeds): 8.4M
• value_embeds: 16.8M
• lm_head: 8.4M
• Transformer matrices: 3.1M

Embeddings are 92% of params at d4 — this flips at d20+ where the transformer dominates.

Hardware

• M4 Max MacBook Pro, 128GB unified memory (~96GB usable for GPU)
• MLX / Metal, fp16/bf16
• Compute, not memory, is the bottleneck at this depth

Utilisation (pretrain)

metric	actual	M4 Max ceiling	headroom
GPU memory	605 MB	~96 GB	~158×
MLX limit	16 GB	~110 GB	~7×
Batch size	512	4096+ @ d4	8×
Seq length	512	2048 model max	4×

At d4, bigger batch/seq just burns more electricity for the same loss — model capacity caps quality, not data throughput.

Training log

1. Data — python -m nanochat_mlx.dataset -n 8

• 8 FineWeb-EDU parquet shards, 776 MB total

2. Tokenizer — python -m scripts.tok_train

• BPE on shard 0, vocab=32,768, max_chars=2B, doc_cap=10k
• 25 seconds

3. Pretrain — python -m scripts.train --depth=4

• 50 default iters (smoke test, not a Chinchilla-sized run)
• batch=512, seq=512, grad_accum=1
• Muon + AdamW multi-optimizer (Muon on transformer matrices, AdamW on embeddings)
• Loss: 10.397 → 8.887 val
• ~17k tokens/sec, peak 1.2 GB MLX memory

Starting loss ≈ ln(vocab_size) = ln(32768) = 10.40. The model first learns to always predict "the", then later learns structure. 50 iters lands well before that transition.

4. SFT — python -m scripts.sft --depth=4

• 256,946 iters on smoltalk (full default run, ~3hr)
• Loss: 8.89 → ~3.5 val
• ~35k tokens/sec (shorter sequences than pretrain)
• 605 MB stable memory

5. Chat — python -m scripts.chat --depth=4 --source=sft

• Pipeline confirmed end-to-end working
• Output incoherent — d4 is too small to be useful, by design

Reproduce

git clone https://github.com/scasella/nanochat-mlx
cd nanochat-mlx
uv sync && source .venv/bin/activate

python -m nanochat_mlx.dataset -n 8
python -m scripts.tok_train
python -m scripts.train --depth=4
python -m scripts.sft --depth=4
python -m scripts.chat --depth=4 --source=sft --interactive

Use these weights

These checkpoints work directly with scasella/nanochat-mlx. Drop into ~/.cache/nanochat/mlx_checkpoints/d4/ and ~/.cache/nanochat/mlx_checkpoints/d4_sft/ with original filenames, then run python -m scripts.chat --depth=4 --source=sft.

Concepts cheat sheet

• BPE tokenizer: greedy merging of frequent byte pairs. ~4 chars/token English. Training your own = vocab matched to data.
• Chinchilla ratio: optimal params:tokens ≈ 1:20. nanochat uses 1:12. For 36.7M params → ~440M tokens for proper training.
• Mode collapse at undertraining: model first learns most-frequent token before learning structure. Garbage at 50 steps is healthy, not bug.
• Muon optimizer: used by Kimi K2. Faster convergence on transformer matrices; embeddings still prefer AdamW.
• Apple Unified Memory: no PCIe copy step between CPU/GPU. MLX lazy eval → Metal fused kernels on demand.

Gotchas

• d4 pretrain defaults to 50 iters (smoke only). Use --num-iterations=2000+ for marginally coherent d4.
• SFT default iter count is ~257k (full smoltalk pass). Pass --num-iterations=500 for a smoke test.
• MLX defaults to 16 GB GPU memory cap. Raise via MLX_METAL_MEMORY_LIMIT_MB for d26+.
• SFT progress display has a denominator-display bug — percentage is correct, the /00000 part is cosmetic.

What's next

Skipping d12 (small quality jump over d4 not worth a day). Pivoting from nanochat scale-up to RL'ing a small Gemma model instead.

Citation

@misc{nanochat-mlx-d4-experiment-2026,
  author = {Worrall, Tj},
  title  = {nanochat-mlx d4 laptop smoke test},
  year   = {2026},
  url    = {https://huggingface.co/tjdoomer/nanochat-mlx-d4-experiment}
}

Upstream:

• Karpathy nanochat: https://github.com/karpathy/nanochat
• MLX port: https://github.com/scasella/nanochat-mlx