Causal LM Architectures

Ever picked a model off the Hugging Face Hub, slapped it into your training script, and realized too late that its architecture silently shaped everything — from memory usage to whether your labels even make sense? That's the hidden trap when you're fine-tuning causal language models. This lesson gives you the mental map and hands-on moves to inspect, compare, and choose the right architecture for your next fine-tuning run, so you stop guessing and start shipping.

The Problem This Lesson Solves

When you fine-tune a causal language model (LM), the architecture is not just a container for weights — it dictates how the model attends to tokens, how autoregressive generation behaves, and how much VRAM you'll burn. If you ignore it, you'll hit mysterious CUDA out of memory errors, see generations that look repetitive and broken, or waste hours training on a model that can't do the task you need.

The reality? Most developers grab the newest base checkpoint from the Hub without second-guessing the design. They then wonder why their chat dataset produces loss curves that plateau, or why the model repeats the same phrase forever. Architecture choices — decoder-only versus encoder-decoder, attention variants, token-to-loss mapping — ripple through every downstream step: data prep, training loop, and even how you evaluate.

By the end of this lesson, you'll be able to look at any causal LM checkpoint and understand its architecture, pick the right one for your use case, and inspect your training setup with targeted probes. No black boxes.

Core Concept / Mental Model

Think of a causal LM as a left-to-right storyteller. It reads your prompt token by token, and at each position, it only looks at what came before — never the future. This is the causal attention mask: a triangular matrix that zeros out future tokens. The transformer blocks process the sequence autoregressively, predicting the next token at each step.

Architecture families differ in how they structure this storytelling:

• Decoder-only: pure storyteller. GPT, Llama, Mistral, Phi. Ideal for generation, chat, and instruction following. Most generative LLMs you see fall here.
• Encoder-decoder: a two-act play. The encoder reads the whole input bidirectionally (like a full reader), then the decoder generates output causally. T5, BART. Great for translation, summarization where the source is fully understood before output.
• Causal decoder with prefix: a hybrid where part of the input is attended bidirectionally (the prefix), and the rest is causal. Used in some summarization and infilling models.

Inside any architecture, you also have configuration knobs: number of layers, heads, hidden size, and attention variant (e.g., flash_attention_2). These aren't just metadata — they determine whether your fine-tune fits on a GPU and how fast it runs.

Pro tip: The model card and config.json are your first stop. They tell you the architecture family, the vocab size, and whether it was trained with a causal mask — no guessing.

How It Works Step by Step

When you instantiate a causal LM with Hugging Face AutoModelForCausalLM, here's what happens under the hood:

1. Config parsing: The library reads config.json and builds a model class matching the architectures field (e.g., LlamaForCausalLM).
2. Embedding: Tokens become vectors — size vocab_size x hidden_size.
3. Positional encoding: Rotary or learned, adds order info.
4. Transformer blocks: Each block runs self-attention (with a causal mask) through a feed-forward network.
5. Output head: The final hidden states map to vocab_size logits.

Crucially, when you fine-tune, the loss is computed only over the target tokens (the ones after your prompt), because every position in the input gets a next-token prediction. If your labels don't align that way, you get silent training issues.

At inference, generation uses this same causal forward pass: you feed the prompt, sample the next token, append it, and repeat. That's why the architecture can't be swapped casually across tasks — the attention pattern is baked in.

Hands-on Walkthrough

Let's put this into practice. You'll inspect a model's architecture, verify its causal nature, and prepare a tiny training example.

1. Inspect the architecture

from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer

model_name = "gpt2"  # classic small causal LM
config = AutoConfig.from_pretrained(model_name)
print(config.architectures)  # -> ["GPT2LMHeadModel"]
print(config.model_type)     # -> "gpt2"
print(config.n_layer, config.n_head, config.n_embd)  # -> 12 12 768

tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name)

