Memory-Constrained Trainer

Implement accumulation, checkpointing, mixed precision, all-reduce, and ZeRO-style optimizer sharding.

7 parts40 individual lessons3.3 estimated hours200 XP available

Memory-Constrained Trainer project artwork

0%0 of 40 complete

No compressed chapters.

Every source step is its own lesson with intuition, concepts, correctly rendered MathJax mathematics, implementation, tests, mistakes, and a checkpoint.

PART 1

MLP Forward and Backward Core

Set up data and parameters, then implement the forward pass, loss, and manual backprop for a two-layer MLP.

0/10

001make synthetic regression batch+5 XP 002init mlp params+5 XP 003linear forward+5 XP 004relu forward+5 XP 005mlp forward+5 XP 006mse loss and grad+5 XP 007linear backward+5 XP 008relu backward+5 XP 009first linear backward+5 XP 010mlp backward+5 XP

PART 2

Gradient Accumulation

Split batches into micro batches and accumulate gradients to emulate a large batch step under a tight memory budget.

0/4

011split into micro batches+5 XP 012accumulate gradients+5 XP 013scale accumulated gradients+5 XP 014grad accumulation step+5 XP

PART 3

Activation Checkpointing

Trade compute for memory by recomputing activations during backward and verify correctness and savings.

0/4

015mlp forward checkpointed+5 XP 016recompute block activations+5 XP 017mlp backward checkpointed+5 XP 018estimate checkpointing memory savings+5 XP

PART 4

Mixed Precision Training

Run forward and backward in half precision while keeping a full precision master copy of the weights, with loss scaling and non-finite detection.

0/6

019cast to half precision+5 XP 020make master params+5 XP 021scale loss+5 XP 022unscale gradients+5 XP 023has non finite gradients+5 XP 024mixed precision step+5 XP

PART 5

Data Parallel Training

Shard data across workers, compute local gradients, and synchronize them with all-reduce, including a ring all-reduce and bucketed communication.

0/6

025shard dataset across workers+5 XP 026compute local gradients+5 XP 027all reduce mean+5 XP 028ring all reduce mean+5 XP 029data parallel train step+5 XP 030bucket gradients+5 XP

PART 6

ZeRO-Style Optimizer Sharding

Partition Adam optimizer state and parameter updates across workers, then all-gather to reconstruct the full model.

0/5

031init adam state+5 XP 032partition optimizer state+5 XP 033local shard adam update+5 XP 034all gather param shards+5 XP 035zero optimizer step+5 XP

PART 7

Memory Accounting and Full Training Loop

Quantify model, optimizer, and activation memory, compare with and without optimizations, and run the end-to-end distributed memory-aware training loop.

0/5

036compute param memory bytes+5 XP 037compute optimizer memory bytes+5 XP 038compute peak activation memory bytes+5 XP 039compare memory with and without optimizations+5 XP 040full distributed training loop+5 XP