agenx-lora-training/train.py

#!/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 bf16 model to CPU
    print(f"\n[INFO] Loading {config['base_model']} bf16 to CPU...")
    model = AutoModelForCausalLM.from_pretrained(
        config["base_model"],
        device_map="cpu",
        torch_dtype=torch.bfloat16,
        trust_remote_code=True,
        low_cpu_mem_usage=True,
    )
    print("✓ Model loaded to CPU (~70GB bf16)")
    
    # Apply PEFT k-bit training preparation
    print("   Applying PEFT k-bit preparation...")
    from peft import prepare_model_for_kbit_training
    model = prepare_model_for_kbit_training(model, use_gradient_checkpointing=False)
    print("   ✓ Model prepared for k-bit training")
    
    # Manually quantize linear layers
    print("   Quantizing linear layers to 4-bit...")
    from bitsandbytes.nn import Linear4bit
    from torch import nn
    
    quantized_count = 0
    for name, module in model.named_modules():
        if isinstance(module, nn.Linear) and 'lm_head' not in name:
            new_module = Linear4bit(
                module.in_features,
                module.out_features,
                bias=module.bias is not None,
            )
            new_module.weight = nn.Parameter(module.weight.data.clone())
            if module.bias is not None:
                new_module.bias = nn.Parameter(module.bias.data.clone())
            
            layers = name.split('.')
            parent = model
            for layer in layers[:-1]:
                parent = getattr(parent, layer)
            setattr(parent, layers[-1], new_module)
            quantized_count += 1
    
    print(f"   ✓ Quantized {quantized_count} linear layers to 4-bit")
    
    # 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()