Pre-training: model additions

The LLM implementation we have built so far (the inference-only model code on GitHub) is not suitable to use for pre-training. To do pre-training, we will use almost the exact same model code, but will need to make a few amendments to it to allow for:

1. Batching (passing multiple inputs in at the same time)
2. Producing next token predictions for every input token

Batching

Our current LLM implementation accepts a single sequence of tokens, like

[“It”, “ was”, “ a”, “ bright”, “ cold”]

which we then use to predict the next token. When pre-training, we could just feed in a single input sequence at a time and adjust weights each time, but this would be very slow. We instead augment our LLM to allow for inputting a whole batch of inputs, and adjust weights after each batch:

Updating our embedding code to allow for batches

Suppose our batched token_ids variable now looks like this:

token_ids = torch.tensor([
    [1026,   373,  257, 6016,  4692],
    [1110,   287, 3035,   11,   290],
    [ 262, 29906,  547, 8871, 28306],
])

Our code for embedding our now becomes

def forward(self, token_ids):
    if token_ids.ndim != 2:
        raise ValueError(
            "token_ids must have shape (batch_size, sequence_length)"
        )

    sequence_length = token_ids.size(1)
    if sequence_length > self.max_sequence_length:
        raise ValueError(
            f"input has {sequence_length} tokens, but the maximum is "
            f"{self.max_sequence_length}"
        )

    positions = torch.arange(
        sequence_length,
        device=token_ids.device,
    )
    token_vectors = self.token_embedding(token_ids)
    position_vectors = self.position_embedding(positions)

    return token_vectors + position_vectors

This method now returns our embeddings:

Updating our attention code to allow for batches

The attention calculation itself is unchanged. The only new concept is choosing the correct axes when splitting token vectors into heads and joining them again:

batch_size, num_tokens, _ = x.shape

queries = self.W_query(x)
keys = self.W_key(x)
values = self.W_value(x)

queries = queries.view(
    batch_size,
    num_tokens,
    self.num_heads,
    self.head_dim,
).transpose(1, 2)
keys = keys.view(
    batch_size,
    num_tokens,
    self.num_heads,
    self.head_dim,
).transpose(1, 2)
values = values.view(
    batch_size,
    num_tokens,
    self.num_heads,
    self.head_dim,
).transpose(1, 2)

attention_scores = queries @ keys.transpose(-2, -1)

mask = torch.triu(
    torch.ones(
        num_tokens,
        num_tokens,
        dtype=torch.bool,
        device=x.device,
    ),
    diagonal=1,
)
attention_scores = attention_scores.masked_fill(mask, float("-inf"))

attention_weights = torch.softmax(
    attention_scores / self.head_dim**0.5,
    dim=-1,
)

context = attention_weights @ values
context = context.transpose(1, 2).reshape(
    batch_size,
    num_tokens,
    self.num_heads * self.head_dim,
)

return self.out_proj(context)

After transpose(1, 2), each sequence and each attention head has its own token-by-token attention calculation. The batch items never attend to one another; they are simply processed in parallel.

The causal mask still has shape (tokens, tokens). PyTorch broadcasts it across every batch item and attention head, so the same look-ahead rule applies throughout the batch.

Producing next token predictions for every input token

Our existing inference-only model selects the final token vector before applying the vocabulary projection. For pre-training, we instead apply the vocabulary projection to every token in the sequence, in order to obtain logits (which can be translated into prediction probabilities) for every token:

The code is very similar:

def forward(self, token_ids):
    if token_ids.ndim != 2:
        raise ValueError(
            "token_ids must have shape (batch_size, sequence_length)"
        )
    if token_ids.size(1) == 0:
        raise ValueError("each sequence must contain at least one token")

    x = self.embedding(token_ids)

    for block in self.transformer_blocks:
        x = block(x)

    x = self.final_norm(x)

    # Training needs a next-token prediction at every input position.
    return self.vocabulary_projection(x)

If we want to use our updated model for inference, we still can; we just select the final prediction after running the model:

with torch.inference_mode():
    all_logits = model(token_ids)
    next_token_logits = all_logits[0, -1]

0 selects the single prompt in the batch, and -1 selects its final position. This lets one model implementation support both batched training and ordinary next-token generation.