251 lines
8.4 KiB
Python
251 lines
8.4 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 accelerate device_map (Test 7 method - DISTRIBUTED)
|
|
print(f"\n[INFO] Loading {config['base_model']}...")
|
|
|
|
# Detect layer names dynamically
|
|
from transformers import AutoConfig
|
|
print(" Detecting layer names...")
|
|
config_model = AutoConfig.from_pretrained(config["base_model"], trust_remote_code=True)
|
|
layer_names = [f"{config_model.model_type.title().replace('_', '')}DecoderLayer"]
|
|
if 'moe' in config_model.model_type.lower():
|
|
layer_names.append(f"{config_model.model_type.title().replace('_', '')}SparseMoeBlock")
|
|
print(f" Detected layers: {layer_names}")
|
|
|
|
# Use accelerate to create device_map
|
|
from accelerate import infer_auto_device_map
|
|
print(" Creating device_map with accelerate...")
|
|
|
|
# First load to CPU to get the model structure
|
|
model_temp = AutoModelForCausalLM.from_pretrained(
|
|
config["base_model"],
|
|
device_map="cpu",
|
|
torch_dtype=torch.float16,
|
|
trust_remote_code=True,
|
|
low_cpu_mem_usage=True,
|
|
)
|
|
|
|
# Create device_map
|
|
device_map = infer_auto_device_map(
|
|
model_temp,
|
|
max_memory={i: "15GB" for i in range(torch.cuda.device_count())},
|
|
no_split_module_classes=layer_names,
|
|
)
|
|
print(f" Created device_map with {len(device_map)} entries")
|
|
|
|
# Reload with device_map
|
|
model = AutoModelForCausalLM.from_pretrained(
|
|
config["base_model"],
|
|
device_map=device_map,
|
|
torch_dtype=torch.float16,
|
|
trust_remote_code=True,
|
|
low_cpu_mem_usage=True,
|
|
)
|
|
print("✓ Success: Model distributed across GPUs via accelerate device_map")
|
|
|
|
# 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 already distributed across GPUs via device_map
|
|
# No FSDP needed - device_map handles distribution
|
|
print("✓ Model already distributed across GPUs (device_map)")
|
|
print(f" GPU 0: {torch.cuda.memory_allocated(0) / 1e9:.2f} GB")
|
|
if torch.cuda.device_count() > 1:
|
|
print(f" GPU 1: {torch.cuda.memory_allocated(1) / 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", "steps"),
|
|
eval_steps=config.get("eval_steps", 100),
|
|
bf16=True,
|
|
gradient_checkpointing=config.get("gradient_checkpointing", True),
|
|
)
|
|
|
|
# SFT Trainer
|
|
from trl import SFTTrainer
|
|
|
|
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()
|