Massed Compute

Launch GPU VM

Create a new VM with GPU support and pre-installed NVIDIA drivers:

• SKU: gpu_1x_l40 or similar
• Image ID: 184 (Ubuntu Server 24.04 w/ Drivers)
• Instance name: lora-train
• Attach your SSH key at launch

Wait for the VM to reach running status and become SSH accessible.

Verify GPU Access

Connect to your VM and confirm the GPU is visible:

ssh ubuntu@YOUR_VM_IP
nvidia-smi --query-gpu=name,memory.total,memory.free --format=csv,noheader

You should see output like:

NVIDIA L40, 45634 MiB, 45628 MiB

Bootstrap Training Environment

Install system dependencies and create the Python environment:

sudo apt-get update
sudo DEBIAN_FRONTEND=noninteractive apt-get install -y \
  python3-venv python3-pip git curl

mkdir -p ~/lora-training/scripts ~/lora-training/outputs
python3 -m venv ~/lora-training-env
source ~/lora-training-env/bin/activate

python -m pip install --upgrade pip setuptools wheel

Install Training Stack

Install the pinned versions of PyTorch and training libraries:

python -m pip install \
  --extra-index-url https://download.pytorch.org/whl/cu124 \
  'torch==2.6.0' \
  'transformers==4.49.0' \
  'datasets==3.3.2' \
  'accelerate==1.4.0' \
  'peft==0.14.0' \
  'trl==0.15.2' \
  'bitsandbytes==0.45.3' \
  sentencepiece protobuf packaging

Create Environment Script

Set up environment activation and configuration:

cat > ~/lora-training/env.sh <<'ENV'
source ~/lora-training-env/bin/activate
export TOKENIZERS_PARALLELISM=false
ENV

source ~/lora-training/env.sh

Verify Installation

Test that all components are working correctly:

python -c "
import torch, transformers, peft, trl, bitsandbytes
print('torch', torch.__version__, 'cuda', torch.cuda.is_available(), torch.cuda.get_device_name(0))
print('transformers', transformers.__version__)
print('peft', peft.__version__)
print('trl', trl.__version__)
print('bitsandbytes', bitsandbytes.__version__)
"

Create Training Script

Create the QLoRA training script with proper configuration:

cat > ~/lora-training/scripts/smoke_lora_train.py <<'PY'
import os
from pathlib import Path
import torch
from datasets import Dataset
from transformers import (
    AutoModelForCausalLM,
    AutoTokenizer,
    BitsAndBytesConfig,
    DataCollatorForLanguageModeling,
    Trainer,
    TrainingArguments,
)
from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training

base_model = os.environ.get('BASE_MODEL', 'Qwen/Qwen2.5-0.5B-Instruct')
output_dir = Path(os.environ.get('OUTPUT_DIR', '~/lora-training/outputs/smoke')).expanduser()
max_seq_length = int(os.environ.get('MAX_SEQ_LENGTH', '512'))
max_steps = int(os.environ.get('SMOKE_MAX_STEPS', '10'))
output_dir.mkdir(parents=True, exist_ok=True)

# Create synthetic training data
texts = [
    '### Instruction: Say hello in one short sentence.\n### Response: Hello from a tiny LoRA smoke test.',
    '### Instruction: Name one GPU vendor.\n### Response: NVIDIA.',
    '### Instruction: What is QLoRA used for?\n### Response: Efficient fine tuning of language models.',
    '### Instruction: Reply with the word adapter.\n### Response: adapter',
] * 16
raw_dataset = Dataset.from_dict({'text': texts})

# Load tokenizer
tokenizer = AutoTokenizer.from_pretrained(base_model, use_fast=True)
if tokenizer.pad_token is None:
    tokenizer.pad_token = tokenizer.eos_token

# Tokenize dataset
def tokenize(batch):
    return tokenizer(batch['text'], truncation=True, max_length=max_seq_length)

train_dataset = raw_dataset.map(tokenize, batched=True, remove_columns=['text'])

# Configure 4-bit quantization
bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type='nf4',
    bnb_4bit_compute_dtype=torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16,
    bnb_4bit_use_double_quant=True,
)

# Load quantized model
model = AutoModelForCausalLM.from_pretrained(
    base_model,
    quantization_config=bnb_config,
    device_map='auto',
    torch_dtype=torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16,
)
model.config.use_cache = False
model = prepare_model_for_kbit_training(model)

# Configure LoRA
peft_config = LoraConfig(
    r=8,
    lora_alpha=16,
    lora_dropout=0.05,
    bias='none',
    task_type='CAUSAL_LM',
    target_modules=['q_proj', 'k_proj', 'v_proj', 'o_proj', 'gate_proj', 'up_proj', 'down_proj'],
)
model = get_peft_model(model, peft_config)
model.print_trainable_parameters()

# Configure training
training_args = TrainingArguments(
    output_dir=str(output_dir),
    per_device_train_batch_size=1,
    gradient_accumulation_steps=1,
    learning_rate=2e-4,
    max_steps=max_steps,
    logging_steps=1,
    save_steps=max_steps,
    save_total_limit=1,
    bf16=torch.cuda.is_bf16_supported(),
    fp16=not torch.cuda.is_bf16_supported(),
    optim='paged_adamw_8bit',
    report_to=[],
    remove_unused_columns=False,
)

collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False)
trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=train_dataset,
    data_collator=collator,
)

# Run training
trainer.train()
trainer.model.save_pretrained(output_dir)
tokenizer.save_pretrained(output_dir)

# Verify output
assert (output_dir / 'adapter_config.json').exists()
print('adapter_config', output_dir / 'adapter_config.json')
print('LORA_SMOKE_OK')
PY

Run Training

Execute the smoke test to validate your setup:

source ~/lora-training/env.sh
rm -rf ~/lora-training/outputs/smoke
export BASE_MODEL="Qwen/Qwen2.5-0.5B-Instruct"
export OUTPUT_DIR="$HOME/lora-training/outputs/smoke"
export MAX_SEQ_LENGTH="512"
export SMOKE_MAX_STEPS="10"
python ~/lora-training/scripts/smoke_lora_train.py

The training should complete in 2-3 minutes and print LORA_SMOKE_OK.

Verify Adapter Files

Check that the LoRA adapter was created successfully:

ls -lah ~/lora-training/outputs/smoke
test -f ~/lora-training/outputs/smoke/adapter_config.json && echo "Adapter created successfully"

You should see files like adapter_config.json and adapter_model.safetensors.