axolotl

Axolotl is a tool designed to streamline the fine-tuning of various AI models, offering support for multiple configurations and architectures.

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 xformer, 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!

Table of Contents Axolotl Table of Contents Quickstart ⚡ Edge Builds Axolotl CLI Usage Badge ❤🏷️ Contributing 🤝 Sponsors 🤝❤ Axolotl supports Advanced Setup Environment Docker Conda/Pip venv Cloud GPU Bare Metal Cloud GPU LambdaLabs GCP Windows Mac Google Colab Launching on public clouds via SkyPilot Launching on public clouds via dstack Dataset Config All Config Options Train Preprocess dataset Multi-GPU DeepSpeed FSDP FSDP + QLoRA Weights & Biases Logging Special Tokens Liger Kernel Inference Playground Merge LORA to base Common Errors 🧰 Tokenization Mismatch b/w Inference & Training Debugging Axolotl Need help? 🙋

Axolotl provides a unified repository for fine-tuning a variety of AI models with ease Go ahead and Axolotl questions!!

Quickstart ⚡

Get started with Axolotl in just a few steps! This quickstart guide will walk you through setting up and running a basic fine-tuning task.

Requirements: Nvidia GPU (Ampere architecture or newer for bf16 and Flash Attention) or AMD GPU, Python >=3.10 and PyTorch >=2.3.1.

pip3 install axolotl[flash-attn,deepspeed]

# download examples and optionally deepspeed configs to the local path
axolotl fetch examples
axolotl fetch deepspeed_configs  # OPTIONAL

# finetune using lora
axolotl train examples/llama-3/lora-1b.yml

Edge Builds 🏎️

If you're looking for the latest features and updates between releases, you'll need to install from source.

git clone https://github.com/axolotl-ai-cloud/axolotl.git
cd axolotl
pip3 install packaging ninja
pip3 install -e '.[flash-attn,deepspeed]'

Axolotl CLI Usage

We now support a new, more streamlined CLI using click.

# preprocess datasets - optional but recommended
CUDA_VISIBLE_DEVICES="0" axolotl preprocess examples/llama-3/lora-1b.yml

# finetune lora
axolotl train examples/llama-3/lora-1b.yml

# inference
axolotl inference examples/llama-3/lora-1b.yml \
    --lora-model-dir="./outputs/lora-out"

# gradio
axolotl inference examples/llama-3/lora-1b.yml \
    --lora-model-dir="./outputs/lora-out" --gradio

# remote yaml files - the yaml config can be hosted on a public URL
# Note: the yaml config must directly link to the **raw** yaml
axolotl train https://raw.githubusercontent.com/axolotl-ai-cloud/axolotl/main/examples/llama-3/lora-1b.yml

We've also added a new command for fetching examples and deepspeed_configs to your local machine. This will come in handy when installing axolotl from PyPI.

# Fetch example YAML files (stores in "examples/" folder)
axolotl fetch examples

# Fetch deepspeed config files (stores in "deepspeed_configs/" folder)
axolotl fetch deepspeed_configs

# Optionally, specify a destination folder
axolotl fetch examples --dest path/to/folder

Legacy Usage

# preprocess datasets - optional but recommended
CUDA_VISIBLE_DEVICES="0" python -m axolotl.cli.preprocess examples/llama-3/lora-1b.yml

# finetune lora
accelerate launch -m axolotl.cli.train examples/llama-3/lora-1b.yml

# inference
accelerate launch -m axolotl.cli.inference examples/llama-3/lora-1b.yml \
    --lora_model_dir="./outputs/lora-out"

# gradio
accelerate launch -m axolotl.cli.inference examples/llama-3/lora-1b.yml \
    --lora_model_dir="./outputs/lora-out" --gradio

# remote yaml files - the yaml config can be hosted on a public URL
# Note: the yaml config must directly link to the **raw** yaml
accelerate launch -m axolotl.cli.train https://raw.githubusercontent.com/axolotl-ai-cloud/axolotl/main/examples/llama-3/lora-1b.yml

Badge ❤🏷️

Building something cool with Axolotl? Consider adding a badge to your model card.

