Axolotl¶
This example shows how use Axolotl
with dstack to fine-tune 4-bit Quantized Llama-4-Scout-17B-16E using FSDP and QLoRA.
Prerequisites
Once dstack is installed, go ahead clone the repo, and run dstack init.
$ git clone https://github.com/dstackai/dstack
$ cd dstack
$ dstack init
Training configuration recipe¶
Axolotl reads the model, QLoRA, and dataset arguments, as well as trainer configuration from a scout-qlora-fsdp1.yaml file. The configuration uses 4-bit axolotl quantized version of meta-llama/Llama-4-Scout-17B-16E, requiring only ~43GB VRAM/GPU with 4K context length.
Single-node training¶
The easiest way to run a training script with dstack is by creating a task configuration file.
This file can be found at examples/fine-tuning/axolotl/.dstack.yml .
type: task
# The name is optional, if not specified, generated randomly
name: axolotl-nvidia-llama-scout-train
# Using the official Axolotl's Docker image
image: axolotlai/axolotl:main-latest
# Required environment variables
env:
- HF_TOKEN
- WANDB_API_KEY
- WANDB_PROJECT
- WANDB_NAME=axolotl-nvidia-llama-scout-train
- HUB_MODEL_ID
# Commands of the task
commands:
- wget https://raw.githubusercontent.com/axolotl-ai-cloud/axolotl/main/examples/llama-4/scout-qlora-fsdp1.yaml
- axolotl train scout-qlora-fsdp1.yaml
--wandb-project $WANDB_PROJECT
--wandb-name $WANDB_NAME
--hub-model-id $HUB_MODEL_ID
resources:
# Two GPU (required by FSDP)
gpu: H100:2
# Shared memory size for inter-process communication
shm_size: 24GB
disk: 500GB..
The task uses Axolotl's Docker image, where Axolotl is already pre-installed.
AMD
The example above uses NVIDIA accelerators. To use it with AMD, check out AMD.
Running a configuration¶
Once the configuration is ready, run dstack apply -f <configuration file>, and dstack will automatically provision the
cloud resources and run the configuration.
$ HF_TOKEN=...
$ WANDB_API_KEY=...
$ WANDB_PROJECT=...
$ WANDB_NAME=axolotl-nvidia-llama-scout-train
$ HUB_MODEL_ID=...
$ dstack apply -f examples/fine-tuning/axolotl/.dstack.yml
Source code¶
The source-code of this example can be found in
examples/fine-tuning/axolotl .