Assemble the Model

Goal

Assemble a small decoder-only PyTorch language model and trace the tensor shape from token IDs to vocabulary logits.

Put the Level 4 and Level 5 pieces on one path

You already know the pieces separately. Model assembly connects them:

token ids                (B, T)
  │ token embedding
  ▼
token vectors            (B, T, C)
  │ + position embedding
  ▼
position-aware vectors   (B, T, C)
  │ transformer blocks
  ▼
hidden states            (B, T, C)
  │ layer norm + lm head
  ▼
logits                    (B, T, V)

B is batch size, T is sequence length, C is the model's channel/embedding width, and V is vocabulary size.

Trace the model from token IDs to vocabulary logits

Select each interface. The batch and time axes stay aligned while the representation changes from token IDs to hidden features and finally vocabulary scores.

Token IDs: (B,T)

→→→

B

batch items

T

token positions

Token IDs: Each position contains one integer vocabulary ID.

A logit is a raw score. At each sequence position, the language-model head produces one score for every token in the vocabulary.

Predict

If batch=2, time=8, and vocab_size=66, what is the logits shape?

(2, 8, 66)(2, 66)(8, 8)(66, 8, 2)

Read the forward pass one boundary at a time

Token and position embeddings have the same channel width, so they can be added:

tok = self.token_embedding(idx)          # (B, T, C)
pos = self.position_embedding(positions) # (T, C)
x = tok + pos                            # broadcasts to (B, T, C)

Then the remaining path is compact:

x = self.token_embedding(idx) + self.position_embedding(positions)
x = self.blocks(x)
x = self.ln_f(x)
logits = self.lm_head(x)

Compact code hides several must-agree shape rules. The transformer stack must preserve (B, T, C), and the final head must replace the channel dimension with V vocabulary scores.

If targets are provided, cross-entropy compares those vocabulary scores with the correct next-token IDs:

loss = F.cross_entropy(
    logits.reshape(-1, vocab_size),
    targets.reshape(-1),
)

Lower cross-entropy means the correct next-token IDs are receiving more probability on average. It does not mean generated text is automatically good.

Trace the Lab

The Lab traces shapes for a batch of 2 sequences, 8 positions each, channel width 48, and vocabulary size 66.

1. Click Run unchanged.
2. Record the four printed shapes: token_ids (2, 8), hidden (2, 8, 48), transformer_output (2, 8, 48), and logits (2, 8, 66).
3. Confirm the path follows (B,T) → (B,T,C) → (B,T,C) → (B,T,V). Only the last step replaces C with V.
4. Find BLOCK_SIZE = 8. The position table has exactly one row for each position 0 to 7.
5. Change trace = shape_trace(2, 8, 48, 66) to trace = shape_trace(2, 9, 48, 66), asking for 9 positions.
6. Click Run. The Lab stops with ValueError: sequence length 9 is longer than block_size 8; no position row exists. A clear error at the boundary is much better than silently reading a position row that does not exist.
7. Press Reset afterward.

Loading lab…

When a shape or cross-entropy error appears, check the earliest boundary that could explain it:

• logits must end in vocab_size;
• targets must contain integer token IDs;
• each target ID must lie between 0 and vocab_size - 1;
• batch/time counts must still match after flattening.

Why this tiny model is enough for the lesson

The tiny model has tens of thousands of parameters rather than billions. That is intentional. The architecture and training mechanics are visible on CPU without changing the core input-to-logits story.

Larger decoder-only models add engineering choices such as different position representations, normalization variants, optimized attention kernels, mixed precision, and distributed execution. Those choices change scale and efficiency, not the basic tensor path you are tracing here.

Quick Check

1. What does the language-model head produce at each position?

One hidden-state feature per residual channel after decodingOne probability per attention head before softmaxOne raw score for every token in the vocabulary

2. Why add position information?

To identify which vocabulary row each token ID selectsTo replace causal masking with a learned ordering ruleToken embeddings alone do not encode sequence order

3. Which dimension changes when the final language-model head maps hidden states to logits?

Time changes from T to VThe final channel dimension changes from C to VBatch size changes from B to V

0 of 3 questions answered.

Key Takeaways

• Token IDs become token and position-aware vectors before entering transformer blocks.
• Transformer blocks preserve (B,T,C).
• The language-model head converts C hidden features into V vocabulary logits.
• Cross-entropy compares those logits with next-token target IDs.
• Debug the earliest wrong shape or ID-range boundary before changing training settings.

Next Lesson

Next, L6.6 — Training Loop uses this assembled model to sample input/target windows, compute loss, backpropagate, update parameters, and inspect the evidence when training fails.

References

• PyTorch, Embedding.
• PyTorch, CrossEntropyLoss.