GELU d=12 Chinchilla (261M) — PyTorch / HuggingFace Transformers

A 261M-parameter nanochat-architecture GPT with GELU MLP, originally trained in JAX/Flax on a TPU v6e-8 and ported to PyTorch for easy inference via the HuggingFace transformers API.

Weights are bit-exact with the Flax checkpoint (mlnomad/gelu-d12-chinchilla-261M) — parity validated at max |Δ logits| = 1.3e-5 on CPU/fp32.

Quick start

pip install torch transformers safetensors

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

model = AutoModelForCausalLM.from_pretrained(
 "mlnomad/gelu-d12-chinchilla-261M-pytorch",
 trust_remote_code=True, # the model class ships in the repo
 dtype=torch.float32,
).eval()

# The model was trained with the Mistral-7B-v0.1 tokenizer (vocab 32,768)
tokenizer = AutoTokenizer.from_pretrained("mistralai/Mistral-7B-v0.1")

prompt = "The meaning of life is"
ids = tokenizer(prompt, return_tensors="pt").input_ids

with torch.no_grad():
 out = model.generate(
 ids,
 max_new_tokens=50,
 do_sample=True,
 temperature=0.8,
 top_p=0.9,
 use_cache=True,
 pad_token_id=tokenizer.eos_token_id or 0,
 )
print(tokenizer.decode(out[0], skip_special_tokens=True))

Greedy completion examples from this checkpoint:

"The meaning of life is" → to live in a place where you can live. The meaning of life is to live in a

"Once upon a time," → the first thing I would do was to make a list of all the things I would like to do

Model details

Architecture features

Full nanochat stack, faithfully ported to PyTorch:

• RoPE (base 100,000), split-half layout
• MHA (n_head = n_kv_head = 12; the code supports GQA via n_kv_head < n_head, but all d=12 models use full MHA)
• QK-norm with 1.2× scaling (after RoPE)
• Parameterless RMSNorm (no learnable gain) post-embedding and per block
• Sliding-window attention with "SSSL" pattern
• Tied embeddings (lm_head = wte.T)
• Value embeddings on alternating layers (ResFormer-style)
• Per-layer learnable residual scalars (resid_lambdas, x0_lambdas)
• Smear — learnable gate on first 24 dims of token embedding mixes in prev token
• Backout — subtract mid-layer residual from late layers
• Logit soft-cap: 15 · tanh(logits / 15)
• No biases in any Linear
• MLP: Linear(n → 4n) → F.gelu(approximate="tanh") → Linear(4n → n)

KV cache

The GeluGPTForCausalLM class implements a smear-aware KV cache for fast autoregressive generation. KV-cache parity vs full forward is validated at max |Δ| < 3e-5. Pass use_cache=True (the default for .generate()).

Files in this repo

.
├── config.json # HF config with auto_map pointing to the classes below
├── generation_config.json
├── model.safetensors # 1.04 GB, fp32 weights + persistent RoPE buffers
├── torch_gpt.py # pure PyTorch GPT module (GELU_GPT)
├── configuration_gelu_gpt.py # PretrainedConfig subclass
├── modeling_gelu_gpt.py # PreTrainedModel + GenerationMixin wrapper with KV cache
└── README.md

Related

• mlnomad/gelu-d12-chinchilla-261M — original JAX/Flax Orbax checkpoint (model + AdamW optimizer state, resumable)
• flaxchat — the JAX/Flax training harness that produced the weights

Wikitext-103 evaluation

Evaluated on ~330K tokens from wikitext-103 test set (model trained on C4 only — this is a zero-shot transfer metric).

License

Apache 2.0.

Base model

mlnomad/gelu-d12-chinchilla-261M

(1)