train.py

#!/usr/bin/env python3
"""
Training script for sparsity predictors using datasets generated by generate_dataset.py.

This script trains predictors to identify which MLP neurons will be most activated
based on the hidden states before each MLP layer. Uses the last token representations
from the generated datasets.

Usage:
    # Start fresh training for single layer
    python train.py \
        --config meta-llama/Llama-2-7b-hf \
        --dataset_dir ./data/c4 \
        --output_dir ./trained_predictors \
        --layer_indices 0 \
        --batch_size 32 \
        --num_epochs 10 \
        --learning_rate 1e-5
    
    # Start fresh training for multiple layers
    python train.py \
        --config meta-llama/Llama-2-7b-hf \
        --dataset_dir ./data/c4 \
        --output_dir ./trained_predictors \
        --layer_indices 0 1 2 3 \
        --batch_size 32 \
        --num_epochs 10 \
        --learning_rate 1e-5
    
    # Train predictors for ALL layers automatically
    python train.py \
        --config meta-llama/Llama-2-7b-hf \
        --dataset_dir ./data/c4 \
        --output_dir ./trained_predictors \
        --layer_indices all \
        --batch_size 32 \
        --num_epochs 10 \
        --learning_rate 1e-5
    
    # Train predictors with hyperparameter grid (different LoRA sizes)
    python train.py \
        --config meta-llama/Llama-2-7b-hf \
        --dataset_dir ./data/c4 \
        --output_dir ./trained_predictors \
        --layer_indices 0 1 2 \
        --lora_sizes 4.0 10.0 20.0 30.0 \
        --batch_size 32 \
        --num_epochs 10 \
        --learning_rate 1e-5
    
    # Resume from latest checkpoint
    python train.py \
        --config meta-llama/Llama-2-7b-hf \
        --dataset_dir ./data/c4 \
        --output_dir ./trained_predictors \
        --layer_indices 0 1 2 \
        --batch_size 32 \
        --num_epochs 10 \
        --learning_rate 1e-5 \
        --resume_from_checkpoint
    
    # Resume from specific checkpoint
    python train.py \
        --config meta-llama/Llama-2-7b-hf \
        --dataset_dir ./data/c4 \
        --output_dir ./trained_predictors \
        --layer_indices 0 \
        --batch_size 32 \
        --num_epochs 10 \
        --learning_rate 1e-5 \
        --resume_from_checkpoint \
        --checkpoint_path ./trained_predictors/checkpoint_layer_0_step_5000.pt

    # Resume only from best-performing lora size per layer
    python train.py \
        --config meta-llama/Llama-2-7b-hf \
        --dataset_dir ./data/c4 \
        --output_dir ./trained_predictors \
        --layer_indices 0 \
        --batch_size 32 \
        --num_epochs 10 \
        --learning_rate 1e-5 \
        --resume_from_checkpoint \
        --load_best_only
"""

import argparse
import logging
import os
import time

import torch

from transformers import AutoConfig
import wandb

from transformers.trainer_utils import set_seed

from src.trainer import MultiLayerPredictorTrainer

# Setup logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


