Axolotl: YAML-Configured LLM Fine-Tuning for Any Dataset Format

Fine-tuning open-weights foundation models has evolved from an academic research endeavor into a core engineering discipline for modern software teams. While cloud-hosted commercial APIs provide rapid prototyping, enterprise applications frequently demand complete data sovereignty, predictable inference latency, custom token distributions, and zero third-party data leakage.

However, writing boilerplate PyTorch training loops with Hugging Face Transformers for every new model architecture - from Llama 3 and Mistral to Gemma 2 and Qwen - introduces acute operational friction. Teams routinely struggle with silent tokenization mismatches, memory fragmentation, complex distributed training hooks, and broken padding masks.

Enter Axolotl, an open-source, configuration-driven LLM post-training framework. By treating model tuning as declarative infrastructure defined in clean YAML manifests, Axolotl standardizes dataset preparation, memory optimizations, quantization, and multi-GPU distributed orchestration into a reproducible GitOps workflow.

At Pristren, our engineering team relies on Axolotl to build domain-specialized language models for enterprise clients who require on-premises and private cloud deployments. In this comprehensive guide, we unpack the architecture of Axolotl, provide production-ready YAML templates, detail sample packing mechanics, and outline the end-to-end pipeline from raw data to quantized inference.

Why Declarative Fine-Tuning Wins: The Configuration as Code Philosophy

Traditional LLM fine-tuning scripts tightly couple model definitions, dataset loaders, optimizer schedulers, and distributed hardware logic. When an engineer updates a tokenizer or switches from LoRA to full parameter tuning, they often have to refactor hundreds of lines of brittle Python code.

Axolotl solves this anti-pattern by enforcing a strict separation between training logic and experiment configuration:

1. Reproducibility: Every model run is governed by a version-controlled YAML manifest. If a model exhibits degraded performance or loss divergence, you can inspect the exact commit and YAML diff to isolate learning rates, LoRA ranks, or batch accumulation shifts.
2. Universal Dataset Parsing: Rather than manually converting JSON files into Hugging Face Dataset objects with complex mapping functions, Axolotl natively ingests ShareGPT, Alpaca, OpenAI JSONL, ChatML, and raw text formats.
3. Turnkey Distributed Hardware Integration: Switching from a single local RTX 4090 to an 8x NVIDIA H100 cluster requires changing a few flags - such as selecting DeepSpeed ZeRO-3 or PyTorch FSDP - without touching a single line of application code.
4. Native Acceleration Kernels: Out of the box, Axolotl integrates FlashAttention-2 and FlashAttention-3, xFormers, and bitsandbytes 4-bit/8-bit optimizers to maximize Tensor Core utilization.

The Core Performance Secret: Sample Packing (Multipack)

The single biggest source of wasted compute in standard fine-tuning pipelines is padding overhead.

In conventional batching, sequences within a batch are padded with dummy zeros to match the length of the longest sequence in that batch. If your context window is 4,096 tokens, but your average user conversation is only 600 tokens, up to 80% of your matrix multiplications are spent computing attention over meaningless padding tokens.

Conventional Batching (VRAM & Compute Waste):
[Example 1: 650 tokens] [PAD: 3446 tokens ..............................] -> 4096
[Example 2: 1200 tokens] [PAD: 2896 tokens .............................] -> 4096
[Example 3: 400 tokens] [PAD: 3696 tokens ..............................] -> 4096

Sample Packing (Multipack - 100% Dense Utilization):
[Ex 1: 650t][EOS][Ex 2: 1200t][EOS][Ex 3: 400t][Ex 4: 1800t][PAD: 46t] -> 4096

Axolotl implements Multipack (Sample Packing). Instead of padding each sequence individually, Axolotl's preprocessor concatenates multiple training examples into a continuous block of 4,096 or 8,192 tokens separated by end-of-sequence (<|eot_id|> or </s>) tokens.

Crucially, Axolotl constructs cumulative position identifiers and block-diagonal attention masks so that tokens in Example 2 never attend to tokens in Example 1. This unlocks:

