| |
| """Merge four independent consecutive three-block training sweeps.""" |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import csv |
| import json |
| import os |
| import shutil |
| from pathlib import Path |
| from typing import Any |
|
|
|
|
| EXPECTED_TRIPLES = [(index, index + 1, index + 2) for index in range(28)] |
| EXPECTED_PROMPTS = {"train": list(range(80)), "val": list(range(80, 100))} |
| GPU_SHARDS = { |
| "0": EXPECTED_TRIPLES[0:7], |
| "2": EXPECTED_TRIPLES[7:14], |
| "6": EXPECTED_TRIPLES[14:21], |
| "7": EXPECTED_TRIPLES[21:28], |
| } |
|
|
|
|
| def parse_args() -> argparse.Namespace: |
| parser = argparse.ArgumentParser(description=__doc__) |
| parser.add_argument("--shard_root", type=Path, required=True) |
| parser.add_argument("--output_dir", type=Path, required=True) |
| return parser.parse_args() |
|
|
|
|
| def experiment_name(triple: tuple[int, int, int]) -> str: |
| return "triple_" + "_".join(f"{layer:02d}" for layer in triple) |
|
|
|
|
| def atomic_json(path: Path, value: Any) -> None: |
| temporary = path.with_suffix(path.suffix + ".tmp") |
| temporary.write_text( |
| json.dumps(value, indent=2, ensure_ascii=False) + "\n", |
| encoding="utf-8", |
| ) |
| os.replace(temporary, path) |
|
|
|
|
| def validate_metrics( |
| metrics: dict[str, Any], triple: tuple[int, int, int] |
| ) -> None: |
| name = experiment_name(triple) |
| expected = { |
| "status": "complete", |
| "name": name, |
| "architecture": "three_block_predictor", |
| "initialization_method": "teacher_full", |
| "source_layers": list(triple), |
| "seed": 0, |
| "max_steps": 1000, |
| "train_prompt_ids": EXPECTED_PROMPTS["train"], |
| "val_prompt_ids": EXPECTED_PROMPTS["val"], |
| } |
| for key, value in expected.items(): |
| if metrics.get(key) != value: |
| raise RuntimeError( |
| f"{name}: expected {key}={value!r}, got {metrics.get(key)!r}" |
| ) |
|
|
|
|
| def result_row(metrics: dict[str, Any]) -> dict[str, Any]: |
| return { |
| "name": metrics["name"], |
| "source_layer_1": metrics["source_layers"][0], |
| "source_layer_2": metrics["source_layers"][1], |
| "source_layer_3": metrics["source_layers"][2], |
| "final_val_flow_mse": metrics["final_val_flow_mse"], |
| "final_val_hidden_mse": metrics["final_val_hidden_mse"], |
| "final_val_total_loss": metrics["final_val_total_loss"], |
| "val_flow_mse_auc": metrics["val_flow_mse_auc"], |
| "val_hidden_mse_auc": metrics["val_hidden_mse_auc"], |
| "val_total_loss_auc": metrics["val_total_loss_auc"], |
| "training_time_s": metrics["training_time_s"], |
| } |
|
|
|
|
| def write_summary(output_dir: Path, rows: list[dict[str, Any]]) -> None: |
| rows.sort(key=lambda row: float(row["final_val_flow_mse"])) |
| destination = output_dir / "summary.csv" |
| temporary = destination.with_suffix(".csv.tmp") |
| with temporary.open("w", encoding="utf-8", newline="") as handle: |
| writer = csv.DictWriter(handle, fieldnames=list(rows[0])) |
| writer.writeheader() |
| writer.writerows(rows) |
| os.replace(temporary, destination) |
| atomic_json(output_dir / "summary.json", rows) |
|
|
|
|
| def main() -> None: |
| args = parse_args() |
| shard_root = args.shard_root.resolve() |
| output_dir = args.output_dir.resolve() |
| output_dir.mkdir(parents=True, exist_ok=True) |
|
|
| rows = [] |
| fingerprints = set() |
| for gpu, triples in GPU_SHARDS.items(): |
| shard_dir = shard_root / f"gpu{gpu}" |
| shard_manifest = json.loads( |
| (shard_dir / "sweep_manifest.json").read_text(encoding="utf-8") |
| ) |
| if shard_manifest.get("status") != "complete": |
| raise RuntimeError(f"GPU {gpu} shard is not complete") |
| if shard_manifest.get("triples") != [list(value) for value in triples]: |
| raise RuntimeError(f"GPU {gpu} shard contains unexpected triples") |
| fingerprints.add(shard_manifest["batch_schedule_sha256"]) |
|
|
| for triple in triples: |
| name = experiment_name(triple) |
| source = shard_dir / name |
| metrics = json.loads( |
| (source / "metrics.json").read_text(encoding="utf-8") |
| ) |
| validate_metrics(metrics, triple) |
| destination = output_dir / name |
| if destination.exists(): |
| shutil.rmtree(destination) |
| shutil.copytree(source, destination) |
| rows.append(result_row(metrics)) |
|
|
| if len(rows) != 28: |
| raise RuntimeError(f"Expected 28 results, found {len(rows)}") |
| if len(fingerprints) != 1: |
| raise RuntimeError(f"Batch schedules differ across shards: {fingerprints}") |
| write_summary(output_dir, rows) |
| atomic_json( |
| output_dir / "sweep_manifest.json", |
| { |
| "status": "complete", |
| "architecture": "three_block_predictor", |
| "initialization_method": "teacher_full", |
| "triples": [list(value) for value in EXPECTED_TRIPLES], |
| "num_triples": 28, |
| "max_steps": 1000, |
| "seed": 0, |
| "train_prompt_ids": EXPECTED_PROMPTS["train"], |
| "val_prompt_ids": EXPECTED_PROMPTS["val"], |
| "batch_schedule_sha256": fingerprints.pop(), |
| "gpu_shards": { |
| gpu: [list(value) for value in triples] |
| for gpu, triples in GPU_SHARDS.items() |
| }, |
| "source_shard_root": str(shard_root), |
| }, |
| ) |
| print( |
| f"[merge] {len(rows)} triples -> {output_dir / 'summary.csv'}", |
| flush=True, |
| ) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
