Module 8 - GPU and TPU Infrastructure
Training and serving large AI models requires understanding the hardware beneath them. GPU utilization metrics, OOM errors, distributed training communication overhead, and hardware costs directly determine what models you can build, how fast you can train them, and how much it costs to serve them.
This module covers the GPU and TPU infrastructure knowledge that distinguishes senior ML engineers from those who just know frameworks.
What You Will Learn
Lessons in This Module
| # | Lesson | Key Concepts |
|---|---|---|
| 01 | GPU Architecture for ML | Tensor Cores, HBM, roofline model, warp execution, NVLink |
| 02 | GPU Memory Management | VRAM limits, activation checkpointing, mixed precision, ZeRO |
| 03 | Distributed Training | DDP, FSDP, tensor/pipeline parallelism, 3D parallelism |
| 04 | Training Infrastructure | InfiniBand, NCCL, Slurm/K8s, fault tolerance, checkpointing |
| 05 | TPU Architecture and Use | Systolic arrays, XLA, JAX, TPU pods, GPU vs TPU comparison |
| 06 | Inference Hardware | Inferentia, L4/L40S, Jetson, hardware-specific quantization |
| 07 | GPU Cost Optimization | Spot instances, MIG, right-sizing, utilization monitoring |
Key Mental Models
The roofline model separates compute-bound from memory-bound. If your operation is memory-bound, buying a GPU with more FLOPS changes nothing - you need more memory bandwidth. If it is compute-bound, you need more TFLOPS. Knowing which regime you are in is the first step in any optimization.
Communication is the bottleneck in distributed training. When you run DDP across 8 GPUs, every gradient synchronization involves an all-reduce across all 8 GPUs. The time for that all-reduce must be overlapped with computation or it becomes the limiting factor. At large scale, the network topology (NVLink vs InfiniBand vs Ethernet) determines whether you spend 30% or 60% of training time waiting for communication.
TPUs are not faster GPUs - they are a different programming model. TPUs require XLA compilation, static shapes, and a different way of thinking about parallelism. The productivity cost of learning JAX and XLA is real. The payoff is only worth it at scale (large models, long training runs) where the cost savings justify the investment.
Inference and training hardware are different. A100 GPUs are excellent for training but wasteful for inference - you pay for enormous VRAM and high-bandwidth compute that is underutilized by inference workloads. L4 and L40S GPUs, AWS Inferentia, and Google Cloud TPU v4 Lite are purpose-built for inference at 3–10× better cost efficiency.
Prerequisites
- Python and PyTorch basics
- Understanding of neural network training (forward/backward pass, gradients)
- Basic Linux system administration