286 lines
10 KiB
Python
286 lines
10 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']}")
|
|
|
|
# Check if model is quantized (has model-XXXXX-of-XXXXX files)
|
|
import glob
|
|
import json
|
|
safetensor_files = glob.glob(f"{config['base_model']}/*.safetensors")
|
|
is_quantized = any("of-" in Path(f).name for f in safetensor_files)
|
|
|
|
if is_quantized:
|
|
# Create proper index file from safetensors metadata
|
|
index_file = Path(config["base_model"]) / "model.safetensors.index.json"
|
|
if not index_file.exists():
|
|
print(f"\n[INFO] Creating model index file from safetensors metadata...")
|
|
shards = sorted([Path(f).name for f in safetensor_files if "of-" in Path(f).name])
|
|
|
|
# Load metadata from first shard to get weight names
|
|
from safetensors import safe_open
|
|
weight_map = {}
|
|
for shard_name in shards:
|
|
shard_path = Path(config["base_model"]) / shard_name
|
|
with safe_open(str(shard_path), framework="pt", device="cpu") as f:
|
|
for key in f.keys():
|
|
weight_map[key] = shard_name
|
|
|
|
# Create index
|
|
index = {
|
|
"metadata": {"total_size": sum((Path(config["base_model"]) / s).stat().st_size for s in shards)},
|
|
"weight_map": weight_map
|
|
}
|
|
with open(index_file, 'w') as f:
|
|
json.dump(index, f)
|
|
print(f"✓ Created index: {index_file} ({len(weight_map)} weights mapped)")
|
|
|
|
print(f"\n[INFO] Loading pre-quantized BnB 4-bit model...")
|
|
model = AutoModelForCausalLM.from_pretrained(
|
|
config["base_model"],
|
|
device_map="cpu",
|
|
torch_dtype=torch.float16,
|
|
trust_remote_code=True,
|
|
low_cpu_mem_usage=True,
|
|
)
|
|
print("✓ Model loaded to CPU (BnB 4-bit)")
|
|
else:
|
|
# Load bf16 with BnB 4-bit quantization
|
|
print(f"\n[INFO] Loading {config['base_model']} with BnB 4-bit to CPU...")
|
|
from transformers import BitsAndBytesConfig
|
|
bnb_config = BitsAndBytesConfig(
|
|
load_in_4bit=True,
|
|
bnb_4bit_quant_type="nf4",
|
|
bnb_4bit_compute_dtype=torch.float16,
|
|
)
|
|
|
|
model = AutoModelForCausalLM.from_pretrained(
|
|
config["base_model"],
|
|
quantization_config=bnb_config,
|
|
device_map="cpu",
|
|
torch_dtype=torch.float16,
|
|
trust_remote_code=True,
|
|
low_cpu_mem_usage=True,
|
|
)
|
|
print("✓ Model loaded to CPU with BnB 4-bit (~17.5GB)")
|
|
|
|
# Move to GPU
|
|
print(" Moving to GPU 0...")
|
|
model = model.to("cuda:0")
|
|
print("✓ Success: Model loaded to GPU 0 (4-bit)")
|
|
print(f" GPU 0: {torch.cuda.memory_allocated(0) / 1e9:.2f} GB")
|
|
print(f" Free VRAM: {torch.cuda.get_device_properties(0).total_memory / 1e9 - torch.cuda.memory_allocated(0) / 1e9:.2f} GB")
|
|
|
|
# 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 path from config
|
|
dataset_path = config["dataset"][0]["path"]
|
|
print(f"Loading dataset from: {dataset_path}")
|
|
|
|
dataset = load_dataset(
|
|
"json",
|
|
data_files={
|
|
"train": dataset_path,
|
|
},
|
|
)
|
|
|
|
# Model is on single GPU
|
|
print("✓ Model loaded to single GPU")
|
|
print(f" GPU 0: {torch.cuda.memory_allocated(0) / 1e9:.2f} GB")
|
|
print(f" Free VRAM: {torch.cuda.get_device_properties(0).total_memory / 1e9 - torch.cuda.memory_allocated(0) / 1e9:.2f} GB")
|
|
|
|
# Training arguments
|
|
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", "no"),
|
|
bf16=True,
|
|
gradient_checkpointing=config.get("gradient_checkpointing", False),
|
|
optim=config["train_params"].get("optim", "adamw_torch"),
|
|
optim_args=config["train_params"].get("optim_args"),
|
|
)
|
|
print(f"Using optimizer: {training_args.optim}")
|
|
|
|
# SFT Trainer
|
|
from trl import SFTTrainer
|
|
|
|
# Get text column from config (default to 'text' if not specified)
|
|
text_column = config["dataset"][0].get("text_column", "text")
|
|
print(f"Using text column: {text_column}")
|
|
|
|
# Rename column to 'text' if needed (SFTTrainer expects 'text')
|
|
if text_column != "text":
|
|
print(f" Renaming '{text_column}' column to 'text'...")
|
|
dataset["train"] = dataset["train"].rename_column(text_column, "text")
|
|
if "test" in dataset and dataset["test"] is not None:
|
|
dataset["test"] = dataset["test"].rename_column(text_column, "text")
|
|
|
|
trainer = SFTTrainer(
|
|
model=model,
|
|
processing_class=tokenizer,
|
|
train_dataset=dataset["train"],
|
|
eval_dataset=dataset.get("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()
|