Module 4 of 6

Pack the bag smarter: single-GPU tricks

Your GPU memory is a backpack. When it overflows, you can buy a bigger bag (expensive), pack smarter (this module), or share the load (Module 5).

An open purple backpack holding books labelled Model Weights, Activations, Gradients and Optimizer States

Quantization

Smaller items: paperbacks instead of hardcovers.

Mixed precision

Mix sizes: full folders only for important documents.

Gradient accumulation

Several smaller trips instead of one overloaded one.

Quantization: keep fewer digits

Four large hardcover books become four small paperbacks

Why? To make a model smaller, so it needs less memory and answers faster. We already do this with π: 3.14159265 becomes 3.14. Quantization stores each number with fewer bits, at a small cost in accuracy.

How it works

Pick a small set of evenly spaced levels, then snap every number to the nearest one. The scale s sets the gap between levels. The zero-point z shifts the levels when the numbers aren't centred on zero.

q = round (\frac{x}{s}) + z

x original numberq stored integers scalez zero-point

To use a number, the GPU turns it back: x ≈ s × (q − z). The small gap left over is the quantization error. For int8 with z = 0, the scale is the largest |x| ÷ 127, so the biggest number lands exactly on 127.

Benefits and trade-offs

Benefits

4-bit weights need about a quarter of the memory of 16-bit ones.
Faster text generation: each new token reads fewer bytes from memory (Module 6).
Bigger models fit on smaller, cheaper GPUs.

Trade-offs

Some quality loss: 8-bit is close to lossless, 4-bit usually loses a little, 3-bit and 2-bit lose a lot.
Small models suffer more than big ones.
Not always faster: it depends on the GPU, the software and the batch size.
You can't train quantized weights directly. QLoRA freezes a 4-bit model and trains small 16-bit adapters instead.

In practice

import torch

torch.manual_seed(0)
w = torch.randn(4096, 4096)  # a weight matrix in float32
print(f"float32: {w.nbytes / 1e6:.0f} MB")

for bits in (8, 4):
    q_max = 2 ** (bits - 1) - 1         # 127 for int8, 7 for int4
    s = w.abs().max() / q_max           # scale: the largest weight maps to q_max
    q = torch.round(w / s)              # quantize: q = round(x / s)
    w_hat = q * s                       # dequantize: x ≈ q × s
    size = w.numel() * bits / 8 / 1e6
    error = (w - w_hat).abs().mean()
    print(f"int{bits}:    {size:.0f} MB, average error {error:.3f}")

Output

float32: 67 MB
int8:    17 MB, average error 0.010
int4:    8 MB, average error 0.189

int4 with one scale for the whole matrix is crude. Real 4-bit formats give each small block of weights (for example 64) its own scale.

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig

name = "Qwen/Qwen2.5-1.5B-Instruct"
tokenizer = AutoTokenizer.from_pretrained(name)

model_16 = AutoModelForCausalLM.from_pretrained(name, dtype=torch.float16, device_map="auto")
model_4 = AutoModelForCausalLM.from_pretrained(
    name,
    device_map="auto",
    quantization_config=BitsAndBytesConfig(
        load_in_4bit=True,                     # store the weights in 4 bits
        bnb_4bit_quant_type="nf4",             # the 4-bit format QLoRA uses
        bnb_4bit_compute_dtype=torch.float16,  # turn back into 16-bit for the maths
    ),
)
print(f"16-bit: {model_16.get_memory_footprint() / 1e9:.1f} GB")
print(f"4-bit:  {model_4.get_memory_footprint() / 1e9:.1f} GB")

prompt = tokenizer("The capital of France is", return_tensors="pt").to(model_4.device)
answer = model_4.generate(**prompt, max_new_tokens=10)
print(tokenizer.decode(answer[0], skip_special_tokens=True))

Prints both sizes and a short answer. The 4-bit model is about a third of the 16-bit size, not a quarter, because the embedding table stays in 16-bit.

from unsloth import FastLanguageModel

# Load the base model in 4 bits. It stays frozen.
model, tokenizer = FastLanguageModel.from_pretrained(
    model_name="unsloth/Qwen3-4B-Instruct-2507",
    max_seq_length=2048,
    load_in_4bit=True,
)

# Add small LoRA adapters. Only these are trained, in 16-bit.
model = FastLanguageModel.get_peft_model(
    model,
    r=16,
    lora_alpha=16,
    target_modules=["q_proj", "k_proj", "v_proj", "o_proj",
                    "gate_proj", "up_proj", "down_proj"],
)
model.print_trainable_parameters()

Prints how many parameters you will train: under 1% of the model. This is QLoRA: a frozen 4-bit model plus small 16-bit adapters.

