File size: 3,938 Bytes
65ed589 d0c87cc 65ed589 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 | #!/usr/bin/env python3
# /// script
# requires-python = ">=3.10"
# dependencies = [
# "trl>=0.12.0",
# "peft>=0.7.0",
# "transformers>=4.44.2",
# "accelerate>=0.24.0",
# "bitsandbytes>=0.41.0",
# "datasets",
# "scipy",
# "hf_transfer",
# "rich",
# "trackio",
# ]
# ///
import torch
import os
import trackio
from datasets import load_dataset
from peft import LoraConfig
from trl import SFTTrainer, SFTConfig
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
# === CONFIGURATION ===
MODEL_NAME = "Qwen/Qwen3-8B" # Base model — fits locally on M2 Pro 16GB after fine-tuning
DATASET_NAME = "ceperaltab/diamond-vision-dataset"
OUTPUT_DIR = "diamond-vision-expert"
HF_USERNAME = "ceperaltab"
def main():
print("=" * 60)
print("Diamond Vision Expert — QLoRA Fine-tuning")
print(f"Base model : {MODEL_NAME}")
print(f"Dataset : {DATASET_NAME}")
print("=" * 60)
# Load dataset
print(f"📦 Loading dataset: {DATASET_NAME}...")
dataset = load_dataset(DATASET_NAME, split="train")
print(f"✅ Dataset loaded: {len(dataset)} examples")
# Train / eval split
print("🔀 Creating train/eval split...")
dataset_split = dataset.train_test_split(test_size=0.05, seed=42)
train_dataset = dataset_split["train"]
eval_dataset = dataset_split["test"]
print(f" Train: {len(train_dataset)} | Eval: {len(eval_dataset)}")
# Training config
config = SFTConfig(
output_dir=OUTPUT_DIR,
push_to_hub=True,
hub_model_id=f"{HF_USERNAME}/{OUTPUT_DIR}",
hub_strategy="every_save",
# Training
num_train_epochs=1,
per_device_train_batch_size=1,
gradient_accumulation_steps=8,
learning_rate=2e-4,
# NOTE: max_seq_length is NOT supported in SFTConfig (trl>=0.12.0) — removed
# Logging & checkpointing
logging_steps=10,
save_strategy="steps",
save_steps=500,
save_total_limit=2,
# Evaluation
eval_strategy="steps",
eval_steps=500,
# Optimization
warmup_ratio=0.03,
lr_scheduler_type="cosine",
gradient_checkpointing=True,
bf16=True, # A10G supports bf16
# Monitoring
report_to="trackio",
project="diamond-vision-training",
run_name="diamond-vision-qwen3-8b-v1",
)
# LoRA
peft_config = LoraConfig(
r=64,
lora_alpha=16,
lora_dropout=0.1,
bias="none",
task_type="CAUSAL_LM",
target_modules=[
"q_proj", "k_proj", "v_proj", "o_proj",
"gate_proj", "up_proj", "down_proj",
],
)
# 4-bit QLoRA quantization
bnb_config = BitsAndBytesConfig(
load_in_4bit=True,
bnb_4bit_quant_type="nf4",
bnb_4bit_compute_dtype=torch.bfloat16,
bnb_4bit_use_double_quant=True,
)
# Load model
print(f"🔄 Loading base model: {MODEL_NAME}...")
model = AutoModelForCausalLM.from_pretrained(
MODEL_NAME,
quantization_config=bnb_config,
device_map="auto",
trust_remote_code=True,
)
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, trust_remote_code=True)
tokenizer.pad_token = tokenizer.eos_token
tokenizer.padding_side = "right"
# Train
print("🎯 Initializing trainer...")
trainer = SFTTrainer(
model=model,
processing_class=tokenizer,
train_dataset=train_dataset,
eval_dataset=eval_dataset,
args=config,
peft_config=peft_config,
)
print("🚀 Starting training...")
trainer.train()
print("💾 Pushing final adapter to Hub...")
trainer.push_to_hub()
trackio.finish()
print("✅ Done! Adapter pushed to:", f"https://huggingface.co/{HF_USERNAME}/{OUTPUT_DIR}")
if __name__ == "__main__":
main()
|