def main():
    parser = argparse.ArgumentParser(description="Train sparsity predictors from pre-generated datasets")
    parser.add_argument("--config", type=str, required=True, help="Path to model config file")
    parser.add_argument("--dataset_dir", type=str, required=True, help="Directory containing dataset.csv and arrays/")
    parser.add_argument("--output_dir", type=str, required=True, help="Output directory for trained models")
    parser.add_argument("--layer_indices", nargs='+', required=True, help="Which layers to train predictors for (can specify multiple layer numbers or 'all' for all layers)")
    parser.add_argument("--batch_size", type=int, default=64, help="Training batch size")
    parser.add_argument("--num_epochs", type=int, default=10, help="Number of training epochs")
    parser.add_argument("--learning_rate", type=float, default=1e-5, help="Learning rate")
    parser.add_argument("--lora_size", type=int, default=None, help="LoRA bottleneck size (default: 4% of intermediate_size)")
    parser.add_argument("--lora_sizes", nargs='+', type=float, default=None, help="LoRA sizes as percentages (e.g., 4.0 10.0 20.0 30.0)")
    parser.add_argument("--val_split", type=float, default=0.1, help="Validation split fraction")
    parser.add_argument("--cache_size", type=int, default=50, help="Number of .npz chunk files to cache in memory")
    parser.add_argument("--load_full_dataset", action="store_true", help="Load full dataset into memory at initialization (faster but uses more memory)")
    parser.add_argument("--checkpoint_save_interval", type=int, default=20000, help="Save checkpoint every N steps")
    parser.add_argument("--resume_from_checkpoint", action="store_true", help="Resume training from the latest checkpoint")
    parser.add_argument("--checkpoint_path", type=str, default=None, help="Specific checkpoint path to resume from (optional)")
    parser.add_argument("--load_best_only", action="store_true", help="Resume training only from best performing lora size for each layer.")
    parser.add_argument("--seed", type=int, default=42, help="Random seed")
    parser.add_argument("--use_wandb", action="store_true", help="Use Weights & Biases logging")
    parser.add_argument("--wandb_project", type=str, default="llama-skip-predictors", help="W&B project name")
    parser.add_argument("--wandb_entity", type=str, default="llama-skip-predictors", help="W&B entity name")
    parser.add_argument("--device", type=str, default="auto", help="Device to use (auto, cpu, cuda)")

    
    args = parser.parse_args()
    
    # Set seed
    set_seed(args.seed)
    
    # Setup device
    if args.device == "auto":
        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    else:
        device = torch.device(args.device)
    
    logger.info(f"Using device: {device}")
    
    # Load config
    config = AutoConfig.from_pretrained(args.config)
    
    # Handle 'all' option for layer_indices
    if len(args.layer_indices) == 1 and args.layer_indices[0] == "all":
        # Get number of layers from config
        num_layers = getattr(config, 'num_hidden_layers', None) or getattr(config, 'n_layer', None) or getattr(config, 'num_layers', None)
        if num_layers is None:
            raise ValueError("Could not determine number of layers from config. Please specify layer indices explicitly.")
        args.layer_indices = list(range(num_layers))
        logger.info(f"Training predictors for all {num_layers} layers: {args.layer_indices}")
    else:
        # Convert to integers
        try:
            args.layer_indices = [int(idx) for idx in args.layer_indices]
        except ValueError:
            raise ValueError("Layer indices must be integers or 'all'")
        logger.info(f"Training predictors for specified layers: {args.layer_indices}")
    
    logger.info(f"Model config: hidden_size={config.hidden_size}, intermediate_size={config.intermediate_size}")
    
    # Initialize wandb for multi-layer training
    if args.use_wandb:
        wandb.init(
            entity=args.wandb_entity,
            project=args.wandb_project,
            config=vars(args),
            name=f"predictor-layers-{'-'.join(map(str, args.layer_indices))}-training-{int(time.time())}"
        )

    # Initialize multi-layer trainer
    trainer = MultiLayerPredictorTrainer(
        config=config,
        layer_indices=args.layer_indices,
        device=device,
        lora_size=args.lora_size,
        lora_sizes=args.lora_sizes
    )
    
    # Train all layers
    trainer.train_all_layers(
        dataset_dir=args.dataset_dir,
        num_epochs=args.num_epochs,
        batch_size=args.batch_size,
        learning_rate=args.learning_rate,
        val_split=args.val_split,
        cache_size=args.cache_size,
        load_full_dataset=args.load_full_dataset,
        use_wandb=args.use_wandb,
        save_dir=args.output_dir,
        save_interval=args.checkpoint_save_interval,
        resume_from_checkpoint=args.resume_from_checkpoint,
        load_best_only= args.load_best_only,
        checkpoint_path=args.checkpoint_path,
        seed=args.seed
    )

    if args.use_wandb:
        wandb.finish()
    
    logger.info(f"Training completed for all layers: {args.layer_indices}")


if __name__ == "__main__":
    main()