Fine-Tune LLMs with Axolotl

vastai search offers \
  "gpu_ram >= 16 num_gpus = 1 cuda_vers >= 12.4 disk_space >= 100 reliability > 0.98 inet_down >= 5000" \
  --order "dph_base" --limit 10

vastai create instance <OFFER_ID> \
  --template_hash 43e16621b7e24ec58a340f33a6afd3ef \
  --disk 100 \
  --ssh --direct

until vastai show instance <CONTRACT_ID> --raw | grep -q '"actual_status": "running"'; do
  echo "Waiting for instance to start..."; sleep 10
done
echo "Instance is running"

SSH_URL=$(vastai ssh-url <CONTRACT_ID>)
SSH_HOST=$(echo "$SSH_URL" | sed -E 's|ssh://root@([^:]+):.*|\1|')
SSH_PORT=$(echo "$SSH_URL" | sed -E 's|.*:||')

base_model: Qwen/Qwen2.5-3B

# Use the model's built-in chat template for formatting conversations
chat_template: tokenizer_default
datasets:
  - path: mlabonne/FineTome-100k
    type: chat_template
    split: train[:10%]  # 10% = ~10K examples, keeps training fast
    field_messages: conversations
    message_property_mappings:
      role: from
      content: value
val_set_size: 0.05
output_dir: ./outputs/qwen25-3b-lora

sequence_len: 2048
sample_packing: true  # Packs multiple examples into each sequence to avoid wasted padding

# LoRA: train small adapter layers instead of the full model
adapter: lora
lora_r: 16
lora_alpha: 32
lora_dropout: 0.05
lora_target_linear: true  # Apply LoRA to all linear layers

gradient_accumulation_steps: 4
micro_batch_size: 2
num_epochs: 1
optimizer: adamw_torch
lr_scheduler: cosine
learning_rate: 0.0002

bf16: auto  # Use 16-bit precision to halve memory vs 32-bit
tf32: true

gradient_checkpointing: true  # Saves ~30% VRAM at the cost of ~20% slower training
gradient_checkpointing_kwargs:
  use_reentrant: false
logging_steps: 1
flash_attention: true

warmup_ratio: 0.1  # Gradually increase learning rate for first 10% of training
evals_per_epoch: 4
saves_per_epoch: 1
weight_decay: 0.0

scp -P "$SSH_PORT" config.yml root@"$SSH_HOST":/workspace/config.yml

ssh -p "$SSH_PORT" root@"$SSH_HOST"
cd /workspace
WANDB_MODE=disabled axolotl train config.yml

trainable params: 29,933,568 || all params: 3,115,872,256 || trainable%: 0.9607

{'loss': '0.82', 'grad_norm': '0.21', 'learning_rate': '0.0',      'epoch': '0.003'}
{'loss': '0.67', 'grad_norm': '0.05', 'learning_rate': '0.000186', 'epoch': '0.254'}
...
{'loss': '0.60', 'grad_norm': '0.05', 'learning_rate': '2.67e-08', 'epoch': '0.994'}

Training completed! Saving trained model to ./outputs/qwen25-3b-lora

from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import PeftModel
import torch

# Load the base model (uses the HuggingFace cache from training — no re-download)
model = AutoModelForCausalLM.from_pretrained(
    "Qwen/Qwen2.5-3B",
    torch_dtype=torch.bfloat16,
    device_map="auto"
)
tokenizer = AutoTokenizer.from_pretrained("./outputs/qwen25-3b-lora")

# Load the LoRA adapter on top of the base model
model = PeftModel.from_pretrained(model, "./outputs/qwen25-3b-lora")

# Generate a response
messages = [{"role": "user", "content": "Write a Python function to check if a number is prime."}]
text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
inputs = tokenizer(text, return_tensors="pt").to(model.device)

with torch.no_grad():
    outputs = model.generate(
        **inputs, max_new_tokens=256,
        do_sample=True, temperature=0.7, top_p=0.9
    )

response = tokenizer.decode(outputs[0][inputs.input_ids.shape[1]:], skip_special_tokens=True)
print(response)

scp -P "$SSH_PORT" test_inference.py root@"$SSH_HOST":/workspace/test_inference.py
ssh -p "$SSH_PORT" root@"$SSH_HOST" "cd /workspace && python test_inference.py"

def is_prime(n: int) -> bool:
    """Check if a number is prime."""
    if n <= 1:
        return False
    if n <= 3:
        return True
    if n % 2 == 0 or n % 3 == 0:
        return False
    ...

scp -P "$SSH_PORT" -r root@"$SSH_HOST":/workspace/outputs/qwen25-3b-lora ./qwen25-3b-lora

vastai destroy instance <CONTRACT_ID>