310 lines
11 KiB
Python
310 lines
11 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
Train LoRA adapter on Cyron summary dataset.
|
|
|
|
Uses Hugging Face TRL for SFT training.
|
|
"""
|
|
|
|
import argparse
|
|
import os
|
|
import yaml
|
|
from pathlib import Path
|
|
|
|
import torch
|
|
|
|
|
|
def train(config_path):
|
|
"""Train LoRA adapter using TRL."""
|
|
|
|
with open(config_path) as f:
|
|
config = yaml.safe_load(f)
|
|
|
|
from transformers import (
|
|
AutoModelForCausalLM,
|
|
AutoTokenizer,
|
|
BitsAndBytesConfig,
|
|
TrainingArguments,
|
|
)
|
|
from peft import (
|
|
LoraConfig,
|
|
get_peft_model,
|
|
)
|
|
from trl import SFTTrainer
|
|
|
|
print(f"Loading model: {config['base_model']}")
|
|
|
|
# Load model with multiple strategies
|
|
print(f"\n[INFO] Loading {config['base_model']}...")
|
|
errors = []
|
|
|
|
# ------------------------------------------------------------------
|
|
# Strategy 1: QLoRA with FSDP (preferred)
|
|
# ------------------------------------------------------------------
|
|
print("\n[1/4] Trying: 4-bit QLoRA (FSDP)...")
|
|
try:
|
|
bnb_config = BitsAndBytesConfig(
|
|
load_in_4bit=True,
|
|
bnb_4bit_quant_type="nf4",
|
|
bnb_4bit_compute_dtype=torch.bfloat16,
|
|
bnb_4bit_use_double_quant=True,
|
|
bnb_4bit_quant_storage=torch.bfloat16, # Enable FSDP sharding of 4-bit weights
|
|
)
|
|
model = AutoModelForCausalLM.from_pretrained(
|
|
config["base_model"],
|
|
quantization_config=bnb_config,
|
|
device_map="cpu", # Load to CPU, FSDP shards later
|
|
trust_remote_code=True,
|
|
low_cpu_mem_usage=True,
|
|
)
|
|
print("✓ Success: QLoRA 4-bit loaded to CPU")
|
|
except Exception as e:
|
|
errors.append(("QLoRA 4-bit", e))
|
|
print(f"✗ Failed: {e}")
|
|
|
|
# --------------------------------------------------------------
|
|
# Strategy 2: BF16 CPU (model too large for single GPU)
|
|
# --------------------------------------------------------------
|
|
print("\n[2/4] Trying: bf16 CPU...")
|
|
try:
|
|
model = AutoModelForCausalLM.from_pretrained(
|
|
config["base_model"],
|
|
torch_dtype=torch.bfloat16,
|
|
device_map="cpu",
|
|
low_cpu_mem_usage=True,
|
|
trust_remote_code=True,
|
|
)
|
|
print("✓ Success: bf16 CPU")
|
|
except Exception as e:
|
|
errors.append(("bf16 CPU", e))
|
|
print(f"✗ Failed: {e}")
|
|
|
|
# ----------------------------------------------------------
|
|
# Strategy 3: FP16 CPU (fallback)
|
|
# ----------------------------------------------------------
|
|
print("\n[3/4] Trying: fp16 CPU...")
|
|
try:
|
|
model = AutoModelForCausalLM.from_pretrained(
|
|
config["base_model"],
|
|
torch_dtype=torch.float16,
|
|
device_map="cpu",
|
|
trust_remote_code=True,
|
|
low_cpu_mem_usage=True,
|
|
)
|
|
print("✓ Success: fp16 CPU")
|
|
except Exception as e:
|
|
errors.append(("fp16 CPU", e))
|
|
print(f"✗ Failed: {e}")
|
|
|
|
# ------------------------------------------------------
|
|
# Strategy 4: 4-bit AS-IS (already quantized)
|
|
# ----------------------------------------------------------
|
|
print("\n[4/4] Trying: 4-bit AS-IS...")
|
|
try:
|
|
model = AutoModelForCausalLM.from_pretrained(
|
|
config["base_model"],
|
|
torch_dtype=torch.float16,
|
|
device_map="cpu",
|
|
trust_remote_code=True,
|
|
low_cpu_mem_usage=True,
|
|
)
|
|
print("✓ Success: 4-bit AS-IS")
|
|
except Exception as e:
|
|
errors.append(("4-bit AS-IS", e))
|
|
print(f"✗ Failed: {e}")
|
|
msg = "\n".join(
|
|
f"{name}: {err}" for name, err in errors
|
|
)
|
|
raise RuntimeError(
|
|
f"All loading strategies failed:\n\n{msg}"
|
|
)
|
|
|
|
# Add LoRA
|
|
lora_config = LoraConfig(
|
|
r=config["lora_r"],
|
|
lora_alpha=config["lora_alpha"],
|
|
lora_dropout=config["lora_dropout"],
|
|
target_modules=config["target_modules"],
|
|
task_type=config.get("lora_task_type", "CAUSAL_LM"),
|
|
)
|
|
model = get_peft_model(model, lora_config)
|
|
model.print_trainable_parameters()
|
|
|
|
# Load tokenizer
|
|
tokenizer = AutoTokenizer.from_pretrained(config["base_model"])
|
|
tokenizer.pad_token = tokenizer.eos_token
|
|
tokenizer.padding_side = "right"
|
|
|
|
# Load dataset
|
|
from datasets import load_dataset
|
|
import os
|
|
|
|
# Get dataset paths - training/data/ is relative to project root
|
|
repo_root = Path(__file__).parent
|
|
train_path = str(repo_root / "training" / "data" / "train.jsonl")
|
|
test_path = str(repo_root / "training" / "data" / "test.jsonl")
|
|
print(f"Looking for dataset at: {train_path}")
|
|
|
|
dataset = load_dataset(
|
|
"json",
|
|
data_files={
|
|
"train": train_path,
|
|
"test": test_path,
|
|
},
|
|
)
|
|
|
|
# Manually wrap model with FSDP before passing to trainer
|
|
from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
|
|
from torch.distributed.fsdp.wrap import transformer_auto_wrap_policy
|
|
|
|
def get_auto_wrap_policy(model):
|
|
from transformers.models.qwen3_5_moe.modeling_qwen3_5_moe import Qwen3_5MoeDecoderLayer
|
|
return transformer_auto_wrap_policy({Qwen3_5MoeDecoderLayer}, model)
|
|
|
|
# Wrap model with FSDP on CPU first
|
|
print("Wrapping model with FSDP on CPU...")
|
|
model = FSDP(
|
|
model,
|
|
auto_wrap_policy=get_auto_wrap_policy(model),
|
|
device_id=None, # Keep on CPU initially
|
|
mixed_precision=None,
|
|
sync_module_states=True,
|
|
use_orig_params=True,
|
|
)
|
|
print("✓ Model wrapped with FSDP")
|
|
|
|
# Training arguments (no FSDP config - we handle it manually)
|
|
training_args = TrainingArguments(
|
|
output_dir=config["train_params"]["output_dir"],
|
|
num_train_epochs=config["train_params"]["num_train_epochs"],
|
|
per_device_train_batch_size=config["train_params"]["per_device_train_batch_size"],
|
|
gradient_accumulation_steps=config["train_params"]["gradient_accumulation_steps"],
|
|
learning_rate=float(config["train_params"]["learning_rate"]),
|
|
lr_scheduler_type=config["train_params"]["lr_scheduler_type"],
|
|
weight_decay=config["train_params"]["weight_decay"],
|
|
warmup_steps=int(config["train_params"]["warmup_ratio"] * config["train_params"]["num_train_epochs"] * len(dataset["train"]) // config["train_params"]["per_device_train_batch_size"]),
|
|
logging_steps=config["train_params"]["logging_steps"],
|
|
save_steps=config["train_params"]["save_steps"],
|
|
save_total_limit=config["train_params"]["save_total_limit"],
|
|
eval_strategy=config.get("eval_strategy", "steps"),
|
|
eval_steps=config.get("eval_steps", 100),
|
|
bf16=True,
|
|
gradient_checkpointing=False, # FSDP handles it
|
|
# No fsdp config - we wrap manually above
|
|
)
|
|
|
|
# SFT Trainer
|
|
from trl import SFTTrainer
|
|
|
|
trainer = SFTTrainer(
|
|
model=model,
|
|
processing_class=tokenizer,
|
|
train_dataset=dataset["train"],
|
|
eval_dataset=dataset["test"],
|
|
args=training_args,
|
|
)
|
|
|
|
# Train
|
|
print("Starting training...")
|
|
trainer.train()
|
|
|
|
# Save
|
|
trainer.save_model(config["train_params"]["output_dir"])
|
|
tokenizer.save_pretrained(config["train_params"]["output_dir"])
|
|
|
|
print(f"Training complete! Model saved to {config['train_params']['output_dir']}")
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(description="Train LoRA adapter")
|
|
parser.add_argument("--config", type=str, default="training/configs/ornith-35b-lora.yaml",
|
|
help="Training configuration file")
|
|
parser.add_argument("--check-only", action="store_true",
|
|
help="Validate config and dependencies without training")
|
|
args = parser.parse_args()
|
|
|
|
if args.check_only:
|
|
check_setup(args.config)
|
|
else:
|
|
train(args.config)
|
|
|
|
|
|
def check_setup(config_path):
|
|
"""Validate config and dependencies without loading model."""
|
|
print("=== Checking Setup ===")
|
|
|
|
# Check config file
|
|
print(f"\n1. Config file: {config_path}")
|
|
if not Path(config_path).exists():
|
|
print(f" ERROR: Config file not found: {config_path}")
|
|
return False
|
|
print(" ✓ Config file exists")
|
|
|
|
# Load and validate config
|
|
with open(config_path) as f:
|
|
config = yaml.safe_load(f)
|
|
print(" ✓ Config file is valid YAML")
|
|
|
|
# Check model path
|
|
print(f"\n2. Model path: {config['base_model']}")
|
|
if Path(config['base_model']).exists():
|
|
print(f" ✓ Model path exists: {config['base_model']}")
|
|
else:
|
|
print(f" ⚠ Model path not found: {config['base_model']}")
|
|
print(" (Will download from HuggingFace during training)")
|
|
|
|
# Check dataset files
|
|
print("\n3. Dataset files:")
|
|
repo_root = Path(__file__).parent
|
|
train_path = repo_root / "training" / "data" / "train.jsonl"
|
|
test_path = repo_root / "training" / "data" / "test.jsonl"
|
|
|
|
if train_path.exists():
|
|
print(f" ✓ Train data: {train_path}")
|
|
else:
|
|
print(f" ✗ Train data missing: {train_path}")
|
|
|
|
if test_path.exists():
|
|
print(f" ✓ Test data: {test_path}")
|
|
else:
|
|
print(f" ✗ Test data missing: {test_path}")
|
|
|
|
# Check GPU
|
|
print("\n4. GPU:")
|
|
import torch
|
|
if torch.cuda.is_available():
|
|
print(f" ✓ GPU available: {torch.cuda.get_device_name(0)}")
|
|
print(f" ✓ VRAM: {torch.cuda.get_device_properties(0).total_memory / 1e9:.1f} GB")
|
|
print(f" ✓ GPU count: {torch.cuda.device_count()}")
|
|
else:
|
|
print(" ✗ No GPU detected!")
|
|
return False
|
|
|
|
# Check required packages
|
|
print("\n5. Required packages:")
|
|
packages = ['transformers', 'datasets', 'trl', 'peft', 'accelerate', 'deepspeed']
|
|
for pkg in packages:
|
|
try:
|
|
__import__(pkg)
|
|
print(f" ✓ {pkg}")
|
|
except ImportError:
|
|
print(f" ✗ {pkg} not installed")
|
|
|
|
# Check DeepSpeed config
|
|
print("\n6. DeepSpeed config:")
|
|
if 'deepspeed_config' in config:
|
|
print(" ✓ DeepSpeed config present")
|
|
ds_config = config['deepspeed_config']
|
|
if 'zero_optimization' in ds_config:
|
|
stage = ds_config['zero_optimization'].get('stage', 'N/A')
|
|
print(f" ✓ ZeRO stage: {stage}")
|
|
else:
|
|
print(" ✗ No DeepSpeed config found")
|
|
|
|
print("\n=== Check Complete ===")
|
|
print("If all checks pass, you can run training with:")
|
|
print(f" bash train-on-this-server.sh")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|