Video summary

Stanford CS336 Language Modeling from Scratch | Spring 2026 | Lecture 2: PyTorch (einops)

Main summary

Key takeaways

Educational

Main ideas and lessons

  • Context & project update

    • The instructor shares results from a “Marine” project:
      • many smaller model runs were used to fit scaling laws
      • an “optimal” compute point was predicted and run
    • The achieved loss matched the forecast closely (within ~0.05).
    • The result is extrapolated to hypothetical GPT-5-level performance (with the caveat that scaling laws may vary).

  • What the course is trying to do (resource accounting)

    • Core goal: train the best model possible given finite resources (primarily compute and memory; data is not treated as a limiting factor here).
    • Objective: maximize computational efficiency.
    • To do that, you must understand compute and memory characteristics of operations and training.

  • Back-of-the-envelope performance questions (practice problems)

    • Example question: how long to train a 70B-parameter model on 15T tokens using 1024 H100 “cell” hardware (clarifies what an H100 “cell” means).
    • Uses a FLOPs estimate of the form:
      • Required FLOPs ≈ 6 × (number of parameters) × (number of tokens) (presented as the “formula we’ll talk about”).
    • Another example: maximum trainable model size using H100s with AdamW, using memory constraints:
      • H100 has 80 GB HBM
      • Memory per parameter involves optimizer/gradient/parameter storage; the instructor gives bytes-per-parameter as a combination (2 + 2 + 4 + 4) (explained later).
      • Important caveat: activations are not included in this rough estimate; activations depend on batch size and sequence length.

  • Training “mechanics” vs systems/resource thinking

    • “Mechanics” today: how PyTorch tensors operate (no ML magic).
    • Mindset: when writing code, continuously consider performance characteristics.
    • Intuition: learn how resources (memory/compute) get spent.

Detailed methodology / instruction-like content

1) Resource accounting workflow (implicit method)

  • Count memory

    • For tensors: memory usage ≈ (number of elements) × (bytes per element).
    • Determine dtype (FP32/FP16/BF16/etc.) to get bytes-per-element.

  • Count compute

    • Use FLOPs to estimate compute cost.
    • Convert FLOPs to time using measured/specified hardware FLOPs/sec.

  • Diagnose bottlenecks

    • Compare:
      • Arithmetic intensity of an algorithm vs
      • Accelerator intensity (hardware FLOPs/sec divided by bytes/sec).
    • Decide whether the operation is:
      • memory-bound (low arithmetic intensity) or
      • compute-bound (high arithmetic intensity)

2) Floating-point precision guidance (training implications)

  • FP32

    • 32-bit float: sign + exponent + mantissa.
    • Safe and straightforward but uses 4 bytes per value.

  • FP16

    • 16-bit float: fewer exponent bits → poor dynamic range → instability (overflow/underflow, NaNs).

  • BF16

    • 16-bit format that preserves exponent range closer to FP32 while reducing mantissa resolution.
    • Presented as a sweet spot for training stability vs efficiency.

  • Mixed precision training (typical practice)

    • Use BF16 for parameters, activations, gradients.
    • Use FP32 for optimizer states (for stability).
    • Use PyTorch AMP:
      • wrap code with AMP so PyTorch casts where safe
      • heuristic mentioned: matmuls are generally safe in lower precision; exponentiation may remain FP32

  • Lower precision notes

    • Mentions FP8 (two versions depending on dynamic range vs resolution); supported via Nvidia “Transformer Engine.”
    • Mentions FP4 (NVFP4):
      • only 4 bits per value, but uses block-wise scaling so effective range is improved
      • example model: NeMo-3 Super trained in FP4
    • Key point: many low-precision formats are handled under the hood by Nvidia/software stacks.

  • Important distinction: training vs inference

    • Very low-bit training (like 1-bit) is described as essentially unrealistic; quantization is typically done after training.

3) Einops / einsum approach to tensor operations (with conceptual rules)

  • Why named-dimension notation
    • Index-heavy code (like transpose with cryptic indices) is error-prone.
    • Einops offers clearer, named-dimension manipulation inspired by Einstein summation notation.

A) Einsum (einsum notation as generalized matmul)

  • General rule

    • Name dimensions in inputs; output names dimensions you keep.
    • Any dimension not in the output is implicitly summed out.

  • Example: matrix multiplication

    • If:
      • X has dims (seq1, hidden)
      • Y has dims (hidden, seq2)
    • Then Z has dims (seq1, seq2), and hidden is summed.

  • More complex/batched example

    • Uses named dimensions to avoid reasoning about transposes manually.
    • Supports ellipsis-style batching:
      • ... represents unspecified batch dims (batch, sequence, head, etc.).

B) Reduce (generalization of sum/mean/max/min)

  • Rule

    • Specify which named dimensions remain; all other included dimensions are reduced using an operation (e.g., sum, mean, max, min).

  • Performance note

    • Considered “syntax sugar” over primitives: no significant speed advantage beyond underlying ops.

C) Rearrange (reshape by splitting/merging dims)

  • Rule
    • Reshape by explicitly mapping one dimension into multiple (or vice versa) using parentheses grouping.
    • Example concept:
      • a dimension of size 8 representing (2 × 4) is split into two dims
      • then you can do a matmul on the structured shape and rearrange back

Compute accounting details

1) FLOPs terminology clarification

  • flops: count of operations (lowercase s)
  • FLOPS / FLOPS/s: hardware throughput measure (per second; instructor clarifies with “/s”)

2) FLOPs for matrix multiplication (core primitive)

  • For multiplying a matrix of shape (B × D) by a matrix of shape (D × K):

    • FLOPs ≈ 2 × B × D × K
      • “2” corresponds to multiply + add per inner-product term.

  • Memory/compute intuition

    • Forward matmul cost scales like:
      • (tokens or data points) × (parameters)
    • This generalizes to transformers and motivates 6 × N × D-style training accounting.

