agenx-lora-training/training/scripts/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,
        TrainingArguments,
    )
    from peft import (
        LoraConfig,
        get_peft_model,
    )
    from trl import SFTTrainer
    
    print(f"Loading model: {config['base_model']}")
    
    # Load model - try 4-bit first, fall back to bf16 with CPU offload
    print(f"Loading model: {config['base_model']}")
    from transformers import BitsAndBytesConfig, AutoConfig
    
    # Try multiple loading strategies in order
    print(f"Loading model: {config['base_model']}")
    
    # Strategy 1: Load 4-bit with BitsAndBytes
    print("\n[1/4] Trying: 4-bit with BitsAndBytes...")
    try:
        from transformers import BitsAndBytesConfig
        bnb_config = BitsAndBytesConfig(
            load_in_4bit=True,
            bnb_4bit_quant_type="nf4",
            bnb_4bit_compute_dtype=torch.bfloat16,
            bnb_4bit_use_double_quant=True,
        )
        model = AutoModelForCausalLM.from_pretrained(
            config["base_model"],
            quantization_config=bnb_config,
            device_map="auto",
            trust_remote_code=True,
        )
        print("✓ Success: 4-bit model loaded")
    except Exception as e:
        print(f"✗ Failed: {e}")
        
        # Strategy 1: 4-bit AS-IS with auto device_map
        print("\n[1/6] Trying: 4-bit AS-IS...")
        try:
            model = AutoModelForCausalLM.from_pretrained(
                config["base_model"],
                torch_dtype=torch.float16,
                device_map="auto",
                trust_remote_code=True,
            )
            print("✓ Success: 4-bit AS-IS")
        except Exception as e:
            print(f"✗ Failed: {e}")
            
            # Strategy 2: 4-bit with bf16 compute
            print("\n[2/6] Trying: 4-bit with bf16 compute...")
            try:
                from transformers import BitsAndBytesConfig
                bnb_config = BitsAndBytesConfig(load_in_4bit=True, bnb_4bit_quant_type="nf4", bnb_4bit_compute_dtype=torch.bfloat16, bnb_4bit_use_double_quant=True)
                model = AutoModelForCausalLM.from_pretrained(config["base_model"], quantization_config=bnb_config, device_map="auto", trust_remote_code=True)
                print("✓ Success: 4-bit with bf16 compute")
            except Exception as e:
                print(f"✗ Failed: {e}")
                
                # Strategy 3: 4-bit to CPU
                print("\n[3/6] Trying: 4-bit to CPU...")
                try:
                    model = AutoModelForCausalLM.from_pretrained(config["base_model"], torch_dtype=torch.float16, device_map="cpu", trust_remote_code=True)
                    print("✓ Success: 4-bit to CPU")
                except Exception as e:
                    print(f"✗ Failed: {e}")
                    
                    # Strategy 4: 4-bit fp32
                    print("\n[4/6] Trying: 4-bit fp32...")
                    try:
                        model = AutoModelForCausalLM.from_pretrained(config["base_model"], torch_dtype=torch.float32, device_map="auto", trust_remote_code=True)
                        print("✓ Success: 4-bit fp32")
                    except Exception as e:
                        print(f"✗ Failed: {e}")
                        
                        # Strategy 5: bf16 with accelerate CPU offload
                        print("\n[5/6] Trying: bf16 with accelerate CPU offload...")
                        try:
                            from accelerate import Accelerator
                            model = AutoModelForCausalLM.from_pretrained(
                                config["base_model"],
                                torch_dtype=torch.bfloat16,
                                device_map="cpu",
                                low_cpu_mem_usage=True,
                                trust_remote_code=True,
                            )
                            accelerator = Accelerator(cpu_offload=True)
                            model = accelerator.prepare(model)
                            print("✓ Success: bf16 with accelerate CPU offload")
                        except Exception as e:
                            print(f"✗ Failed: {e}")
                            
                            # Strategy 6: bf16 to CPU with gradient checkpointing
                            print("\n[6/6] Trying: bf16 to CPU with gradient checkpointing...")
                            try:
                                model = AutoModelForCausalLM.from_pretrained(
                                    config["base_model"],
                                    torch_dtype=torch.bfloat16,
                                    device_map="cpu",
                                    low_cpu_mem_usage=True,
                                    trust_remote_code=True,
                                )
                                model.gradient_checkpointing_enable()
                                print("✓ Success: bf16 to CPU with gradient checkpointing")
                            except Exception as e:
                                print(f"✗ Failed: {e}")
                                raise RuntimeError("All loading strategies failed!")
    
    # Prepare model for k-bit training
    from peft import prepare_model_for_kbit_training
    model = prepare_model_for_kbit_training(model)
    print("Model prepared for k-bit training.")
    
    # 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 paths - dataset is in training/data/ relative to repo root
    repo_root = Path(__file__).parent.parent.parent
    train_path = str(repo_root / "training" / "data" / "train.jsonl")
    test_path = str(repo_root / "training" / "data" / "test.jsonl")
    
    dataset = load_dataset(
        "json",
        data_files={
            "train": train_path,
            "test": test_path,
        },
    )
    
    # Training arguments (max_seq_length removed - passed to SFTTrainer instead)
    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=config.get("mixed_precision", "bf16") == "bf16",
        fp16=config.get("mixed_precision", "bf16") == "fp16",
        gradient_checkpointing=config.get("gradient_checkpointing", True),
    )
    
    # SFT Trainer (DeepSpeed handles distributed training via config)
    from trl import SFTTrainer
    
    trainer = SFTTrainer(
        model=model,
        processing_class=tokenizer,
        train_dataset=dataset["train"],
        eval_dataset=dataset["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="configs/llama2-7b-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.parent.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()