Llama 2 Reference Implementation — Administrator & Developer Guide

This guide is the operator and developer companion to the Llama 2 reference inference repository at github.com/meta-llama/llama (commit 689c7f2, 18 source files, ~47.5 KB of code). It documents the architecture, the configuration surface, deployment topology, security posture (license, weights provenance, multi-tenancy), and the developer internals that you will need to extend, port, or audit the codebase.

Important

Status: maintenance. Active Meta development of Llama models has moved to llama-models (model definitions, including Llama 3 family) and llama-stack (inference, agents, evaluations, safety). This repository is the canonical Llama 2 reference and remains useful for understanding the architecture, but new production work should not start here.

1. Architecture Overview

The repository is a thin, end-to-end PyTorch implementation of a decoder-only transformer matching the Llama 2 architecture. Eight files do all the work — four Python modules under llama/, two example scripts at the root, one shell script for weight download, and one setup.py.

1.1 Component diagram

                       Operator
                          |
                          v
            torchrun --nproc_per_node N
            [WORLD_SIZE, RANK, LOCAL_RANK,
             MASTER_ADDR, MASTER_PORT set]
                          |
                          v
        +----------------------------------+
        |  example_text_completion.py /    |
        |  example_chat_completion.py      |    fire.Fire(main)
        |  (parses CLI flags via fire)     |
        +----------------+-----------------+
                         |
                         v
        +----------------------------------+
        |        Llama.build()             |
        |  - torch.distributed init (nccl) |
        |  - fairscale.initialize_model_   |
        |    parallel(WORLD_SIZE)          |
        |  - torch.cuda.set_device(rank)   |
        |  - load consolidated.NN.pth shard|
        |  - read params.json -> ModelArgs |
        |  - build Tokenizer(model_path)   |
        |  - build Transformer(model_args) |
        |  - load_state_dict(checkpoint)   |
        +----------------+-----------------+
                         |
                         v
        +----------------------------------+
        |    Llama (generator instance)    |
        |  text_completion / chat_compl. / |
        |  generate (token-level loop)     |
        +----------------+-----------------+
                         |
                         v
        +----------------------------------+
        |        Transformer.forward       |
        |  ParallelEmbedding ->            |
        |  [TransformerBlock] x n_layers   |
        |    RMSNorm -> Attention(KV cache)|
        |    -> RMSNorm -> FeedForward     |
        |  -> RMSNorm -> ColumnParallel    |
        |  Linear (output logits)          |
        +----------------+-----------------+
                         |
                         v
        +----------------------------------+
        |  sample_top_p / argmax sampler   |
        |  Tokenizer.decode(out_tokens)    |
        +----------------------------------+

1.2 Component inventory

1.3 Technology stack

Note

The model-parallelism is tensor parallel only (column / row parallel linear layers and a parallel embedding). There is no pipeline parallel, no FSDP / ZeRO, and no automatic sharding for arbitrary GPU counts. The shard count is hardcoded by the published checkpoints: 7B = 1, 13B = 2, 70B = 8.

1.4 Data flow — single inference call

1. torchrun launches N Python processes; each gets RANK, WORLD_SIZE, LOCAL_RANK, MASTER_ADDR, MASTER_PORT.
2. Llama.build:
   • Initializes torch.distributed with NCCL backend.
   • Calls fairscale.initialize_model_parallel(WORLD_SIZE).
   • Pins each rank to cuda:LOCAL_RANK.
   • Each rank reads exactly one consolidated.NN.pth shard (its get_model_parallel_rank()).
   • Loads params.json once on every rank to build ModelArgs.
   • Sets default tensor type to torch.cuda.HalfTensor (FP16 weights).
   • Constructs Transformer(model_args) and load_state_dict(checkpoint, strict=False).
3. text_completion or chat_completion is called.
4. Tokenizer encodes prompts (BOS prepended, EOS only on chat templating internals).
5. generate runs an autoregressive loop:
   • Initializes a (bsz, total_len) token tensor padded with pad_id.
   • For each new position, calls Transformer.forward(tokens[:, prev_pos:cur_pos], prev_pos).
   • Applies softmax(logits / temperature) + sample_top_p (or argmax if temperature == 0).
   • Stops at EOS or max_gen_len.
