Memory Optimization

~22 min · bf16, gradient-checkpointing, flash-attention, memory

Level 0Observer

0 XP0/43 lessons0/11 achievements

0/120 XP to next level120 XP to go0% complete

The memory toolbox

Mixed precision (bf16/fp16)

Training in 16-bit instead of 32-bit halves memory and speeds up training on modern GPUs.

bf16 vs fp16: bf16 has the same range as fp32 (just less precision), making it more numerically stable. Always prefer bf16 if your hardware supports it (Ampere+ : A100, RTX 3090+).

Gradient checkpointing

Trades compute for memory — recomputes activations during backward pass instead of storing them. Reduces VRAM by 30–50% at the cost of ~20% slower training.

Flash Attention 2

A faster, more memory-efficient attention implementation. Reduces attention memory from O(n²) to O(n) and significantly speeds up training.

Memory checklist (in order of impact)

Use QLoRA (4-bit) — biggest single win for large models.
Enable gradient checkpointing.
Use bf16 (not fp32 or fp16 if avoidable).
Enable Flash Attention 2.
Reduce per_device batch and increase gradient accumulation.
Reduce max_seq_length.

Code

All four memory tricks combined·python

from transformers import AutoModelForCausalLM, BitsAndBytesConfig
from trl import SFTConfig
import torch

# 1. 4-bit QLoRA load
bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.bfloat16,
    bnb_4bit_use_double_quant=True,
)

# 2. Flash Attention 2 + bf16
model = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Llama-3.1-8B-Instruct",
    quantization_config=bnb_config,
    attn_implementation="flash_attention_2",
    torch_dtype=torch.bfloat16,
    device_map="auto",
)

# 3. Gradient checkpointing + bf16 mixed precision
args = SFTConfig(
    bf16=True,                                        # mixed precision
    gradient_checkpointing=True,                      # checkpointing
    gradient_checkpointing_kwargs={"use_reentrant": False},
    per_device_train_batch_size=2,                    # small per-device
    gradient_accumulation_steps=8,                    # effective batch 16
    max_seq_length=2048,                              # tune for memory
)

External links

Exercise

Set up a Llama 3.1 8B QLoRA training run on a 24 GB GPU. Apply the four memory tricks one at a time and record peak VRAM after each: baseline (fp16) → +bf16 → +grad checkpoint → +Flash Attention 2 → +QLoRA. Map the cumulative VRAM savings.