342 lines
11 KiB
Python
342 lines
11 KiB
Python
"""
|
|
Opus-Candid v2 — Qwen 3 Fine-Tuning Script
|
|
============================================
|
|
Trains 8B and 32B sequentially on RTX PRO 6000 Blackwell (96GB VRAM)
|
|
8B: Full 16-bit LoRA (fits in VRAM easily)
|
|
32B: 4-bit QLoRA
|
|
|
|
Usage:
|
|
python train_qwen3.py
|
|
"""
|
|
|
|
import json
|
|
import os
|
|
import sys
|
|
import torch
|
|
import gc
|
|
from datetime import datetime
|
|
from pathlib import Path
|
|
|
|
# ============================================================
|
|
# CONFIG
|
|
# ============================================================
|
|
|
|
DATASET_PATH = "/workspace/opus_candid_full_dataset.json"
|
|
OUTPUT_BASE = "/workspace/opus_candid_v2"
|
|
HF_ORG = "Verdugie"
|
|
|
|
MODELS = [
|
|
{
|
|
"name": "Opus-Candid-8B-v2",
|
|
"base_model": "Qwen/Qwen3-8B",
|
|
"load_in_4bit": False, # 16-bit — fits on 96GB easily
|
|
"lora_r": 16,
|
|
"lora_alpha": 32,
|
|
"per_device_batch_size": 4,
|
|
"gradient_accumulation": 4, # effective batch = 16
|
|
"num_epochs": 3,
|
|
"learning_rate": 2e-4,
|
|
"max_seq_length": 4096,
|
|
"warmup_ratio": 0.05,
|
|
},
|
|
{
|
|
"name": "Opus-Candid-32B-v2",
|
|
"base_model": "Qwen/Qwen3-32B",
|
|
"load_in_4bit": True, # QLoRA — 32B needs quantization
|
|
"lora_r": 16,
|
|
"lora_alpha": 32,
|
|
"per_device_batch_size": 2,
|
|
"gradient_accumulation": 8, # effective batch = 16
|
|
"num_epochs": 3,
|
|
"learning_rate": 1e-4, # slightly lower for larger model
|
|
"max_seq_length": 4096,
|
|
"warmup_ratio": 0.05,
|
|
},
|
|
]
|
|
|
|
LORA_TARGET_MODULES = [
|
|
"q_proj", "k_proj", "v_proj", "o_proj",
|
|
"gate_proj", "up_proj", "down_proj",
|
|
]
|
|
|
|
# ============================================================
|
|
# DATASET LOADING
|
|
# ============================================================
|
|
|
|
def load_and_prepare_dataset(tokenizer, max_seq_length):
|
|
"""Load ShareGPT-format dataset and convert to ChatML for Qwen 3."""
|
|
from datasets import Dataset
|
|
|
|
print(f" Loading dataset from {DATASET_PATH}...")
|
|
with open(DATASET_PATH, "r", encoding="utf-8") as f:
|
|
raw_data = json.load(f)
|
|
|
|
print(f" Total conversations: {len(raw_data)}")
|
|
|
|
# Convert ShareGPT format (from/value) to ChatML (role/content)
|
|
role_map = {"human": "user", "gpt": "assistant", "system": "system"}
|
|
|
|
formatted = []
|
|
skipped = 0
|
|
for conv in raw_data:
|
|
messages = []
|
|
for turn in conv["conversations"]:
|
|
role = role_map.get(turn["from"], turn["from"])
|
|
messages.append({"role": role, "content": turn["value"]})
|
|
|
|
# Apply chat template WITHOUT thinking mode
|
|
text = tokenizer.apply_chat_template(
|
|
messages,
|
|
tokenize=False,
|
|
add_generation_prompt=False,
|
|
enable_thinking=False, # Critical: disable thinking for Opus-Candid
|
|
)
|
|
|
|
# Check length
|
|
token_count = len(tokenizer.encode(text))
|
|
if token_count > max_seq_length:
|
|
skipped += 1
|
|
continue
|
|
|
|
formatted.append({"text": text})
|
|
|
|
print(f" Formatted: {len(formatted)} | Skipped (too long): {skipped}")
|
|
|
|
# Train/eval split
|
|
dataset = Dataset.from_list(formatted)
|
|
split = dataset.train_test_split(test_size=0.05, seed=42)
|
|
|
|
print(f" Train: {len(split['train'])} | Eval: {len(split['test'])}")
|
|
return split
|
|
|
|
|
|
# ============================================================
|
|
# TRAINING
|
|
# ============================================================
|
|
|
|
def train_model(config):
|
|
"""Train a single model with the given config."""
|
|
|
|
print("\n" + "=" * 60)
|
|
print(f"TRAINING: {config['name']}")
|
|
print(f"Base model: {config['base_model']}")
|
|
print(f"{'QLoRA (4-bit)' if config['load_in_4bit'] else 'LoRA (16-bit)'}")
|
|
print("=" * 60)
|
|
|
|
start_time = datetime.now()
|
|
output_dir = os.path.join(OUTPUT_BASE, config["name"])
|
|
os.makedirs(output_dir, exist_ok=True)
|
|
|
|
# --- Load model ---
|
|
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
|
|
from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training, TaskType
|
|
from trl import SFTTrainer, SFTConfig
|
|
|
|
print("\n Loading tokenizer...")
|
|
tokenizer = AutoTokenizer.from_pretrained(
|
|
config["base_model"],
|
|
trust_remote_code=True,
|
|
)
|
|
|
|
# Qwen 3 has no pad token — use eos
|
|
if tokenizer.pad_token is None:
|
|
tokenizer.pad_token = tokenizer.eos_token
|
|
tokenizer.pad_token_id = tokenizer.eos_token_id
|
|
|
|
print(f" EOS token: {tokenizer.eos_token} (id: {tokenizer.eos_token_id})")
|
|
print(f" PAD token: {tokenizer.pad_token} (id: {tokenizer.pad_token_id})")
|
|
|
|
# --- Prepare dataset ---
|
|
dataset = load_and_prepare_dataset(tokenizer, config["max_seq_length"])
|
|
|
|
# --- Load model ---
|
|
print(f"\n Loading model: {config['base_model']}...")
|
|
|
|
model_kwargs = {
|
|
"trust_remote_code": True,
|
|
"torch_dtype": torch.bfloat16,
|
|
"device_map": "auto",
|
|
"attn_implementation": "flash_attention_2",
|
|
}
|
|
|
|
if config["load_in_4bit"]:
|
|
bnb_config = BitsAndBytesConfig(
|
|
load_in_4bit=True,
|
|
bnb_4bit_quant_type="nf4",
|
|
bnb_4bit_compute_dtype=torch.bfloat16,
|
|
bnb_4bit_use_double_quant=True,
|
|
)
|
|
model_kwargs["quantization_config"] = bnb_config
|
|
|
|
model = AutoModelForCausalLM.from_pretrained(
|
|
config["base_model"],
|
|
**model_kwargs,
|
|
)
|
|
|
|
if config["load_in_4bit"]:
|
|
model = prepare_model_for_kbit_training(model)
|
|
|
|
# --- LoRA ---
|
|
print(" Applying LoRA config...")
|
|
lora_config = LoraConfig(
|
|
r=config["lora_r"],
|
|
lora_alpha=config["lora_alpha"],
|
|
lora_dropout=0.05,
|
|
target_modules=LORA_TARGET_MODULES,
|
|
bias="none",
|
|
task_type=TaskType.CAUSAL_LM,
|
|
)
|
|
|
|
model = get_peft_model(model, lora_config)
|
|
model.print_trainable_parameters()
|
|
|
|
# --- Training args ---
|
|
print(" Configuring trainer...")
|
|
|
|
effective_batch = config["per_device_batch_size"] * config["gradient_accumulation"]
|
|
steps_per_epoch = len(dataset["train"]) // effective_batch
|
|
total_steps = steps_per_epoch * config["num_epochs"]
|
|
eval_steps = max(steps_per_epoch // 4, 10) # eval ~4x per epoch
|
|
save_steps = steps_per_epoch # save each epoch
|
|
|
|
print(f" Effective batch size: {effective_batch}")
|
|
print(f" Steps per epoch: {steps_per_epoch}")
|
|
print(f" Total steps: {total_steps}")
|
|
print(f" Eval every: {eval_steps} steps")
|
|
|
|
training_args = SFTConfig(
|
|
output_dir=output_dir,
|
|
num_train_epochs=config["num_epochs"],
|
|
per_device_train_batch_size=config["per_device_batch_size"],
|
|
per_device_eval_batch_size=1,
|
|
gradient_accumulation_steps=config["gradient_accumulation"],
|
|
learning_rate=config["learning_rate"],
|
|
lr_scheduler_type="cosine",
|
|
warmup_ratio=config["warmup_ratio"],
|
|
weight_decay=0.01,
|
|
bf16=True,
|
|
logging_steps=10,
|
|
eval_strategy="steps",
|
|
eval_steps=eval_steps,
|
|
save_strategy="steps",
|
|
save_steps=save_steps,
|
|
save_total_limit=2,
|
|
max_seq_length=config["max_seq_length"],
|
|
dataset_text_field="text",
|
|
gradient_checkpointing=True,
|
|
gradient_checkpointing_kwargs={"use_reentrant": False},
|
|
optim="adamw_torch_fused",
|
|
report_to="none",
|
|
padding="right", # Qwen 3 prefers right padding
|
|
)
|
|
|
|
# --- Train ---
|
|
trainer = SFTTrainer(
|
|
model=model,
|
|
args=training_args,
|
|
train_dataset=dataset["train"],
|
|
eval_dataset=dataset["test"],
|
|
processing_class=tokenizer,
|
|
)
|
|
|
|
print(f"\n 🚀 Training {config['name']}...\n")
|
|
train_result = trainer.train()
|
|
|
|
# --- Save ---
|
|
print(f"\n Saving LoRA adapters to {output_dir}...")
|
|
trainer.save_model(output_dir)
|
|
tokenizer.save_pretrained(output_dir)
|
|
|
|
# Save training stats
|
|
stats = {
|
|
"model": config["name"],
|
|
"base_model": config["base_model"],
|
|
"quantization": "4-bit QLoRA" if config["load_in_4bit"] else "16-bit LoRA",
|
|
"dataset_size": len(dataset["train"]),
|
|
"eval_size": len(dataset["test"]),
|
|
"epochs": config["num_epochs"],
|
|
"learning_rate": config["learning_rate"],
|
|
"lora_r": config["lora_r"],
|
|
"lora_alpha": config["lora_alpha"],
|
|
"effective_batch_size": effective_batch,
|
|
"max_seq_length": config["max_seq_length"],
|
|
"train_loss": train_result.training_loss,
|
|
"train_runtime_seconds": train_result.metrics["train_runtime"],
|
|
"train_samples_per_second": train_result.metrics["train_samples_per_second"],
|
|
"timestamp": datetime.now().isoformat(),
|
|
}
|
|
|
|
with open(os.path.join(output_dir, "training_stats.json"), "w") as f:
|
|
json.dump(stats, f, indent=2)
|
|
|
|
elapsed = datetime.now() - start_time
|
|
print(f"\n ✅ {config['name']} complete!")
|
|
print(f" Training loss: {train_result.training_loss:.4f}")
|
|
print(f" Time: {elapsed}")
|
|
print(f" Adapters saved: {output_dir}")
|
|
|
|
# --- Cleanup ---
|
|
del model, trainer
|
|
gc.collect()
|
|
torch.cuda.empty_cache()
|
|
|
|
return stats
|
|
|
|
|
|
# ============================================================
|
|
# MAIN
|
|
# ============================================================
|
|
|
|
def main():
|
|
print("=" * 60)
|
|
print("OPUS-CANDID v2 — Qwen 3 Training Pipeline")
|
|
print(f"GPU: {torch.cuda.get_device_name(0)}")
|
|
print(f"VRAM: {torch.cuda.get_device_properties(0).total_mem / 1e9:.1f} GB")
|
|
print(f"Models to train: {', '.join(m['name'] for m in MODELS)}")
|
|
print("=" * 60)
|
|
|
|
# Verify dataset exists
|
|
if not os.path.exists(DATASET_PATH):
|
|
print(f"\n❌ Dataset not found: {DATASET_PATH}")
|
|
sys.exit(1)
|
|
|
|
with open(DATASET_PATH, "r", encoding="utf-8") as f:
|
|
data = json.load(f)
|
|
print(f"\nDataset: {len(data)} conversations")
|
|
|
|
all_stats = []
|
|
|
|
for config in MODELS:
|
|
try:
|
|
stats = train_model(config)
|
|
all_stats.append(stats)
|
|
except Exception as e:
|
|
print(f"\n❌ FAILED: {config['name']}")
|
|
print(f" Error: {e}")
|
|
import traceback
|
|
traceback.print_exc()
|
|
# Continue to next model
|
|
gc.collect()
|
|
torch.cuda.empty_cache()
|
|
continue
|
|
|
|
# Summary
|
|
print("\n" + "=" * 60)
|
|
print("TRAINING COMPLETE — SUMMARY")
|
|
print("=" * 60)
|
|
for s in all_stats:
|
|
print(f"\n {s['model']}:")
|
|
print(f" Loss: {s['train_loss']:.4f}")
|
|
print(f" Time: {s['train_runtime_seconds']:.0f}s")
|
|
print(f" Saved: {OUTPUT_BASE}/{s['model']}")
|
|
|
|
print(f"\nNext steps:")
|
|
print(f" 1. Merge adapters into base models")
|
|
print(f" 2. Convert to GGUF")
|
|
print(f" 3. Quantize Q4_K_M and Q8_0")
|
|
print(f" 4. Upload to HuggingFace")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|