PART 1
0/6Model, Forward & Loss
Initialize parameters and implement the forward pass, activations, softmax, cross-entropy loss, and backprop that every worker will run locally.
DiLoCo Distributed Training
DiLoCo Distributed Training
Train workers locally, aggregate pseudo-gradients with an outer optimizer, and quantify communication savings.
Every source step is its own lesson with intuition, concepts, correctly rendered MathJax mathematics, implementation, tests, mistakes, and a checkpoint.
PART 2
0/5Inner Optimizer: AdamW
Build the per-worker AdamW optimizer: state init, moment updates, bias correction, the parameter step, and decoupled weight decay.
PART 3
0/4Parameter Arithmetic Utilities
Implement the pytree-style parameter operations (clone, scale, subtract, average) that DiLoCo uses to compute pseudo-gradients and merge worker replicas.
PART 4
0/5Data Sharding & Local Worker Training
Partition the dataset across workers (IID and non-IID), sample per-worker batches, and run local training steps that emulate a single DiLoCo inner worker.
PART 5
0/4Outer Optimizer with Nesterov Momentum
Implement the server-side outer optimizer: state init, momentum buffer updates, Nesterov parameter updates, and computing the outer pseudo-gradient from worker deltas.
PART 6
0/3DiLoCo Training Loop & Baseline
Compose everything into a full DiLoCo communication round, the outer training loop across rounds, and a synchronous SGD baseline for comparison.
PART 7
0/3Evaluation & Communication Analysis
Measure held-out loss, classification accuracy, and quantify DiLoCo's communication savings versus the synchronous baseline.