6. Decoded text is returned. chat_completion overwrites with UNSAFE_ERROR if the input contained chat templating tokens.

2. System Requirements

2.1 Hardware

Note

The host-RAM number is the transient peak during torch.load(..., map_location="cpu"), before weights move to GPU. Running on a host with less RAM than the on-disk shard is the most common cause of exit code -9 (SIGKILL by the OOM killer, GitHub #928, #960).

2.2 Software

2.3 Network

• Outbound HTTPS to *.meta.com / *.fbcdn.net during download.sh to fetch checkpoint shards. Sustained throughput matters: 70B is over 130 GB.
• NCCL collectives between ranks — for multi-host runs, low-latency high-bandwidth networking (InfiniBand or 100 Gbps+ Ethernet) is effectively required to avoid bottlenecking on the all-reduce in every layer's tensor-parallel forward.

2.4 Filesystem layout (post-download)

llama/
├─ download.sh
├─ example_chat_completion.py
├─ example_text_completion.py
├─ requirements.txt
├─ setup.py
├─ llama/
│   ├─ __init__.py
│   ├─ generation.py
│   ├─ model.py
│   └─ tokenizer.py
├─ tokenizer.model           # SentencePiece model
├─ tokenizer_checklist.chk
├─ LICENSE
├─ USE_POLICY.md
├─ llama-2-7b/               # one of these per downloaded variant
│   ├─ consolidated.00.pth
│   ├─ params.json
│   └─ checklist.chk
├─ llama-2-13b/
│   ├─ consolidated.00.pth
│   ├─ consolidated.01.pth
│   ├─ params.json
│   └─ checklist.chk
└─ llama-2-70b-chat/
    ├─ consolidated.00.pth ... consolidated.07.pth
    ├─ params.json
    └─ checklist.chk

3. Configuration Reference

The repository has no configuration file. All configuration lives at three layers: distributed launcher env vars (read by Llama.build), build-time arguments (passed to Llama.build), and per-call sampling parameters.

3.1 Environment variables consumed

Important

No other environment variables are read by application code. Anything beyond this set comes from the underlying PyTorch / NCCL / fairscale stack. Do not invent app-level env vars — there is no plumbing for them.

3.2 Llama.build arguments

@staticmethod
def build(
    ckpt_dir: str,
    tokenizer_path: str,
    max_seq_len: int,
    max_batch_size: int,
    model_parallel_size: Optional[int] = None,
    seed: int = 1,
) -> "Llama":

3.3 ModelArgs (per-checkpoint structural config)

ModelArgs is a frozen @dataclass with the following defaults; the actual per-checkpoint values override these from params.json:

Note

Llama 2 base checkpoints use dim=4096, n_layers=32, n_heads=32 for 7B; dim=5120, n_layers=40, n_heads=40 for 13B; dim=8192, n_layers=80, n_heads=64, n_kv_heads=8 for 70B. The exact values come from the params.json file inside each variant's checkpoint directory.

3.4 Per-call sampling parameters

3.5 CLI flags exposed by example scripts

Both example scripts expose every main argument as a --flag via fire.Fire(main):

example_text_completion.py:
  --ckpt_dir <str>           required
  --tokenizer_path <str>     required
  --temperature <float>      default 0.6
  --top_p <float>            default 0.9
  --max_seq_len <int>        default 128
  --max_gen_len <int>        default 64
  --max_batch_size <int>     default 4

example_chat_completion.py:
  --ckpt_dir <str>           required
  --tokenizer_path <str>     required
  --temperature <float>      default 0.6
  --top_p <float>            default 0.9
  --max_seq_len <int>        default 512
  --max_batch_size <int>     default 8
  --max_gen_len <int>        default None (=> model.params.max_seq_len-1)

There are no other CLI flags. The download.sh script takes no flags — it is interactive (read -p prompts for the presigned URL and the comma-separated list of variants).

4. Deployment & Scaling

4.1 Single-host deployment (typical research workstation)

# 7B base on a single GPU
torchrun --nproc_per_node 1 example_text_completion.py \
    --ckpt_dir llama-2-7b/ --tokenizer_path tokenizer.model \
    --max_seq_len 1024 --max_batch_size 4

Operationally:

• One torchrun invocation launches one Python process.
• That process initializes a single-rank distributed group (still required — Llama.build calls init_process_group("nccl") unconditionally).
• All weights live on cuda:0.

4.2 Single-host, multi-GPU (13B and 70B)

# 70B chat on a single 8-GPU host
torchrun --nproc_per_node 8 example_chat_completion.py \
    --ckpt_dir llama-2-70b-chat/ --tokenizer_path tokenizer.model \
    --max_seq_len 2048 --max_batch_size 6

torchrun spawns 8 Python processes. Each rank loads exactly one consolidated.NN.pth shard. Forward passes synchronize via NCCL all- reduce inside each RowParallelLinear and the column-parallel embedding.

4.3 Multi-host deployment (rare; only for 70B on smaller hosts)

# Two hosts, 4 GPUs each, MP=8
# Host 0
torchrun --nproc_per_node 4 --nnodes 2 --node_rank 0 \
    --rdzv_endpoint <host0_ip>:29500 \
    example_chat_completion.py \
    --ckpt_dir llama-2-70b-chat/ --tokenizer_path tokenizer.model
# Host 1
torchrun --nproc_per_node 4 --nnodes 2 --node_rank 1 \
    --rdzv_endpoint <host0_ip>:29500 \
    example_chat_completion.py \
    --ckpt_dir llama-2-70b-chat/ --tokenizer_path tokenizer.model

Warning

Multi-host runs with this code are bandwidth-sensitive. Without InfiniBand or 100+ Gbps Ethernet, all-reduce latency dominates per-token generation time. For production-grade serving, use vLLM or llama-stack instead — they implement pipeline parallelism and efficient batching that this reference does not.

4.4 Containerization sketch

There is no Dockerfile in this repository. A minimal container looks like:

FROM nvcr.io/nvidia/pytorch:24.05-py3
WORKDIR /opt/llama
COPY . /opt/llama
RUN pip install -e .
ENTRYPOINT ["torchrun", "--nproc_per_node", "1", \
    "example_text_completion.py"]

Mount the checkpoint directory and tokenizer at runtime:

docker run --rm --gpus all \
    -v /data/llama-2-7b:/weights \
    -v /data/tokenizer.model:/weights/tokenizer.model:ro \
    llama-ref \
    --ckpt_dir /weights --tokenizer_path /weights/tokenizer.model

4.5 Scaling characteristics

This codebase does not scale beyond one in-flight request per process. There is no:

• Continuous batching
• Paged KV cache
• Asynchronous request queue
• HTTP server
• Streaming response API

The generate loop processes one batch from start to EOS before returning. For higher throughput, port to vLLM, TensorRT-LLM, or llama-stack.

5. Operations

5.1 Standard run procedure

• [x] Confirm GPU visibility: nvidia-smi
• [x] Confirm CUDA + NCCL: python -c "import torch; print(torch.cuda.is_available(), torch.distributed.is_nccl_available())"
• [x] Verify checkpoint shards: ls llama-2-7b/consolidated.*.pth
• [x] Verify checksums: (cd llama-2-7b && md5sum -c checklist.chk)
• [x] Verify tokenizer checksum: md5sum -c tokenizer_checklist.chk
• [x] Launch under torchrun with the right --nproc_per_node
• [x] Confirm the Loaded in <N> seconds line prints before any generation

Tip

Wrap the launch in a small shell script that exports NCCL_DEBUG=INFO and TORCH_DISTRIBUTED_DEBUG=DETAIL for the first run on any new host. They surface NCCL transport problems (P2P, IB, sockets) early, before they become silent hangs.

5.2 Logging and observability

This repository has no structured logging. The Tokenizer class uses Python's logging.getLogger() at module level to log SentencePiece load info; everything else uses print(...). Specifically:

• Llama.build prints Loaded in <N> seconds.
• For ranks where local_rank > 0, sys.stdout = open(os.devnull, "w") is set after device pinning, so only rank 0 stdout reaches the operator.

Note

Rank 0 silencing is intentional — it prevents N copies of the same generation output. If you need per-rank diagnostics, log to per-rank files keyed by int(os.environ['LOCAL_RANK']) before Llama.build redirects stdout.

5.3 Health checks

There is no in-process health endpoint. To verify a deployment is healthy:

1. Run a short text completion (max_gen_len=8, max_batch_size=1).
2. Compare the output to a known-good golden output for the same seed.
3. Time the call; spike-vs-baseline catches GPU thermal throttling and slow PCIe.

5.4 Common operational alerts

6. Performance

6.1 Hot path

Per generation step (i.e. per generated token), the hot path is:

Transformer.forward(tokens[:, prev_pos:cur_pos], prev_pos)
   for layer in self.layers:           # n_layers iterations
       attention_norm = RMSNorm(x)     # cheap
       attention_out  = Attention(...) # KV cache update + scaled dot-prod
                                       # contains an NCCL all-reduce in
                                       # RowParallelLinear self.wo
       ffn_norm  = RMSNorm(...)
       ffn_out   = FeedForward(...)    # SwiGLU; another all-reduce in
                                       # RowParallelLinear self.w2
   norm_out  = RMSNorm(...)
   logits    = ColumnParallelLinear(self.output, ...)
softmax + sample_top_p (or argmax)

Per-token cost is dominated by:

1. Two NCCL all-reduces per layer (one in attention output, one in FFN output). For 80-layer 70B that's 160 collectives per token.
2. KV cache append at every layer.
3. Softmax over the full vocab (32k tokens for Llama 2) at the final logits.

6.2 KV cache footprint

Allocated once per Attention instance at Llama.build time:

cache_k, cache_v shape:
  (max_batch_size, max_seq_len, n_local_kv_heads, head_dim)

For 7B (n_kv_heads=32, head_dim=128, FP16):

• max_batch_size=4, max_seq_len=2048 → 4 × 2048 × 32 × 128 × 2 B ≈ 64 MB per cache tensor, per layer. Across 32 layers: ~4 GB.

For 70B (n_kv_heads=8, n_local_kv_heads=1 with MP=8, head_dim=128):

• max_batch_size=6, max_seq_len=2048 → 6 × 2048 × 1 × 128 × 2 B ≈ 3 MB per cache tensor per layer per rank. Across 80 layers: ~240 MB per rank.

Tip

The single biggest knob for memory pressure is max_seq_len. Drop it to the smallest value that fits your prompts plus expected generation length — KV cache memory scales linearly with it.

6.3 Sampling cost

sample_top_p performs a full sort of the vocab probabilities per token via torch.sort(probs, dim=-1, descending=True). For a 32k vocab this is small relative to the transformer forward, but it is not free. When temperature == 0, the code path skips the softmax entirely and goes straight to torch.argmax(logits[:, -1], dim=-1) — measurably cheaper on long generations.

6.4 What this code does not optimize

• No FlashAttention — Attention.forward uses an explicit torch.matmul(xq, keys.transpose(2, 3)) / sqrt(head_dim) followed by softmax then matmul(scores, values). Modern stacks replace this with F.scaled_dot_product_attention or FlashAttention v2/v3.
• No paged KV cache — caches are dense, full-max_seq_len tensors.
• No continuous batching — one prompt batch per generate call.
• No quantization — weights stay FP16 throughout.
• No torch.compile — graph capture is left to the user.

For production throughput, port to vLLM or use llama-stack.

Note

The pre-allocated KV cache is the reason Llama.build requires max_batch_size and max_seq_len up front: every Attention layer sizes its caches at construction time, with .cuda() allocations.

7. Security

Caution

License-bound weights. The weights downloaded by download.sh are governed by the Llama 2 Community License Agreement and the Acceptable Use Policy that ship alongside them. Do not redistribute the weights, do not share your presigned URL, and do not bundle the weights into a public repository or container image. A breach of the license is a legal exposure for the operator, not just a policy violation.

7.1 License acceptance flow

1. Operator visits ai.meta.com/resources/models-and-libraries/llama-downloads.
2. Submits the request and accepts the Llama 2 Community License.
3. Meta emails a single-tenant, time-limited presigned URL.
4. Operator runs ./download.sh, pastes the URL, picks variants.
5. The script downloads LICENSE and USE_POLICY.md first, before any model files.

7.2 Weights provenance and integrity

download.sh does three things to defend against silent corruption:

• Uses wget --continue so partial downloads can resume without re-fetching from byte zero.
• Downloads a tokenizer_checklist.chk and verifies the tokenizer with md5sum -c (or md5 on macOS arm64).
• Downloads a per-variant checklist.chk inside each model folder and verifies every shard:

(cd ${MODEL_PATH} && md5sum -c checklist.chk)

Warning

If a checksum fails, stop — the most likely cause is a truncated download (the URL expired mid-transfer or the connection dropped). Re-run ./download.sh and let wget --continue complete the missing bytes. Do NOT silently proceed; a partially-corrupted weight tensor can produce plausible-looking but degenerate outputs.

7.3 Acceptable Use Policy (USE_POLICY.md)

The downloaded USE_POLICY.md enumerates uses that are explicitly prohibited under the license. The operator is responsible for ensuring their deployment complies. Categories include (verbatim from the policy that ships with the weights — read the file for the authoritative text):

• Violence, harm, or illegal activity
• Child exploitation
• Human trafficking, exploitation, sexual violence
• Discriminatory practices including unlawful targeting
• Unauthorized practice of regulated professions (medical, legal, financial advice, etc.)
• Privacy violations (collection of PII without consent)
• Generation of malware, exploits, or other code intended to harm
• Misrepresentation as a human in deceptive contexts

The model also implements a prompt-injection guardrail at the chat template layer: if a user message contains [INST], [/INST], <<SYS>>, or <</SYS>>, the assistant's content is overwritten with UNSAFE_ERROR = "Error: special tags are not allowed as part of the prompt." (see llama/generation.py lines ~57 and ~388).

7.4 Multi-tenancy and adversarial environments

Caution

Do not run this code in shared / multi-tenant environments without additional isolation. Specifically:

• The generate loop has no rate limiting, no per-tenant resource quotas, and no input length validation beyond a single assert max_prompt_len <= params.max_seq_len.
• The unsafe-tag check is a brittle string match — sophisticated prompt injection bypasses it. Treat the model output as untrusted.
• The chat_completion template concatenates user content directly into the chat string. There is no filtering for harmful content beyond the four chat tokens.
• The model is fully trusted by the runtime — it can emit any string in its vocabulary, including tokens that look like control sequences, shell commands, or credentials.

For adversarial deployment, layer:

• An input filter that strips or escapes [INST], [/INST], <<SYS>>, <</SYS>>.
• An output filter (Llama Guard, classifier, or rule-based).
• Per-tenant rate limits and concurrency caps at the launcher.
• A network sandbox that disallows the inference process from making outbound connections.

7.5 Things you must not do

• Do not commit weights to git — even private repos have leaked.
• Do not bake weights into a public container image.
• Do not share the presigned download URL. It carries your license acceptance.
• Do not modify LICENSE or USE_POLICY.md before redistribution.
• Do not use the model for purposes prohibited by USE_POLICY.md.
• Do not redistribute the tokenizer — it is part of the licensed artifact bundle.

7.6 Code-level secrets and surface

The repository contains no hardcoded credentials, API keys, or remote endpoints in source. The only network call in code paths is the wget-based fetch in download.sh, parameterized by the operator- supplied URL.

There is no AuthN, no AuthZ, no TLS termination, no input validation beyond shape assertions. The reference is for trusted environments only.

8. Developer Internals

8.1 Repository structure (post-install)

llama/                         Python package (importable as `llama`)
├─ __init__.py                 Re-exports public surface
├─ generation.py               Llama class, sample_top_p, chat templating
├─ model.py                    Transformer architecture, RoPE, RMSNorm, FFN
└─ tokenizer.py                SentencePiece wrapper

example_text_completion.py     fire-wrapped main, four built-in prompts
example_chat_completion.py     fire-wrapped main, six demo dialogs
download.sh                    Bash; presigned-URL prompted weight downloader
setup.py                       setuptools; reads requirements.txt
requirements.txt               torch, fairscale, fire, sentencepiece

LICENSE                        Llama 2 Community License Agreement
USE_POLICY.md                  Acceptable Use Policy
MODEL_CARD.md                  Llama 2 model card
Responsible-Use-Guide.pdf      Meta's Responsible Use Guide
README.md                      Top-level repo readme
UPDATES.md                     Repo update log
CONTRIBUTING.md                Contribution guidelines
CODE_OF_CONDUCT.md             Code of conduct

8.2 Class inventory

Type aliases:

• Role = Literal["system", "user", "assistant"] (generation.py:32)
• Dialog = List[Message] (generation.py:52)

8.3 Function inventory

8.4 The model architecture in detail

ModelArgs is the structural config. Key fields and what they control:

• dim — model embedding width (4096 / 5120 / 8192 for 7B / 13B / 70B).
• n_layers — number of TransformerBlocks.
• n_heads — number of attention query heads.
• n_kv_heads — number of key/value heads. When < n_heads, this enables grouped-query attention (GQA), used by 70B (n_heads=64, n_kv_heads=8). For 7B and 13B it is None (defaults to n_heads).
• multiple_of — rounding factor for the SwiGLU hidden dim (hidden_dim is rounded up to a multiple of this; default 256).
• ffn_dim_multiplier — extra multiplier on the SwiGLU hidden dim (used by 70B for a wider FFN).
• norm_eps — epsilon inside RMSNorm.

RMSNorm (Root Mean Square Layer Norm, model.py:34):

def _norm(self, x):
    return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
def forward(self, x):
    output = self._norm(x.float()).type_as(x)
    return output * self.weight

A drop-in replacement for LayerNorm that drops the mean-subtraction and the bias. The .float() cast is intentional — RMSNorm is computed in FP32 for stability, then cast back to the input dtype (FP16) before multiplying by the learnable weight.

Rotary Positional Embedding (RoPE) lives in three helpers:

• precompute_freqs_cis(dim, end, theta=10000.0) — builds a complex freqs_cis tensor of shape (end, dim // 2) containing e^{i * t * freq} for each position t and each frequency. The theta=10000.0 matches the Llama 2 convention.
• reshape_for_broadcast(freqs_cis, x) — reshapes for elementwise multiply with (bsz, seqlen, n_heads, head_dim) tensors.
• apply_rotary_emb(xq, xk, freqs_cis) — views queries and keys as pairs of reals, multiplies as complex numbers, returns reals. This is the actual position injection.

Transformer.__init__ precomputes the RoPE table for max_seq_len * 2 positions (the * 2 reservation accommodates Llama 2's 4096-token training context — the comment in source explains the rationale).

Attention (model.py:176):

self.wq = ColumnParallelLinear(dim, n_heads * head_dim, ...)
self.wk = ColumnParallelLinear(dim, n_kv_heads * head_dim, ...)
self.wv = ColumnParallelLinear(dim, n_kv_heads * head_dim, ...)
self.wo = RowParallelLinear(n_heads * head_dim, dim, ...)

self.cache_k = torch.zeros((max_batch, max_seq, n_local_kv_heads, head_dim)).cuda()
self.cache_v = torch.zeros((max_batch, max_seq, n_local_kv_heads, head_dim)).cuda()

Forward pass (simplified):

1. Q/K/V projections via column-parallel linears (no all-reduce — output is sharded along the head dim).
2. Reshape to (bsz, seqlen, n_local_heads, head_dim).
3. Apply RoPE to xq and xk.
4. Append xk, xv to the persistent KV cache at `start_pos:start_pos
   • seqlen`.
5. Read the full cached keys, values up to start_pos + seqlen.
6. repeat_kv(keys, n_rep) if GQA — repeat KV heads to match query heads.
7. Standard scaled dot-product attention: softmax(Q @ K^T / sqrt(head_dim) + mask) @ V.
8. wo (row-parallel linear) — this is where the NCCL all-reduce happens; outputs are gathered across ranks here.

FeedForward is the SwiGLU variant:

def forward(self, x):
    return self.w2(F.silu(self.w1(x)) * self.w3(x))

with w1, w3 column-parallel and w2 row-parallel. The hidden dim is computed as int(2 * hidden_dim / 3), optionally multiplied by ffn_dim_multiplier, then rounded up to the next multiple of multiple_of. The default in TransformerBlock.__init__ passes hidden_dim=4 * args.dim.

TransformerBlock is pre-norm:

h   = x + self.attention(self.attention_norm(x), start_pos, freqs_cis, mask)
out = h + self.feed_forward(self.ffn_norm(h))

Transformer wires it all together:

self.tok_embeddings = ParallelEmbedding(vocab_size, dim, ...)
self.layers         = ModuleList([TransformerBlock(i, params) for i in range(n_layers)])
self.norm           = RMSNorm(dim, eps=norm_eps)
self.output         = ColumnParallelLinear(dim, vocab_size, ...)

The forward pass:

1. Embed tokens.
2. Slice the precomputed freqs_cis for the current position window.
3. Build a causal mask if seqlen > 1 (only on the prefill / prompt pass; per-token decode steps have seqlen=1 and skip the mask).
4. Run all n_layers transformer blocks.
5. Final RMSNorm + output linear (column-parallel — output is gathered across ranks).
6. Return logits cast to FP32.

Note

The mask construction in Transformer.forward builds a (seqlen, start_pos + seqlen) mask so attention is masked over both the new tokens and the cached prefix, with -inf above the diagonal in the new-token quadrant. This is what makes streaming-style decode possible with the persistent KV cache.

8.5 The generation loop in detail

Llama.generate is one tight loop in generation.py:130:

for cur_pos in range(min_prompt_len, total_len):
    logits = self.model.forward(tokens[:, prev_pos:cur_pos], prev_pos)
    if temperature > 0:
        probs = torch.softmax(logits[:, -1] / temperature, dim=-1)
        next_token = sample_top_p(probs, top_p)
    else:
        next_token = torch.argmax(logits[:, -1], dim=-1)

    next_token = next_token.reshape(-1)
    next_token = torch.where(
        input_text_mask[:, cur_pos], tokens[:, cur_pos], next_token
    )
    tokens[:, cur_pos] = next_token
    eos_reached |= (~input_text_mask[:, cur_pos]) & (next_token == self.tokenizer.eos_id)
    prev_pos = cur_pos
    if all(eos_reached):
        break

Key invariants:

• tokens is allocated (bsz, total_len) upfront, padded with pad_id.
• input_text_mask = tokens != pad_id — true for prompt positions.
• The torch.where step preserves prompt tokens, only overwriting the positions that the model is asked to predict.
• eos_reached short-circuits the loop when every batch row has hit EOS (or, optionally, never — there is no min_gen_len).

When min_prompt_len == total_len (no generation slots), the loop is skipped and the function only computes log-probs over the prompt.

8.6 Chat templating in chat_completion

generation.py:284 builds the prompt for each dialog:

1. If dialog[0]["role"] == "system", fold the system content into the first user message: <<SYS>>\n{system}\n<</SYS>>\n\n{user_content}.
2. Assert the remaining roles strictly alternate user / assistant.
3. For each (user_msg, assistant_msg) pair, encode f"{B_INST} {user.strip()} {E_INST} {assistant.strip()} " with bos=True, eos=True. The trailing space and EOS are deliberate — they delimit dialog turns inside the token stream.
4. Encode the final user turn as f"{B_INST} {final_user.strip()} {E_INST}" with bos=True, eos=False.
5. The unsafe-tag scan happens before this — if any tag was found in the input, unsafe_requests[i] = True and the assistant content gets replaced with UNSAFE_ERROR after generation.

Constants:

B_INST, E_INST = "[INST]", "[/INST]"
B_SYS, E_SYS   = "<<SYS>>\n", "\n<</SYS>>\n\n"
SPECIAL_TAGS   = [B_INST, E_INST, "<<SYS>>", "<</SYS>>"]
UNSAFE_ERROR   = "Error: special tags are not allowed as part of the prompt."

8.7 The tokenizer

Tokenizer (llama/tokenizer.py) is a 50-line wrapper around sentencepiece.SentencePieceProcessor:

• __init__(model_path) — loads tokenizer.model, asserts the file exists, exposes n_words, bos_id, eos_id, pad_id.
• encode(s, bos, eos) — returns List[int]. Optionally prepends bos_id and/or appends eos_id. Asserts type(s) is str.
• decode(t) — returns str from List[int].

There is no streaming decode and no batch decode — both example scripts decode entire generated sequences at once after the loop ends.

8.8 Code patterns and conventions

• Public surface is tiny. llama/__init__.py re-exports five names: Llama, Dialog, ModelArgs, Transformer, Tokenizer. Anything else is internal.
• No factory pattern. Llama.build is a static factory — call it, don't subclass it.
• Inference-only. @torch.inference_mode() on both Llama.generate and Transformer.forward — no autograd, no backward, no training.
• All assertions are loud. Failure modes use assert with human-readable messages, not custom exception classes. Error-handling callers should catch AssertionError.
• Constants live where they are used. Chat templating tokens (B_INST, etc.) are module-level in generation.py; nothing is configurable.

8.9 Testing strategy

There are no unit tests in this repository. Validation is end-to-end: run example_text_completion.py and example_chat_completion.py against a known checkpoint. For a contributor adding tests, the natural seams are:

• Tokenizer.encode / decode round-trip on a known-good tokenizer.model (does not require GPU or NCCL).
• sample_top_p correctness on known probability distributions (greedy threshold, threshold below the largest single mass, etc.).
• apply_rotary_emb numeric agreement with a reference implementation.

The architecture itself is best validated against a short reference checkpoint by comparing generate(prompt_tokens, ...) outputs at fixed seed and temperature=0 against captured golden tokens.

8.10 Potential improvements / known limitations

• No FlashAttention path. Modernizing Attention.forward to use F.scaled_dot_product_attention would yield a 2–4× speedup on long contexts.
• Dense KV cache only. Paged attention would unlock continuous batching.
• No streaming output. The generate loop returns only after every batch row finishes. A streaming variant would yield tokens as they are produced.
• No Llama 3 support. The tokenizer abstraction is hardcoded to SentencePiece; Llama 3's tiktoken-based tokenizer would need a new class.
• No quantization. All weights stay FP16. Adding INT8 / FP8 / NF4 quantization is a separate branch of work.
• No safety filtering beyond the four chat tokens. Harden user inputs and outputs externally.
• requirements.txt is unpinned. Pin torch, fairscale, fire, sentencepiece to specific versions for reproducible builds.

8.11 Contributing guide

• The repo follows standard CONTRIBUTING.md conventions for Meta open source.
• Active development has moved upstream to llama-models and llama-stack. Bug fixes for the reference are welcome; large features (Llama 3 support, paged attention, etc.) should target the successor repos.

9. Integration Points

9.1 Importable Python surface

from llama import Llama, Dialog
from llama.model import ModelArgs, Transformer
from llama.tokenizer import Tokenizer

9.2 Process-level integration

9.3 Filesystem contract

9.4 What this code does NOT integrate with

• No HTTP server. No FastAPI, Flask, or REST surface ships with this repo. Bring your own.
• No streaming protocol. No SSE, WebSocket, or gRPC handlers.
• No metrics export. No Prometheus, OpenTelemetry, statsd hooks.
• No queue integration. No Celery, RabbitMQ, Redis, or worker pool.
• No model registry. Checkpoints are addressed by filesystem path; there is no MLflow / W&B / Hugging Face Hub plumbing.
• No fine-tuning hooks. Training is out of scope.

9.5 Successor integration paths

For higher-level integration — REST APIs, agent runtimes, evaluation harnesses, safety filtering — use:

• llama-stack — ships an inference server, an agents API, evaluations, and the llama-stack-client CLI. This is the supported way to expose Llama inference behind an API.
• vLLM — production- grade inference server with paged attention and continuous batching.
• Hugging Face transformers — community Llama implementation with broader tokenizer / quantization / training support, at the cost of being a much larger codebase.