Deep Dive into nanochat: Source Code, Model Size Calculations, and Optimization Techniques

Model Size and Parameter Calculations

The nanochat d20 configuration uses vocab_size=65536, n_embd=1280, n_layer=20, and n_head=10. Using the formula Size = VE + N*12E (where V is vocabulary size and E is embedding dimension) yields 560,988,160 ≈ 561M parameters, matching the reported size.

For a Transformer layer the generic size equation is

Size = EmbeddingLayer + Layers*(MultiheadAttention + FeedForwardLayer + LayerNormalization) + FinalLayerNormalization

. Assuming H = 4E and ignoring position‑encoding and norm layers simplifies to Size ≈ VE + N*12E.

GPT vs. Llama Decoder Differences

Both GPT‑style and Llama‑style Transformers share core blocks, but GPT includes explicit position encoding while Llama does not. Llama’s MLP uses three parameter groups (up, gate, down) with a hidden‑to‑embedding ratio of 2.66–3.5, compared to GPT’s two‑group (fc and proj) design where H = 4E.

Key Architectural Features

nanochat adopts bias‑free linear layers, parameter‑free RMSNorm, FlashNorm (2025) for efficient normalization, and the ReLU² activation (Tsinghua 2024) which provides high sparsity with minimal performance loss. Weight tying is deliberately omitted, keeping the embedding and LM head separate.

Memory Estimation

Training memory consists of model parameters (2N bytes for FP16), gradients (2N), optimizer states (8N for Adam/AdamW), activations (20‑40% of total), and temporary buffers (10‑20%). Approximate peak memory is ≈ 12N * 1.2 * 1.1 ≈ 16‑20N bytes, i.e., 3‑4 × the parameter count for a 561 M model.

Inference memory adds KV‑cache storage:

2 × batch_size × seq_len × num_layers × hidden_size × 2 bytes

, growing linearly with generated token length. FP16 inference typically uses 2.4‑2.8N bytes, while INT4 quantization can reduce it to 0.7‑0.9N bytes.

Scaling Laws

The article reviews Kaplan (2020), Chinchilla (2022), and Inference‑Time Scaling (2024) laws, emphasizing that model size matters more than data volume when compute budget is fixed, and presenting the power‑law relationships L = f(N, D, S, C).

Optimizer Innovations

nanochat uses a hybrid optimizer setup: embedding and LM‑head parameters are updated with AdamW, while decoder parameters use the Muon optimizer, which inserts a Newton‑Schulz orthogonalization step between momentum and parameter update. Muon’s 5‑step Newton‑Schulz iteration approximates matrix inverse‑square‑root efficiently.

Distributed Training Architecture

Training runs via torchrun with NCCL backend. Parameter groups are split: group 1 (embedding) + group 3 (LM head) → AdamW; group 2 (decoder) → Muon. Gradient and optimizer‑state sharding follow ZeRO‑1 and ZeRO‑2 patterns, using reduce_scatter and all_gather for all‑reduce. Both parameter‑server (master‑slave) and ring‑based all‑reduce topologies are illustrated.

All steps—from model initialization, data loading, scheduler setup, periodic evaluation (BPB, CORE, sample generation), to gradient update and checkpointing—are visualized with flow diagrams.

Training Stages

Four stages are described: pre‑training (next‑token prediction), mid‑training (dialogue format adaptation), supervised fine‑tuning (SFT), and reinforcement learning (GRPO/PPO). Each stage reuses the same core training loop but varies data, loss, and evaluation metrics.

Finally, the article lists reference links to papers and blogs on optimizers, scaling laws, and LLM memory calculations.