GPUs, memory, and quantization

Bytes per parameter and formats

The memory the weights take is a multiplication:

with $N$ parameters and $b$ bytes per parameter. For Llama 3.1 70B in FP16: $70\times10^{9}\cdot 2 = 1.4\times10^{11}$ bytes $=140$ GB.

Decimal formats (the IEEE 754 standard) store a sign, an exponent (the range), and a mantissa (the digits):

format	bytes	sign · exponent · mantissa	decimal digits	largest value
FP32	4	1 · 8 · 23	$\log_{10}2^{24}\approx 7.2$	$\approx 3.4\times10^{38}$
FP16	2	1 · 5 · 10	$\log_{10}2^{11}\approx 3.3$	65,504
BF16	2	1 · 8 · 7	$\log_{10}2^{8}\approx 2.4$	same as FP32

BF16 (Kalamkar et al., 2019) gives up digits to keep FP32’s range; Llama 3 was trained in BF16. INT8 and INT4 have no exponent: they store an integer plus a shared scale (see below).

A detail about units: datasheets say “80 GB” of VRAM, but that is $80\cdot 2^{30}$ bytes (GiB), 7% more than $80\times10^{9}$. In this lesson, 1 GB = $10^{9}$ bytes, as with bandwidth; the difference changes no conclusion.

Training: sixteen bytes per parameter

With Adam in mixed precision (Micikevicius et al., 2018), each parameter needs the weight and gradient in FP16 and, in FP32, a copy of the weight and Adam’s two averages (Kingma and Ba, 2015). ZeRO (Rajbhandari et al., 2020, §3.1) sums it up:

The FP32 copy exists because each step’s update (gradient times learning rate) is often thousands of times smaller than the weight: added in FP16, with its 10 mantissa bits, it would be lost to rounding. We measured it on an RTX 3060 (measure-memory.py, the Pythia-160M architecture with random weights): 16.19 bytes per parameter in full FP32 and 16.05 in mixed precision, after AdamW’s first step.

On top of that come the activations, the intermediate results backpropagation needs. Per layer of a 16-bit transformer with linear-memory attention they are about $34\,s\,b\,h$ bytes ($s$ = tokens per sequence, $b$ = sequences per batch, $h$ = width; Korthikanti et al., 2022). On the same RTX 3060 we measured 0.949 times that formula. They can be reduced by recomputing them on the way back, at the cost of more compute.

ZeRO and its PyTorch version, FSDP (the one Meta used), split those 16 bytes across the data copies, so each GPU stores only part of the states.

Bandwidth: the ceiling when writing

When generating one token at a time, every step reads all the weights:

An H100 SXM moves 3,350 GB/s; an 8-billion model in FP16 (16 GB) stays under about 209 tokens per second, however fast it computes. Its dense FP16 compute peak is about 989 trillion operations per second (the datasheet says 1,979 “with sparsity”): $989\times10^{12}/3.35\times10^{12}\approx 295$ operations per byte, about 590 per 2-byte parameter, which a matrix–vector product uses for only 2 (multiply and add).

Measured on an RTX 3060 (measure-memory.py): copying 1 GiB, 333 GB/s; multiplying a $16384\times16384$ FP16 matrix by one vector, 1.57 ms, 341 GB/s, and just 0.34 trillion operations per second; by 4,096 vectors at once, 27.7 trillion: 81 times more on the same chip, because each weight read is reused. Batching requests is the way out, which we will see in lesson 26.

Parallelism

The number of GPUs is the product of the three splits: $G = T\cdot P\cdot D$ (tensor, pipeline, data).

