Machine Learning Systems, Part 3 — Motivating GPU Hardware

Given our understanding of dense computation and tensors and how they relate to machine learning, we'll watch how the GPU architecture naturally emerges as a good candidate on which to do computational tasks in deep learning.

To get there, rather than describing the GPU architecture, let's see how we can design a piece of hardware from first principles, and notice how it resembles modern GPUs.

Our Goal: Parallel Operations

We have a good sense of the building-block ingredients for dense computation in deep learning. To start mapping things to hardware, let's remind ourselves of the matrix multiplication operation:

What are the constituent arithmetic operations here? We can start by breaking down matrix multiplication into sets of dot products:

Clearly, each element of the output matrix is an dot product (or inner product). Without loss of generality, let's pick one of those and explicitly enumerate the multiplication and addition operations therein:

Rewriting in a particular way:

From this rewrite, one thing is clear: we can do the multiplication operations in each constituent inner product in parallel! If we had $n$ circuits, we could do all multiplications in a single step, then sum all of the results together in $O(n)$ steps (although we can do better than this).

As we've seen, there are many other operations in deep learning computation that have numerous parallel arithmetic operations. We can also consider the Hadamard Product ($\odot$) which involves element-wise multiplication of matrices:

Here, it's clear that the element-wise multiplication operations for each element of the output matrix can be performed in parallel; there's no dependence between individual multiplication operations.

Designing Hardware for Parallel Arithmetic

Our desiderata are clear: we want hardware that can perform numerous arithmetic operations in parallel, with the ability to combine values with operations like reductions. Let's begin with these goals.

We'll start with a von Neumann architecture including a memory unit, registers, and an arithmetic logic unit with a loop that writes data back into registers, which can read and write from memory.

In this loop, data is loadable from main memory into registers, where it then flows into an arithmetic-logic unit (ALU), which is capable of performing arithmetic operations (e.g. addition, multiplication, logical AND), the result of which is then written back to registers and can be stored back into main memory on subsequent cycles.

With a single loop, each "cycle" denotes a roundtrip from registers (or from main memory) through the ALU and back. Given our original formulation of the inner product above, we have $n$ multiplication operations to perform. With the above circuit, this requires $n$ cycles, with subsequent cycles required to add (i.e. reduce) the resulting products together.

What's our bottleneck? If we were able to perform multiple multiplication operations simultaneously, we could perform all $n$ at once! We're operating with a blank canvas; let's create more loops!

If we create $n$ loops, we can do all $n$ multiplications in a single step if we have reach of the needed operands in memory! There are two problems here:

1. If we have operand tensors contiguously in memory, how do we get each pair of elements in the input tensors into memory for each loop?
2. Once we have the results of each multiplication, how do we combine them if we need to access all of the intermediate results at once?

To solve this, we'll introduce the notion of "shared memory" — that is, memory that's shared between loops! Revising our circuit diagram, we have the following:

In this setup, we can load and store data from anywhere in shared memory into registers in each loop, perform operations, then write back, all at once. Now that ALUs can operate on shared data, we can start with the tensors contiguously in memory, assign each loop parallel arithmetic tasks, execute them, and write back to a shared location.

Putting it all Together

There are a few basic components we'll add to fill in the gaps between our toy model and a real GPU¹:

• An instruction control unit which controls sets of the memory-register-ALU loops
• Caches which sit in between shared memory and sets of registers.

Reassigning proper terminology to the abstractions we've already created:

• Streaming processor (SP) — the register-ALU abstraction that, along with other proximal SPs, share a control unit and instruction cache.²
• Streaming mutliprocessor (SM) — a collection of SPs, connected to common shared memory and caches. SMs are the hardware building blocks on top of which parallel computing abstractions are based.

Modern GPUs typically feature hundreds of streaming processors per streaming multiprocessor.³. Given the above, a full streaming multiprocessor looks like the following:

Modern datacenter GPUs typically contain dozens⁴ of SMs. The shared memory in each SM is connected to an intermediate cache (usually labeled level 2, or L2), which is connected to device memory (also called VRAM, or "Video Random-Access Memory" on GPUs). Combining the above together, our revised picture of a GPU is as follows:

Take notice of several artifacts of how busses (high-bandwidth connections between components) are connected:

• Direct VRAM to shared memory loads — new instructions enable direct I/O from GPU VRAM into shared memory in an SM, thus bypassing caches and control units.
• SMs can write to an L2 cache — SMs can cache arbitrary data, which is faster to read/write from cache than from VRAM.

This diagram is an approximation of how a real GPU is laid out. In sections that follow, we'll revisit this schematic in the context of hierarchical data transfers and storage as it relates to bottlenecks and performance as well as in the context of distributed computation in deep learning.

Citation

@misc{kahn2024mlsystems,
  title        = {Machine Learning Systems},
  author       = {Jacob Kahn},
  year         = {2024},
  url          = {https://jacobkahn.me/writing/tag/machine_learning_systems/},
}