GPUs and compute

Feel the cost grow

The code below times three matrix multiplies. Each uses square matrices, so all three sizes are the same: $m = k = n$. The work is then $2n^3$ FLOPs. Here $n^3$ means $n$ times $n$ times $n$.

The code reports speed in GFLOP/s, which means billions of FLOPs per second. Your timings will differ from anyone else's.

import time
import numpy as np

for n in [128, 256, 512]:
    A, B = np.random.rand(n, n).astype(np.float32), np.random.rand(n, n).astype(np.float32)
    start = time.perf_counter()
    for _ in range(5):
        A @ B
    seconds = (time.perf_counter() - start) / 5
    flops = 2 * n ** 3
    print(f"{n} × {n}: {flops / 1e6:7.1f} million FLOPs in {seconds * 1000:6.1f} ms  →  {flops / seconds / 1e9:5.1f} GFLOP/s")

Doubling $n$ multiplies the work by 8, because 2 × 2 × 2 = 8. This browser tab manages a few GFLOP/s at best.

A modern data-center GPU does about $10^{15}$ FLOPs per second in low precision (explained next). That's a 1 followed by 15 zeros, called a petaFLOP. It's hundreds of thousands of times faster than this tab.

Precision: how many bytes per number

Every number in a model takes up memory. A byte is the basic unit of computer memory. One byte is 8 bits, and a bit is a single 0 or 1.

More bytes per number give more exact numbers. It's like keeping more decimal places on a price. Fewer bytes give rougher numbers that take less room. That choice is called precision.

The number in each name is its size in bits. So float32 uses 32 bits, which is 4 bytes. "Inference" means running a trained model to get answers.

Halving the bytes halves the memory. On hardware that supports the smaller format, it also roughly doubles the speed. The loss in quality is usually small. The quantization lesson in Phase 6 covers the trade-offs.

Memory to run a model

To run a model, its weights must sit in the GPU's memory. The weights are also called the model's parameters: the numbers it learned in training. The rule is one multiplication:

memory = parameters × bytes per parameter

Take a model with 7 billion parameters, called "7B" for short. In bfloat16 that's 7 billion × 2 bytes, or 14 billion bytes. A gigabyte (GB) is a billion bytes, so that's 14 GB. At 4 bits, which is half a byte, it's 7 billion × 0.5 = 3.5 GB. That fits on a laptop.

params = 7e9                          # 7 billion; 7e9 means 7 × 10**9
print(params * 2 / 1e9, "GB in bfloat16")
print(params * 0.5 / 1e9, "GB at 4 bits")

It prints 14.0 and 3.5. Dividing by 1e9 turns bytes into gigabytes.

You need some extra room on top. The model stores activations, the in-between results of each layer. When it writes text, it also keeps an attention cache. That cache holds notes about the earlier words, so the model doesn't redo that work.

Memory to train a model: about 16 bytes per parameter

Training needs far more memory than running. The usual setup has two parts. Mixed precision keeps some numbers in 2 bytes and some in 4. Adam is the popular update rule from the optimizers lesson. With both, the standard estimate is about 16 bytes per parameter:

• 2 bytes: the weights in bfloat16, used for the forward and backward passes.
• 2 bytes: the gradients in bfloat16. A gradient says which way to nudge each weight.
• 4 bytes: a float32 "master copy" of the weights. Tiny updates are added here, so rounding doesn't wipe them out.
• 8 bytes: Adam's two running averages for each weight, 4 bytes each in float32.

Activations come on top of all this.

In plain words: to train, multiply the number of parameters by 16. That gives the bytes for weights, gradients and Adam's notes, before activations.

So a 7B model needs about 7 billion × 16 = 112 billion bytes. That's 112 GB. It's more than a typical 80 GB data-center GPU holds, and activations push it higher still.

print(2 + 2 + 4 + 8, "bytes per parameter")
print(7e9 * 16 / 1e9, "GB to train a 7B model")

It prints 16 and 112.0. That's why big training runs are split across many GPUs. It's also why parameter-efficient fine-tuning is popular. One method, LoRA (Phase 6), trains a small add-on instead of every weight.

print(3e9 * 16 / 1e9)

Compute: the 6ND rule

Memory tells you which GPU you need. Compute tells you how long it will run. For a transformer, a short formula gives the total work of training.

It uses two symbols. $N$ is the number of parameters in the model. $D$ is the number of training tokens. A token is a small chunk of text, often a whole word or part of one.

Here $C$ is the total compute, counted in FLOPs.

In plain words: each training token costs about 6 FLOPs for every parameter. About 2 are for the forward pass, which makes a prediction. About 4 are for the backward pass, which works out the gradients.

Here's a toy case. A model with 1,000 parameters trains on 500 tokens. That's 6 × 1,000 × 500 = 3,000,000 FLOPs. It sounds like a lot, but a laptop does it in a tiny fraction of a second.

N = 1000          # parameters
D = 500           # training tokens
print(6 * N * D)  # training FLOPs

Running a trained model costs far less. It's about $2N$ FLOPs per generated token, because it needs only the forward pass.

params = 100e6                    # N: 100 million parameters
tokens = 3e9                      # D: 3 billion tokens
flops = 6 * params * tokens
seconds = flops / (1e15 * 0.5)    # peak FLOP/s × utilization
hours = seconds / 3600
print(f"{flops:.1e} FLOPs, {seconds:,.0f} seconds, {hours:.1f} GPU-hours")
print(f"cost: about {hours * 3:.0f} dollars at 3 dollars per GPU-hour")

Real training runs

Now the same recipe for three real-sized jobs. The code uses the peak speeds of two common NVIDIA data-center GPUs: the A100 and the newer H100. It assumes 40% utilization.

def training_flops(params, tokens):
    return 6 * params * tokens

def gpu_hours(flops, flops_per_second, utilization=0.4):
    return flops / (flops_per_second * utilization) / 3600

A100_BF16 = 312e12        # NVIDIA A100 peak dense bfloat16 FLOP/s
H100_BF16 = 989e12        # NVIDIA H100 (SXM) peak dense bfloat16 FLOP/s

runs = {"GPT-3 (175B params, 300B tokens)": (175e9, 300e9),
        "a 7B model on 2T tokens": (7e9, 2e12),
        "fine-tuning 7B on 50M tokens": (7e9, 50e6)}
for name, (n, d) in runs.items():
    c = training_flops(n, d)
    print(f"{name:<34} {c:.1e} FLOPs = {gpu_hours(c, A100_BF16):>9,.0f} A100-hours = {gpu_hours(c, H100_BF16):>9,.0f} H100-hours")

Read the first line. GPT-3 has 175 billion parameters and trained on 300 billion tokens. The rule gives about $3 \times 10^{23}$ FLOPs. That matches what the GPT-3 paper reports.

Multiply the GPU-hours by a cloud price per GPU-hour, and you have a budget. The price is a few dollars at the time of writing, but check current rates. At that price, even a GPT-3-sized run costs on the order of a million dollars of compute. Today's frontier models cost far more.

The last line is the hopeful one. Fine-tuning a 7B model on 50 million tokens takes only about 1 to 5 GPU-hours.

Where the course's code runs

On your own machine, torch.cuda.is_available() tells you whether PyTorch can see an NVIDIA GPU. On a Mac with Apple silicon, use torch.backends.mps.is_available() instead.

Tiny models like the ones in this phase train fine on a CPU. Anything the size of a transformer wants a GPU. Cloud notebooks such as Colab and Kaggle give you one for free, in small amounts.

In this browser, everything runs on the CPU. It uses the course's own small version of PyTorch, built on NumPy. That's why the models here stay tiny.