• Data: each copy sees another slice of the batch; at the end of the step the gradients are averaged (a sum across all of them, an all-reduce).
• Pipeline (GPipe, Huang et al., 2019): $K$ stages of layers; the batch is cut into $M$ micro-batches so the stages work at the same time. The idle fraction (the “bubble”) is on the order of $\frac{K-1}{M+K-1}$; GPipe found it negligible with $M \ge 4K$.
• Tensor (Megatron-LM, Shoeybi et al., 2019): in each block’s network, the first table is split by columns and the second by rows, so one sum across GPUs per block in each direction is enough. Since it happens at every layer, it stays inside the server, where the GPUs are joined by NVLink (900 GB/s on the H100).

Llama 3.1 405B (Llama Team, 2024, table 4): 8 tensor × 16 pipeline × 128 data = 16,384 H100s, using 41% of their compute; the long-context stage added a fourth split, by context (8 × 16 context × 16 × 8). And to use it: in BF16 it does not fit in one 8-H100 server (§6.1), which is why Meta serves it split across two servers or quantized to FP8.

Group quantization

The lesson’s rule (symmetric round-to-nearest, no calibration), for a group $g$ of 32 weights and $b$ bits:

With 8 bits there are 255 marks (from −127 to 127); with 4, 15. Each group also stores its scale in FP16: $16/32 = 0.5$ extra bits per weight, so “INT4” is really 4.5 bits (the 70-billion model takes about 39 GB, not 35).

Measured on Pythia-1.4B (build-data.py: the 96 matrices of its 24 layers, 1.208 billion weights; loss in nats per token over 287,744 tokens of WikiText-2):

bits	FP16	8	6	5	4	3	2
one scale per group of 32	2.5805	2.5806	2.5858	2.5975	2.6565	3.1589	11.045
one scale per matrix	2.5805	2.6002	3.2172	5.3017	10.692	10.445	21.347

With a single scale per matrix, the model already breaks at 4 bits: a few outlier weights, up to 13 times the standard deviation in the median matrix, stretch the ruler, and in that matrix two out of three weights land on the zero mark at 4 bits. That is why small groups are used. Production methods go further: LLM.int8() (Dettmers et al., 2022) separates the dimensions with outliers and keeps them in 16 bits; GPTQ (Frantar et al., 2023) corrects the rounding with second-order information and brings 175-billion models to 3 or 4 bits with negligible accuracy loss. Caveat: Pythia saw Wikipedia during training, so the absolute level of the loss is optimistic; what is compared is the difference, measured on the same text.

In code

import numpy as np

def gb(params, bytes_per_param):
    return params * bytes_per_param / 1e9

for fmt, b in [("FP32", 4), ("FP16", 2), ("INT8", 1), ("INT4", 0.5)]:
    print(f"70B in {fmt}: {gb(70e9, b):5.0f} GB")
print(f"training with Adam (16 B): {gb(70e9, 16):.0f} GB")
print(f"writing ceiling, 8B in FP16 on an H100: {3350 / gb(8e9, 2):.0f} tokens/s")

# quantize a group of 32 weights: a ruler with 2^(b-1) - 1 marks per side
rng = np.random.default_rng(0)
w = rng.normal(0, 0.02, 32)
for bits in (8, 4, 2):
    qmax = 2 ** (bits - 1) - 1
    s = np.abs(w).max() / qmax                     # the group's scale
    q = np.clip(np.round(w / s), -qmax, qmax)      # the integers that are stored
    err = np.linalg.norm(w - q * s) / np.linalg.norm(w)
    print(f"{bits} bits: relative error {err:.3f}, weights at zero {np.mean(q == 0):.0%}")

Output:

70B in FP32:   280 GB
70B in FP16:   140 GB
70B in INT8:    70 GB
70B in INT4:    35 GB
training with Adam (16 B): 1120 GB
writing ceiling, 8B in FP16 on an H100: 209 tokens/s
8 bits: relative error 0.006, weights at zero 0%
4 bits: relative error 0.100, weights at zero 19%
2 bits: relative error 0.792, weights at zero 84%

What comes next

When the model is used, another memory appears, the attention cache, which grows with every token of the conversation, along with the way to make the most of bandwidth: serving many people at once. That is lesson 26, inference.