Learn TPUs and Google Cloud

Who this path is for

Model builders who want to provision TPUs, launch single-host and multi-host experiments, and connect sequence models to accelerator engineering.

Preparation

Shared foundations and basic machine learning

Your first useful artifact

A verified TPU runtime receipt, a checkpoint-resumed TPU experiment log, and a memory/precision/XProf decision dossier.

How to study

Complete the 3-part Cloud TPU mini-course (19: TPU Setup, 20: TPU Workflows, 21: TPU Systems) immediately after 00: Setup & first steps so you know how to get a TPU running, stage scripts with gcloud compute tpus tpu-vm scp, launch single-host and --worker=all jobs, and capture XProf traces before starting any TPU-dependent course. Then build the tiny Transformer, verify causal masking and held-out token loss, restore the complete state in a fresh process, and run the distributed, pretraining, post-training, and operations experiments on your TPU VM in parallel with the course.

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. TPU Setup: Provisioning & Runtime Verification

   Step 1 of the 3-part Cloud TPU mini-course. Choose the right TPU slice, provision a Cloud TPU VM (or Colab TPU), bootstrap ~/jax-tpu-lab/.venv with jax[tpu], verify a synchronized pmap runtime receipt, and practice clean teardown.

   • Select a TPU slice and provision your Cloud TPU VM
   • Bootstrap .venv and verify a TPU runtime receipt
3. TPU Workflows: Launching Jobs & Checkpointed Experiments

   Step 2 of the 3-part Cloud TPU mini-course. Stage scripts with gcloud compute tpus tpu-vm scp, separate XLA compilation warmup from steady-state step timing, recover interrupted runs from atomic checkpoints, and launch single-host or --worker=all Pod slice experiments.

   • Stage and launch checkpointed TPU jobs
   • Run course scripts and multi-host TPU Pod slices
4. TPU Systems: Generations, Memory, Precision & XProf

   Step 3 of the 3-part Cloud TPU mini-course. Compare TPU v5e, v5p, and v6e specifications and ridge points, budget training and KV-cache HBM before provisioning, measure BF16 and INT8 numerical error, and capture warmed XProf traces.

   • TPU generations and HBM memory budgeting
   • bfloat16, int8 precision, and XProf profiling on TPU
5. 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
6. 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
7. 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
8. 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
9. 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
10. 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
11. Transformers & language models

    Build attention and causal masking from small arrays, then connect tokenization, mixed precision, training and generation.

    • Attention from small pieces
    • Masking, tokenization, and sequence packing
    • A Transformer block and mixed precision
    • Train, checkpoint, and generate
12. Self-supervised pretraining: masked and contrastive learning

    Train small masked-token, hidden-patch and paired-view models. Audit information leakage, reduction counts and pair identity before interpreting a falling objective.

    • Masked language modeling: predict hidden tokens
    • Masked image modeling: reconstruct missing patches
    • Contrastive learning: views, positives and negatives
13. Reinforcement learning

    Build a functional environment, collect randomized trajectories, train policy-gradient and clipped-policy agents, and evaluate independent seeds.

    • Write a functional environment
    • Batch environments and scan rollouts
    • Policy gradients and PPO
    • Evaluate agents across seeds
14. Post-training: SFT, LoRA, reward models and RLHF

    Connect supervised demonstrations, frozen-base LoRA, reward modeling, sampled PPO and direct preference optimization. Keep the reference, behavior policy, reward model and dataset identities separate.

    • Supervised fine-tuning with response-only token loss
    • LoRA: adapt, save and merge low-rank updates
    • Learn a reward model from pairwise preferences
    • RLHF mechanics: a frozen reward and PPO policy updates
    • Direct preference optimization and reference-corrected margins
15. 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
16. 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
17. Deployment, interoperability & edge AI

    Adapt a trained model, verify framework and export contracts, compare precision policies and measure local inference before planning server or edge capacity. Package a model service in a reproducible container and verify readiness, artifact identity and request contracts.

    • Adapt a pretrained model and choose a post-training objective
    • Move models between JAX, Keras, TensorFlow and PyTorch
    • Convert PyTorch weights to Flax and locate numerical errors
    • Export a computation and verify its serving contract
    • Choose weight, activation and accumulation precision
    • Containerize a model service and verify its boundary
    • Inference capacity, batching, and autoscaling
    • Deploy at the edge: conversion, budgets and device checks
18. Workload operations

    Launch and observe real local jobs, detect failures and stalls, restore complete training state, verify artifact quality and rehearse rollback before planning capacity. Connect MLOps data/release gates, LLMOps evaluation/traces and ModelOps ownership/approval to the workload runbook.

    • Accelerator jobs and runtime lifecycle
    • Health, logs, and workload observability
    • Failure recovery and operational runbooks
    • Changes, rollback, and artifact provenance
    • Capacity, utilization, and operating cost
    • MLOps: data contracts, CI gates and monitoring
    • LLMOps: version prompts, evaluate behavior and trace requests
    • ModelOps: own, approve, roll back and retire a model

Causal Transformer training-system capstone

A guided project is available for this route.

Train and recover a byte Transformer, verify cache precision and exported inference, and benchmark on CPU and Cloud TPU VM.

• Independent causal-attention and valid-token loss checks on disjoint data
• Fresh-process recovery of parameters, both Adam moments, dropout keys and document position
• FP32/BF16/INT8 cache parity, stored bytes and numerical-error comparisons
• Six serialized prefill/decode endpoints with fresh-process reload and artifact rejection
• A genuine profile and synchronized CPU measurements, plus separate actual-target qualification before any TPU claim
• Cloud resource identity, resolved environment, actual backend events, retrieved checkpoints and verified resource cleanup for any claimed TPU run
• Generation and topology choice, explicit weight/activation/accumulator/cache policy, memory budget and one tested profile hypothesis with retained quality evidence

Read the portfolio project guide

Work through the synthesis assessment

Use this path in your learning workspace