init: add LoRA training infrastructure and 20k dataset

training/scripts/train.py

@@ -0,0 +1,127 @@
+#!/usr/bin/env python3
+"""
+Train LoRA adapter on Cyron summary dataset.
+
+Uses Hugging Face TRL for SFT training with QLoRA.
+"""
+
+import argparse
+import yaml
+from pathlib import Path
+
+
+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 quantization
+    bnb_config = BitsAndBytesConfig(
+        load_in_4bit=config.get("load_in_4bit", True),
+        bnb_4bit_compute_dtype=config.get("bnb_4bit_compute_dtype", "bfloat16"),
+        bnb_4bit_quant_type=config.get("bnb_4bit_quant_type", "nf4"),
+        use_nested_quant=config.get("use_nested_quant", False),
+    )
+    
+    model = AutoModelForCausalLM.from_pretrained(
+        config["base_model"],
+        quantization_config=bnb_config,
+        device_map="auto",
+    )
+    model = prepare_model_for_kbit_training(model)
+    
+    # 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),
+    )
+    
+    # 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"],
+    )
+    
+    # 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()