Profile first
nvidia-smi dmon -s um -d 5 # utilisation and memory every 5 s
| Observation | Likely bottleneck |
|---|---|
| GPU util jumps between 0% and 100% | Data loading / CPU preprocessing |
| GPU util high, memory low | Batch could be larger |
| Out of memory | Precision, batch size, sequence length, model size |
| Multi-GPU scaling poor | Communication, small per-GPU batches |
The PyTorch profiler or framework logs (samples/sec, tokens/sec) give more detail.
Levers
| Lever | Effect |
|---|---|
| bf16 mixed precision | Roughly halves activation memory, faster on modern GPUs |
| Gradient accumulation | Large effective batch with small per-step memory |
| Gradient checkpointing | Much less activation memory, ~20–30% more compute |
| LoRA / QLoRA (PEFT) | Train small adapters; base model frozen (QLoRA: 4-bit base) |
| Sequence packing | Fewer padding tokens per batch |
| Fused optimizers / FlashAttention kernels | Faster steps where supported |
| FSDP / DeepSpeed ZeRO | Shard optimizer state and parameters across GPUs |
from transformers import TrainingArguments
args = TrainingArguments(
output_dir="out",
bf16=True,
per_device_train_batch_size=8,
gradient_accumulation_steps=8, # effective batch 64 per device
gradient_checkpointing=True,
dataloader_num_workers=8,
dataloader_pin_memory=True,
save_strategy="steps", save_steps=500, save_total_limit=3,
logging_steps=50,
)
Data loading
- Pre-tokenize datasets once and store them (Arrow/Parquet) instead of tokenizing every epoch.
- Enough
num_workersand pinned memory. - Local NVMe or cached datasets rather than reading small files from remote storage per step.
Checkpoints and cheaper capacity
- Save checkpoints regularly (and to durable storage) so jobs survive preemption.
- With reliable resume, spot/low-priority GPUs become an option for long jobs.
- Keep only the last few checkpoints to save storage.
Measure cost per result
Track GPU hours, tokens/sec and final evaluation score per run. A faster configuration only helps if model quality stays the same.
Tip
If GPU utilisation is below ~70% for most of a run, fix data loading before buying bigger GPUs. It's the most common waste in fine-tuning jobs.
Checklist
- Utilisation and throughput profiled
- bf16 mixed precision on
- Batch size + gradient accumulation tuned
- Checkpointing/PEFT used when memory is tight
- Data pre-tokenized and loaders tuned
- Regular durable checkpoints; spot capacity considered
- Cost per run and model quality tracked