Axolotl is a tool designed to streamline post-training for various AI models. Post-training refers to any modifications or additional training performed on pre-trained models - including full model fine-tuning, parameter-efficient tuning (like LoRA and QLoRA), supervised fine-tuning (SFT), instruction tuning, and alignment techniques. With support for multiple model architectures and training configurations, Axolotl makes it easy to get started with these techniques.
Axolotl is designed to work with YAML config files that contain everything you need to preprocess a dataset, train or fine-tune a model, run model inference or evaluation, and much more.
Features:
- Train various Huggingface models such as llama, pythia, falcon, mpt
- Supports fullfinetune, lora, qlora, relora, and gptq
- Customize configurations using a simple yaml file or CLI overwrite
- Load different dataset formats, use custom formats, or bring your own tokenized datasets
- Integrated with xformers, flash attention, liger kernel, rope scaling, and multipacking
- Works with single GPU or multiple GPUs via FSDP or Deepspeed
- Easily run with Docker locally or on the cloud
- Log results and optionally checkpoints to wandb, mlflow or Comet
- And more!
Requirements:
- NVIDIA GPU (Ampere or newer for
bf16
and Flash Attention) or AMD GPU - Python ≥3.10
- PyTorch ≥2.4.1
pip3 install --no-build-isolation axolotl[flash-attn,deepspeed]
# Download example axolotl configs, deepspeed configs
axolotl fetch examples
axolotl fetch deepspeed_configs # OPTIONAL
Other installation approaches are described here.
# Fetch axolotl examples
axolotl fetch examples
# Or, specify a custom path
axolotl fetch examples --dest path/to/folder
# Train a model using LoRA
axolotl train examples/llama-3/lora-1b.yml
That's it! Check out our Getting Started Guide for a more detailed walkthrough.
- Multiple Model Support: Train various models like LLaMA, Mistral, Mixtral, Pythia, and more
- Training Methods: Full fine-tuning, LoRA, QLoRA, and more
- Easy Configuration: Simple YAML files to control your training setup
- Performance Optimizations: Flash Attention, xformers, multi-GPU training
- Flexible Dataset Handling: Use various formats and custom datasets
- Cloud Ready: Run on cloud platforms or local hardware
- Installation Options - Detailed setup instructions for different environments
- Configuration Guide - Full configuration options and examples
- Dataset Guide - Supported formats and how to use them
- Multi-GPU Training
- Multi-Node Training
- Multipacking
- FAQ - Frequently asked questions
- Join our Discord community for support
- Check out our Examples directory
- Read our Debugging Guide
- Need dedicated support? Please contact ✉️[email protected] for options
Contributions are welcome! Please see our Contributing Guide for details.
fp16/fp32 | lora | qlora | gptq | gptq w/flash attn | flash attn | xformers attn | |
---|---|---|---|---|---|---|---|
llama | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ |
Mistral | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ |
Mixtral-MoE | ✅ | ✅ | ✅ | ❓ | ❓ | ❓ | ❓ |
Mixtral8X22 | ✅ | ✅ | ✅ | ❓ | ❓ | ❓ | ❓ |
Pythia | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ | ❓ |
cerebras | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ | ❓ |
btlm | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ | ❓ |
mpt | ✅ | ❌ | ❓ | ❌ | ❌ | ❌ | ❓ |
falcon | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ | ❓ |
gpt-j | ✅ | ✅ | ✅ | ❌ | ❌ | ❓ | ❓ |
XGen | ✅ | ❓ | ✅ | ❓ | ❓ | ❓ | ✅ |
phi | ✅ | ✅ | ✅ | ❓ | ❓ | ❓ | ❓ |
RWKV | ✅ | ❓ | ❓ | ❓ | ❓ | ❓ | ❓ |
Qwen | ✅ | ✅ | ✅ | ❓ | ❓ | ❓ | ❓ |
Gemma | ✅ | ✅ | ✅ | ❓ | ❓ | ✅ | ❓ |
Jamba | ✅ | ✅ | ✅ | ❓ | ❓ | ✅ | ❓ |
✅: supported ❌: not supported ❓: untested
Thank you to our sponsors who help make Axolotl possible:
- Modal - Modal lets you run jobs in the cloud, by just writing a few lines of Python. Customers use Modal to deploy Gen AI models at large scale, fine-tune large language models, run protein folding simulations, and much more.
Interested in sponsoring? Contact us at [email protected]
This project is licensed under the Apache 2.0 License - see the LICENSE file for details.