Home / AI & GenAI

AI & GenAIHow-toAdvanced

Optimization: getting more from GPUs when fine-tuning - mixed precision, batch sizing, data loading and checkpoints

Fine-tuning jobs often leave GPUs idle while waiting for data or wasting memory on full precision. Profile utilisation first, use bf16 mixed precision, size batches with gradient accumulation, enable gradient checkpointing and parameter-efficient methods when memory is tight, fix the data loader, and checkpoint so spot or preempted capacity is usable.

Deepak Bhardwaj2 min read

Applies toPyTorch fine-tuning (Hugging Face Transformers/TRL, PEFT/LoRA, DeepSpeed/FSDP) on NVIDIA GPUs in cloud VMs, Kubernetes, Azure ML, SageMaker or Databricks

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_workers and 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