run_experiments.py

"""
Automated LoRA hyperparameter sweep pipeline.

Runs a grid search over (rank, alpha, target_modules) combinations,
using LLaMA-Factory's training backend. Each experiment is logged to WandB
and outputs are saved for later analysis.
"""

import os
import sys
import json
import yaml
import itertools
import subprocess
from pathlib import Path
from concurrent.futures import ProcessPoolExecutor, as_completed


def load_config(config_path: str) -> dict:
    """Load YAML experiment configuration."""
    with open(config_path, "r") as f:
        return yaml.safe_load(f)


def build_grid(config: dict) -> list[dict]:
    """Expand parameter grid into a list of individual experiment configs."""
    grid = config["lora_grid"]
    experiments = []

    for rank, alpha_factor, target_mods in itertools.product(
        grid["rank"], grid["alpha"], grid["target_modules"]
    ):
        alpha = int(alpha_factor * rank)
        experiments.append({
            "rank": rank,
            "alpha": alpha,
            "alpha_factor": alpha_factor,
            "target_modules": target_mods,
        })

    return experiments


def run_single_experiment(
    exp_config: dict,
    base_config: dict,
    dataset_name: str,
    output_dir: str,
    gpu_id: int,
) -> dict:
    """
    Run one LoRA fine-tuning experiment.

    Uses LLaMA-Factory CLI: llamafactory-cli train with a temporary YAML config.
    """
    exp_name = (
        f"r{exp_config['rank']}_a{exp_config['alpha']}"
        f"_af{exp_config['alpha_factor']}_"
        f"{exp_config['target_modules'].split(',')[0]}_{dataset_name.replace('/', '_')}"
    )
    exp_output = os.path.join(output_dir, exp_name)
    os.makedirs(exp_output, exist_ok=True)

    # Build LLaMA-Factory training config
    train_config = {
        "model_name_or_path": base_config["model"]["base_model"],
        "dataset": dataset_name,
        "output_dir": exp_output,
        "finetuning_type": "lora",
        "lora_rank": exp_config["rank"],
        "lora_alpha": exp_config["alpha"],
        "lora_target": exp_config["target_modules"],
        "lora_dropout": 0.05,
        "num_train_epochs": base_config["training"]["num_epochs"],
        "per_device_train_batch_size": base_config["training"]["per_device_batch_size"],
        "gradient_accumulation_steps": base_config["training"]["gradient_accumulation_steps"],
        "learning_rate": base_config["training"]["learning_rate"],
        "lr_scheduler_type": base_config["training"]["lr_scheduler"],
        "warmup_ratio": base_config["training"]["warmup_ratio"],
        "bf16": base_config["training"]["bf16"],
        "logging_steps": base_config["training"]["logging_steps"],
        "save_strategy": base_config["training"]["save_strategy"],
        "per_device_eval_batch_size": base_config["evaluation"]["per_device_eval_batch_size"],
        "val_size": 0.05,
        "report_to": "wandb",
        "run_name": exp_name,
    }

    # Write temp config
    config_path = os.path.join(exp_output, "train_config.yaml")
    with open(config_path, "w") as f:
        yaml.dump(train_config, f)

    # Launch training
    env = os.environ.copy()
    env["CUDA_VISIBLE_DEVICES"] = str(gpu_id)

    cmd = [
        sys.executable, "-m", "llamafactory.cli", "train", config_path
    ]

    result = subprocess.run(
        cmd, env=env, capture_output=True, text=True, cwd="LLaMA-Factory"
    )

    success = result.returncode == 0
    return {
        "experiment": exp_name,
        "success": success,
        "config": exp_config,
        "dataset": dataset_name,
        "output_dir": exp_output,
    }


def main():
    config_path = sys.argv[2] if "--config" in sys.argv else "configs/experiment_config.yaml"
    config = load_config(config_path)

    experiments = build_grid(config)
    total = len(config["datasets"]) * len(experiments)
    print(f"Total experiments: {total} ({len(experiments)} configs x {len(config['datasets'])} datasets)")

    output_dir = config["output_dir"]
    os.makedirs(output_dir, exist_ok=True)

    # Run experiments with GPU parallelism
    results = []
    max_parallel = config.get("max_parallel_jobs", 1)

    with ProcessPoolExecutor(max_workers=max_parallel) as executor:
        futures = {}
        gpu_id = 0
        for dataset_name in config["datasets"]:
            for exp_config in experiments:
                future = executor.submit(
                    run_single_experiment,
                    exp_config,
                    config,
                    dataset_name["name"],
                    output_dir,
                    gpu_id % max_parallel,
                )
                futures[future] = f"{exp_config['rank']}_{dataset_name['name']}"
                gpu_id += 1

        for future in as_completed(futures):
            try:
                result = future.result()
                status = "OK" if result["success"] else "FAILED"
                print(f"[{status}] {futures[future]}")
                results.append(result)
            except Exception as e:
                print(f"[ERROR] {futures[future]}: {e}")

    # Save results summary
    summary_path = os.path.join(output_dir, "experiment_summary.json")
    with open(summary_path, "w") as f:
        json.dump(results, f, indent=2)

    print(f"\nDone! {sum(1 for r in results if r['success'])}/{len(results)} succeeded.")
    print(f"Summary: {summary_path}")


if __name__ == "__main__":
    main()