Why is the softmax + cross-entropy at the LM head a training bottleneck?
hardAnswer
- For a vocab V=128k and hidden d=4k, the logit matrix (batch × seq × V) can dominate activations — often the single biggest tensor in the graph.
- Materializing it in fp32 wastes memory and compute.
- Modern trick: 'fused CE' — compute logits and cross-entropy in a single fused CUDA kernel that never materializes the full V-dim tensor.
- Liger Kernel and (2024) do this — 2-5x memory reduction on the LM head, faster training.
Check yourself — multiple choice
- LM head is negligible
- Huge V·d logit tensor dominates activations → fused CE kernels save 2-5× memory
- Only in inference
- Only for MoE
LM head cross-entropy is a memory hog → fused CE kernels cut activation memory 2-5×.
#architecture#pretraining
Practise LLMs & GenAI
214 interview questions in this topic.
Related questions
- How do you prevent MoE routing collapse?
- What parallelism strategies are combined to train a 70B model?
- What is embedding tying and when do LLMs use it?
- What is the context window and what tricks extend it?
- What is a Mixture-of-Experts (MoE) LLM?
- Walk through how one training step of an autoregressive LLM works.