run_train_cli.py

#!/usr/bin/env python
"""
Local training script for CLI model.
This script runs the training locally without cluster job submission.

Usage:
    python run_train_cli.py [--num_machines N] [--num_gpus G] [--config CONFIG]
"""

import argparse
import os
import subprocess
import sys


def parse_args():
    parser = argparse.ArgumentParser(description="Local CLI Training Script")
    parser.add_argument("--num_machines", type=int, default=1, help="Number of machines")
    parser.add_argument("--num_gpus", type=int, default=8, help="Number of GPUs per machine")
    parser.add_argument("--pretrain_vlm_model", type=str,
                        default="llava-onevision-qwen2-0.5b-mid-stage-a4",
                        help="Pretrained VLM model name")
    parser.add_argument("--vlm_exp_layers", type=str, default="RANGE-1-24-4",
                        help="VLM experiment layers")
    parser.add_argument("--vision_exp_layers", type=str, default="RANGE-1-28-4",
                        help="Vision experiment layers")
    parser.add_argument("--projector_name", type=str, default="mlp2x_gelu",
                        help="Projector name")
    parser.add_argument("--output_path", type=str, default="./output/CLI/",
                        help="Output path for saving models")
    parser.add_argument("--data_path", type=str,
                        default="./scripts/train/single_image_ov_cli.yaml",
                        help="Training data config path")
    parser.add_argument("--init_env", action="store_true",
                        help="Run environment initialization script")
    return parser.parse_args()


def build_exp_name(pretrain_vlm_model, projector_name,
                   vlm_exp_layers, vision_exp_layers):
    """Build experiment name based on configuration."""
    exp_name = (
        f"{pretrain_vlm_model}-{projector_name}-Layer{vlm_exp_layers}-bs256"
    )
    if vision_exp_layers:
        exp_name = (
            f"{pretrain_vlm_model}-{projector_name}-Layer{vlm_exp_layers}-Vision{vision_exp_layers}-bs256"
        )
    return exp_name


def main():
    args = parse_args()

    # Build experiment name
    exp_name = build_exp_name(
        args.pretrain_vlm_model,
        args.projector_name,
        args.vlm_exp_layers,
        args.vision_exp_layers,
    )

    save_path = os.path.join(args.output_path, "finetune", exp_name)
    os.makedirs(save_path, exist_ok=True)

    # Set environment variables
    env = os.environ.copy()
    env.update({
        "NUM_MACHINES": str(args.num_machines),
        "NUM_GPUS": str(args.num_gpus),
        "VLM_EXP_LAYERS": args.vlm_exp_layers,
        "EXP_NAME": exp_name,
        "SAVE_NAME": exp_name,
        "PRETRAIN_VLM_MODEL": args.pretrain_vlm_model,
        "PROJECTOR_NAME": args.projector_name,
        "VISION_EXP_LAYERS": args.vision_exp_layers,
        "DATA_PATH": args.data_path,
    })

    print("=" * 60)
    print("CLI Training Configuration")
    print("=" * 60)
    print(f"Experiment Name: {exp_name}")
    print(f"Save Path: {save_path}")
    print(f"Save Name: {exp_name}")
    print(f"Num Machines: {args.num_machines}")
    print(f"Num GPUs: {args.num_gpus}")
    print(f"Pretrain VLM Model: {args.pretrain_vlm_model}")
    print(f"VLM Exp Layers: {args.vlm_exp_layers}")
    print(f"Vision Exp Layers: {args.vision_exp_layers}")
    print(f"Data Path: {args.data_path}")
    print("=" * 60)

    # Run training script
    script_dir = os.path.dirname(os.path.abspath(__file__))
    train_script = os.path.join(script_dir, "scripts/train/finetune_cli_max9.sh")

    if not os.path.exists(train_script):
        print(f"Error: Training script not found at {train_script}")
        sys.exit(1)

    # Optionally run init script
    if args.init_env:
        init_script = os.path.join(script_dir, "scripts/env/init_env.sh")
        if os.path.exists(init_script):
            print("Running environment initialization...")
            subprocess.run(["bash", init_script], env=env, check=True)

    print("Starting training...")
    result = subprocess.run(["bash", train_script], env=env)

    if result.returncode == 0:
        print(f"Training completed successfully. Model saved to: {save_path}")
    else:
        print(f"Training failed with return code: {result.returncode}")
        sys.exit(result.returncode)


if __name__ == "__main__":
    main()