Make it fast. Make it scale.

Who this path is for

Engineers comfortable with training loops who want to reason about memory, communication and throughput.

Preparation

Shared foundations and a working training loop

Your first useful artifact

A synchronized timing report and a checked array placement on logical CPU devices.

How to study

Keep one correct reference workload throughout this path. Measure completed work, inspect its trace and relate the bottleneck to memory traffic or computation. Verify partitioning and collectives on logical devices, then audit a Pallas kernel against independent references. Run the separate guarded TPU measurement only on supported target hardware before claiming a device or system speedup.

Course phases and lessons

1. Setup & first steps

   Run a saved program and explain its environment, array output and device placement. Start on one CPU, practice sharding on four logical CPU devices, and learn how to get a TPU running, launch a TPU job, and run course experiments on TPU in parallel from Phase 00 onward.

   • Set up your learning workspace
   • Meet your arrays and devices
   • Practice with four virtual CPU devices
   • Move your experiment to a TPU
2. Arrays & pure functions

   Treat every axis as part of a contract. Work through normalization, broadcasting and immutable updates, then diagnose mistakes that still return plausible numbers.

   • From NumPy to jax.numpy
   • Shapes, broadcasting, and dtypes
   • Pure functions and explicit inputs
   • Immutable updates and indexing
3. JAX transformations

   Begin with a scalar function, then differentiate it, batch it and compile it. Compare each transformed result with a simple reference so you can explain what changed.

   • Your first gradient
   • Losses and value_and_grad
   • Batch a function with vmap
   • Compile a function with jit
   • Tracing, static arguments, and recompilation
4. State, randomness & control flow

   Make randomness, parameters and loop state explicit. Learn how to repeat an experiment without accidentally reusing randomness or mutating a caller’s state.

   • Random keys without surprises
   • Pytrees and structured parameters
   • Compiled loops with lax.scan
   • Branches with lax.cond
5. Math & optimization

   Connect geometry and probability to the updates in a training loop. Begin with linear algebra, gradient checks and gradient descent, then investigate conditioning, curvature, likelihoods, regularization and optimizer behavior.

   • Vectors, norms, and projections
   • Linear algebra and loss intuition
   • Least squares, rank, and conditioning
   • The chain rule, Jacobians, and directions
   • Gradient checking and numerical accuracy
   • Curvature, Hessians, and learning rates
   • Write gradient descent yourself
   • Probability, likelihood, and stable losses
   • Regularization and honest validation
   • Minibatches, expectation, and gradient noise
   • Optimize with Optax
   • Adam, clipping, and learning-rate schedules
6. Neural network training

   Build the same ideas into small neural networks. Understand hidden layers, parameter state, stable objectives and evaluation before moving to a small image classifier.

   • Build a tiny multilayer perceptron
   • Model and state with Flax NNX
   • A compiled train and evaluation step
   • Train a CNN on a small dataset
   • Debug unstable learning
7. Data & checkpoint recovery

   Follow examples from the dataset into a training update, then save enough information to reproduce the next update after interruption. Track experiment configuration, data identity and metric histories with MLflow alongside complete checkpoints.

   • Design an input pipeline
   • Load, batch, and checkpoint order with Grain
   • Save model, optimizer, and random state
   • Recover data position and resume
   • Asynchronous checkpoints and failure boundaries
   • Track experiments and lineage with MLflow
8. Performance diagnosis

   Measure completed work, inspect a real profiler trace and relate compiler output to memory traffic. Compare rematerialization choices while separating saved residuals, compiler memory estimates and observed latency.

   • Benchmark asynchronous work correctly
   • Diagnose recompilation and host synchronization
   • Read an XProf trace
   • Reason with HLO and the roofline model
   • Memory, rematerialization, and optimization tradeoffs
9. Distributed training

   Partition arrays and training updates, compare all-reduce with reduce-scatter, and restore a complete sharded state across an epoch boundary. Verify CPU placement and arithmetic before separate cluster validation.

   • Arrays, meshes, and sharding
   • A sharded training step
   • Communication-efficient algorithms
   • Resilient distributed training
10. Pallas kernels

    Write actual Pallas grids and buffered pipelines, verify boundaries and precision on CPU, and prepare guarded target measurements before claiming an optimization win.

    • Pallas grids and BlockSpecs
    • A first TPU kernel
    • Tiling, memory, and pipelining
    • Iterate with correctness and performance evidence

A checked kernel and a measurement-backed optimization argument

A guided project is available for this route.

Defend the profile, global objective, partitioning and kernel correctness; quantify a whole-workload opportunity. CPU checks establish semantics. Actual accelerator execution and performance remain a separate required qualification for device claims.

• Reference agreement across more than one shape and parameter setting
• Synchronized compile and execution measurements with a trace of the bottleneck
• Array partitioning, collective communication and a before/after comparison on the stated hardware

Read the portfolio project guide

Work through the synthesis assessment

Use this path in your learning workspace