nanogpt-mlx: Training a GPT-2 Class Model from Scratch on Apple Silicon

I built nanogpt-mlx, a port of Andrej Karpathy's nanoGPT to Apple's MLX framework. It trains a GPT-2 class model (173M parameters, similar to GPT-2 Small) from scratch on Apple Silicon in about 2.5 days.

The interesting parts: adapting PyTorch patterns to MLX's functional style, and integrating modern architectural improvements that weren't in the original nanoGPT.

What Changed from PyTorch to MLX

MLX is Apple's answer to PyTorch for their M-series chips. Similar API, but different enough to require real porting work.

Automatic Differentiation

PyTorch uses loss.backward() then optimizer.step(). MLX uses a functional pattern with value_and_grad:

# PyTorch
loss = model(x, y)
loss.backward()
optimizer.step()

# MLX
def lossfn(model, x, y):
    _, loss = model(x, y)
    return loss

loss_and_grad_fn = nn.value_and_grad(model, lossfn)
loss, grads = loss_and_grad_fn(model, x, y)
optimizer.update(model, grads)
mx.eval(model.parameters(), optimizer.state)

The mx.eval() call is crucial. MLX uses lazy evaluation, so computations don't actually run until you evaluate. This enables graph optimization but trips you up if you forget it.

Parameter Counting

PyTorch's model.parameters() returns an iterable of tensors. MLX returns a nested dict/list structure. Had to write a recursive counter:

def count_params(tree):
    total = 0
    if isinstance(tree, dict):
        for v in tree.values():
            total += count_params(v)
    elif isinstance(tree, list):
        for v in tree:
            total += count_params(v)
    elif isinstance(tree, mx.array):
        total += tree.size
    return total

Sequential Modules

PyTorch's nn.Sequential doesn't translate directly. I switched to a plain list and explicit iteration:

# Instead of nn.Sequential
self.blocks = [Block(...) for _ in range(layers)]

def __call__(self, x):
    for block in self.blocks:
        x = block(x)
    return x

Architectural Improvements

The original nanoGPT uses the GPT-2 architecture faithfully. I added several modern improvements from recent LLM research:

RoPE (Rotary Position Embeddings)

Replaced learned positional embeddings with RoPE. Instead of adding a position vector to token embeddings, RoPE rotates the query and key vectors based on position. Better extrapolation to longer sequences:

self.rope = nn.RoPE(head_size, traditional=False)

def __call__(self, x):
    q = self.query(x)
    k = self.key(x)
    q = self.rope(q)
    k = self.rope(k)

RMSNorm

Replaced LayerNorm with RMSNorm. Simpler (no mean subtraction), faster, and no learnable parameters:

def rms_norm(x, eps=1e-5):
    rms = mx.sqrt(mx.mean(x * x, axis=-1, keepdims=True) + eps)
    return x / rms

ReLU² Activation

Squared ReLU in the feed-forward blocks. Better gradient flow than plain ReLU:

x = mx.maximum(x, 0)  # ReLU
x = x * x             # Square it

QK Normalization

Normalize queries and keys after RoPE. Stabilizes attention scores during training:

q = self.rope(q)
k = self.rope(k)
q = rms_norm(q)
k = rms_norm(k)

Zero-Init Residual

Initialize output projections to zero so residual branches start as identity. Helps training stability in deep networks:

def _init_weights(self):
    self.lm_head.weight = mx.zeros(self.lm_head.weight.shape)
    for block in self.blocks:
        block.sa.projection.weight = mx.zeros(...)
        block.feed_forward.fc2.weight = mx.zeros(...)

Training Pipeline

Three distinct stages, each with different objectives:

Stage 1: Base Pretraining

Train on raw text to learn language patterns. 50k iterations on FineWeb-Edu dataset (high-quality educational web pages).

  • Data: 1 billion tokens from ~50 parquet shards
  • Batch: 32 sequences × 512 tokens = 16,384 tokens/step
  • Learning rate: 3e-4
  • Time: ~52 hours

Loss drops from 10.3 to 3.5. The model learns grammar, sentence structure, and some factual knowledge. But it has no concept of instruction-following—it just predicts the next token.

Stage 2: Instruction Fine-Tuning

Train on conversation examples to learn instruction-following behavior. 20k iterations on SmolTalk dataset (460k conversations).

The key is masked loss—only compute loss on assistant responses:

def format_conversation(messages, tokenizer):
    tokens = []
    mask = []

    for msg in messages:
        if role == "user":
            user_tokens = tokenizer.encode(f"\n\nUser: {content}\n\nAssistant:")
            tokens.extend(user_tokens)
            mask.extend([0] * len(user_tokens))  # Don't train
        elif role == "assistant":
            asst_tokens = tokenizer.encode(f" {content}")
            tokens.extend(asst_tokens)
            mask.extend([1] * len(asst_tokens))  # Train on this

    return tokens, mask
  • Learning rate: 1e-5 (10x lower to preserve base knowledge)
  • Time: ~9 hours

Loss drops from ~2.5 to 1.8. The model learns to respond to questions rather than just continuing text.

Stage 3: Testing

The base model (after stage 1) produces grammatically correct but semantically drifting text:

"The future of artificial intelligence is the way that we've witnessed various types of DNA sequencing, specifically at Morton Stanley..."

The fine-tuned model (after stage 2) attempts to answer questions:

"Explain photosynthesis refers to the process of turning carbon dioxide into glucose for use in cells. In plants, this process is known as primary photosynthesis..."

Better, but still hallucinates ("Dr. Suzuki at Wenwei University") and loses coherence. The fundamental limitation: 1B tokens isn't enough base knowledge, and 173M parameters isn't enough capacity.

Tokenization

BPE tokenizer trained on FineWeb-Edu using HuggingFace's tokenizers library:

GPT_SPLIT_PATTERN = r"""'(?i:[sdmt]|ll|ve|re)|[^\r\n\p{L}\p{N}]?+\p{L}+|\p{N}{1,2}| ?[^\s\p{L}\p{N}]++[\r\n]*|\s*[\r\n]|\s+(?!\S)|\s+"""

32,768 token vocabulary. Achieves 2.67x compression vs character-level. Training on 200k documents takes ~10 minutes.

Performance

On M-series Macs:

  • Step time: ~3.7 seconds (16k tokens/step)
  • Throughput: ~4,400 tokens/second
  • Memory: ~8GB unified memory
  • Total training: ~62 hours (2.5 days)

MLX's unified memory architecture means no CPU-GPU transfer overhead. The model and data share the same memory pool.

Future Work

The 2.5-day run is educational but produces mediocre results. The path to better quality:

  1. More pretraining: 7-10B tokens instead of 1B. Industry standard for GPT-2 Small is 10B+.
  2. Better instruction data: OpenHermes-2.5 instead of SmolTalk. Higher quality conversations with better reasoning.
  3. Multi-phase fine-tuning: General instruction tuning, then specialized (reasoning, code, etc).
  4. Scale up: GPT-2 Medium (355M) or Large (774M) if you have weeks to spare.

I'm running a 7B token experiment now. Two weeks for base pretraining, then higher-quality instruction tuning. Should produce significantly better results.

Try It

git clone https://github.com/dzoba/nanogpt-mlx
cd nanogpt-mlx
chmod +x speedrun.sh
./speedrun.sh

Come back in 2.5 days. Or run steps individually if you want to inspect intermediate results.

The speedrun downloads data, trains the tokenizer, runs base pretraining, and does instruction fine-tuning. Everything from scratch, no pretrained weights.