Why CausalLM labels equal input_ids in Llama_cookbook (no shift needed)

In the Llama_cookbook finetuning example, labels are set equal to input_ids for the summary tokens (not shifted by one position) because the Hugging Face LlamaForCausalLM model internally shifts the logits before computing the loss. The model takes the full sequence of input tokens, generates logits for each position, then shifts those logits left by one token so that position i predicts token i+1. The labels therefore align with the original input_ids without any manual offset. The prompt tokens are masked with -100 so the loss function ignores them, focusing the training only on the summary portion.

The Full Answer

How the dataset is prepared in the cookbook

The cookbook example uses LlamaForCausalLM with the samsum dataset (dialog and summary pairs). The preprocessing code builds each sample like this:

prompt = tokenizer.encode(tokenizer.bos_token + sample["prompt"], add_special_tokens=False)
summary = tokenizer.encode(sample["summary"] + tokenizer.eos_token, add_special_tokens=False)

sample = {
    "input_ids": prompt + summary,
    "attention_mask" : [1] * (len(prompt) + len(summary)),
    "labels": [-100] * len(prompt) + summary,
}

The prompt tokens get -100 as labels, and the summary tokens get the same values as input_ids. This is the pattern that confuses many readers.

Why the labels are not shifted

The accepted answer from the Stack Overflow thread explains that the CausalLM paradigm predicts token i based on tokens 0 through i-1. When you look at the whole sequence, each labels[i] should equal input_ids[i] because the model internally handles the shift. The confusion often comes from zero-based indexing: input_ids[0:i] does not include the ith element itself.

Hugging Face's LlamaForCausalLM (and all *ForCausalLM models) implement this shift in the forward pass. The model computes logits for every position in the input sequence, then shifts the logits so that logits at position i correspond to predicting token i+1. The loss is then computed between these shifted logits and the labels. Since the labels are aligned with the original input_ids, no manual shift is needed.

How the loss calculation works

The loss function (usually CrossEntropyLoss) receives the shifted logits and the labels. The -100 values in the labels act as an ignore index. CrossEntropyLoss in PyTorch and Hugging Face automatically skips positions where the target is -100. This means the prompt portion contributes nothing to the gradient, and only the summary tokens drive the training.

What happens during generation

During training, the model generates the full sequence (prompt + summary) in one forward pass. The loss is computed on the entire sequence at once. If the generated sequence is shorter than the labels, it gets padded. If it is longer, it gets truncated. The shift ensures that the model learns to predict each summary token given all previous tokens (including the prompt).

Verifying with the dataloader output

The questioner printed a batch from the dataloader and saw:

batch = next(iter(train_dataloader))
print(batch['input_ids'][0][35:40])
print(batch['labels'][0][35:40])

# Output:
# tensor([19791, 512, 32, 36645, 41778])
# tensor([ -100, -100, 32, 36645, 41778])

Positions 35 and 36 are still in the prompt (labels are -100), while positions 37 through 39 are the first summary tokens. The labels match the input_ids for the summary part, confirming no shift was applied.