Pretraining Objective

A decoder-only language model consumes tokens from the left context and predicts the next token at each position. The common implementation pattern is to produce logits with shape [batch, time, vocab], then compare logits[:, :-1] with labels[:, 1:]. Hugging Face's causal language modeling guide notes that its data collator can copy input_ids into labels because the alignment shift is handled by the model path. In hand-written PyTorch, make the shift explicit so tests can catch off-by-one errors.

def shift_for_causal_lm(logits, labels):
    # logits: [batch, time, vocab], labels: [batch, time]
    return logits[:, :-1, :].contiguous(), labels[:, 1:].contiguous()

1. One-token shift: position t predicts the label at t + 1; the final logit has no next token target in the same sequence.
2. Do not double shift: if a framework model shifts internally, pass unshifted labels; if you own the loss function, shift exactly once.
3. Teacher forcing: training sees the true previous tokens, while generation feeds back sampled tokens one step at a time.

PyTorch CrossEntropyLoss supports ignore_index: ignored targets do not contribute to the loss or input gradients. For language modeling, that is the clean way to exclude padding tokens and any prompt positions that should condition the model without being scored. Keep the attention mask and loss mask separate: attention controls what the model can read, while ignore_index controls which predictions count toward the objective.

IGNORE_INDEX = -100

labels = input_ids.clone()
labels[attention_mask == 0] = IGNORE_INDEX
labels[:, :prompt_len] = IGNORE_INDEX  # Optional for completion-only SFT.

1. Mask padding: padded positions should not teach the model to emit pad tokens unless that is a deliberate objective.
2. Mask prompts deliberately: for instruction tuning, score only the assistant/completion span when the prompt is context rather than target text.
3. Average honestly: report loss over non-ignored targets, not over padded sequence length.

Packing improves token utilization by concatenating multiple short examples into one fixed-length training sequence. The contract is stricter than simple concatenation: labels must still predict the next token within the intended example, and the attention mask must prevent a later packed sample from using tokens from an unrelated earlier sample unless that boundary was deliberately bridged with an EOS objective. Hugging Face's CLM guide shows concatenating text into blocks for next-token training, while TRL and NVIDIA NeMo document packing as an efficiency feature that needs explicit packing metadata and masks.

import torch

def packed_causal_keep_mask(segment_ids):
    # segment_ids: [time], same id means tokens belong to one packed example.
    time = segment_ids.numel()
    causal = torch.ones(time, time, dtype=torch.bool).tril()
    same_segment = segment_ids[:, None] == segment_ids[None, :]
    return causal & same_segment

segment_ids = torch.tensor([0, 0, 0, 1, 1, 2])
keep = packed_causal_keep_mask(segment_ids)
assert keep[4, 3]        # sample 1 can read its own previous token
assert not keep[4, 2]    # sample 1 cannot read sample 0

1. Track segment ids: keep enough metadata to know which packed tokens came from the same source example or conversation.
2. Mask attention boundaries: a plain lower-triangular mask allows cross-example leakage inside the same packed row.
3. Mask loss boundaries: ignore boundary targets when the next token would come from a different example and should not be learned as a continuation.
4. Test one tiny pack: assert allowed and blocked positions before trusting a fused attention backend or trainer packing flag.

Sources: Hugging Face causal language modeling guide, Hugging Face TRL SFTTrainer packing docs, NVIDIA NeMo Megatron Bridge packed sequences, Hugging Face sequence packing note.

The interview-safe version is small but strict: assert shapes, shift once, flatten only after alignment, and pass ignore_index through to cross entropy. This mirrors the math while keeping the batch and time dimensions inspectable until the final loss call.

import torch
import torch.nn.functional as F

IGNORE_INDEX = -100

def causal_lm_loss(logits, labels, ignore_index=IGNORE_INDEX):
    if logits.ndim != 3 or labels.ndim != 2:
        raise ValueError("expected logits [batch, time, vocab] and labels [batch, time]")
    if logits.shape[:2] != labels.shape:
        raise ValueError("logits and labels must share batch/time dimensions")

    shift_logits = logits[:, :-1, :].contiguous()
    shift_labels = labels[:, 1:].contiguous()
    vocab = shift_logits.size(-1)

    return F.cross_entropy(
        shift_logits.view(-1, vocab),
        shift_labels.view(-1),
        ignore_index=ignore_index,
    )

def test_causal_lm_loss_ignores_padding():
    torch.manual_seed(0)
    logits = torch.randn(2, 5, 11, requires_grad=True)
    labels = torch.randint(0, 11, (2, 5))
    labels[0, -2:] = IGNORE_INDEX

    loss = causal_lm_loss(logits, labels)
    assert torch.isfinite(loss)
    loss.backward()
    assert torch.isfinite(logits.grad).all()

1. Test the alignment: create a tiny batch where the correct next token is obvious and prove the loss uses labels[:, 1:].
2. Test ignored rows: include padding or prompt-masked labels and confirm gradients remain finite.
3. Test framework parity: if you switch to AutoModelForCausalLM, verify whether the model shifts labels internally before reusing this helper.

Sources: Hugging Face causal language modeling guide, Hugging Face LLM course causal LM note, PyTorch CrossEntropyLoss reference.