3) FLOPs → time: benchmarking and MFU

  • GPU timing considerations

    • Use CUDA synchronization before and after to avoid falsely short timings from async execution.
    • Repeat and average timings.

  • MFU (Model FLOPs Utilization)

    • Definition:
      • MFU = (actual achieved FLOPs/sec) / (promised/spec FLOPs/sec)
      • (approx ignoring comms/overhead)
    • Typical expectation:
      • modern models: MFU around 0.5 is “pretty happy”
      • pure matmul could be higher (e.g., 0.8), but real models often lower

Arithmetic intensity and roofline-style diagnosis (system-level methodology)

1) Hardware “intensity”

  • Accelerator intensity ≈ (FLOPs/sec) / (bytes/sec) from the spec sheet.
  • Example value given for H100: around ~295 FLOPs per byte.

2) Algorithm arithmetic intensity

  • Arithmetic intensity = (algorithm FLOPs) / (bytes moved).

  • Examples

    • ReLU
      • moves ~4N bytes total and does ~N FLOPs → intensity ~ 0.25
      • concluded memory-bound
    • GELU
      • more math per element, but still memory-bound
    • Dot product
      • intensity ~ 0.5 → memory-bound
    • Matrix-vector product
      • still memory-bound (intensity only slightly better)
    • Matrix multiplication
      • intensity ~ ~N/3 (roughly)
      • for sufficiently large matrices, becomes compute-bound
      • key conclusion: large matmuls are compute-bound; smaller ones may not fully utilize compute

3) Bottleneck criterion

  • If arithmetic intensity < accelerator intensity → memory-bound
  • If arithmetic intensity > accelerator intensity → compute-bound

4) Roofline plot interpretation

  • X-axis: arithmetic intensity
  • Y-axis: realized FLOPs/sec
  • As intensity increases, performance rises toward peak until becoming compute-bound; cannot exceed peak FLOPs.

Training compute/memory accounting

1) Where the “6 × tokens × parameters” comes from

  • For a deep network built from repeated matmuls + elementwise activations:

    • Forward pass FLOPs ≈ ~2 × (tokens/data points) × (parameters)
    • Backward pass ≈ 2 × forward
      • computes:
        • gradients w.r.t. parameters
        • gradients w.r.t. inputs
    • Total training FLOPs ≈ 6 × tokens × parameters

  • Extension note

    • Approximates transformers when context length isn’t too large
    • If context length is large, attention adds extra ~context² FLOPs not captured in the simpler formula

2) Optimizer cost and memory footprint (AdamW-style notes)

  • Optimizer state is a major source of memory.
  • Optimizer states often stored in FP32 for stability even when model/activations are FP16/BF16.
  • Note: optimizer memory usually limits fit, but isn’t usually the compute-speed bottleneck (compute dominates elsewhere).

3) Memory usage in training (and scaling intuition)

  • Activations scale with:
    • batch size
    • number of layers
  • Training stores activations for backprop; inference generally stores fewer activations because it doesn’t need gradients.

Two systems strategies to reduce memory

1) Gradient accumulation (microbatches)

  • Goal: emulate a larger batch size without storing activations for the entire large batch at once.
  • Mechanism:
    • split a large batch into microbatches
    • compute gradients per microbatch
    • accumulate gradients (don’t zero gradients each microstep)
    • once enough microsteps complete to reach the effective batch size:
      • perform an optimizer update
      • then zero gradients
  • Benefits: saves memory by reducing per-step activation footprint.
  • Tradeoff: more steps/overhead (not quantified in the subtitles).

2) Activation checkpointing / gradient checkpointing (rematerialization)

  • Key idea:
    • don’t store activations for every layer during forward
    • store activations only at checkpoints
    • during backward, recompute missing activations from the last checkpoint
  • PyTorch mechanism:
    • use torch.utils.checkpoint around blocks/layers
  • Example consequence:
    • checkpointing blocks containing linear + ReLU can save roughly about half activation memory (as described)
  • Extreme tradeoff:
    • store no activations → maximum memory saving, but recomputation cost grows toward L²
  • Suggested balance:
    • store checkpoints at about √L layers to balance memory and recompute overhead (both scale ~√L)

Lecture wrap-up / takeaways

  • Everything in training comes down to:
    • tensors (parameters/gradients/activations/data)
    • optimizer states
    • their memory + compute costs
  • Introduced Einops as a clearer mental model for tensor transformations.
  • Demystified 6 × (tokens) × (parameters) as training FLOPs accounting (forward + backward).
  • Introduced arithmetic intensity + roofline analysis to determine:
    • matrix multiplications are often compute-bound
    • most other ops are often memory-bound
  • Practical memory-reduction techniques:
    • gradient accumulation
    • activation checkpointing
  • Conclusion: reducing memory enables larger batch sizes and makes training more feasible.

Speakers / sources featured

  • Speaker: Instructor for Stanford CS336 (no name given in subtitles)
  • Referenced sources/products:
    • Marine project (internal project mentioned by instructor)
    • DeepSeek 3.2 model (example for tensor sizes)
    • PyTorch (tensors, AMP, torch.utils.checkpoint)
    • Einops library (einsum/reduce/rearrange concepts)
    • NVIDIA H100 (hardware specs discussion)
    • NVIDIA Transformer Engine (FP8 support mention)
    • NeMo-3 Super (FP4 training example)
    • Adagrad (2011) (optimizer mentioned)
    • Adam/AdamW (mentioned; assignment reference)
    • Jensen (referenced humorously as “tell Jensen” about hardware design)

Original video