LM head
The final layer that turns the last position's vector into one logit per vocabulary token.
Interview questions about LM head
- Your LM head is too big for one device. How do you shard it by vocabulary or by hidden width, combine the logits, and compute an exact top-k across shards?
Vocabulary sharding: each device holds a slice of rows and computes logits for its slice only.
- What are logits, and how does a model go from them to a whole sentence of output?
Logits are the raw scores the LM head gives every vocabulary token, 50,257 for GPT-2, from the last position's vector.