Practice — The PyTorch Training Loop from Scratch (6 questions)

You are asked to train a named-entity tagger. Each example is a sentence of 3 to 80 token ids with one tag id (0-8) per token. The vocabulary has 30,000 entries and id 0 is reserved for padding. You have an nn.Module that maps a (B, T) batch of token ids to (B, T, 9) logits.

1. Write the collate_fn that turns a list of (tokens, tags) pairs into a batch. What value do you pad the tags with, and why not 0?
2. Write the loss computation for one batch so that padding positions contribute nothing and the loss is a per-token mean. Why does the normalization matter when batches have different amounts of padding?
3. A 7-sentence toy sample has lengths 5, 12, 3, 40, 7, 9 and 22 and you use batch_size=3 without shuffling. How many real tokens and how many padded cells does each batch contain, and what does that suggest for throughput?
4. How do you compute token accuracy for the whole validation set so that it is correct?

A teammate switches a vision model from fp32 to fp16 autocast on a GPU without bf16 support to double the batch size. The loss goes down for the first epoch, then plateaus; the pre-clip gradient norm is exactly 0.0 on many steps and some layers' .grad tensors are all zeros. They propose "just raise the learning rate". A second teammate says "add GradScaler and the problem disappears; also we should call clip_grad_norm_ right after loss.backward() as before".

1. Explain, with the actual number ranges of fp16, why small gradients vanish and what loss scaling does about it. Use a true gradient of 10^{-8} and an initial scale of 65{,}536.
2. A gradient whose true value is 2.0 appears in the same step. What happens at scale 65{,}536, and what does the scaler do next? Trace the scale until the step succeeds.
3. Write the correct step (autocast, scaler, clipping, scheduler) and say what is wrong with clipping right after backward().
4. Would you rather switch to bf16 if the hardware supported it? What do you give up?
5. Does fp16 autocast reduce the memory of weights and optimizer state for a 100M-parameter AdamW model? What does it reduce?