gelu-d12-chinchilla-261M

GELU d=12 Chinchilla (261M)

A 261M-parameter nanochat-architecture GPT with GELU activation in the MLP. Trained on English C4 to Chinchilla-optimal token budget (20× params ≈ 5.22 B tokens) on a single TPU v6e-8.

This repository hosts the seed 0 checkpoint (model + optimizer state, resumable). Training was also run for seeds 1 and 2 under identical compute to quantify seed variance — summary below.

Final result — seed 0 (this checkpoint)

3-seed variance (same arch, same data, same compute)

Architecture

Full nanochat stack with GELU MLP — everything else identical to the default nanochat GPT:

• RoPE (base 100K), no learned positional embeddings
• MHA (n_head = n_kv_head = 12, head_dim = 64; GQA-capable but all d=12 models use full MHA)
• QK norm with 1.2× scaling
• Sliding window attention with "SSSL" pattern
• Tied embeddings (wte ↔ lm_head)
• Parameterless RMSNorm post-embedding and per block
• Value embeddings (ResFormer-style, alternating layers)
• Per-layer learnable residual scalars
• Smear (learnable bigram gate on first 24 dims of token embedding)
• Backout (subtract mid-layer residual)
• Logit soft-capping via tanh(x/15)·15
• No biases in any Linear (attention Q/K/V/proj, MLP fc/proj)
• MLP FFN: Linear(n → 4n) → jax.nn.gelu → Linear(4n → n)
  • Applied via the --mlp gelu flag in scripts/train_d12_chinchilla.py which monkey-patches MLP.__call__ in flaxchat.gpt

Training

The optimizer state is preserved → checkpoint is resumable for further training.

Contents

This repository contains the full Orbax checkpoint (model weights + AdamW optimizer state), the frozen config, and code snapshots sufficient to load the model:

.
├── 19920/                          # final Orbax checkpoint (~2.86 GB total)
│   ├── _CHECKPOINT_METADATA
│   ├── metadata/                   # JSON: {"final_loss": 3.07, "smooth": 3.11}
│   ├── model/                      # nnx.Param state — architecture weights
│   └── optimizer/                  # optax AdamW state (Adam m/v moments + step + LR state)
├── config.json
├── README.md
└── code/                           # snapshots of flaxchat code (source of truth: github)
    ├── gpt.py
    ├── checkpoint.py
    ├── config.py
    ├── train_d12_chinchilla.py     # the exact training script (use --mlp gelu)
    └── load_model.py

Loading

git clone https://github.com/mlnomadpy/flaxchat && cd flaxchat
pixi install   # or pip install with the deps in code/load_model.py
huggingface-cli download mlnomad/gelu-d12-chinchilla-261M --local-dir ./hf_model

# Inside your flaxchat clone — adapt code/load_model.py for GELU:
from flax import nnx
import jax, jax.numpy as jnp
from flaxchat.gpt import GPT, GPTConfig, MLP
from flaxchat.checkpoint import restore_model_from_checkpoint

# Patch MLP to use GELU instead of ReLU² (matches training)
_orig = MLP.__call__
def _gelu_call(self, x):
    x = self.c_fc(x)
    x = jax.nn.gelu(x)
    x = self.c_proj(x)
    return x
MLP.__call__ = _gelu_call

config = GPTConfig(
    sequence_len=1024, vocab_size=32768,
    n_layer=12, n_head=12, n_kv_head=12, n_embd=768,
    window_pattern="SSSL", tie_embeddings=True,
)
model = GPT(config, rngs=nnx.Rngs(0))
meta = restore_model_from_checkpoint(model, "./hf_model/19920")
print(f"Loaded — final loss: {meta['final_loss']:.4f}")  # expect 3.0745

Weights and optimizer state

Both the model parameters and the full AdamW optimizer state (Adam m/v moments, step counter, LR schedule state) are stored at:

• HuggingFace Hub: mlnomad/gelu-d12-chinchilla-261M
  • 19920/model/ — nnx.Param state
  • 19920/optimizer/ — optax AdamW state (resumable)

The checkpoint format is Orbax (OCDBT PyTree). flaxchat.checkpoint.restore_model_from_checkpoint(model, ckpt_path, optimizer=optimizer) restores both in one call.

License

Apache 2.0.