Llama 3 Fine-tuning

This example fine-tunes Llama3-8B with a code instruction dataset, illustrating how VESSL AI offloads the infrastructural challenges of large-scale AI workloads and help you train multi-billion-parameter models in hours, not weeks.

This is the most compute-intensive workload yet but you will see how VESSL AI’s training stack enables you to seamlessly scale and execute multi-node training. For a more in-depth introduction, refer to our blog post.

​

What you will do

• Fine-tune an LLM with zero-to-minimum setup
• Mount a custom dataset
• Store and export model artifacts

​

Writing the YAML

Let’s fill in the llama3_fine-tuning.yml file.

name: llama-3-finetuning
description: Fine-tune Llama3-8B with instruction datasets
resources:
  cluster: vessl-gcp-oregon
  preset: gpu-l4-small-spot
image: quay.io/vessl-ai/torch:2.2.0-cuda12.3-r4

name: llama-3-finetuning
description: Fine-tune Llama3-8B with instruction datasets
resources:
  cluster: vessl-gcp-oregon
  preset: gpu-l4-small-spot
image: quay.io/vessl-ai/torch:2.2.0-cuda12.3-r4
import:
  /code/:
    git:
  	  url: https://github.com/vessl-ai/examples
  	  ref: main
/dataset/: hf://huggingface.co/datasets/VESSL/code_instructions_small_alpaca
env:
  HF_TOKEN: YOUR_HF_TOKEN

name: llama-3-finetuning
description: A template Run to finetune Llama 3 with code instruction alpaca dataset
resources:
  cluster: vessl-gcp-oregon
  preset: gpu-l4-small-spot
image: quay.io/vessl-ai/torch:2.2.0-cuda12.3-r4
import:
  /code/:
    git:
	  url: https://github.com/vessl-ai/examples
	  ref: main
/dataset/: hf://huggingface.co/datasets/VESSL/code_instructions_small_alpaca
run:
  - command: |-
      pip install -r requirements.txt
      python main.py --dataset /dataset --output-dir /artifacts --batch-size 4 --max-seq-length 512
    workdir: /code/llama-3-finetuning
env:
  HF_TOKEN: YOUR_HF_TOKEN

name: llama-3-finetuning
description: A template Run to finetune Llama 3 with code instruction alpaca dataset
resources:
  cluster: vessl-gcp-oregon
  preset: gpu-l4-small-spot
image: quay.io/vessl-ai/torch:2.2.0-cuda12.3-r4
import:
  /code/:
    git:
	  url: https://github.com/vessl-ai/examples
	  ref: main
/dataset/: hf://huggingface.co/datasets/VESSL/code_instructions_small_alpaca
run:
  - command: |-
      pip install -r requirements.txt
      python main.py --dataset /dataset --output-dir /artifacts --batch-size 4 --max-seq-length 512
    workdir: /code/llama-3-finetuning
export:
  /artifacts/: vessl-artifact://
env:
  HF_TOKEN: YOUR_HF_TOKEN

​

Running the workload

Once the workload is completed, you can follow the link in the terminal to get the output files including the model checkpoints under Files.

vessl run create -f llama3_fine-tuning.yml

​

Behind the scenes

With VESSL AI, you can launch a full-scale LLM fine-tuning workload on any cloud, at any scale, without worrying about these underlying system backends.

• Model checkpointing — VESSL AI stores .pt files to mounted volumes or model registry and ensures seamless checkpointing of fine-tuning progress.
• GPU failovers — VESSL AI can autonomously detect GPU failures, recover failed containers, and automatically re-assign workload to other GPUs.
• Spot instances — Spot instance on VESSL AI works with model checkpointing and export volumes, saving and resuming the progress of interrupted workloads safely.
• Distributed training — VESSL AI comes with native support for PyTorch DistributedDataParallel and simplifies the process for setting up multi-cluster, multi-node distributed training.
• Autoscaling — As more GPUs are released from other tasks, you can dedicate more GPUs to fine-tuning workloads. You can do this on VESSL AI by adding the following to your existing fine-tuning YAML.

​

Tips & tricks

In addition to the model checkpoints, you can track key metrics and parameters with vessl.log Python SDK. Here’s a snippet from main.py.

class VesslLogCallback(TrainerCallback):
    def on_log(
        self, args, state, control, logs: dict[str, float] = None, **kwargs
    ) -> None:
        # log metrics in the main process only
        if state.is_world_process_zero:
            # if the current step is the max step, and the summary keys are subset of log keys, just skip logging.
            # It is summary log and it will be handled in on_train_end logs
            log_keys = set(logs.keys())
            if state.max_steps != state.global_step or not SUMMARY_KEYS.issubset(
                log_keys
            ):
                vessl.log(payload=logs, step=state.global_step) ``

​

Using our web interface

You can repeat the same process on the web. Head over to your Organization, select a project, and create a New run.

​

What’s next?

Next, let’s see how you can serve and deploy your fine-tuning Llama 3 to the cloud and create a text-generation API endpoint.