Lectures on ML Systems

Machine learning models are pervasively used to solve problems in varied fields such as vision, robotics, NLP, and scientific discovery. The increased capabilities of these models has corresponded with increase in their size and compute requirements. Besides, use of these models in real-world applications demands strict requirements on performance parameters, such as latency, throughput, and hardware usage efficiency.

The focus of this course is on exploring these systems-related challenges during training and serving of large language models (LLMs) with special emphasis on Transformer architecture. Topics include GPU architecture and hardware-aware algorithms, ML frameworks and compilers, techniques to parallelize LLMs over multiple GPUs, and reduction of computational complexity and memory footprint.

Course Overview

Info not available

Learning Objectives

Info not available

Learning Outcomes

Info not available

How to Scale Your Model by Austin, J, Douglas, S, Frostig, R, Levskaya, A, Chen, C, Vikram, S, Lebron, F, Choy, P, Ramasesh, V, Webson, A, & Pope, R (2025).

Lectures

Topic 1: Introduction

Slides, Video lecture

Topic 2: Automatic Differentiation

Slides, Video lecture

Practical 1: Automatic Differentiation Implementation

Video

We do a code walkthrough on how reverse-mode automatic differentiation is implemented in a modern ML framework. We choose Needle, an educational framework with a similar interface as PyTorch, developed at CMU.

Topic 3: Understanding GPU Bottlenecks for ML

Slides, Video lecture

Topic 4: GPU Programming Model

Slides, Video lecture

Example code: Vector addition, Simple matrix multiplication, Tiled matrix multiplication

Practical 2: GPU Profiling

Video

We run Nvidia Nsight Systems and Nsight Compute profilers on our naive and tiled matrix multiplication examples. We identify the key information to look for when using these tools and see how they’re valuable in identifying bottlenecks.

Topic 5: Transformer FLOPs Math

Slides, Video lecture

Topic 6: Introduction to LLM Inference

Slides, Video lecture

Additional Readings

Topic 1: Introduction

  1. Gholami, Amir, et al. “AI and Memory Wall.” IEEE Micro Journal (arxiv link)
    This very readable paper presents the interplay between throughput, bandwidth, and end-to-end runtime via a case study on Transformer models.
  2. Kaplan, Jared, et al. “Scaling laws for neural language models.” arXiv:2001.08361 (2020) (arxiv link)
    An empirical study quantifying the improvement in LLM performance with model size, training dataset size, and amount of compute. Also known as Kaplan’s scaling law.
  3. Jacob, Austin, et al. How to Scale Your Model (Part 0: Intro)
    A useful summary of why study of ML Systems is important.

Topic 2: Automatic Differentiation

  1. Kevin Clark, “Computing Neural Network Gradients” (link)
    A concise refresher on analytical gradient computation for neural networks in terms of matrices and vectors.
  2. Roger Grosse, “CSC321 Lecture 10: Automatic Differentiation” (Slides)
    A presentation of automatic differentiation in the context of the Autograd library. It’s nice to see how Autograd builds computation graph in a different manner than PyTorch or Needle.
  3. Chen, Tianqi, et al. “Training Deep Nets with Sublinear Memory Cost” (arxiv link)
    Seminal paper that introduced the idea of gradient checkpointing.

Topic 3: Understanding GPU Bottlenecks for ML

  1. Stephen Jones, “How GPU Computing Works”, GTC 2021 (Video)
    Excellent introduction to principles behind GPU design and CUDA architecture. Specially good is the discussion of co-design of GPU hardware and CUDA.
  2. Nvidia Docs, “GPU Performance Background User’s Guide” (Link)
    A useful discussion of how different deep learning operations get limited by either compute or memory of GPU.
  3. Horace He, “Making Deep Learning Go Brrrr From First Principles” (Blog post)
    Nice diagrams and figures to compactly illustrate the 3 components where our ML program spends time - compute, memory, and overhead.

Topic 4: GPU Programming Model

  1. Mark Harris, “An Even Easier Introduction to CUDA” (Nvidia Technical Blog post)
    Introduces GPU execution model by vector addition example. One useful concept used is Unified Memory - a pool of managed memory shared by GPU and CPU, thereby simplifying memory management for the programmer. Also, Mark introduces a new CUDA programming pattern called grid-stride loop.
  2. Jeremy Howard, “Getting Started With CUDA for Python Programmers” (Video)
    Jeremy shows how to call a CUDA kernel from PyTorch using its cpp_extension. He introduces kernels to do RGB to grayscale conversion and matrix multiplication in a more Pythonic way.
  3. Sasha Rush, “GPU Puzzles” (Link)
    A nice collection of puzzles to implement CUDA kernels in Python. Uses Numba Python JIT Compiler.

Topic 5: Transformer FLOPs Math

  1. Einops tutorial (Link), Einsum tutorial (Link)
    Einsum is a notation to concisely express complex tensor operations like multiplication and summation. Einops is a library to reshape and manipulate tensors such as rearranging axes and reducing dimensions. One doesn’t have to remember all the different PyTorch functions for tensor manipulations if using einsum and einops.

Topic 6: Introduction to LLM Inference

  1. Niels Rogge, “How a Transformer works at inference vs training time” (Video)
  2. Jacob, Austin, et al. How to Scale Your Model (Part 4: Transformers, Part 7: Inference)