Stanford CS336 Language Modeling from Scratch | Spring 2026 | Lecture 2: PyTorch (einops)
TL;DR
This lecture covers resource accounting fundamentals for training large language models, including FLOPs calculations and memory estimation, explores numerical precision trade-offs from FP32 down to FP4, and introduces einops as a readable alternative to PyTorch tensor operations using named dimensions.
💻 Resource Accounting & Training Efficiency 3 insights
Training FLOPs calculation formula
Estimate total compute required using 6 × parameters × tokens, which represents the floating point operations needed for a full training run.
Hardware utilization reality check
Training a 70B parameter model on 15 trillion tokens requires approximately 143 days on 1,024 H100 GPUs assuming 50% Model FLOPs Utilization (MFU).
Memory-constrained capacity limits
Eight H100s with 80GB memory can theoretically train approximately 53 billion parameters using AdamW, calculated from 12 bytes per parameter (weights, gradients, and two optimizer states).
🔢 Numerical Precision Trade-offs 3 insights
BF16 as the practical training standard
Brain Float 16 maintains the same dynamic range as FP32 with half the memory footprint, avoiding the underflow and overflow instability that plagues FP16 training.
Mixed precision training strategy
Use BF16 for parameters, activations, and gradients while keeping optimizer states in FP32 for numerical stability, typically managed via PyTorch's AMP library.
Frontier quantization formats
FP8 and FP4 formats with block scaling enable extreme memory reduction, with Nemotron 3 Super demonstrating successful training at 4-bit precision using NVIDIA's transformer engine.
🧮 Einops for Tensor Operations 3 insights
Named dimensions replace integer indices
Einops uses explicit dimension names instead of cryptic axis numbers like -1 or -2, eliminating transpose errors and making tensor shapes self-documenting.
Einstein summation simplified
Operations follow input-to-output notation where dimensions listed only on the input side are automatically summed, such as 'seq hidden, hidden seq2 -> seq seq2' for matrix multiplication.
Arbitrary batch dimension handling
The ellipsis (...) syntax accommodates any number of leading batch dimensions without explicit enumeration, crucial for language modeling with variable batch and sequence structures.
Bottom Line
Always perform back-of-the-envelope resource accounting before training, default to BF16 mixed precision for optimal efficiency, and use einops to write readable tensor operations that prevent shape errors.
More from Stanford Online
View all
Stanford CME296 Diffusion & Large Vision Models | Spring 2026 | Lecture 7 - Evaluation
This Stanford lecture establishes aesthetics and prompt adherence as the dual pillars for evaluating text-to-image models, compares human evaluation methods from noisy absolute ratings to reliable pairwise comparisons, and details the ELO rating system for robust model benchmarking before addressing the scalability crisis that necessitates automated metrics.
Stanford CS153 Frontier Systems | The Road Ahead: Resilience Required
Former federal prosecutor and tech security chief Joe Sullivan recounts his journey from prosecuting cybercrime to leading security at eBay, Facebook, Uber, and Cloudflare, sharing hard-won lessons on the critical importance of transparency in security incidents through the lens of his personal prosecution for the 2016 Uber data breach cover-up.
Stanford CS336 Language Modeling from Scratch | Spring 2026 | Lecture 16: Post-Training - RLVR
This lecture explains why RLHF hits overoptimization limits with learned reward models, and how RLVR (Reinforcement Learning from Verifiable Rewards) enables unlimited compute scaling on verifiable tasks like math and coding through simpler algorithms like GRPO.
Stanford CS336 Language Modeling from Scratch | Spring 2026 | Lecture 15: Mid/Post-Training
This lecture explains how post-training transforms raw pre-trained models like GPT-3 into instruction-following systems like ChatGPT through supervised fine-tuning and reinforcement learning, emphasizing that high-quality data curation matters more than algorithmic sophistication.