[<img height="32" width="200" alt="Built with Axolotl" src="https://raw.githubusercontent.com/axolotl-ai-cloud/axolotl/main/image/axolotl-badge-web.png">](https://github.com/axolotl-ai-cloud/axolotl)

Sponsors 🤝❤

If you love axolotl, consider sponsoring the project by reaching out directly to wing@axolotl.ai.

• Modal Modal lets you run data/AI 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 LLM models, run protein folding simulations, and much more.

Contributing 🤝

Please read the contributing guide

Bugs? Please check the open issues else create a new Issue.

PRs are greatly welcome!

Please run the quickstart instructions followed by the below to setup env:

pip3 install -r requirements-dev.txt -r requirements-tests.txt
pre-commit install

# test
pytest tests/

# optional: run against all files
pre-commit run --all-files

Thanks to all of our contributors to date. Help drive open source AI progress forward by contributing to Axolotl.

Axolotl supports

✅: supported ❌: not supported ❓: untested

Advanced Setup

Environment

Docker

docker run --gpus '"all"' --rm -it axolotlai/axolotl:main-latest

Or run on the current files for development:

docker compose up -d

>[!Tip] > If you want to debug axolotl or prefer to use Docker as your development environment, see the debugging guide's section on Docker.

docker run --privileged --gpus '"all"' --shm-size 10g --rm -it --name axolotl --ipc=host --ulimit memlock=-1 --ulimit stack=67108864 --mount type=bind,src="${PWD}",target=/workspace/axolotl -v ${HOME}/.cache/huggingface:/root/.cache/huggingface axolotlai/axolotl:main-latest

Conda/Pip venv

1. Install python >=3.10

2. Install pytorch stable https://pytorch.org/get-started/locally/

3. Install Axolotl along with python dependencies

   pip3 install packaging
   pip3 install -e '.[flash-attn,deepspeed]'
   

4. (Optional) Login to Huggingface to use gated models/datasets.

   huggingface-cli login
   

   Get the token at huggingface.co/settings/tokens

Cloud GPU

For cloud GPU providers that support docker images, use axolotlai/axolotl-cloud:main-latest

• on Latitude.sh use this direct link
• on JarvisLabs.ai use this direct link
• on RunPod use this direct link

Bare Metal Cloud GPU

LambdaLabs

sudo apt update
sudo apt install -y python3.10

sudo update-alternatives --install /usr/bin/python python /usr/bin/python3.10 1
sudo update-alternatives --config python # pick 3.10 if given option
python -V # should be 3.10

wget https://bootstrap.pypa.io/get-pip.py
python get-pip.py

pip3 install protobuf==3.20.3
pip3 install -U --ignore-installed requests Pillow psutil scipy

export LD_LIBRARY_PATH=/usr/lib/x86_64-linux-gnu:$LD_LIBRARY_PATH

GCP

pip uninstall -y torch_xla[tpu]

Windows

Please use WSL or Docker!

Mac

Use the below instead of the install method in QuickStart.

pip3 install -e '.'

More info: mac.md

Google Colab

Please use this example notebook.

Launching on public clouds via SkyPilot

To launch on GPU instances (both on-demand and spot instances) on 7+ clouds (GCP, AWS, Azure, OCI, and more), you can use SkyPilot:

pip install "skypilot-nightly[gcp,aws,azure,oci,lambda,kubernetes,ibm,scp]"  # choose your clouds
sky check

Get the example YAMLs of using Axolotl to finetune mistralai/Mistral-7B-v0.1:

git clone https://github.com/skypilot-org/skypilot.git
cd skypilot/llm/axolotl

Use one command to launch:

# On-demand
HF_TOKEN=xx sky launch axolotl.yaml --env HF_TOKEN

# Managed spot (auto-recovery on preemption)
HF_TOKEN=xx BUCKET= sky spot launch axolotl-spot.yaml --env HF_TOKEN --env BUCKET

Launching on public clouds via dstack

To launch on GPU instance (both on-demand and spot instances) on public clouds (GCP, AWS, Azure, Lambda Labs, TensorDock, Vast.ai, and CUDO), you can use dstack.

Write a job description in YAML as below:

# dstack.yaml
type: task

image: axolotlai/axolotl-cloud:main-latest

env:
  - HUGGING_FACE_HUB_TOKEN
  - WANDB_API_KEY

commands:
  - accelerate launch -m axolotl.cli.train config.yaml

ports:
  - 6006

resources:
  gpu:
    memory: 24GB..
    count: 2

then, simply run the job with dstack run command. Append --spot option if you want spot instance. dstack run command will show you the instance with cheapest price across multi cloud services:

pip install dstack
HUGGING_FACE_HUB_TOKEN=xxx WANDB_API_KEY=xxx dstack run . -f dstack.yaml # --spot

For further and fine-grained use cases, please refer to the official dstack documents and the detailed description of axolotl example on the official repository.

Dataset

Axolotl supports a variety of dataset formats. It is recommended to use a JSONL. The schema of the JSONL depends upon the task and the prompt template you wish to use. Instead of a JSONL, you can also use a HuggingFace dataset with columns for each JSONL field.

See the documentation for more information on how to use different dataset formats.

Config

See examples for quick start. It is recommended to duplicate and modify to your needs. The most important options are:

• model

  base_model: ./llama-7b-hf # local or huggingface repo
  

  Note: The code will load the right architecture.

• dataset

  datasets:
      # huggingface repo
    - path: vicgalle/alpaca-gpt4
      type: alpaca
  
      # huggingface repo with specific configuration/subset
    - path: EleutherAI/pile
      name: enron_emails
      type: completion # format from earlier
      field: text # Optional[str] default: text, field to use for completion data
  
      # huggingface repo with multiple named configurations/subsets
    - path: bigcode/commitpackft
      name:
        - ruby
        - python
        - typescript
      type: ... # unimplemented custom format
  
      # chat_template https://axolotl-ai-cloud.github.io/axolotl/docs/dataset-formats/conversation.html#chat_template
    - path: ...
      type: chat_template
      chat_template: chatml # defaults to tokenizer's chat_template
  
      # local
    - path: data.jsonl # or json
      ds_type: json # see other options below
      type: alpaca
  
      # dataset with splits, but no train split
    - path: knowrohit07/know_sql
      type: context_qa.load_v2
      train_on_split: validation
  
      # loading from s3 or gcs
      # s3 creds will be loaded from the system default and gcs only supports public access
    - path: s3://path_to_ds # Accepts folder with arrow/parquet or file path like above. Supports s3, gcs.
      ...
  
      # Loading Data From a Public URL
      # - The file format is `json` (which includes `jsonl`) by default. For different formats, adjust the `ds_type` option accordingly.
    - path: https://some.url.com/yourdata.jsonl # The URL should be a direct link to the file you wish to load. URLs must use HTTPS protocol, not HTTP.
      ds_type: json # this is the default, see other options below.
  

• loading

  load_in_4bit: true
  load_in_8bit: true
  
  bf16: auto # require >=ampere, auto will detect if your GPU supports this and choose automatically.
  fp16: # leave empty to use fp16 when bf16 is 'auto'. set to false if you want to fallback to fp32
  tf32: true # require >=ampere
  
  bfloat16: true # require >=ampere, use instead of bf16 when you don't want AMP (automatic mixed precision)
  float16: true # use instead of fp16 when you don't want AMP
  

  Note: Repo does not do 4-bit quantization.

• lora

  adapter: lora # 'qlora' or leave blank for full finetune
  lora_r: 8
  lora_alpha: 16
  lora_dropout: 0.05
  lora_target_modules:
    - q_proj
    - v_proj
  

All Config Options

See these docs for all config options.

Train

Run

accelerate launch -m axolotl.cli.train your_config.yml

> [!TIP] > You can also reference a config file that is hosted on a public URL, for example accelerate launch -m axolotl.cli.train https://yourdomain.com/your_config.yml

Preprocess dataset

You can optionally pre-tokenize dataset with the following before finetuning. This is recommended for large datasets.

• Set dataset_prepared_path: to a local folder for saving and loading pre-tokenized dataset.
• (Optional): Set push_dataset_to_hub: hf_user/repo to push it to Huggingface.
• (Optional): Use --debug to see preprocessed examples.

python -m axolotl.cli.preprocess your_config.yml

Multi-GPU

Below are the options available in axolotl for training with multiple GPUs. Note that DeepSpeed is the recommended multi-GPU option currently because FSDP may experience loss instability.

DeepSpeed

Deepspeed is an optimization suite for multi-gpu systems allowing you to train much larger models than you might typically be able to fit into your GPU's VRAM. More information about the various optimization types for deepspeed is available at https://huggingface.co/docs/accelerate/main/en/usage_guides/deepspeed#what-is-integrated

We provide several default deepspeed JSON configurations for ZeRO stage 1, 2, and 3.

deepspeed: deepspeed_configs/zero1.json

accelerate launch -m axolotl.cli.train examples/llama-2/config.yml --deepspeed deepspeed_configs/zero1.json

FSDP

• llama FSDP

fsdp:
  - full_shard
  - auto_wrap
fsdp_config:
  fsdp_offload_params: true
  fsdp_state_dict_type: FULL_STATE_DICT
  fsdp_transformer_layer_cls_to_wrap: LlamaDecoderLayer

FSDP + QLoRA

Axolotl supports training with FSDP and QLoRA, see these docs for more information.

Weights & Biases Logging

Make sure your WANDB_API_KEY environment variable is set (recommended) or you login to wandb with wandb login.

• wandb options

wandb_mode:
wandb_project:
wandb_entity:
wandb_watch:
wandb_name:
wandb_log_model:

Comet Logging

Make sure your COMET_API_KEY environment variable is set (recommended) or you login to wandb with comet login.

• wandb options

use_comet:
comet_api_key:
comet_workspace:
comet_project_name:
comet_experiment_key:
comet_mode:
comet_online:
comet_experiment_config:

Special Tokens

It is important to have special tokens like delimiters, end-of-sequence, beginning-of-sequence in your tokenizer's vocabulary. This will help you avoid tokenization issues and help your model train better. You can do this in axolotl like this:

special_tokens:
  bos_token: "<s>"
  eos_token: "</s>"
  unk_token: ""
tokens: # these are delimiters
  - "&lt;|im_start|&gt;"
  - "&lt;|im_end|&gt;"

When you include these tokens in your axolotl config, axolotl adds these tokens to the tokenizer's vocabulary.

Liger Kernel

Liger Kernel: Efficient Triton Kernels for LLM Training

https://github.com/linkedin/Liger-Kernel

Liger (LinkedIn GPU Efficient Runtime) Kernel is a collection of Triton kernels designed specifically for LLM training. It can effectively increase multi-GPU training throughput by 20% and reduces memory usage by 60%. The Liger Kernel composes well and is compatible with both FSDP and Deepspeed.

plugins:
  - axolotl.integrations.liger.LigerPlugin
liger_rope: true
liger_rms_norm: true
liger_glu_activation: true
liger_layer_norm: true
liger_fused_linear_cross_entropy: true

Inference Playground

Axolotl allows you to load your model in an interactive terminal playground for quick experimentation. The config file is the same config file used for training.

Pass the appropriate flag to the inference command, depending upon what kind of model was trained:

• Pretrained LORA:

  python -m axolotl.cli.inference examples/your_config.yml --lora_model_dir="./lora-output-dir"
  

• Full weights finetune:

  python -m axolotl.cli.inference examples/your_config.yml --base_model="./completed-model"
  

• Full weights finetune w/ a prompt from a text file:

  cat /tmp/prompt.txt | python -m axolotl.cli.inference examples/your_config.yml \
    --base_model="./completed-model" --prompter=None --load_in_8bit=True
  

-- With gradio hosting

python -m axolotl.cli.inference examples/your_config.yml --gradio

Please use --sample_packing False if you have it on and receive the error similar to below:

> RuntimeError: stack expects each tensor to be equal size, but got [1, 32, 1, 128] at entry 0 and [1, 32, 8, 128] at entry 1

Merge LORA to base

The following command will merge your LORA adapater with your base model. You can optionally pass the argument --lora_model_dir to specify the directory where your LORA adapter was saved, otherwhise, this will be inferred from output_dir in your axolotl config file. The merged model is saved in the sub-directory {lora_model_dir}/merged.

python3 -m axolotl.cli.merge_lora your_config.yml --lora_model_dir="./completed-model"

You may need to use the gpu_memory_limit and/or lora_on_cpu config options to avoid running out of memory. If you still run out of CUDA memory, you can try to merge in system RAM with

CUDA_VISIBLE_DEVICES="" python3 -m axolotl.cli.merge_lora ...

although this will be very slow, and using the config options above are recommended instead.

Common Errors 🧰

See also the FAQ's and debugging guide.

> If you encounter a 'Cuda out of memory' error, it means your GPU ran out of memory during the training process. Here's how to resolve it:

Please reduce any below

• micro_batch_size
• eval_batch_size
• gradient_accumulation_steps
• sequence_len

If it does not help, try running without deepspeed and without accelerate (replace "accelerate launch" with "python") in the command.

Using adamw_bnb_8bit might also save you some memory.

> failed (exitcode: -9)

Usually means your system has run out of system memory. Similarly, you should consider reducing the same settings as when you run out of VRAM. Additionally, look into upgrading your system RAM which should be simpler than GPU upgrades.

> RuntimeError: expected scalar type Float but found Half

Try set fp16: true

> NotImplementedError: No operator found for memory_efficient_attention_forward ...

Try to turn off xformers.

> accelerate config missing

It's safe to ignore it.

> NCCL Timeouts during training

See the NCCL guide.

Tokenization Mismatch b/w Inference & Training

For many formats, Axolotl constructs prompts by concatenating token ids after tokenizing strings. The reason for concatenating token ids rather than operating on strings is to maintain precise accounting for attention masks.

If you decode a prompt constructed by axolotl, you might see spaces between tokens (or lack thereof) that you do not expect, especially around delimiters and special tokens. When you are starting out with a new format, you should always do the following:

1. Materialize some data using python -m axolotl.cli.preprocess your_config.yml --debug, and then decode the first few rows with your model's tokenizer.
2. During inference, right before you pass a tensor of token ids to your model, decode these tokens back into a string.
3. Make sure the inference string from #2 looks exactly like the data you fine tuned on from #1, including spaces and new lines. If they aren't the same, adjust your inference server accordingly.
4. As an additional troubleshooting step, you can look at the token ids between 1 and 2 to make sure they are identical.

Having misalignment between your prompts during training and inference can cause models to perform very poorly, so it is worth checking this. See this blog post for a concrete example.

Debugging Axolotl

See this debugging guide for tips on debugging Axolotl, along with an example configuration for debugging with VSCode.

Need help? 🙋

Join our Discord server where our community members can help you.

Need dedicated support? Please contact us at ✉️wing@axolotl.ai for dedicated support options.