simplecloud/VidChain-exercise

1 2import os3import sys4import argparse5import json6import torch7from dataclasses import dataclass, field8from typing import Optional, Dict, Any9 10# Add project root to path11current_dir = os.getcwd()12if 'VidChain' in current_dir:13    root_dir = os.path.join(current_dir, "..", "..")14else:15    root_dir = current_dir16sys.path.append(root_dir)17 18try:19    import transformers20    from transformers import HfArgumentParser21    from vtimellm.train.train_dpo import (22        ModelArguments, 23        TrainingArguments, 24        DataArguments,25        train26    )27    from vtimellm.train.dataset import make_supervised_dpo_data_module28    from vtimellm.model import VTimeLLMLlamaForCausalLM29    from vtimellm.model.builder import load_lora30    from vtimellm.mm_utils import print_trainable_parameters31    from trl.trl.trainer import DPOTrainer32    print("✓ Training modules imported successfully")33except ImportError as e:34    print(f"✗ Error importing training modules: {e}")35    print("Make sure you're in the correct environment and all dependencies are installed")36    sys.exit(1)37 38 39@dataclass40class DPOConfig:41    """Configuration class for DPO training demo"""42    43    # Model paths44    model_name_or_path: str = "checkpoints/vicuna-7b-v1.5"45    stage2_path: str = "checkpoints/vtimellm-vicuna-v1-5-7b-stage2"46    stage3_path: str = "checkpoints/vtimellm-vicuna-v1-5-7b-stage3"47    stage4_path: str = "checkpoints/vtimellm-vicuna-v1-5-7b-activitynet-stage4"48    pretrain_mm_mlp_adapter: str = "checkpoints/vtimellm-vicuna-v1-5-7b-stage1/mm_projector.bin"49    50    # Data paths51    data_path: str = "data/activitynet/mdpo-train.json"52    data_folder: str = "data/activitynet/videos/train"53    feat_folder: str = "data/activitynet/clipvitl14-vtimellm.pth"54    55    # Training parameters56    output_dir: str = "outputs/vtimellm-dpo-demo"57    num_train_epochs: int = 158    per_device_train_batch_size: int = 259    gradient_accumulation_steps: int = 460    learning_rate: float = 1e-661    model_max_length: int = 204862    63    # DPO parameters64    beta: float = 0.5  # DPO beta parameter (controls preference strength)65    dpo_alpha: float = 1.0  # DPO alpha parameter66    gamma: float = 0.0  # Additional regularization parameter67    68    # LoRA parameters69    lora_enable: bool = True70    lora_r: int = 6471    lora_alpha: int = 12872    lora_dropout: float = 0.0573    74    # Other settings75    bf16: bool = True76    gradient_checkpointing: bool = True77    dataloader_num_workers: int = 478    logging_steps: int = 1079    save_steps: int = 50080    save_total_limit: int = 381    warmup_ratio: float = 0.182    weight_decay: float = 0.083    freeze_mm_mlp_adapter: bool = True84 85 86def check_training_requirements(config: DPOConfig) -> bool:87    """Check if all requirements for training are met"""88    print("🔍 Checking training requirements...")89    90    # Check model files91    model_files = [92        config.model_name_or_path,93        config.stage2_path,94        config.stage3_path,95        config.stage4_path,96        config.pretrain_mm_mlp_adapter97    ]98    99    missing_files = []100    for file_path in model_files:101        if not os.path.exists(file_path):102            missing_files.append(file_path)103    104    if missing_files:105        print("❌ Missing model files:")106        for file_path in missing_files:107            print(f"  - {file_path}")108        print("\nPlease download the required model checkpoints.")109        return False110    111    # Check data files112    if not os.path.exists(config.data_path):113        print(f"❌ Training data not found: {config.data_path}")114        print("Please prepare your DPO training data.")115        return False116    117    # Check GPU118    if not torch.cuda.is_available():119        print("⚠ No CUDA GPU available. Training will be very slow on CPU.")120    else:121        gpu_memory = torch.cuda.get_device_properties(0).total_memory / 1024**3122        print(f"✓ GPU available: {torch.cuda.get_device_name(0)} ({gpu_memory:.1f} GB)")123        124        if gpu_memory < 16:125            print("⚠ Warning: GPU memory is less than 16GB. Consider reducing batch size.")126    127    # Check output directory128    os.makedirs(config.output_dir, exist_ok=True)129    print(f"✓ Output directory: {config.output_dir}")130    131    print("✅ Training requirements check completed")132    return True133 134 135def create_training_arguments(config: DPOConfig) -> tuple:136    """Create training arguments from config"""137    138    # Model arguments139    model_args = ModelArguments(140        model_name_or_path=config.model_name_or_path,141        stage2_path=config.stage2_path,142        stage3_path=config.stage3_path,143        stage4_path=config.stage4_path,144        pretrain_mm_mlp_adapter=config.pretrain_mm_mlp_adapter,145        version="v1"146    )147    148    # Data arguments149    data_args = DataArguments(150        data_path=config.data_path,151        data_folder=config.data_folder,152        feat_folder=config.feat_folder153    )154    155    # Training arguments156    training_args = TrainingArguments(157        output_dir=config.output_dir,158        num_train_epochs=config.num_train_epochs,159        per_device_train_batch_size=config.per_device_train_batch_size,160        gradient_accumulation_steps=config.gradient_accumulation_steps,161        learning_rate=config.learning_rate,162        model_max_length=config.model_max_length,163        bf16=config.bf16,164        gradient_checkpointing=config.gradient_checkpointing,165        dataloader_num_workers=config.dataloader_num_workers,166        logging_steps=config.logging_steps,167        save_steps=config.save_steps,168        save_total_limit=config.save_total_limit,169        warmup_ratio=config.warmup_ratio,170        weight_decay=config.weight_decay,171        freeze_mm_mlp_adapter=config.freeze_mm_mlp_adapter,172        173        # DPO specific174        beta=config.beta,175        dpo_alpha=config.dpo_alpha,176        gamma=config.gamma,177        train4dpo=True,178        179        # LoRA specific180        lora_enable=config.lora_enable,181        lora_r=config.lora_r,182        lora_alpha=config.lora_alpha,183        lora_dropout=config.lora_dropout,184        185        # Single GPU training settings186        no_cuda=False,  # Keep CUDA enabled187        local_rank=-1,  # Force single GPU by setting local_rank to -1188        dataloader_pin_memory=False,  # Disable pin memory for single GPU189        190        # Other settings191        training_stage=3,192        finetuning=True,193        evaluation_strategy="no",194        save_strategy="steps",195        lr_scheduler_type="cosine",196        tf32=True,197        report_to="none",  # Disable wandb for demo198        remove_unused_columns=False199    )200    201    return model_args, data_args, training_args202 203 204def load_and_prepare_model(model_args, training_args):205    """Load and prepare the model for DPO training"""206    print("🤖 Loading and preparing model...")207    208    # Load base model209    model = VTimeLLMLlamaForCausalLM.from_pretrained(210        model_args.model_name_or_path,211        cache_dir=training_args.cache_dir,212        torch_dtype=torch.bfloat16 if training_args.bf16 else torch.float16213    )214    model.config.use_cache = False215    216    # Load tokenizer217    tokenizer = transformers.AutoTokenizer.from_pretrained(218        model_args.model_name_or_path,219        cache_dir=training_args.cache_dir,220        model_max_length=training_args.model_max_length,221        padding_side="right",222        use_fast=False,223    )224    tokenizer.pad_token = tokenizer.unk_token225    226    # Initialize vision modules227    model.get_model().initialize_vision_modules(model_args)228    model.cuda()229    230    # Load stage 2 weights231    print("📥 Loading stage 2 weights...")232    model = load_lora(model, model_args.stage2_path)233    model = model.merge_and_unload()234    235    # Load stage 3 weights236    print("📥 Loading stage 3 weights...")237    model = load_lora(model, model_args.stage3_path)238    model = model.merge_and_unload()239    240    # Load stage 4 weights (for DPO training)241    print("📥 Loading stage 4 weights...")242    model = load_lora(model, model_args.stage4_path)243    model = model.merge_and_unload()244    245    # Add LoRA adapters for DPO training246    if training_args.lora_enable:247        from peft import LoraConfig, get_peft_model248        249        # Find target modules for LoRA250        def find_all_linear_explicit_names(model):251            cls = torch.nn.Linear252            lora_module_names = set()253            multimodal_keywords = ['mm_projector', 'vision_tower', 'vision_resampler']254            for name, module in model.named_modules():255                if any(mm_keyword in name for mm_keyword in multimodal_keywords):256                    continue257                if isinstance(module, cls):258                    lora_module_names.add(name)259            if 'lm_head' in lora_module_names:260                lora_module_names.remove('lm_head')261            return list(lora_module_names)262        263        explicit_linear = find_all_linear_explicit_names(model)264        lora_config = LoraConfig(265            r=training_args.lora_r,266            lora_alpha=training_args.lora_alpha,267            target_modules=explicit_linear,268            lora_dropout=training_args.lora_dropout,269            bias="none",270            task_type="CAUSAL_LM",271        )272        273        print("🔧 Adding LoRA adapters...")274        model = get_peft_model(model, lora_config)275        print_trainable_parameters(model)276    277    # Move to GPU278    model = model.cuda()279    280    return model, tokenizer281 282 283def save_training_results(model, trainer, training_args, config):284    """Save training results and model"""285    print("💾 Saving training results...")286    287    # Save model288    if training_args.lora_enable:289        from vtimellm.train.train_dpo import get_peft_state_maybe_zero_3, get_peft_state_non_lora_maybe_zero_3290        291        state_dict = get_peft_state_maybe_zero_3(292            model.named_parameters(), training_args.lora_bias293        )294        non_lora_state_dict = get_peft_state_non_lora_maybe_zero_3(295            model.named_parameters()296        )297        298        if training_args.local_rank in [0, -1]:299            model.config.save_pretrained(training_args.output_dir)300            model.save_pretrained(training_args.output_dir, state_dict=state_dict)301            torch.save(non_lora_state_dict, os.path.join(training_args.output_dir, 'non_lora_trainables.bin'))302    303    # Save training config304    config_dict = {305        "model_paths": {306            "model_name_or_path": config.model_name_or_path,307            "stage2_path": config.stage2_path,308            "stage3_path": config.stage3_path,309            "stage4_path": config.stage4_path,310        },311        "training_params": {312            "num_train_epochs": config.num_train_epochs,313            "per_device_train_batch_size": config.per_device_train_batch_size,314            "learning_rate": config.learning_rate,315            "beta": config.beta,316            "dpo_alpha": config.dpo_alpha,317        },318        "lora_params": {319            "lora_r": config.lora_r,320            "lora_alpha": config.lora_alpha,321            "lora_dropout": config.lora_dropout,322        }323    }324    325    with open(os.path.join(training_args.output_dir, "training_config.json"), "w") as f:326        json.dump(config_dict, f, indent=2)327    328    print(f"✅ Results saved to: {training_args.output_dir}")329 330 331def demo_dpo_training(config: DPOConfig):332    """Main function for DPO training demo"""333    print("🎬 VTimeLLM DPO Training Demo")334    print("=" * 50)335    336    # Force single GPU usage337    print("🔧 Setting up single GPU training...")338    if torch.cuda.is_available():339        # Set CUDA device to 0 (first GPU)340        torch.cuda.set_device(0)341        print(f"✓ Using GPU: {torch.cuda.get_device_name(0)}")342        343        # Set environment variables for single GPU344        os.environ['CUDA_VISIBLE_DEVICES'] = '0'345        os.environ['OMPI_COMM_WORLD_SIZE'] = '1'346        print("✓ Environment variables set for single GPU")347    else:348        print("⚠ No CUDA GPU available")349    350    # Check requirements351    if not check_training_requirements(config):352        print("❌ Training requirements not met. Please fix the issues above.")353        return False354    355    try:356        # Create training arguments357        model_args, data_args, training_args = create_training_arguments(config)358        359        # Load and prepare model360        model, tokenizer = load_and_prepare_model(model_args, training_args)361        362        # Create trainer363        trainer = create_dpo_trainer(model, tokenizer, data_args, training_args)364        365        # Train model366        train_model(trainer, training_args)367        368        # Save results369        save_training_results(model, trainer, training_args, config)370        371        print("\n🎉 DPO Training Demo Completed Successfully!")372        print(f"📁 Model saved to: {config.output_dir}")373        print("\nYou can now use the trained model for inference!")374        375        return True376        377    except Exception as e:378        print(f"❌ Training failed with error: {e}")379        import traceback380        traceback.print_exc()381        return False382 383 384def create_sample_config():385    """Create a sample configuration for students"""386    return DPOConfig(387        # Use smaller batch size for demo388        per_device_train_batch_size=1,389        gradient_accumulation_steps=2,390        num_train_epochs=1,391        logging_steps=5,392        save_steps=100,393        save_total_limit=2394    )395 396 397def main():398    """Main function with command line interface"""399    # Use default config400    config = DPOConfig()401    print("⚙️ Using default configuration")402    check_only = False403    if check_only:404        # Only check requirements405        check_training_requirements(config)406    else:407        # Run full training demo408        demo_dpo_training(config)409 410 411if __name__ == "__main__":412    main()413