137 lines
4.2 KiB
Python
Executable File
137 lines
4.2 KiB
Python
Executable File
#!/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 (
|
|
prepare_model_for_kbit_training,
|
|
LoraConfig,
|
|
get_peft_model,
|
|
)
|
|
from trl import SFTTrainer, SFTConfig
|
|
|
|
print(f"Loading model: {config['base_model']}")
|
|
|
|
# Load model with distributed training support
|
|
# Use FSDP for multi-GPU training
|
|
from accelerate import Accelerator
|
|
accelerator = Accelerator()
|
|
|
|
# Load model on CPU first, then distribute
|
|
model = AutoModelForCausalLM.from_pretrained(
|
|
config["base_model"],
|
|
torch_dtype=torch.bfloat16,
|
|
device_map="cpu", # Load on CPU first
|
|
)
|
|
|
|
# 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
|
|
|
|
dataset = load_dataset(
|
|
"json",
|
|
data_files={
|
|
"train": config["dataset"][0]["path"].replace("../", ""),
|
|
"test": config["dataset"][0]["path"].replace("../", "").replace("combined_20k.jsonl", "test.jsonl"),
|
|
},
|
|
)
|
|
|
|
# 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=config["train_params"]["learning_rate"],
|
|
lr_scheduler_type=config["train_params"]["lr_scheduler_type"],
|
|
weight_decay=config["train_params"]["weight_decay"],
|
|
warmup_ratio=config["train_params"]["warmup_ratio"],
|
|
max_seq_length=config["train_params"]["max_seq_length"],
|
|
logging_steps=config["train_params"]["logging_steps"],
|
|
save_steps=config["train_params"]["save_steps"],
|
|
save_total_limit=config["train_params"]["save_total_limit"],
|
|
evaluation_strategy=config.get("eval_strategy", "steps"),
|
|
eval_steps=config.get("eval_steps", 100),
|
|
mixed_precision=config.get("mixed_precision", "bf16"),
|
|
gradient_checkpointing=config.get("gradient_checkpointing", True),
|
|
)
|
|
|
|
# Use FSDP for multi-GPU training
|
|
from trl import SFTTrainer
|
|
|
|
# Prepare model for FSDP
|
|
model = accelerator.prepare(model)
|
|
|
|
# SFT Trainer
|
|
trainer = SFTTrainer(
|
|
model=model,
|
|
tokenizer=tokenizer,
|
|
train_dataset=dataset["train"],
|
|
eval_dataset=dataset["test"],
|
|
args=training_args,
|
|
max_seq_length=config["train_params"]["max_seq_length"],
|
|
)
|
|
|
|
# Prepare trainer for distributed training
|
|
trainer = accelerator.prepare(trainer)
|
|
|
|
# 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="configs/llama2-7b-lora.yaml",
|
|
help="Training configuration file")
|
|
args = parser.parse_args()
|
|
|
|
train(args.config)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|