File size: 4,472 Bytes
f325414 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 | """Single LoRA-task run via any optimizer.
Usage::
python -m ahdcma.cli.run_task \\
--algo ahdcma --task cifar100_vit --seed 0 \\
--pop 8 --gens 5 --num-steps 50 --output outputs/runs/lora/
Writes ``{output}/{run_id}/result.json`` with the best LoRA config,
final fitness, history, and config snapshot. Each fitness call is a
short LoRA fine-tune (``num_steps`` optimiser steps) — the wall-time
per call dominates total runtime.
"""
from __future__ import annotations
import argparse
import json
import time
from pathlib import Path
from typing import Any
import numpy as np
import yaml
from ahdcma.algorithms.base import SearchSpace
from ahdcma.cli.run_benchmark import ALGO_REGISTRY
from ahdcma.fitness.lora_finetune import TaskConfig, lora_fitness
from ahdcma.search_space.encoder import DIM, DIM_NAMES
from ahdcma.utils.logging import make_run_id, setup_logging
from ahdcma.utils.seed import set_global_seed
def _load_task_config(task_name: str) -> TaskConfig:
cfg_path = Path(__file__).resolve().parents[3] / "configs" / "tasks" / f"{task_name}.yaml"
with cfg_path.open() as f:
raw = yaml.safe_load(f)
return TaskConfig(**raw)
def _load_algo_config(algo: str, *, seed: int, pop: int, gens: int) -> dict[str, Any]:
cfg_path = Path(__file__).resolve().parents[3] / "configs" / "algo" / f"{algo}.yaml"
with cfg_path.open() as f:
cfg = yaml.safe_load(f)
cfg["seed"] = seed
cfg["population_size"] = pop
cfg["max_generations"] = gens
return dict(cfg)
def run_lora_task(
algo: str,
task_name: str,
seed: int,
*,
pop: int = 8,
gens: int = 10,
num_steps: int = 100,
output_dir: Path | str = "outputs/runs/lora",
) -> dict[str, Any]:
"""Run one (algo, task, seed) LoRA tuning job and write result.json."""
if algo not in ALGO_REGISTRY:
raise KeyError(f"unknown algorithm {algo!r}")
set_global_seed(seed)
run_id = make_run_id(algo, task_name, seed)
out_root = Path(output_dir) / run_id
out_root.mkdir(parents=True, exist_ok=True)
setup_logging(run_id, log_dir=out_root / "logs")
task = _load_task_config(task_name)
algo_cfg = _load_algo_config(algo, seed=seed, pop=pop, gens=gens)
search_space = SearchSpace.unit_cube(DIM, names=DIM_NAMES)
eval_log: list[dict[str, Any]] = []
def fitness(x: np.ndarray[Any, np.dtype[np.float64]]) -> float:
f = lora_fitness(x, task, num_steps=num_steps, seed=seed)
eval_log.append({"x": x.tolist(), "f": f})
return f
cls = ALGO_REGISTRY[algo]
opt = cls(algo_cfg, fitness, search_space, run_id=run_id)
t0 = time.time()
result = opt.optimize()
wall = time.time() - t0
summary: dict[str, Any] = {
"run_id": run_id,
"algo": algo,
"task": task_name,
"seed": seed,
"pop": pop,
"gens": gens,
"num_steps": num_steps,
"best_f": float(result.best_f),
"best_x": result.best_x.tolist(),
"best_accuracy": float(-result.best_f),
"wall_time": wall,
"n_generations": len(result.history),
"n_evals": len(eval_log),
"config": algo_cfg,
"best_fitness_curve": list(result.history.best_fitness),
"mode_curve": list(result.history.mode_per_gen),
}
(out_root / "result.json").write_text(json.dumps(summary, indent=2))
(out_root / "evaluations.json").write_text(json.dumps(eval_log, indent=2))
return summary
def main() -> None:
p = argparse.ArgumentParser(description="Run a single LoRA tuning job.")
p.add_argument("--algo", required=True, choices=sorted(ALGO_REGISTRY))
p.add_argument("--task", required=True)
p.add_argument("--seed", type=int, default=0)
p.add_argument("--pop", type=int, default=8)
p.add_argument("--gens", type=int, default=10)
p.add_argument("--num-steps", type=int, default=100)
p.add_argument("--output", default="outputs/runs/lora")
args = p.parse_args()
summary = run_lora_task(
args.algo,
args.task,
args.seed,
pop=args.pop,
gens=args.gens,
num_steps=args.num_steps,
output_dir=args.output,
)
print(
json.dumps(
{
"run_id": summary["run_id"],
"best_accuracy": summary["best_accuracy"],
"wall_time": summary["wall_time"],
}
)
)
if __name__ == "__main__":
main()
|