Stanford CS336 from-scratch LLM notes and code released: DDP sync was 54% of step time, GRPO prompt tricks
stanfordnlp · x · 2026-09-22
A student who completed all five Stanford CS336 assignments shared detailed notes and open-sourced the full from-scratch implementation (raw PyTorch + Triton, no Hugging Face): BPE tokenizer, Transformer, Triton kernels, distributed training, scaling laws, and GRPO.
Key experimental findings:
- FFN vs Attention: at 1024 seq length the FFN dominates (a key reason MoE lives on FFN); attention takes over around 16k context.
- Distributed training: naive DDP on 2 GPUs spent 54% of each step on gradient sync; overlapping comms with backward cut it to 3%; ZeRO-1 sharding reduced peak memory 68.7 → 47.7 GB.
- GRPO: with the r1-zero prompt a base model scored 0/1300 on GSM8K, but 224/1300 with a 3-shot prompt — RL works best amplifying signal that's already there.
Course materials and the writeup are publicly available, ideal for anyone who wants to understand LLM internals hands-on.
More from Companies & People
- Altman Calls for AI Standards: Outsiders Need Real Say, Power Must Not Concentrate — patience_cave · 2026-09-22
- Snorkel AI raises $350M Series E at $3.5B valuation, passes $375M ARR — ajratner · 2026-09-22
- Ex-Apple designer: taste isn't born, it's drilled in painfully — floguo · 2026-09-22
- MIT Tech Review's 35 Innovators Under 35 Includes Agent-Planning Founder Danijar Hafner — thegautamkamath · 2026-09-22
- Dreamforce talk: AI generates the work, but you still have to earn the trust — i_am_dy · 2026-09-22
- Rippling AI Lab goes live in Bengaluru, targeting trusted enterprise AI — SuB8u · 2026-09-22