SFT — Agent Reference
Supervised fine-tuning pipeline reference. For config templates and dataset format examples, see getting-started.qmd and dataset-formats/.
Architecture
YAML Config → axolotl train config.yaml
1. Load base model (+ quantization if QLoRA/8-bit)
2. Apply adapter layers (LoRA/QLoRA) if configured
3. Load + tokenize dataset(s)
- Apply prompt template (chat_template / alpaca / custom)
- Mask inputs (train_on_inputs: false)
- Pack samples into sequences (sample_packing: true)
4. Training loop (HuggingFace Trainer)
- forward → loss → backward → optimizer step → lr scheduler step
5. Save model / adapter weights + tokenizer
Multi-GPU: FSDP or DeepSpeed shards model across GPUs automatically.
Components Required
- A YAML config — model, dataset(s), adapter settings, hyperparameters
- A dataset — HuggingFace Hub, local JSONL/JSON/Parquet, or S3/GCS path
- (Optional) A custom prompt strategy — for non-standard dataset formats
No external server processes needed (unlike GRPO which requires vLLM).
Dataset Format Decision Tree
Is your data in chat/message format?
├─ YES: OpenAI message format (role/content)?
│ ├─ YES ──────────────────────> type: chat_template (recommended)
│ └─ NO (custom field names) ──> type: chat_template + message_property_mappings
└─ NO: Instruction/response pairs?
├─ YES ──> type: alpaca (instruction, input, output)
└─ NO: Raw text?
├─ YES with segments ─────> type: input_output (template-free masking)
└─ YES continuous ────────> type: completion (pretraining-style)
Full format specs: dataset-formats/
Model Size to Adapter Choice
| Model Size | LoRA | QLoRA (4-bit) | Full Fine-Tune | VRAM (approx) |
|---|---|---|---|---|
| 1-3B | Preferred | Low-budget option | Single GPU OK | 8-16 GB (LoRA) |
| 7-8B | Preferred | Good balance | Needs multi-GPU | 16-24 GB (LoRA) |
| 13-14B | Preferred | Good balance | Multi-GPU required | 24-40 GB (LoRA) |
| 30-70B | LoRA or QLoRA | Preferred for single GPU | Multi-node | 40-80 GB (QLoRA) |
Hyperparameter Ranges
| Parameter | LoRA | QLoRA | Full FT |
|---|---|---|---|
learning_rate |
1e-4 to 3e-4 | 1e-4 to 3e-4 | 1e-5 to 5e-5 |
lora_r |
16-64 | 16-64 | N/A |
lora_alpha |
1-2x lora_r |
1-2x lora_r |
N/A |
micro_batch_size |
2-8 | 2-4 | 1-2 |
gradient_accumulation_steps |
2-8 | 4-16 | 4-16 |
num_epochs |
1-3 | 1-3 | 1-3 |
optimizer |
adamw_8bit |
adamw_bnb_8bit |
adamw_torch_fused |
Effective batch = micro_batch * grad_accum * num_gpus. Lower LR for larger models.
Healthy Training Indicators
| Metric | Healthy | Problem |
|---|---|---|
train_loss |
Decreasing, starting ~2-4 for chat models | Flat or increasing from step 1 — data or LR issue |
eval_loss |
Decreasing, tracks train_loss | Increasing while train_loss decreases — overfitting |
grad_norm |
0.1-10, relatively stable | Spikes >100 — instability. 0.0 — frozen weights |
learning_rate |
Follows scheduler curve | Flat or NaN — config issue |
Watch for: loss never decreasing (check train_on_inputs, dataset, LR), loss goes to 0 quickly (overfitting), eval_loss diverging (reduce epochs, add regularization). See training_stability.qmd.
Known Issues
| Issue | Fix |
|---|---|
| OOM during training | Reduce micro_batch_size, enable gradient_checkpointing, reduce sequence_len |
sample_packing + SDPA + bf16 = 0.0 loss |
Use flash_attention: true or disable sample_packing |
| Missing chat template error | Set chat_template: chatml explicitly |
| Label masking wrong | Run axolotl preprocess config.yaml --debug and inspect labels |
| Loss NaN | Use bf16: auto, lower LR, check data for empty samples |
| Tokenizer pad token / infinite loss | Set special_tokens: pad_token: "<\|end_of_text\|>" |
| FSDP save hangs | Use fsdp_state_dict_type: FULL_STATE_DICT |
| DeepSpeed CheckpointError | Set use_reentrant: true in gradient_checkpointing_kwargs |
Full troubleshooting: training_stability.qmd, debugging.qmd
File Map
src/axolotl/
cli/train.py # Entry point for `axolotl train`
cli/preprocess.py # Entry point for `axolotl preprocess`
core/builders/causal.py # HFCausalTrainerBuilder — wires config → SFT trainer
core/trainers/base.py # AxolotlTrainer — base trainer class
core/trainers/mixins/ # Packing, optimizer, scheduler, checkpoints
prompt_strategies/ # Format handlers: chat_template, alpaca, completion, input_output
utils/schemas/config.py # AxolotlInputConfig — main config schema
utils/schemas/datasets.py # SFTDataset, DatasetConfig
utils/schemas/peft.py # LoraConfig — LoRA parameters
integrations/liger/ # Liger kernel plugin
examples/llama-3/ # LoRA, QLoRA, full FT example configs
docs/getting-started.qmd # Quickstart with config templates
docs/optimizations.qmd # Flash attention, gradient checkpointing, sample packing
docs/multi-gpu.qmd # FSDP and DeepSpeed setup