""" Modal GPU pipeline for research/finetune.py + slm-lm-eval. Skill-matrix pipeline: train -> eval -> gate -> publish. Each job in experiments.yaml fine-tunes one QLoRA adapter for a skill (math, science, coding, reasoning, teaching, ...), evaluates it against the matching slm-lm-eval profile vs. a per-profile baseline, checks the result against `goals`, and (only if the gate passes) publishes the adapter to the Hugging Face Hub. Run from repo root: modal run research/modal/finetune_app.py modal run research/modal/finetune_app.py --eval-only modal run research/modal/finetune_app.py --job math-lora --max-steps 20 modal run research/modal/finetune_app.py --category science modal run research/modal/finetune_app.py --no-publish --no-pull modal run research/modal/finetune_app.py::publish_only --job math-lora modal run research/modal/finetune_app.py::pull --category math """ from __future__ import annotations import json import subprocess import sys from pathlib import Path from typing import Any import modal # Make `_common` importable both locally (sibling file) and in the Modal # container, where the entrypoint lands at /root but the repo is baked into the # image at /repo (see add_local_dir in _common.py). for _candidate in (Path(__file__).resolve().parent, Path("/repo/research/modal")): if _candidate.is_dir() and str(_candidate) not in sys.path: sys.path.insert(0, str(_candidate)) from _common import ( # noqa: E402 FINETUNE_VOL_PATH, HF_CACHE_PATH, LM_EVAL_OUTPUT, baseline_profiles_for_jobs, build_finetune_cmd, build_lm_eval_cmd, check_gate_files, check_publish_gate_files, commit_volumes, config_for_profile, discover_cached_baselines, eval_paths, finetune_vol, general_eval_profile, general_goals_for_job, hf_cache_vol, hf_secret, image, job_gpu, job_plan_rows, parse_json_object, prepare_jobs, profiles_needing_baseline_run, resolve_base_model_id, split_csv, publish_adapter_files, pull_artifacts, reload_finetune_volume, reload_volumes, repo_env, baseline_experiment_name, ) APP_NAME = "slm-finetune-benchmark" app = modal.App(APP_NAME, image=image) @app.function( gpu="A10G", volumes={ HF_CACHE_PATH: hf_cache_vol, FINETUNE_VOL_PATH: finetune_vol, }, secrets=[hf_secret], timeout=7200, ) def finetune_one(job: dict[str, Any]) -> dict[str, Any]: """Fine-tune one dataset job; persist adapter to Modal Volume.""" name = job["name"] out_dir = f"{FINETUNE_VOL_PATH}/{name}" Path(out_dir).mkdir(parents=True, exist_ok=True) cmd = build_finetune_cmd(job, out_dir) print("Running:", " ".join(cmd)) subprocess.run(cmd, cwd="/repo", check=True, env=repo_env()) commit_volumes() results_path = Path(out_dir) / "training_results.json" payload = json.loads(results_path.read_text()) payload["job_name"] = name return payload @app.function( gpu="A10G", volumes={ HF_CACHE_PATH: hf_cache_vol, FINETUNE_VOL_PATH: finetune_vol, }, secrets=[hf_secret], timeout=3600, ) def run_lm_eval( *, experiment_name: str, config: str = "research/evals/configs/lm_eval_smoke.yaml", preset: str | None = None, model_path: str | None = None, adapter_path: str | None = None, compare_to: str | None = None, tasks: list[str] | None = None, limit: int | None = None, num_fewshot: int | None = None, batch_size: str | None = None, device: str | None = None, dtype: str | None = None, seed: int | None = None, ) -> dict[str, Any]: """Run slm-lm-eval on base model or finetuned checkpoint.""" reload_finetune_volume() if adapter_path: adapter_dir = Path(adapter_path) adapter_cfg = adapter_dir / "adapter_config.json" if not adapter_cfg.is_file(): raise FileNotFoundError( f"LoRA adapter not visible at {adapter_path} " f"(missing {adapter_cfg.name}). " "If training just finished, retry after volume commit/reload." ) cmd = build_lm_eval_cmd( experiment_name=experiment_name, config=config, preset=preset, model_path=model_path, adapter_path=adapter_path, compare_to=compare_to, tasks=tasks, limit=limit, num_fewshot=num_fewshot, batch_size=batch_size, device=device, dtype=dtype, seed=seed, ) print("Running:", " ".join(cmd)) proc = subprocess.run(cmd, cwd="/repo", check=False, env=repo_env()) commit_volumes() out_root = Path(LM_EVAL_OUTPUT) / experiment_name results_json = out_root / "results.json" summary_md = out_root / "summary.md" comparison_md = out_root / "comparison.md" return { "experiment_name": experiment_name, "config": config, "preset": preset, "model_path": model_path, "adapter_path": adapter_path, "compare_to": compare_to, "tasks": tasks, "limit": limit, "num_fewshot": num_fewshot, "batch_size": batch_size, "device": device, "dtype": dtype, "seed": seed, "results_json": str(results_json), "summary_md": str(summary_md), "comparison_md": str(comparison_md) if comparison_md.is_file() else None, "exit_code": proc.returncode, "ok": proc.returncode == 0 and results_json.is_file(), } @app.function(volumes={FINETUNE_VOL_PATH: finetune_vol}, timeout=300) def check_gate( *, candidate_results_path: str, baseline_results_path: str | None, goals: dict[str, Any], general_candidate_results_path: str | None = None, general_baseline_results_path: str | None = None, general_goals: dict[str, Any] | None = None, ) -> dict[str, Any]: """Check skill + general lm-eval results against publish goals.""" reload_finetune_volume() if general_goals: return check_publish_gate_files( skill_candidate_path=candidate_results_path, skill_baseline_path=baseline_results_path, skill_goals=goals, general_candidate_path=general_candidate_results_path, general_baseline_path=general_baseline_results_path, general_goals=general_goals, ) return check_gate_files( candidate_results_path=candidate_results_path, baseline_results_path=baseline_results_path, goals=goals, ) @app.function( volumes={FINETUNE_VOL_PATH: finetune_vol}, secrets=[hf_secret], timeout=900, ) def publish_adapter( *, job: dict[str, Any], adapter_dir: str, gate_result: dict[str, Any], candidate_results_path: str, baseline_results_path: str | None, ) -> dict[str, Any]: """Write a model card and push the adapter to the Hub, but only if the gate passed.""" reload_finetune_volume() return publish_adapter_files( job=job, adapter_dir=adapter_dir, gate_result=gate_result, candidate_results_path=candidate_results_path, baseline_results_path=baseline_results_path, ) def _print_summary(rows: list[dict[str, Any]]) -> None: print("\n--- summary ---") print(f"{'skill':<18} {'category':<12} {'gate':<6} {'published':<10} hub_repo") for row in rows: gate = "PASS" if row.get("gate_passed") else "fail" published = "yes" if row.get("published") else "no" print( f"{row['name']:<18} {row.get('category') or '-':<12} {gate:<6} " f"{published:<10} {row.get('hub_repo') or '-'}" ) @app.local_entrypoint() def main( train: bool = True, eval_only: bool = False, parallel: bool = False, job: str | None = None, category: str | None = None, sector: str | None = None, usecase: str | None = None, profiles: str | None = None, max_steps: int | None = None, max_samples: int | None = None, finetune_args_json: str | None = None, publish: bool = True, pull: bool = True, plan: bool = False, skip_baseline: bool = False, eval_tasks: str | None = None, eval_limit: int | None = None, eval_num_fewshot: int | None = None, eval_batch_size: str | None = None, eval_device: str | None = None, eval_dtype: str | None = None, eval_seed: int | None = None, ): """ Skill-matrix pipeline: per-profile baselines -> train -> eval -> gate -> publish -> pull. Examples: modal run research/modal/finetune_app.py modal run research/modal/finetune_app.py --job math-lora --max-steps 20 modal run research/modal/finetune_app.py --category science modal run research/modal/finetune_app.py --eval-only --job math-lora modal run research/modal/finetune_app.py --no-publish --no-pull """ defaults, prepared = prepare_jobs( job=job, category=category, sector=sector, usecase=usecase, profiles=split_csv(profiles), max_steps=max_steps, max_samples=max_samples, finetune_overrides=parse_json_object( finetune_args_json, flag="--finetune-args-json" ), ) if not prepared: raise SystemExit("No matching jobs; check --job/--category and experiments.yaml") preset = defaults.get("preset", "minicpm5-1b") plan_rows = job_plan_rows(prepared) if plan: print(json.dumps({"preset": preset, "jobs": plan_rows}, indent=2)) return profile_names = baseline_profiles_for_jobs(prepared, defaults) eval_task_list = split_csv(eval_tasks) baselines_ok = discover_cached_baselines( profile_names, preset=preset, eval_tasks=eval_task_list, eval_limit=eval_limit, eval_num_fewshot=eval_num_fewshot, eval_seed=eval_seed, ) missing_baselines = profiles_needing_baseline_run( profile_names, baselines_ok, skip_baseline=skip_baseline ) if missing_baselines: print(f"--- base-model baselines ({', '.join(missing_baselines)}) ---") for profile in missing_baselines: exp = baseline_experiment_name(preset, profile) result = run_lm_eval.remote( experiment_name=exp, config=config_for_profile(profile), preset=preset, tasks=eval_task_list, limit=eval_limit, num_fewshot=eval_num_fewshot, batch_size=eval_batch_size, device=eval_device, dtype=eval_dtype, seed=eval_seed, ) print(json.dumps(result, indent=2)) baselines_ok[profile] = bool(result.get("ok")) elif any(baselines_ok.values()): cached = [p for p in profile_names if baselines_ok.get(p)] print(f"--- base-model baselines: reusing cached ({', '.join(cached)}) ---") train_results: dict[str, dict[str, Any]] = {} if train and not eval_only: print(f"--- finetune ({len(prepared)} job(s), parallel={parallel}) ---") if parallel: handles = { j["name"]: finetune_one.with_options(gpu=job_gpu(j)).spawn(j) for j in prepared } for name, handle in handles.items(): train_results[name] = handle.get() print(json.dumps(train_results[name], indent=2)) else: for j in prepared: print(f"Training {j['name']}...") result = finetune_one.with_options(gpu=job_gpu(j)).remote(j) train_results[j["name"]] = result print(json.dumps(result, indent=2)) print("--- post-train lm-eval / gate / publish ---") summary: list[dict[str, Any]] = [] gen_profile = general_eval_profile(defaults) for j in prepared: job_name = j["name"] profile = j.get("eval_profile", "compare_study") train_payload = train_results.get(job_name) adapter_path = ( train_payload["output_dir"] if train_payload else f"{FINETUNE_VOL_PATH}/{job_name}" ) baseline_path = f"{LM_EVAL_OUTPUT}/{baseline_experiment_name(preset, profile)}/results.json" compare_to = baseline_path if baselines_ok.get(profile) else None base_model_id = resolve_base_model_id(j, defaults) exp_name = f"{job_name}__{profile}" eval_result = run_lm_eval.remote( experiment_name=exp_name, config=config_for_profile(profile), model_path=base_model_id, adapter_path=adapter_path, compare_to=compare_to, tasks=eval_task_list, limit=eval_limit, num_fewshot=eval_num_fewshot, batch_size=eval_batch_size, device=eval_device, dtype=eval_dtype, seed=eval_seed, ) print(json.dumps(eval_result, indent=2)) general_goals = general_goals_for_job(j, defaults) general_eval_result: dict[str, Any] | None = None general_candidate_path: str | None = None general_baseline_path: str | None = None if general_goals: general_baseline_path = ( f"{LM_EVAL_OUTPUT}/{baseline_experiment_name(preset, gen_profile)}/results.json" ) gen_compare_to = ( general_baseline_path if baselines_ok.get(gen_profile) else None ) gen_exp_name = f"{job_name}__{gen_profile}" general_eval_result = run_lm_eval.remote( experiment_name=gen_exp_name, config=config_for_profile(gen_profile), model_path=base_model_id, adapter_path=adapter_path, compare_to=gen_compare_to, tasks=eval_task_list, limit=eval_limit, num_fewshot=eval_num_fewshot, batch_size=eval_batch_size, device=eval_device, dtype=eval_dtype, seed=eval_seed, ) print(json.dumps(general_eval_result, indent=2)) general_candidate_path = general_eval_result["results_json"] row: dict[str, Any] = { "name": job_name, "category": j.get("category"), "profile": profile, "general_profile": gen_profile if general_goals else None, "plan": next((p for p in plan_rows if p["name"] == job_name), None), } gate_result: dict[str, Any] | None = None if j.get("goals"): skill_ok = bool(eval_result.get("ok")) general_ok = ( not general_goals or bool(general_eval_result and general_eval_result.get("ok")) ) if skill_ok and general_ok: gate_result = check_gate.remote( candidate_results_path=eval_result["results_json"], baseline_results_path=baseline_path, goals=j["goals"], general_candidate_results_path=general_candidate_path, general_baseline_results_path=general_baseline_path, general_goals=general_goals, ) print(json.dumps(gate_result, indent=2)) row["gate_passed"] = bool(gate_result and gate_result.get("passed")) if j.get("publish"): row["hub_repo"] = j["publish"].get("hub_repo") if publish and gate_result is not None: publish_result = publish_adapter.remote( job=j, adapter_dir=adapter_path, gate_result=gate_result, candidate_results_path=eval_result["results_json"], baseline_results_path=baseline_path, ) print(json.dumps(publish_result, indent=2)) row["published"] = publish_result.get("published") summary.append(row) if pull: pull_artifacts(job_name, exp_name) if general_goals and general_eval_result: pull_artifacts(job_name, f"{job_name}__{gen_profile}", dest="models/finetuned") _print_summary(summary) @app.local_entrypoint() def publish_only(job: str): """Re-run the gate and Hub publish for a job using already-computed results (no train/eval).""" defaults, prepared = prepare_jobs(job=job) j = prepared[0] if not j.get("goals"): raise SystemExit(f"Job {job!r} has no `goals`; nothing to gate on") if not j.get("publish"): raise SystemExit(f"Job {job!r} has no `publish` config") preset = defaults.get("preset", "minicpm5-1b") profile = j.get("eval_profile", "compare_study") gen_profile = general_eval_profile(defaults) general_goals = general_goals_for_job(j, defaults) adapter_path = f"{FINETUNE_VOL_PATH}/{job}" candidate_results_path, baseline_results_path, _ = eval_paths( job_name=job, preset=preset, profile=profile ) general_candidate_path = None general_baseline_path = None if general_goals: general_candidate_path, general_baseline_path, _ = eval_paths( job_name=job, preset=preset, profile=gen_profile ) gate_result = check_gate.remote( candidate_results_path=candidate_results_path, baseline_results_path=baseline_results_path, goals=j["goals"], general_candidate_results_path=general_candidate_path, general_baseline_results_path=general_baseline_path, general_goals=general_goals, ) print(json.dumps(gate_result, indent=2)) publish_result = publish_adapter.remote( job=j, adapter_dir=adapter_path, gate_result=gate_result, candidate_results_path=candidate_results_path, baseline_results_path=baseline_results_path, ) print(json.dumps(publish_result, indent=2)) @app.local_entrypoint() def pull(job: str | None = None, category: str | None = None, dest: str = "models/finetuned"): """Download adapters and their lm-eval results from the `slm-finetune` Volume.""" _, prepared = prepare_jobs(job=job, category=category) if not prepared: raise SystemExit("No matching jobs; pass --job or --category") for j in prepared: profile = j.get("eval_profile", "compare_study") pull_artifacts(j["name"], f"{j['name']}__{profile}", dest)