looped-nanochat-d24

looped-nanochat-d24

A two‑pass, shared‑weight routing‑loop Transformer built on karpathy/nanochat @ 92d63d4, trained at the d24 (~1.38 B param) scale. This is the looped model from a controlled comparison against the stock nanochat baseline (kz919/baseline-nanochat-d24).

Code: https://github.com/kyleliang919/looped_nanochat

What is the routing loop?

The Transformer trunk is run twice per forward:

1. Pass 1 runs the ordinary causal Transformer and captures the normalized, RoPE‑applied Q/K features from selected deeper layers.

2. Pass 2 re‑runs the trunk from the token embeddings. Each pass‑2 layer adds a causal attention‑logit prior from one pass‑1 layer:

   attn_logits = q2·k2ᵀ/√D + g · q1·k1ᵀ/√D
   

   with a per‑layer learned gate g = softplus(gate_logit). The additive bias is computed via augmented‑width Q/K + ordinary attention (no N×N matrix materialized), and the training path runs it through a fused, compiled FlexAttention kernel. Everything is causal and differentiable.

Extra parameters over baseline: routing_gate_logits (per layer) and pass_embeddings (2×n_embd).

Results (vs. baseline, identical training budget)

The looped model wins or ties every metric, at ~2.4× the training compute (inherent to the two passes). See the full report.

Model details

Training: pretrained for 5,568 steps on ~5.84 B tokens (--target-param-data-ratio=8, 1,048,576‑token batches), then chat‑SFT on SmolTalk + MMLU(×3) + GSM8K(×4). This checkpoint is the SFT model (step 467).

Files

• model_000467.pt — model weights (state_dict)
• meta_000467.json — model_config (incl. routing fields) + training config
• tokenizer.pkl, token_bytes.pt — nanochat tokenizer

Usage

This uses the nanochat runtime (not 🤗 Transformers). Clone the repo, drop the files into ~/.cache/nanochat/chatsft_checkpoints/looped-d24/ (weights/meta) and ~/.cache/nanochat/tokenizer/, then:

python -m scripts.chat_cli -g looped-d24 -i sft -p "Why is the sky blue?"

The routing config is stored in meta_000467.json, so the runtime rebuilds the two‑pass model automatically.

Limitations

A ~$100‑tier micro‑model: it hallucinates, is weak at math/code, and is meant for research on architecture, not production use.

License

MIT (inherited from nanochat).