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 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 paths - training/data/ is relative to project root
    repo_root = Path(__file__).parent
    train_path = str(repo_root / "training" / "data" / "train.jsonl")
    test_path = str(repo_root / "training" / "data" / "test.jsonl")
    print(f"Looking for dataset at: {train_path}")
    
    dataset = load_dataset(
        "json",
        data_files={
            "train": train_path,
            "test": test_path,
        },
    )
    
    # Manually wrap model with FSDP before passing to trainer
    try:
        from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
        from torch.distributed.fsdp.wrap import transformer_auto_wrap_policy
        from functools import partial
        import torch.distributed as dist
        
        # Initialize distributed process group (required for FSDP)
        if not dist.is_initialized():
            dist.init_process_group(backend="nccl")
            print("✓ Distributed process group initialized")
        
        def get_auto_wrap_policy(model):
            from transformers.models.qwen3_5_moe.modeling_qwen3_5_moe import Qwen3_5MoeDecoderLayer
            return partial(
                transformer_auto_wrap_policy,
                transformer_layer_cls={Qwen3_5MoeDecoderLayer},
            )
        
        # Wrap model with FSDP on GPU
        print("Wrapping model with FSDP on GPU...")
        model = FSDP(
            model,
            auto_wrap_policy=get_auto_wrap_policy(model),
            device_id=torch.cuda.current_device(),  # Keep on current GPU
            mixed_precision=None,
            sync_module_states=False,  # Model is already on GPU
            use_orig_params=True,
        )
        print("✓ Model wrapped with FSDP (will be sharded across GPUs during training)")
        
        # Training arguments (no FSDP config - we handle it manually)
        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=False,  # FSDP handles it
            # No fsdp config - we wrap manually above
        )
        
        # SFT Trainer
        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']}")
        
        # Cleanup distributed process group
        if dist.is_initialized():
            dist.destroy_process_group()
        
        # Return early - success
        return
    except Exception as e:
        errors.append(("QLoRA 4-bit FSDP training", e))
        print(f"✗ Failed during FSDP training: {e}")
        # Cleanup distributed process group if initialized
        if dist.is_initialized():
            dist.destroy_process_group()
        # Fall through to next strategy


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()