This example shows how use Axolotl :material-arrow-top-right-thin:{ .external }{:target="_blank"}
with dstack
to fine-tune Llama3 8B using FSDP and QLoRA.
??? info "Prerequisites"
Once dstack
is installed, go ahead clone the repo, and run dstack init
.
<div class="termy">
```shell
$ git clone https://github.com/dstackai/dstack
$ cd dstack
$ dstack init
```
</div>
Axolotl reads the model, LoRA, and dataset arguments, as well as trainer configuration from a YAML file. This file can
be found at examples/fine-tuning/axolotl/config.yaml
:material-arrow-top-right-thin:{ .external }{:target="_blank"}.
You can modify it as needed.
Before you proceed with training, make sure to update the
hub_model_id
inexamples/fine-tuning/axolotl/config.yaml
:material-arrow-top-right-thin:{ .external }{:target="_blank"} with your HuggingFace username.
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/train.dstack.yml
:material-arrow-top-right-thin:{ .external }{:target="_blank"}. Below is its content:
type: task
# The name is optional, if not specified, generated randomly
name: axolotl-train
# Using the official Axolotl's Docker image
image: winglian/axolotl-cloud:main-20240429-py3.11-cu121-2.2.1
# Required environment variables
env:
- HUGGING_FACE_HUB_TOKEN
- WANDB_API_KEY
# Commands of the task
commands:
- accelerate launch -m axolotl.cli.train examples/fine-tuning/axolotl/config.yaml
resources:
gpu:
# 24GB or more vRAM
memory: 24GB..
# Two or more GPU
count: 2..
The task uses Axolotl's Docker image, where Axolotl is already pre-installed.
To run the task, use dstack apply
:
$ HUGGING_FACE_HUB_TOKEN=...
$ WANDB_API_KEY=...
$ dstack apply -f examples/fine-tuning/axolotl/train.dstack.yml
By default,
dstack run
reusesidle
instances from one of the existing fleets. If noidle
instances meet the requirements, it creates a new fleet using one of the configured backends.
The example folder includes a fleet configuration:
examples/fine-tuning/axolotl/fleet.dstack.yml
:material-arrow-top-right-thin:{ .external } {:target="_blank"}
(a single node with a 24GB
GPU).
You can update the fleet configuration to change the vRAM size, GPU model, number of GPUs per node, or number of nodes.
A fleet can be provisioned with dstack apply
:
dstack apply -f examples/fine-tuning/axolotl/fleet.dstack.yml
Once provisioned, the fleet can run dev environments and fine-tuning tasks.
To delete the fleet, use dstack fleet delete
.
To ensure
dstack apply
always reuses an existing fleet, pass--reuse
todstack apply
(or setcreation_policy
toreuse
in the task configuration). The default policy isreuse_or_create
.
If you'd like to play with the example using a dev environment, run
.dstack.yml
:material-arrow-top-right-thin:{ .external }{:target="_blank"} via dstack apply
:
dstack apply -f examples/fine-tuning/axolotl/.dstack.yaml
The source-code of this example can be found in
examples/fine-tuning/axolotl
:material-arrow-top-right-thin:{ .external }{:target="_blank"}.
- Check dev environments, tasks, services, and fleets.
- Browse Axolotl :material-arrow-top-right-thin:{ .external }{:target="_blank"}.