Transformer: Understanding Training and Inference

Transformer: Understanding Training and Inference

Using a Transformer Model: From Training to Inference

If you have implemented a transformer model in PyTorch, you can use the same code for both training and inference, but in very different ways. During training, you typically process a batch of fixed-length token sequences and update the model's weights. During inference, the weights are fixed, and the model generates new tokens one by one.

This difference significantly impacts performance. Training is dominated by large matrix multiplications and backpropagation. Inference is dominated by repeated forward passes, memory movement, and the need to keep previous attention keys and values available for the next token.

In this chapter, you will learn about:

• The autoregressive generation loop
• The difference between pre-filling and decoding
• Why caching keys and values is necessary
• How to implement a simple KV cache
• How to reason about the memory used by the cache

Autoregressive Generation Loop

A decoder-only transformer model predicts the next token from the tokens that precede it. The strict requirement to use only previous tokens is enforced by the causal attention mechanism. If the input tokens are:

The cat sat on the

The model returns a probability distribution over the vocabulary for the next token. A likely next token might be "mat", but the model does not directly return a word. It returns logits, which are unnormalized scores for each token in the vocabulary.

The generation loop is therefore simple:

• Tokenize the prompt.
• Run the model to get the logits for the next token.
• Choose a token from the logits.
• Add this token to the input.
• Repeat until a stopping rule is reached.

This is called autoregressive generation because each new token depends on the tokens generated previously. The model cannot generate the tenth output token before knowing the first nine output tokens.

A very small greedy decoding loop can be written as follows:

import torch
@torch.no_grad()
def greedy_decode(model, input_ids, max_new_tokens):
    output_ids = input_ids.clone()
    for _ in range(max_new_tokens):
        logits = model(output_ids)
        next_token_logits = logits[:, -1, :]
        next_token = next_token_logits.argmax(dim=-1, keepdim=True)
        output_ids = torch.cat([output_ids, next_token], dim=1)
    return output_ids

In the code above, model is a PyTorch model, max_new_tokens is a positive integer, and all other variables are PyTorch tensors. The for loop iterates max_new_tokens times, and at each iteration, it feeds the entire sequence back to the model to get the logits for the next token. The argmax() function selects the token with the highest score. The cat() function is used to concatenate the new token to the output sequence, which will be used in the next iteration until the stopping rule is reached.

This code is easy to understand, but it is inefficient. At each iteration, it feeds the entire sequence back to the model. If the prompt has 1,000 tokens and you generate 100 new tokens, the model repeatedly recomputes the hidden states for the same prompt tokens. The model processes O(N²) tokens in this function, for a prompt of length N.

The actual time complexity of the code is even worse. Without caching, each forward pass recomputes attention for all tokens in the growing sequence. If the sequence length is N, self-attention has a score computation complexity of O(N²). For generation, this means you repeat a large amount of work. (Specifically, if the output sequence length is N=P+G with the prompt length P and the number of generated tokens G, the naive computation complexity should be O(P²G + PG² + G³). With caching, we can reduce it to O(P² + PG).)

Inference systems mitigate this by splitting generation into two phases: pre-filling and decoding.

Pre-filling and Decoding

Generation typically starts with a prompt. The prompt is known before the generation begins. The model can process all tokens of the prompt in a single forward pass. This is called the pre-filling phase.

During pre-filling, the model computes the hidden states for all tokens of the prompt and produces logits for the next token. It also computes the keys and values for all attention layers. These keys and values can be saved as they will be needed for each future token.

After selecting the first new token, generation enters the decoding phase. In this phase, the model receives only the most recent token. It computes the query, key, and value for this token, adds the new key and value to the cache, and attends to the new query over all cached keys and values.

This changes the cost of a decoding step. Instead of recomputing attention for the entire sequence, the model computes attention for a single new query against all previous keys. The attention cost per token drops from O(N²) to O(N) for a sequence of length N. The pre-filling step is still O(N²), but it is only performed once for the prompt.

