How to Debug an LLM Fine-Tuning Config

For ML engineers debugging fine-tuning configs · Based on AI Anytime LLM Fine-Tuning End-to-End Framework

// TL;DR

ML engineers debugging a fine-tuning run can use this framework to systematically diagnose the most common failures: out-of-memory errors, flash-attention version conflicts, incorrect LoRA target module names, and memorisation versus generalisation. It provides a decision tree for tuning lora_r, batch size, gradient accumulation, learning rate, and dataset diversity based on observed symptoms. Whether your loss won't converge, your GPU keeps OOMing, or your model overfits, this structured approach maps symptoms to root causes and concrete YAML changes so you can iterate quickly on Axolotl configs.

How do you diagnose out-of-memory (OOM) errors?

OOM errors mean your effective batch size or model precision exceeds available VRAM. The fix hierarchy is clear: first, reduce micro_batch_size to 1; second, increase gradient_accumulation_steps to 8 or 16 to preserve the effective batch size (which equals micro_batch_size × gradient_accumulation_steps × num_GPUs); third, enable QLoRA with load_in_4bit: true if you're not already quantising; fourth, reduce lora_r. Gradient accumulation is the memory fix that lets you keep large effective batches without the VRAM cost — apply it before sacrificing rank or model size.

Why does training fail immediately with flash-attention errors?

Flash-attention version 1.x causes errors in Axolotl — this is one of the most common and misleading setup failures. Always install flash-attention 2.0 or higher via pip install -e '.[flash-attn,deepspeed]', and verify the installed version. If you inherited a config or environment that half-works then crashes during attention computation, the version is the first thing to check. Misleading stack traces often obscure this simple root cause.

Why doesn't LoRA seem to be adapting anything?

If your model isn't learning despite training completing, your lora_target_modules names may not exist for this model family. Naming conventions like q_proj and v_proj differ across architectures, so a config copied from a Llama example may silently match nothing on another model. Always inspect exact layer names with [name for name in model.state_dict().keys()] before setting targets. Confirm you're at minimum targeting q_proj and v_proj, and add k_proj, o_proj, and MLP modules for complex tasks.

How do you tell memorisation from underfitting?

Symptoms diverge. If the model repeats training examples verbatim but fails on new phrasings, it's memorising: reduce epochs, add dropout and weight decay, and increase dataset diversity. If it performs poorly even on training-like inputs, it's underfitting: increase lora_r to 16 or 32, improve dataset quality and size, and add epochs. Always validate with a held-out split and prompts similar to — but not identical to — training data to distinguish the two.

How do you fix a loss curve that won't converge?

A non-converging or unstable loss usually points to the learning rate. Too small and you get slow convergence or a stall in a local minimum; too large and training becomes unstable with spikes. Start around 2e-4 for LoRA and adjust in factors of 2-3. Also confirm your data is correctly formatted as instruction-response pairs — malformed tokenisation can make loss meaningless. Enable Weights & Biases to watch train_loss, tokens per second, and steps per second in real time.

How do you tune lora_r and lora_alpha deliberately?

Think of lora_r as capacity and lora_alpha as the strength of the new signal. Start at r=8 (paper default); raise to 16 or 32 when the dataset is complex and the model underfits; lower it when you hit memory limits. Set lora_alpha to roughly 2× the rank — a lower alpha preserves more base knowledge, a higher one shifts the model toward new data. Don't set rank too low for a rich dataset, a classic pitfall that caps model quality.

What's the systematic iteration loop?

Map symptom to cause to change. Memorising → fewer epochs, more diversity, dropout. Underperforming on new inputs → higher lora_r, better data, more epochs. OOM → lower micro_batch_size, higher gradient accumulation, enable QLoRA, lower rank. Slow convergence → raise learning rate. Unstable training → lower learning rate. Change one variable at a time so you can attribute effects, and keep each config version-controlled.

Next step: Pull up your failing qlora.yml, run the model.state_dict().keys() check to confirm target module names and flash-attention version, then apply one symptom-mapped change from the iteration loop and re-run with Weights & Biases logging enabled.

// FREQUENTLY ASKED QUESTIONS

How do I fix out-of-memory errors without reducing dataset size?

Reduce micro_batch_size to 1 and increase gradient_accumulation_steps to 8 or 16 to keep the same effective batch size while cutting peak memory. If that's not enough, enable QLoRA with load_in_4bit: true and lower lora_r. Dataset size doesn't drive per-step memory — batch size and model precision do — so you rarely need to shrink your data to fit the GPU.

Why does my config work on Llama but fail on another model?

LoRA target module names differ across model families, so q_proj and v_proj in a Llama config may not exist under the same names elsewhere. Run [name for name in model.state_dict().keys()] to get the exact layer names for your model before setting lora_target_modules. Also check that flash-attention is version 2.0+, since 1.x causes errors regardless of model.

How do I know if I should increase or decrease LoRA rank?

Increase lora_r (to 16 or 32) when the model underfits a complex dataset or performs poorly on training-like inputs — higher rank adds capacity. Decrease it when you hit OOM errors or the model overfits a small dataset. Start at the paper default of 8 and adjust based on whether your symptom is underfitting (raise) or memory/overfitting (lower).