Fine-tuning an open model used to mean writing a training script: dataset glue, tokenizer surgery, trainer arguments, sharding flags, and a prayer that the checkpoint resumes. Axolotl replaces the script with a single YAML file and a codebase whose whole architecture exists to translate that YAML into a correct, memory-efficient training run. The result is one of the most widely used fine-tuning frameworks in the open-source world, and its source is a masterclass in how to layer configuration, data, and distributed training without letting any of them tangle. Version 0.21.1 under an Apache-2.0 license, it currently spans roughly 750 Python modules under its src tree, covering everything from classic Alpaca prompt formats to NVFP4 MoE LoRA kernels and asynchronous GRPO rollouts.
What makes the codebase worth reading is that it is honest about the messy parts. Every performance trick you have heard about - sample packing, flash attention swaps, FSDP2 with 4-bit bases, fused Liger kernels - lives here as an actual patch or plugin with clear seams, not as a blog diagram. You can trace one YAML key from the schema that validates it, through the loader that acts on it, to the trainer mixin that consumes it. As always in this series, what follows is an educational tour of published source code.
Overview: the CLI and config layer feeds dataset machinery and model loading; the trainer core consumes packed batches while patches and plugins reshape the model in between.
Reading the overview from left to right:
- cli/main.py is the dispatcher whose main at line 529 wires every subcommand from train to preprocess to merge to cloud launches.
- utils/config/init.py at 636 lines reads the YAML, applies defaults and plugin additions, and hands it to validation.
- utils/schemas/config.py defines AxolotlInputConfig at line 196, a 2240-line pydantic schema that is effectively the framework’s contract with users.
- utils/data/sft.py starts at prepare_datasets line 49 and assembles raw, tokenized, and shuffled datasets for supervised runs.
- prompt_strategies/chat_template.py implements ChatTemplateStrategy at line 308, turning chat rows into masked label tensors.
- utils/samplers/multipack.py provides MultipackBatchSampler at line 252, the best-fit-decreasing bin packer that keeps sequences together.
- loaders/model.py is ModelLoader at line 209, 1468 lines handling quantization configs, device maps, and attention impl selection.
- train.py closes the loop with train at line 703: setup, execute, save, model card, plugin teardown.
Why You Need This
If you have read our tours of peft, accelerate, or deepspeed, you know the individual parts: adapters, sharding, ZeRO stages. Axolotl is the orchestration layer that binds them behind one config file, and reading it shows you what integration actually costs. The repo does not just call PEFT’s get_peft_model; it patches PEFT’s parameter wrappers, seeds adapter initialization deterministically, and validates that a saved adapter matches the patterns you asked for. Those are the failure modes you only discover after your third night-long run dies at step 4000.
It is also the best available map of the 2025-2026 fine-tuning frontier in executable form. QLoRA on FSDP2 without a full-precision replica per GPU, NVFP4 MoE adapters via custom Triton kernels, selective activation checkpointing with CPU offload, asynchronous GRPO with a rollout producer - each landed as a self-contained module you can read in isolation. When a technique confuses you in a paper, the odds are good this repo contains the production version of it. Our torchtitan tour showed a clean-room trainer; axolotl shows the same ideas bolted onto the Hugging Face ecosystem where most practitioners actually live.
Finally, the plugin architecture is a lesson in extension design. Liger kernels, the ScatterMoE LoRA library, evaluation hooks, and even cloud launchers all enter through the same BasePlugin contract, getting called at pre-config-load, pre-model-load, post-model-load, and post-train phases. Nothing core knows Liger exists; the patch manager just asks registered plugins to do their worst at defined moments. If you maintain any training tool and want to stop merging every user request into core, this is the pattern to steal.
Detailed view: entry and config, dataset machinery, model and adapter loading, the trainer core, RL and preference trainers, and the patch and kernel layer.
Start at the entry, because the config is the product. cli/train.py defines do_cli at line 88 and do_train at line 37; the config path lands in utils/config/init.py which merges plugin-added keys before utils/schemas/config.py validates against AxolotlInputConfig at line 196. The schema is worth an evening: every knob from sequence_len to flash_attention is declared with types and defaults, and a capability-aware subclass AxolotlConfigWCapabilities at line 1939 cross-references the declared attention implementation against the installed hardware before a single GPU is touched. When validation fails you get a pydantic error pointing at the exact key, not a CUDA OOM nine hours later.
The dataset machinery is the second act. utils/data/sft.py loads raw datasets at line 319, applies the selected strategy through prompt_strategies/chat_template.py whose ChatTemplateStrategy at line 308 builds input and label tensors from any Jinja chat template, and persists the tokenized result so runs after the first skip straight to loading at line 268. The strategy ABCs live in prompt_tokenizers.py with PromptTokenizingStrategy at line 43, and the legacy Alpaca family at line 162 shows how train-on-input masking works. For RL runs the parallel module utils/data/rl.py at 528 lines prepares preference and prompt datasets instead.
The packing story is where axolotl earns its throughput. utils/samplers/multipack.py implements MultipackBatchSampler at line 252: sequence lengths come in, ffd_check at line 33 validates the bin estimate, pack_group at line 69 does best-fit-decreasing, and generate_batches at line 378 emits index batches whose padding waste is reported by efficiency at line 597. The sampler plugs into AxolotlTrainer through _create_multipack_sampler at line 363 of core/trainers/base.py, and the matching collator V2BatchSamplerDataCollatorForSeq2Seq at line 213 of utils/collators/batching.py pads to the batch max rather than a global max. A monkeypatch teaches the stock HF Trainer about these variable-length batches so loss normalization counts real tokens, a detail that silently skews results if ignored.
Model loading is a 1468-line negotiation with reality. ModelLoader.load at line 313 of loaders/model.py runs pre-model-load patches, picks device map via _set_device_map_config at line 922, quantization via _set_quantization_config at line 1006, and attention via _set_attention_config at line 1105, then _build_model at line 1243 actually constructs the model and _build_adapters at line 727 attaches LoRA through loaders/adapter.py whose load_lora at line 414 and find_all_linear_names at line 60 decide what gets adapters. The deepest magic is monkeypatch/accelerate/fsdp2.py: fsdp2_load_full_state_dict at line 183 shards a full checkpoint across ranks while keeping 4-bit quantized parameters intact, fsdp2_prepare_model at line 583 applies activation checkpointing, and the QLoRA companion monkeypatch/fsdp2_qlora.py makes adapters trainable on the sharded quantized base - the CPU-staged loading that lets you fine-tune models bigger than any single GPU.
On top sit the trainers and plugins. AxolotlTrainer at line 114 of core/trainers/base.py overrides compute_loss at line 771 for proper token-level normalization and composes mixins from core/trainers/mixins/packing.py. RL lives in core/trainers/grpo/trainer.py with AxolotlGRPOTrainer at line 57 and a context-parallel variant at line 89 whose _generate_and_score_completions at line 297 does the rollout-scoring dance, while core/trainers/grpo/async_trainer.py builds a full producer-consumer system - AsyncDataProducer at line 368, RolloutDataset at line 529, AsyncGRPOTrainer at line 832 - so generation overlaps training. Plugins enter through integrations/base.py; LigerPlugin at line 34 of integrations/liger/plugin.py swaps fused kernels in pre_model_load at line 114, and KernelsPlugin at line 146 of integrations/kernels/plugin.py routes MoE adapters to the ScatterMoE library’s grouped kernels while registering house Triton ops like the swiglu in kernels/swiglu.py.
What happens after the last step matters just as much, and the post-training exit ramps are built in. train at line 703 of train.py finishes with save_trained_model at line 274, which handles the FSDP merge-and-rename path (a helper at line 252 renames merged shards to adapter form), writes the model card at line 488, and calls PluginManager.post_train so plugins can unload or upload. From there the CLI takes over: cli/merge_lora.py collapses adapters into base weights with a memory-efficient variant at line 119, and the export command drives utils/gguf.py whose export_gguf at line 149 resolves llama.cpp binaries, preflights the vocab and template at line 89, and quantizes straight to a deployable GGUF for llama.cpp or Ollama. A fine-tune that ends in a serving-ready file without leaving the tool is the actual deliverable most users want.
From Install to Fine-Tuned Model
The README’s uv-first install and one-command train:
curl -LsSf https://astral.sh/uv/install.sh | sh
export UV_TORCH_BACKEND=cu130
uv venv --python 3.12
uv pip install torch==2.14.0 torchvision
uv pip install --no-build-isolation axolotl[deepspeed]
axolotl fetch examples
axolotl train examples/llama-3/lora-1b.yml
The same path programmatically, which is how notebooks and sweeps drive it:
from axolotl.cli.train import do_cli
do_cli(config="examples/llama-3/lora-1b.yml")
Change a few keys - base_model, learning_rate, lora_target_modules, flash_attention: true - and the identical command fine-tunes a different architecture; add rl: grpo with a reward function and the same entrypoint switches to the GRPO trainer.
Honest limits: axolotl’s breadth is also its surface area - the monkeypatch and model_support trees track upstream Hugging Face and torch releases closely, so pinning versions matters more here than in smaller tools. The config schema documents most keys, but exotic combinations (NVFP4 MoE plus context parallelism, say) carry caveats that only the docs’ support matrix admits to. And like every YAML-driven trainer, it will happily train a mediocre model on a mediocre dataset; no framework validates that your data was worth the GPU hours. Enjoyed this post? Never miss out on future posts by following us