When Meta released Llama 2 in July 2023, the weights grabbed the headlines, but the small Python package that shipped alongside them — meta-llama/llama — quietly became one of the most consequential pieces of inference code in modern machine learning. The llama/ directory is barely two thousand lines across three modules, yet it defined the template that dozens of later open-source serving stacks would copy: a dataclass of hyperparameters, a decoder-only Transformer, a SentencePiece tokenizer wrapper, and an autoregressive loop with temperature and top-p sampling. If you have ever read a newer inference framework and felt a strange sense of déjà vu, this is almost certainly the source you were remembering.

This post is a source tour of that reference implementation as it stands on the main branch today. The repository is now flagged as deprecated in its own README in favor of the consolidated Llama Stack repositories, which makes it an unusually clean time capsule: nothing in it is chasing new features, so every line exists to demonstrate the Llama 2 architecture — grouped-query attention, rotary position embeddings, RMSNorm, SwiGLU — in the most direct PyTorch imaginable. There is no serving layer, no quantization, no CUDA kernels of its own. That minimalism is precisely why it rewards a close reading.

Two diagrams anchor the tour. The first shows the four-load-bearing-piece view: the model, the tokenizer, the generation facade, and the example scripts that tie them together. The second zooms in to individual symbols and the wiring between them, including the checkpoint loader and the chat formatting rules.

Overview of meta-llama/llama: llama/model.py builds the Transformer; llama/tokenizer.py wraps SentencePiece; llama/generation.py's Llama facade loads checkpoints and runs generation; example_chat_completion.py is the torchrun entry point.

Reading the overview from left to right:

Each node in the diagram corresponds to one real symbol in the repository. Follow the links to see them in place:

• ModelArgs, L20 — a plain dataclass holding dim, n_layers, n_heads, and n_kv_heads. That last field defaults to None, and when it does, the code falls back to n_heads; setting it smaller is the entire grouped-query attention mechanism, one line of configuration instead of a new class.
• Transformer, L413 — the top-level nn.Module. It owns a ParallelEmbedding, a ModuleList of n_layers blocks, a final RMSNorm, a ColumnParallelLinear projection to vocab_size, and a precomputed RoPE table sized to max_seq_len * 2.
• Attention, L176 — multi-head attention built on fairscale’s ColumnParallelLinear and RowParallelLinear, with four projections wq, wk, wv, wo and a preallocated KV cache.
• FeedForward, L307 — the SwiGLU MLP. Its entire forward pass is one line: w2(F.silu(w1(x)) * w3(x)), with the hidden dimension rounded up to a multiple of 256.
• RMSNorm, L34 — root-mean-square normalization, x * rsqrt(mean(x²) + eps) with a learned scale, computed in float32 and cast back. No mean subtraction, no bias — that is the whole layer.
• Tokenizer, L14 — a thin wrapper over SentencePieceProcessor that records bos_id, eos_id, and pad_id at load time and exposes exactly two methods.
• Llama.build, L53 — the static factory that initializes the NCCL process group, globs *.pth checkpoint shards, reads params.json, and loads the state dict into a freshly constructed Transformer.
• generate, L130 — the decode loop: pad the batch to a rectangle, call the model once per position window, sample the next token, and track which rows have hit EOS.
• chat_completion, L284 — the chat API that validates role-alternating dialogs, folds system prompts into <<SYS>> blocks, renders the [INST] template, and then delegates to generate.

One model.py, Four Ideas

The elegance of llama/model.py is that it contains only the ideas the Llama 2 paper actually requires. The TransformerBlock, L351 is a textbook pre-norm residual: h = x + attention(norm(x)) followed by out = h + feed_forward(norm(h)). There is no dropout, no bias term anywhere in the linear layers, and no separate embedding scaling. The block is instantiated n_layers times (32 for 7B) and stacked.

Grouped-query attention is the most cited architectural change from Llama 1, and in this codebase it is disarmingly small. The attention module computes n_rep = n_local_heads // n_local_kv_heads, and a helper called repeat_kv, L164 expands each key/value head n_rep times along the head axis using expand and reshape — no copied memory, just a view trick. The 70B model uses 8 KV heads against 64 query heads, which shrinks the KV cache by a factor of 8 relative to full multi-head attention. Because the cache tensors cache_k and cache_v, L236 are allocated once in __init__ at (max_batch_size, max_seq_len, n_local_kv_heads, head_dim), that memory saving applies for the entire process lifetime, not per request.

Positional information arrives through rotary embeddings. precompute_freqs_cis, L80 builds a complex64 table of unit-magnitude complex exponentials once at construction, and apply_rotary_emb, L132 reinterprets the query and key tensors as complex numbers and multiplies them by the broadcast table. The comment above the call site in the Transformer constructor is worth reading: the table is sized to max_seq_len * 2 because the Llama 2 token limit is 4096 and the authors wanted headroom for fine-tuning at longer lengths.

The Tokenizer and the Checkpoint Loader