Libraries: bitsandbytes with Hugging Face Transformers (load any model in 8 or 4 bits), Unsloth (fine-tune in 4 bits with QLoRA), torchao (PyTorch's own quantization), llm-compressor and GPTQModel (make GPTQ or AWQ checkpoints), and vLLM, TensorRT-LLM, llama.cpp and Ollama (serve quantized models).

Quantization trades a little accuracy for a lot of memory. It shines when serving models. For training, use mixed precision, or QLoRA to fine-tune.

Mixed precision: fast where we can, precise where we must

Four big folders become two big folders for important documents and two compact ones

Why? Training in float32 is slow and memory-hungry. GPUs do 16-bit maths many times faster on their Tensor Cores, and 16-bit activations take half the memory. But a few steps really need float32's precision.

How it works

Keep the weights in float32. This is the master copy.
Run the forward and backward passes in mixed precision: matrix multiplications in bfloat16, and sensitive steps like softmax, normalization and the loss in float32. PyTorch's autocast picks for you.
The gradients arrive in float32, and the optimizer updates the float32 weights, so tiny updates aren't rounded away.
With float16 only: multiply the loss by a big number before the backward pass, so tiny gradients don't round to zero, then divide them back before the update. This is loss scaling. bfloat16 has float32's range, so it doesn't need it.

Why must the update stay in float32? Try it:

w \leftarrow w - η \cdot \nabla_{w} L

w master weight, kept in float32η learning rate∇L gradient, from the bfloat16 maths

Benefits and trade-offs

Benefits

Faster maths: an H100 does about 990 TFLOPS in bfloat16, against 67 in float32. Real training often runs 2 to 3 times faster.
Activations take half the memory, so bigger batches fit.
Little or no loss in accuracy.

Trade-offs

The model state doesn't shrink: weights, gradients and Adam still need about 16 bytes per parameter.
bfloat16 needs an Ampere or newer GPU (A100, H100, RTX 30 series and up). On a T4, like free Colab, use float16.
float16 overflows above 65,504 and loses tiny values, so it needs loss scaling.

In practice

import torch

device = "cuda" if torch.cuda.is_available() else "cpu"
model = torch.nn.Linear(512, 512).to(device)
x = torch.randn(64, 512, device=device)
y = torch.randn(64, 512, device=device)

with torch.autocast(device_type=device, dtype=torch.bfloat16):
    out = model(x)                              # matrix multiply: bfloat16
    loss = torch.nn.functional.mse_loss(out, y)  # loss: float32
loss.backward()

print("weights: ", model.weight.dtype)       # the master copy
print("output:  ", out.dtype)
print("loss:    ", loss.dtype)
print("gradient:", model.weight.grad.dtype)  # matches the weights

Output

weights:  torch.float32
output:   torch.bfloat16
loss:     torch.float32
gradient: torch.float32

The weights and gradients stay in float32. Only the matrix multiply ran in bfloat16.

import time
import torch

a = torch.randn(4096, 4096, device="cuda")
b = torch.randn(4096, 4096, device="cuda")

for dtype in (torch.float32, torch.float16):
    x, y = a.to(dtype), b.to(dtype)
    x @ y                    # warm up
    torch.cuda.synchronize()
    start = time.perf_counter()
    for _ in range(50):
        x @ y
    torch.cuda.synchronize()  # wait for the GPU to finish
    ms = (time.perf_counter() - start) / 50 * 1000
    print(f"{dtype}: {ms:.2f} ms per matrix multiply")

Prints the time for each. On NVIDIA GPUs with Tensor Cores, including the T4, float16 is several times faster.

import torch

# bfloat16 on Ampere or newer (A100, H100, RTX 30 series and up). float16 on older GPUs like the T4.
use_bf16 = torch.cuda.is_bf16_supported(including_emulation=False)
dtype = torch.bfloat16 if use_bf16 else torch.float16
scaler = torch.amp.GradScaler("cuda", enabled=(dtype == torch.float16))  # loss scaling: float16 only

for step, (x, y) in enumerate(loader):
    with torch.autocast(device_type="cuda", dtype=dtype):  # forward pass in mixed precision
        loss = loss_fn(model(x), y)
    scaler.scale(loss).backward()  # backward pass, outside autocast
    scaler.step(optimizer)         # update the float32 weights
    scaler.update()
    optimizer.zero_grad()
    if step % 100 == 0:
        print(f"step {step}: loss {loss.item():.3f}, loss scale {scaler.get_scale():.0f}")

With float16, the loss scale starts at 65,536 and may drop when a step overflows: GradScaler skips that step and tries a smaller scale. That is normal. With bfloat16 it stays at 1.

from transformers import TrainingArguments

args = TrainingArguments(output_dir="out", bf16=True)  # fp16=True on a T4 or older GPU

# Hugging Face Accelerate, in your own loop:  Accelerator(mixed_precision="bf16")
# PyTorch Lightning:                          Trainer(precision="bf16-mixed")

One setting. Load the model in float32 as usual; the library handles autocast and loss scaling.

Libraries: PyTorch AMP (torch.autocast and torch.amp.GradScaler), Hugging Face Trainer and Accelerate, PyTorch Lightning, Keras (mixed precision policies), and NVIDIA Transformer Engine for 8-bit (FP8) training on H100 and newer GPUs.

Compute in 16-bit, keep the weights and the update in float32. Training gets faster and activations take half the memory.

Gradient accumulation: several trips

One full backpack becomes four smaller backpack trips

Why? Bigger batches give smoother, more stable training. But activation memory grows with the batch size, so the batch you want may not fit on your GPU.

How it works

Split the batch you want into micro-batches that fit: a batch of 64 becomes 4 micro-batches of 16.
Run the forward and backward pass on each micro-batch. PyTorch adds the gradients up in .grad. Divide each loss by 4, so the sum becomes the average.
After the 4th micro-batch, update the weights once and reset the gradients to zero.

Effective batch = micro-batch \times accumulation steps \times GPUs

what fits in memorymicro-batches per updatedata-parallel copies (Module 5)

Benefits and trade-offs

Benefits

Train with a big effective batch on a small GPU.
The same gradient as the big batch, up to tiny rounding differences.
One setting in most libraries.

Trade-offs

Slower than a real big batch: the maths is the same, but small micro-batches keep the GPU less busy.
Only activation memory shrinks. Weights, gradients and optimizer states stay the same size.
Fewer updates per epoch, so tune the learning rate for the effective batch size.

In practice

import torch

torch.manual_seed(0)
model = torch.nn.Linear(10, 1)
x, y = torch.randn(64, 10), torch.randn(64, 1)
loss_fn = torch.nn.MSELoss()

# One big batch of 64
loss_fn(model(x), y).backward()
big = model.weight.grad.clone()
model.zero_grad()

# Four micro-batches of 16: gradients add up in .grad
steps = 4
for xb, yb in zip(x.chunk(steps), y.chunk(steps)):
    loss = loss_fn(model(xb), yb) / steps  # divide so the sum is an average
    loss.backward()

print("same gradient:", torch.allclose(big, model.weight.grad))
print(f"largest difference: {(big - model.weight.grad).abs().max():.1e}")

Output

same gradient: True
largest difference: 6.7e-08

The tiny difference is floating-point rounding: the gradient is the same as the big batch.

accumulation_steps = 4  # 4 micro-batches of 16 = an effective batch of 64

for step, (x, y) in enumerate(loader):  # the loader yields micro-batches of 16
    loss = loss_fn(model(x), y) / accumulation_steps
    loss.backward()                     # gradients add up in .grad

    if (step + 1) % accumulation_steps == 0:
        optimizer.step()                # one update per 64 examples
        optimizer.zero_grad()
        print(f"update {(step + 1) // accumulation_steps}: loss {loss.item() * accumulation_steps:.3f}")

Prints the loss once per update, every 4 micro-batches.

from transformers import Trainer, TrainingArguments

args = TrainingArguments(
    output_dir="out",
    per_device_train_batch_size=16,  # the micro-batch: what fits in memory
    gradient_accumulation_steps=4,   # effective batch: 16 × 4 = 64 per GPU
    logging_steps=10,                # print the loss every 10 updates
)
trainer = Trainer(model=model, args=args, train_dataset=train_dataset)
trainer.train()

One setting. The Trainer divides the loss, steps the optimizer every 4 micro-batches and logs the loss.

Libraries: plain PyTorch (a few lines), Hugging Face Trainer (gradient_accumulation_steps) and Accelerate (accelerator.accumulate), PyTorch Lightning (accumulate_grad_batches) and DeepSpeed (gradient_accumulation_steps).

Gradient accumulation gives big-batch stability with small-batch memory. The cost is time: the micro-batches run one after another.

Comparing the three tricks

Technique	What changes?	Main benefit	Main trade-off
Quantization	How numbers are stored	Less memory, faster inference	Some quality loss
Mixed precision	The precision of the maths	Faster training, less activation memory	Numerical care: loss scaling and the right GPU
Gradient accumulation	When the weights are updated	Big batches on a small GPU	Slower than a real big batch

Optimization is a trade-off between memory, speed, quality and cost.

So there's no single best trick. Instead of asking "which optimization is best?", I should ask: what is my bottleneck, and what am I willing to trade to remove it?

See slides 57 to 68

Previous: GPU Bottlenecks Mark done and continue

Pack the bag smarter: single-GPU tricks

Quantization

Mixed precision

Gradient accumulation

Quantization: keep fewer digits

How it works

Quantize eight numbers

Benefits and trade-offs

Benefits

Trade-offs

Same model, four formats

In practice

Run it yourself

Mixed precision: fast where we can, precise where we must

How it works

Can a tiny update survive?

Benefits and trade-offs

Benefits

Trade-offs

In practice

Run it yourself

Gradient accumulation: several trips

How it works

Batch of 64 on a GPU that fits 16

Benefits and trade-offs

Benefits

Trade-offs

In practice

Run it yourself

Comparing the three tricks