# Confirm the causal mask in action: logits depend only on past tokens.
import torch
inputs = tokenizer("The capital of France is", return_tensors="pt")
with torch.no_grad():
    outputs = model(**inputs)
logits = outputs.logits  # shape (1, seq_len, vocab_size)

# For each position, the logits predict the *next* token.
next_token_id = logits[0, -1].argmax().item()
print(tokenizer.decode(next_token_id))  # e.g., " Paris"

Expected output:

["GPT2LMHeadModel"]
gpt2
12 12 768
 Paris

2. Build a causal training example

For fine-tuning, the model expects labels that are shifted right by one token.** In practice, you pass the same input as labels, and the library masks them internally.

from transformers import Trainer, TrainingArguments

# Minimal dataset with prompt/completion structure
train_data = [
    {"text": "Question: What is 2+2?\nAnswer: 4"},
    {"text": "Question: What is the sky?\nAnswer: Blue"},
]

def tokenize_function(examples):
    return tokenizer(examples["text"], truncation=True, max_length=128)

tokenized = tokenize_function({"text": [d["text"] for d in train_data]})
tokenized["labels"] = tokenized["input_ids"].copy()

# In the Trainer, each batch's labels are shifted by -1 automatically.

training_args = TrainingArguments(
    output_dir="./gpt2-finetuned",
    num_train_epochs=1,
    per_device_train_batch_size=2,
    learning_rate=5e-5,
    logging_dir="./logs",
)

trainer = Trainer(model=model, args=training_args, train_dataset=tokenized)

# train() would work, but we'll keep it instructional
trainer.train()

Note: If you use AutoModelForCausalLM, the loss is automatically computed only over the shifted labels — every position in the sequence gets a next-token prediction, and the loss is the mean negative log-likelihood over all positions.

3. Compare memory footprint across architectures

import torch

def count_params(model):
    return sum(p.numel() for p in model.parameters())

models = {
    "gpt2": AutoModelForCausalLM.from_pretrained("gpt2"),
    "facebook/opt-125m": AutoModelForCausalLM.from_pretrained("facebook/opt-125m"),
}

for name, model in models.items():
    print(f"{name}: {count_params(model)/1e6:.1f}M params")

Expected: roughly 124M for GPT-2, 125M for OPT — similar size, but different positional encodings (learned vs. learned, but different) and tokenizers.

Compare Options / When to Choose What

When to choose what

• Chat / assistant: decoder-only. It's the natural fit for autoregressive response generation.
• Document summarization: you might get better results with an encoder-decoder that reads the whole document first. But many strong modern summarization models are decoder-only (e.g., GPT-3.5 summary).
• Code completion: decoder-only, because you generate token-by-token.

Troubleshooting & Edge Cases

1. "ValueError: Unrecognized model in transformers"

You picked a model with a custom architecture not in your transformers version. Fix: upgrade transformers or pick a standard one.

2. Loss doesn't decrease

If your labels include padding tokens, the model is forced to predict them. Mask padding with -100 in labels:

labels = input_ids.clone()
labels[labels == tokenizer.pad_token_id] = -100

3. Generating repetitive loops

Causal decode can easily repeat. Adjust no_repeat_ngram_size or temperature.

4. CUDA out of memory on larger models

Architecture size matters. Switch to a smaller variant or use parameter-efficient fine-tuning (LoRA) — but that's for a later lesson.

What You Learned & What's Next

You've demystified explore model architectures for causal lm. You can now inspect a checkpoint's architecture, understand how the causal mask shapes training and generation, and choose between decoder-only and encoder-decoder families. You also built a simple fine-tuning data pipeline and spotted classic pitfalls like padding-token loss.

This is the foundation for the next lesson: choosing a base model for your fine-tuning task. There, you'll learn how to map your task requirements (domain, language, context length) to the best pretrained checkpoint — using the architecture knowledge you just gained.

Keep this mental model handy: always check config.json, know your attention mask, and align labels carefully. Your future fine-tunes will thank you.