llama/tokenizer.py is the smallest module and the easiest to underestimate. Tokenizer.encode, L38 piece-encodes a string and optionally prepends BOS or appends EOS; Tokenizer.decode, L58 goes the other way. The constructor asserts that the SentencePiece vocab size matches the piece count and logs the special-token IDs. Everything chat-related later — stopping conditions, prompt template boundaries — leans on the bos_id/eos_id pair captured here.

The loader inside Llama.build, L53 encodes several assumptions worth noticing. It globs *.pth files in the checkpoint directory and asserts that the shard count equals model_parallel_size — the 7B model ships one shard (model-parallel degree 1), 13B ships two, and 70B ships eight, per the README’s table. Each rank reads only its own shard with torch.load(..., map_location="cpu"), then params.json hydrates a ModelArgs instance, the tokenizer’s word count overrides vocab_size, and model.load_state_dict(checkpoint, strict=False) fills the freshly built module. The strict=False is there because the checkpoint contains no freqs_cis buffer — that tensor is computed, not trained. Finally, torch.set_default_tensor_type(torch.cuda.HalfTensor) flips every subsequent allocation to CUDA fp16, a now-deprecated idiom that dates the file precisely.

From Install to Your First Model

The README’s quick start is honest about the one thing you cannot get from GitHub: the weights. Installing the package itself is trivial because a setup.py exists at the repo root and declares the contents of requirements.txt (torch, fairscale, fire, sentencepiece):

pip install -e .

The weights are gated. You request access on the Meta website, accept the license, and receive a signed, expiring URL by email. That URL is fed to ./download.sh in the repository root, a wget/md5sum script that fetches the .pth shards and the tokenizer.model file for whichever models you selected. (Hugging Face mirrors exist under the same license terms.) Once the checkpoints are on disk, the README’s canonical invocation is:

torchrun --nproc_per_node 1 example_chat_completion.py \
    --ckpt_dir llama-2-7b-chat/ \
    --tokenizer_path tokenizer.model \
    --max_seq_len 512 --max_batch_size 6

example_chat_completion.py, L11 is a Fire CLI that calls Llama.build and then runs six hard-coded dialogs through chat_completion, including one deliberately unsafe prompt containing [INST] tags to demonstrate the special-tag rejection path. For base (non-chat) models, example_text_completion.py, L9 follows the identical pattern with text_completion and a few-shot translation prompt.

Sampling, Stopping, and the Chat Template

The decode loop in generate deserves attention because so many later frameworks reimplemented it line for line. Prompts of different lengths are right-padded into one rectangle; then, for each position from the shortest prompt to total_len, the model runs a forward pass over the not-yet-consumed slice. When temperature > 0, the logits are softmaxed with the temperature and passed to sample_top_p, L398, which sorts probabilities descending, zeroes tokens beyond the cumulative top-p mass, renormalizes, and draws from a multinomial. At temperature = 0 it is a plain argmax. Two subtleties stand out: tokens already belonging to the prompt are never overwritten (the input_text_mask guard), and generation stops per row, with the loop breaking only when every row has produced EOS.

The chat format lives in the same file, just above the Llama class: B_INST, E_INST = "[INST]", "[/INST]" and B_SYS, E_SYS = "<<SYS>>\n", "\n<</SYS>>\n\n". chat_completion, L284 enforces strict role alternation with assertions, merges a leading system message into the first user turn, then encodes each exchange with the template and BOS. Any user content containing the special tags is flagged and answered with an error string rather than being forwarded to the model — a tiny, instructive piece of in-band prompt-injection defense.

The detail graph: how ModelArgs configures the Transformer, how the Attention block consumes RoPE tables and the preallocated KV cache, and how Llama.build, generate, and the chat template connect the modules to the torchrun examples and download.sh.

Honest Limits

The reference code is a teaching artifact, and reading it with production eyes reveals its boundaries quickly. Every rank assumes exactly one CUDA device — torch.cuda.set_device(local_rank) runs unconditionally, the KV caches are allocated with .cuda() in __init__, and the process group is NCCL, so there is no CPU or Apple Silicon path in this code. The chat template is implemented only for the chat-finetuned checkpoints; feeding a base model through chat_completion will run, but the [INST] wrapping is a no-op semantically and will degrade output quality, which is why the separate text_completion API exists.

The KV cache, though preallocated efficiently, is not reusable across requests: each generate call writes keys and values from position zero, and nothing is carried between calls, so serving many short conversations pays full prompt processing each time. State-of-the-art serving engines rebuilt this part first. Finally, the repository itself is superseded — the README’s deprecation notice points to llama-models for Llama 3-era weights and the wider Llama Stack repos, so treat this main branch as frozen history. If you want to see where the ideas landed next, the llama.cpp project re-implemented the same architecture in pure C++ with quantization, vLLM redesigned the KV cache for continuous batching, and the PyTorch APIs used here — load_state_dict, ModuleList, torch.inference_mode — remain the substrate everything above rests on. For the tokenization side, the SentencePiece approach contrasted here with byte-level BPE is covered in our tokenizers tour.

Enjoyed this post? Never miss out on future posts by following us