This distinction is significant enough that service systems generally measure pre-filling and decoding separately:

• Pre-filling affects the time to the first token. A slow pre-filling increases the time to the first token.
• Decoding affects the speed of streaming output tokens. A slow decoding reduces the rate at which output tokens are streamed.

A short prompt with a long response emphasizes decoding. A long prompt with a short response emphasizes pre-filling. A chat application with a long conversation history emphasizes both.

A Simple KV Cache

The KV cache is where the model stores the attention keys and values produced by previous tokens. To see how this works, you do not need a large model. The following code builds a small transformer-like model with a cache.

This model is not intended to produce useful text. Its purpose is to show how the cache is created during pre-filling and extended during decoding.

import math
import torch
import torch.nn as nn
import torch.nn.functional as F

class SelfAttention(nn.Module):
    def __init__(self, hidden_size, num_heads):
        super().__init__()
        assert hidden_size % num_heads == 0
        self.num_heads = num_heads
        self.head_dim = hidden_size // num_heads
        self.qkv = nn.Linear(hidden_size, 3 * hidden_size)
        self.out = nn.Linear(hidden_size, hidden_size)

    def forward(self, x, past_kv=None):
        batch_size, seq_len, hidden_size = x.shape
        qkv = self.qkv(x)
        qkv = qkv.view(batch_size, seq_len, 3, self.num_heads, self.head_dim)
        qkv = qkv.permute(2, 0, 3, 1, 4)
        q, k, v = qkv[0], qkv[1], qkv[2]
        if past_kv is not None:
            past_k, past_v = past_kv
            k = torch.cat([past_k, k], dim=2)
            v = torch.cat([past_v, v], dim=2)
        total_len = k.size(2)
        past_len = total_len - seq_len
        scores = q @ k.transpose(-2, -1)
        scores = scores / math.sqrt(self.head_dim)
        causal_mask = torch.ones(seq_len, total_len, device=x.device, dtype=torch.bool)
        causal_mask = torch.tril(causal_mask, diagonal=past_len)
        scores = scores.masked_fill(~causal_mask, float("-inf"))
        attn = F.softmax(scores, dim=-1)
        y = attn @ v
        y = y.transpose(1, 2).contiguous().view(batch_size, seq_len, hidden_size)
        return self.out(y), (k, v)

class Block(nn.Module):
    def __init__(self, hidden_size, num_heads):
        super().__init__()
        self.attn_norm = nn.LayerNorm(hidden_size)
        self.attn = SelfAttention(hidden_size, num_heads)
        self.ffn_norm = nn.LayerNorm(hidden_size)
        self.ffn = nn.Sequential(
            nn.Linear(hidden_size, 4 * hidden_size),
            nn.GELU(),
            nn.Linear(4 * hidden_size, hidden_size),
        )

    def forward(self, x, past_kv=None):
        attn_out, new_kv = self.attn(self.attn_norm(x), past_kv=past_kv)
        x = x + attn_out
        x = x + self.ffn(self.ffn_norm(x))
        return x, new_kv

class TinyCausalLM(nn.Module):
    def __init__(self, vocab_size=128, hidden_size=64, num_heads=4, num_layers=2):
        super().__init__()
        self.token_emb = nn.Embedding(vocab_size, hidden_size)
        self.blocks = nn.ModuleList([
            Block(hidden_size, num_heads) for _ in range(num_layers)
        ])
        self.norm = nn.LayerNorm(hidden_size)
        self.lm_head = nn.Linear(hidden_size, vocab_size, bias=False)

    def forward(self, input_ids, past_kv=None):
        x = self.token_emb(input_ids)
        new_cache = []
        if past_kv is None:
            past_kv = [None] * len(self.blocks)
        for block, layer_past in zip(self.blocks, past_kv):
            x, layer_cache = block(x, past_kv=layer_past)
            new_cache.append(layer_cache)
        logits = self.lm_head(self.norm(x))
        return logits, new_cache