TopolearnOut-of-memory errors and slow epochs are the main practical obstacles in fine-tuning.
Training memory has four parts: weights, gradients, optimizer state (AdamW keeps two extra values per parameter) and activations. TrainingArguments exposes the standard levers:
- Mixed precision:
bf16=True(Ampere or newer GPUs) runs most math in bfloat16; faster and lighter.fp16=Truefor older GPUs uses loss scaling automatically. - Gradient accumulation:
gradient_accumulation_steps=8withper_device_train_batch_size=2gives an effective batch of 16 per device without the activation memory of 16. - Gradient checkpointing:
gradient_checkpointing=Truerecomputes activations in the backward pass, cutting activation memory substantially for roughly 20 to 30% more compute. - Optimizers:
optim="adamw_torch_fused"is faster;"adamw_bnb_8bit"or"paged_adamw_8bit"shrink optimizer state;"adafactor"uses even less. - Throughput:
train_sampling_strategy="group_by_length",dataloader_num_workers,torch_compile=True, and padding-free packing for LLM training remove wasted compute on pad tokens.
A practical OOM checklist: enable bf16, halve the batch and double accumulation, turn on gradient checkpointing, shorten max_length, switch to an 8-bit optimizer, then move to LoRA or QLoRA.
Use auto_find_batch_size=True to let Trainer retry with smaller batches on OOM. Remember that the effective batch size changes the best learning rate.
Example
from transformers import TrainingArguments
args = TrainingArguments(
output_dir="out",
per_device_train_batch_size=4,
gradient_accumulation_steps=8, # effective batch 32 per device
bf16=True, # mixed precision (Ampere+ GPUs)
gradient_checkpointing=True, # trade compute for activation memory
optim="adamw_torch_fused",
learning_rate=2e-5,
warmup_steps=100,
lr_scheduler_type="cosine",
train_sampling_strategy="group_by_length", # less padding per batch
dataloader_num_workers=4,
report_to="none",
)
eff = args.per_device_train_batch_size * args.gradient_accumulation_steps
print(eff) # 32Remember
- Weights, gradients, optimizer state and activations all consume memory
- Gradient accumulation raises effective batch size without more activation memory
- Gradient checkpointing and 8-bit optimizers are the next levers after bf16
Watch out for
- Raising gradient_accumulation_steps without adjusting the learning rate, then wondering why training dynamics changed.