• 3x to 5x higher training throughput: You process vastly more training tokens per GPU second.
• Lower total cloud spend: Compute runs that normally take 18 hours finish in 4 to 5 hours.
• Stable loss curves: Eliminating zero-padded tokens prevents gradient distortion during backward passes.

Anatomy of a Production-Grade Axolotl YAML Configuration

Below is a complete, battle-tested Axolotl configuration for fine-tuning a Llama-3-8B-Instruct model on custom multi-turn domain conversations using QLoRA (4-bit quantization). This setup is optimized for an environment with single or dual 24 GB/48 GB GPUs (such as an RTX 4090 or A10G/A6000).

# ==========================================================
# Pristren Engineering: Production Axolotl Configuration
# Model: Meta-Llama-3-8B-Instruct
# Technique: QLoRA (4-bit) with FlashAttention-2 & Sample Packing
# ==========================================================

# 1. Base Model & Tokenizer
base_model: meta-llama/Meta-Llama-3-8B-Instruct
model_type: AutoModelForCausalLM
tokenizer_type: AutoTokenizer
is_llama_derived_model: true

load_in_8bit: false
load_in_4bit: true
strict: false

# 2. Dataset Definitions & Schemas
datasets:
  - path: ./data/enterprise_support_sharegpt.json
    type: sharegpt
    conversation: llama3
    train_on_split: train
  - path: ./data/domain_regulations_alpaca.json
    type: alpaca
    train_on_split: train

dataset_prepared_path: ./prepared_data_llama3
val_set_size: 0.05
output_dir: ./lora_out_llama3_8b

# 3. Context Length & Sample Packing
sequence_len: 4096
sample_packing: true
pad_to_sequence_len: true

# 4. LoRA Adapter Hyperparameters
adapter: qlora
lora_r: 32
lora_alpha: 64
lora_dropout: 0.05
lora_target_modules:
  - q_proj
  - k_proj
  - v_proj
  - o_proj
  - gate_proj
  - up_proj
  - down_proj

# 5. Training Dynamics & Optimizer
gradient_accumulation_steps: 4
micro_batch_size: 2
num_epochs: 3
optimizer: adamw_bnb_8bit
lr_scheduler: cosine
learning_rate: 0.0002
warmup_ratio: 0.05
weight_decay: 0.01
max_grad_norm: 1.0

# 6. Memory & Kernel Optimizations
flash_attention: true
gradient_checkpointing: true
early_stopping_patience:
resume_from_checkpoint:

# 7. Training on Responses Only (Crucial for Instruction Quality)
train_on_inputs: false
group_by_length: false

# 8. Logging, Evaluation & Checkpointing
logging_steps: 10
eval_steps: 100
save_strategy: steps
save_steps: 200
save_total_limit: 3
wandb_project: pristren-llm-finetuning
wandb_entity: pristren-ai
wandb_watch: gradients

Key Parameter Breakdown

• adapter: qlora: Loads base model weights in 4-bit NormalFloat (NF4) via Tim Dettmers' QLoRA method, while backpropagating 16-bit LoRA adapter weights. This slashes base model VRAM footprint from 16 GB down to ~5.5 GB.
• lora_target_modules: Notice that we target all linear projection layers (q, k, v, o, gate, up, down). Early LoRA research only targeted q_proj and v_proj, but modern consensus shows that adapting all attention and MLP projection layers is vital for deep domain adaptation.
• train_on_inputs: false: One of the most important flags in modern alignment. When set to false, the cross-entropy loss is masked out for user prompts and instructions. The model is only graded and updated on its own generated responses, preventing it from memorizing user queries.
• lora_r: 32 and lora_alpha: 64: Using a rank (r) of 32 provides sufficient representation capacity for complex formatting and domain vocabularies, while setting alpha = 2 * r maintains optimal scaling stability for gradient updates.

Data Engineering: Formatting Datasets for Axolotl

Axolotl supports virtually every common LLM training dataset format without requiring custom data prep scripts. The two primary formats used in enterprise deployments are ShareGPT (for multi-turn dialogue) and Alpaca (for single-turn instruction-output tasks).

