Mixed Precision, Gradient Accumulation, and Memory

Mixed precision: bf16 default in 2026

fp32 weights, bf16 activations + gradients during forward/backward, master copy of weights in fp32 for the optimizer step. bf16=True in TrainingArguments handles all of it. On hardware that supports it (Ampere+), there's almost no quality cost and roughly 2x speed.

Gradient accumulation: free virtual batch size

If your GPU fits 4 examples but the recipe wants batch=64, set per_device_train_batch_size=4 and gradient_accumulation_steps=16. The optimizer sees a 64-example gradient by accumulating 16 microbatches. The cost: 16x compute time per optimizer step. The benefit: you can train any recipe on any GPU, just slower.

Three more memory tricks

gradient_checkpointing=True — recomputes activations during backward instead of storing them. ~50% memory reduction at ~30% speed cost.
optim="adamw_8bit" (bitsandbytes) — 8-bit Adam state. Saves a lot when training big models.
torch.compile(model) — PyTorch 2.x's graph compiler. 1.2-1.5x speedup on stable shapes; sometimes hurts on dynamic shapes.

Code

Memory-efficient TrainingArguments·python

from transformers import TrainingArguments

args = TrainingArguments(
    output_dir="./out",
    per_device_train_batch_size=2,
    gradient_accumulation_steps=8,   # effective batch = 16 (single GPU)
    bf16=True,
    gradient_checkpointing=True,
    optim="adamw_8bit",              # requires bitsandbytes
    num_train_epochs=3,
    learning_rate=2e-4,
    warmup_ratio=0.03,
    lr_scheduler_type="cosine",
    save_steps=500, eval_steps=500, logging_steps=50,
    torch_compile=False,             # try True after baseline works
)

Inspect peak memory during training·python

import torch

# Reset peak counter
torch.cuda.reset_peak_memory_stats()

# ... run trainer.train() for some steps ...

print(f"peak: {torch.cuda.max_memory_allocated() / 1e9:.2f} GB")

Mixed Precision, Gradient Accumulation, and Memory

Mixed precision: bf16 default in 2026

Gradient accumulation: free virtual batch size

Three more memory tricks

Code

External links

Exercise

Progress

Comments 0