axolotl/docs/agents/sft.md

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

1. A YAML config — model, dataset(s), adapter settings, hyperparameters
2. A dataset — HuggingFace Hub, local JSONL/JSON/Parquet, or S3/GCS path
3. (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 attn_implementation: flash_attention_2 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

Profiling

To profile training and identify optimization opportunities:

# Profile steps 3-7 (after warmup/autotuning settles)
profiler_steps_start: 3
profiler_steps: 5

This produces profiler_trace.json (Chrome trace) and snapshot.pickle (memory snapshot) in output_dir. View the Chrome trace at chrome://tracing.

To programmatically inspect the trace:

python scripts/analyze_profile.py output_dir/

The trace shows per-kernel CUDA times, memory allocations, and operator-level breakdown. Look for:

• Large matmul kernels: candidates for fusion or quantization
• Memory copies (H2D/D2H): unnecessary data movement
• Small frequent kernels: candidates for kernel fusion
• Gaps between kernels: pipeline bubbles from CPU overhead

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