1. Multi-Turn Dialogue (ShareGPT Format)

For interactive chatbots, complex reasoning chains, or tool-calling workflows, ShareGPT is the industry standard:

[
  {
    "conversations": [
      {
        "from": "human",
        "value": "What are the compliance requirements for handling biometric data under GDPR Article 9?"
      },
      {
        "from": "gpt",
        "value": "Under GDPR Article 9, processing biometric data for the purpose of uniquely identifying a natural person is prohibited by default, unless one of the specific exceptions applies: 1) Explicit consent is obtained, 2) Processing is necessary for employment or social security obligations, or 3) Substantial public interest is established."
      },
      {
        "from": "human",
        "value": "Does automated facial recognition in private retail stores qualify under substantial public interest?"
      },
      {
        "from": "gpt",
        "value": "In general, no. European Data Protection Authorities (such as the CNIL and EDPS) have repeatedly ruled that commercial loss prevention does not meet the proportionality threshold for 'substantial public interest' under Article 9(2)(g)."
      }
    ]
  }
]

In your YAML configuration, you simply specify:

datasets:
  - path: ./data/compliance_dialogue.json
    type: sharegpt
    conversation: llama3

Axolotl automatically injects the appropriate special tokens (<|begin_of_text|>, <|start_header_id|>user<|end_header_id|>, <|start_header_id|>assistant<|end_header_id|>) matching the model's native chat template.

2. Single-Turn Instruction (Alpaca Format)

For structured tasks such as JSON entity extraction, classification, or code synthesis:

[
  {
    "instruction": "Extract the enterprise legal entity and invoice total from the provided raw text.",
    "input": "RECEIPT #8491 - Vendor: Northern Analytics Inc. Total Balance Due: $14,250.00 USD by October 31, 2026.",
    "output": "{"vendor": "Northern Analytics Inc.", "total_amount": 14250.00, "currency": "USD", "due_date": "2026-10-31"}"
  }
]

Hardware Sizing & Memory Optimization Matrix

Selecting the right hardware architecture prevents out-of-memory (OOM) crashes and budget overruns. Below is Pristren's field-tested hardware sizing matrix for modern open models:

Core Memory Reduction Tactics

1. Gradient Checkpointing: Recomputes activation tensors during the backward pass instead of caching them in VRAM. This reduces activation memory consumption by up to 60% at the cost of a minor 20% slowdown in forward/backward throughput.
2. FlashAttention-2: Replaces quadratic standard attention memory with tiled, SRAM-optimized fused matrix multiplications. VRAM usage scales linearly with sequence length rather than quadratically.
3. 8-bit Optimizers (adamw_bnb_8bit): Standard AdamW tracks first and second gradient moments in full 32-bit floating point, consuming 8 bytes per trainable parameter. Bitsandbytes quantizes optimizer states to 8-bit, saving 75% of optimizer VRAM with zero loss in convergence stability.

Step-by-Step Training Execution & CLI Workflow

Axolotl is best executed inside an isolated Docker container or a dedicated virtual environment with CUDA 12.x support.

Step 1: Preprocess and Tokenize the Dataset

Before launching a multi-hour GPU training run, validate your dataset schemas and pre-tokenize:

# Pre-tokenize, calculate sample packing, and cache tensors to disk
accelerate launch -m axolotl.cli.preprocess config_llama3.yaml

This command reads your raw JSON datasets, applies the chat templates, packs sequences up to sequence_len: 4096, and saves the resulting arrow tables to ./prepared_data_llama3. If your data contains malformed JSON or unclosed tags, the preprocessor will flag the offending line immediately without wasting cloud GPU compute.

Step 2: Launch Distributed Training

For single-GPU machines:

accelerate launch -m axolotl.cli.train config_llama3.yaml

For multi-GPU setups (e.g., 4x A100 80GB), create or configure an Accelerate configuration file:

# Configure PyTorch Distributed / FSDP
accelerate config

# Launch distributed training across all available GPUs
accelerate launch --num_processes=4 -m axolotl.cli.train config_llama3.yaml

