# Modern Architectural Choices in NanoChat's GPT Model: Efficiency Meets State-of-the-Art

> Explore modern architectural choices like Rotary Embeddings, GQA, and FlashAttention-3 in NanoChat's GPT model. Achieve efficient inference on commodity GPUs.

- Repository: [Andrej/nanochat](https://github.com/karpathy/nanochat)
- Tags: architecture
- Published: 2026-03-10

---

**NanoChat's GPT model combines Rotary Positional Embeddings, Group-Query Attention (GQA), FlashAttention-3, RMSNorm, sliding-window attention, and per-layer residual scaling to deliver a memory-efficient, training-stable transformer that runs inference on commodity GPUs.**

NanoChat, developed by Andrej Karpathy, implements a compact yet powerful GPT architecture that incorporates several cutting-edge design decisions from modern large language model research. This article examines the specific modern architectural choices used in nanochat's GPT model, tracing each implementation detail directly to its source code in the `karpathy/nanochat` repository.

## Efficient Attention Mechanisms

### Group-Query Attention (GQA)

**Group-Query Attention** reduces memory bandwidth during inference by sharing key and value heads across multiple query heads. In [`nanochat/gpt.py`](https://github.com/karpathy/nanochat/blob/main/nanochat/gpt.py), the `CausalSelfAttention` class defines separate counts for `self.n_head` (query heads) and `self.n_kv_head` (key/value heads), with linear projections `c_q`, `c_k`, and `c_v` implemented in the constructor (lines 69-77). This design maintains query expressiveness while decreasing memory footprint, making it ideal for inference on consumer GPUs.

### FlashAttention-3 Integration

The model leverages **FlashAttention-3** (FA3) through direct calls to `flash_attn.flash_attn_func` and `flash_attn.flash_attn_with_kvcache` within the `CausalSelfAttention.forward` method (lines 104-118). This integration yields greater than 2× speed-up over standard attention implementations while maintaining O(n) memory complexity, particularly effective on Hopper-generation GPUs.

### Sliding-Window Attention Pattern

To support long contexts with constant per-token memory, **sliding-window attention** restricts most layers to attending only to a limited left window, while deeper layers see progressively larger contexts. The `_compute_window_sizes` function builds per-layer window tuples from a `window_pattern` configuration (lines 70-96), enabling efficient processing of extended sequences.

### Rotary Positional Embeddings (RoPE)

Rather than using absolute positional embeddings, the model implements **Rotary Positional Embeddings** via the `apply_rotary_emb` function in [`nanochat/gpt.py`](https://github.com/karpathy/nanochat/blob/main/nanochat/gpt.py). This approach pre-computes `cos` and `sin` buffers (lines 57-64) to provide relative position information with negligible memory cost, working naturally with cached key-value stores during inference (lines 92-95).

## Normalization and Stability Improvements

### RMSNorm for Layer Normalization

The architecture replaces traditional LayerNorm with **RMSNorm** (root-mean-square normalization), implemented in the `norm` function using `F.rms_norm` (lines 42-44). This variant removes learnable scale parameters, operates faster in mixed-precision training, and maintains training stability without the bias terms found in standard LayerNorm.

### Q-K Normalization

Following the application of rotary embeddings, the model applies **Q-K normalization** by explicitly normalizing query and key tensors before computing attention scores: `q, k = norm(q), norm(k)` (lines 100-101). This technique stabilizes attention scores when operating with lower-precision dtypes like bfloat16 or float16.

### Per-Layer Residual and Embedding-Skip Scaling

Derived from "modded-nanogpt" research, the implementation introduces **per-layer residual scaling** (`resid_lambdas`) and **embedding-skip scaling** (`x0_lambdas`). These parameters, defined in `GPT.__init__` (lines 76-80) and applied in the forward loop (lines 150-156), grant the optimizer fine-grained control over gradient flow through residual connections and the original embedding stream.

## Embedding and Projection Innovations

### Untied Token Embeddings and Language Model Head

Unlike earlier GPT architectures that tied input and output embeddings, NanoChat uses **untied embeddings**, allowing `self.transformer["wte"]` and `self.lm_head` to remain separate modules (lines 71-75). This pattern, common in modern LLMs, improves training dynamics by decoupling the input representation space from the output projection.

### ResFormer-Style Value Embeddings

The model implements **ResFormer-style value embeddings (VE)** through the `has_ve` flag and `self.ve_gate` parameter (lines 53-56). During the forward pass, a per-token "value-only" embedding is gated into the attention values via `v = v + gate * ve` (lines 79-96), improving expressivity without adding parameters to the main residual stream.

### Bias-Free Linear Layers

All projection layers in the attention mechanism use **bias-free linear layers**. The `Linear` class overrides standard `nn.Linear` to set `bias=False` (lines 45-50), saving parameters and removing unnecessary bias terms when preceding normalization layers already center activations.

## Training and Initialization Optimizations

### Meta-Device Initialization

To support rapid model construction on multi-GPU setups without immediate memory allocation, the weights are created on a **meta device**. The `GPT.__init__` method establishes this pattern (lines 56-61), with real tensor initialization deferred until `init_weights()` is explicitly called (lines 96-100).

### Mixed-Precision Compute Handling

The model handles **mixed-precision compute** by casting embeddings and value embeddings to a specified `COMPUTE_DTYPE` while maintaining master weights in FP32 for optimizer stability. This casting occurs in `init_weights` for `self.transformer.wte` and related modules (lines 48-52).

### MuonAdamW Optimizer

For optimization, NanoChat implements a custom **MuonAdamW** optimizer (defined in [`nanochat/optim.py`](https://github.com/karpathy/nanochat/blob/main/nanochat/optim.py)). The `setup_optimizer` function (lines 84-90) groups parameters by shape, enabling more efficient fused kernel updates compared to standard AdamW implementations, particularly accelerating training on modern hardware.

## Practical Implementation Examples

### Instantiating the Model

```python
from nanochat.gpt import GPT, GPTConfig

# Default configuration: 12 layers, 6 heads, 768-d model, 2048 context length

cfg = GPTConfig()
model = GPT(cfg)

# Initialize weights (meta-device to real tensors)

model.init_weights()
print(f"Model has {sum(p.numel() for p in model.parameters())/1e6:.2f} M parameters")

```

### Training Forward Pass

```python
import torch
from nanochat.gpt import GPT, GPTConfig

cfg = GPTConfig()
model = GPT(cfg)
model.init_weights()

# Dummy batch: batch size 4, sequence length 128

batch_tokens = torch.randint(0, cfg.vocab_size, (4, 128), dtype=torch.long)
targets = torch.clone(batch_tokens)  # Next-token prediction

loss = model(batch_tokens, targets)  # Returns scalar loss

loss.backward()
print(f"Training loss: {loss.item():.4f}")

```

### Autoregressive Generation

```python
from nanochat.gpt import GPT, GPTConfig

cfg = GPTConfig()
model = GPT(cfg)
model.init_weights()

prompt = [1, 502, 23]  # Token IDs from tokenizer

generator = model.generate(prompt, max_tokens=50, temperature=0.8, top_k=40)

generated = list(generator)
print("Generated tokens:", generated)

```

## Key Source Files

- **[`nanochat/gpt.py`](https://github.com/karpathy/nanochat/blob/main/nanochat/gpt.py)**: Core model implementation containing the `GPT` class, `CausalSelfAttention`, rotary embeddings, GQA, FlashAttention-3 integration, sliding-window logic, and value embeddings.
- **[`nanochat/flash_attention.py`](https://github.com/karpathy/nanochat/blob/main/nanochat/flash_attention.py)**: Thin wrapper handling FlashAttention-3 selection on Hopper+ GPUs with fallback to PyTorch SDPA.
- **[`nanochat/optim.py`](https://github.com/karpathy/nanochat/blob/main/nanochat/optim.py)**: Custom `MuonAdamW` and `DistMuonAdamW` optimizers for grouped-by-shape parameter updates.
- **[`nanochat/tokenizer.py`](https://github.com/karpathy/nanochat/blob/main/nanochat/tokenizer.py)**: BPE tokenizer utilities for encoding text to model-compatible token IDs.
- **[`scripts/chat_cli.py`](https://github.com/karpathy/nanochat/blob/main/scripts/chat_cli.py)**: Interactive chat interface demonstrating model loading and streaming generation.

## Summary

- **Group-Query Attention** and **FlashAttention-3** reduce memory bandwidth and increase throughput during both training and inference.
- **RMSNorm** and **Q-K normalization** stabilize training in mixed-precision environments without learnable affine parameters.
- **Rotary Positional Embeddings** and **sliding-window attention** enable efficient long-context modeling with cached key-value stores.
- **ResFormer value embeddings**, **untied input/output projections**, and **bias-free layers** represent modern embedding strategies that improve expressivity and parameter efficiency.
- **Meta-device initialization**, **per-layer residual scaling**, and **MuonAdamW** optimize the training workflow for multi-GPU setups and modern accelerator hardware.

## Frequently Asked Questions

### What is Group-Query Attention and why does NanoChat use it?

**Group-Query Attention (GQA)** is an attention mechanism that uses fewer key and value heads than query heads, grouping multiple queries to share the same key-value pairs. NanoChat implements GQA in [`nanochat/gpt.py`](https://github.com/karpathy/nanochat/blob/main/nanochat/gpt.py) (lines 69-77) to reduce memory consumption during inference while maintaining model quality, making it feasible to run the model on consumer-grade GPUs with limited VRAM.

### How does RMSNorm differ from standard LayerNorm?

**RMSNorm** (root-mean-square normalization) differs from LayerNorm by removing the learnable gain parameter and the bias correction, instead normalizing solely by the root-mean-square of the inputs. As implemented in NanoChat's `norm` function (lines 42-44), RMSNorm operates faster in mixed-precision training and achieves comparable or better stability than traditional LayerNorm without additional parameters.

### Why are the token embeddings and LM head separate in NanoChat?

NanoChat uses **untied embeddings**, meaning `self.transformer["wte"]` (input embeddings) and `self.lm_head` (output projection) are initialized as separate modules (lines 71-75). This modern architectural choice decouples the input and output representation spaces, often leading to improved training dynamics and final model performance compared to weight-tying schemes used in earlier GPT architectures.

### What is the MuonAdamW optimizer and how does it improve training?

**MuonAdamW** is a custom optimizer defined in [`nanochat/optim.py`](https://github.com/karpathy/nanochat/blob/main/nanochat/optim.py) that groups parameters by their tensor shape and applies Adam-style updates more efficiently than standard implementations. Used in NanoChat's `setup_optimizer` function (lines 84-90), it leverages fused kernel operations for grouped updates, reducing overhead and accelerating training convergence on modern hardware compared to conventional AdamW optimizers.