During execution, Axolotl logs real-time step loss, learning rate decay, and gradient norms directly to Weights & Biases or TensorBoard.

[Step 100/1200] Loss: 1.124 | Grad Norm: 0.84 | LR: 1.95e-4 | VRAM: 14.8 GB / 24.0 GB
[Step 200/1200] Loss: 0.942 | Grad Norm: 0.72 | LR: 1.82e-4 | VRAM: 14.8 GB / 24.0 GB
[Step 300/1200] Loss: 0.815 | Grad Norm: 0.65 | LR: 1.63e-4 | VRAM: 14.8 GB / 24.0 GB

Post-Training: Merging Adapters and Exporting to Production

Once training completes, your output directory will contain adapter weights (adapter_model.safetensors and adapter_config.json), not a standalone base model. To deploy this model into high-throughput production inference engines like vLLM or Ollama, you must merge the low-rank delta weights back into the foundation model.

1. Merging LoRA Weights into 16-Bit Precision

Create a lightweight merge configuration or run the Axolotl merge utility:

python -m axolotl.cli.merge_lora config_llama3.yaml \
  --lora_model_dir="./lora_out_llama3_8b" \
  --output_dir="./merged_llama3_8b_enterprise"

This script:

1. Loads the original base model in 16-bit bfloat16 precision.
2. Mathematically multiplies the low-rank matrices ($W = W_0 + rac{alpha}{r}BA$).
3. Adds the adapter deltas to the base weights.
4. Serializes a complete, standalone Hugging Face model directory ready for any standard serving stack.

2. High-Throughput Serving with vLLM

You can now immediately serve your fine-tuned model with paged attention and continuous batching:

python3 -m vllm.entrypoints.openai.api_server \
  --model ./merged_llama3_8b_enterprise \
  --served-model-name "enterprise-llama3-8b" \
  --tensor-parallel-size 1 \
  --gpu-memory-utilization 0.90 \
  --max-model-len 4096

Your custom model is now accessible via a standard OpenAI-compatible REST endpoint at http://localhost:8000/v1/chat/completions.

Real-World Case Study: Pristren Enterprise Domain Adaptation

At Pristren, a multinational logistics client required an automated document parsing and customs classification agent capable of operating strictly on private infrastructure under zero-retention regulatory mandates.

The Challenge

Closed commercial APIs were prohibited due to strict cross-border data protection treaties. Standard off-the-shelf open models (Llama-3-8B-Instruct) scored only 61% accuracy on the client's proprietary customs tariff schedules and failed to consistently format output responses in validated JSON.

The Pristren Solution

1. Curated Instruction Corpus: We generated 14,000 synthetic and verified customs classification records using multi-agent verification loops.
2. Axolotl YAML Pipeline: Using Axolotl with sample packing and QLoRA on a dual-GPU node, we fine-tuned Llama-3-8B over 3 epochs with train_on_inputs: false.
3. Outcome:
   • 97.8% JSON Schema Compliance: Zero parser failures across 100,000 live production invoices.
   • Accuracy Improvement: Tariff classification accuracy increased from 61% to 94.2%.
   • Cost Reduction: Replaced an estimated $18,000/month in third-party API token costs with a private on-premises server consuming less than $600/month in electricity and cloud host depreciation.

Explore more about our enterprise architectural engineering in our Complete Business Solutions or view our past technical milestones in Machine Learning.

Key Takeaways

1. Declarative Architecture: Managing LLM post-training via version-controlled YAML manifests prevents operational entropy and guarantees experiment reproducibility.
2. Eliminate Padding Waste: Always enable sample_packing: true in Axolotl to pack training sequences, tripling GPU throughput and cutting training time.
3. Train on Outputs Only: Set train_on_inputs: false to prevent the model from expending representational capacity memorizing user queries.
4. Target All Linear Layers: When employing QLoRA or LoRA, target all projection modules (q, k, v, o, gate, up, down) for maximum representation transfer.
5. Decouple Training from Serving: Train with low-rank adapters for agility, then merge weights into standalone 16-bit precision for high-throughput serving with vLLM or TensorRT-LLM.