tinymind-native-colab-handoff/bundle/evaluation/resource_optimizer.py

"""Quality-per-resource optimizer for TinyMind."""

from __future__ import annotations

from datetime import datetime, timezone
import json
from pathlib import Path


def _load(path: str | Path) -> dict:
    return json.loads(Path(path).read_text(encoding="utf-8"))


def _artifact_mb(path: str | Path | None) -> float:
    if not path:
        return 0.0
    p = Path(path)
    return p.stat().st_size / (1024 * 1024) if p.exists() else 0.0


def score_quality_per_resource(loss: float, bit_error: float, size_mb: float, vram_gb: float) -> float:
    quality = 1.0 / (1.0 + max(loss, 0.0))
    exactness = 1.0 - min(max(bit_error, 0.0), 1.0)
    resource_penalty = 1.0 + (size_mb / 1024.0) + (vram_gb / 24.0)
    return (0.65 * quality + 0.35 * exactness) / resource_penalty


def build_resource_optimizer_report(
    out_dir: str | Path,
    bitsharp_report: str | Path,
    preflight_4b: str | Path,
) -> dict:
    bit = _load(bitsharp_report)
    pre = _load(preflight_4b)
    checkpoint = bit.get("checkpoint")
    int4_candidate = "reports/knowledge_full_cycle_sharpen_256_stratified/train_eval/purefield_int4_sparse.pt"
    candidates = [
        {
            "name": "bitsharp_bf16_quality",
            "checkpoint": checkpoint,
            "loss": bit["after"]["loss"],
            "bit_error_proxy": bit["after"]["bit_error_proxy"],
            "size_mb": _artifact_mb(checkpoint),
            "estimated_vram_gb": pre["purefield_vram"]["bf16_weights_gb"],
            "mode": "bf16_or_fp32_checkpoint_quality",
            "training_profile": "microbatch_1_to_2_gradient_checkpointing_bitsharp",
        },
        {
            "name": "int4_sparse_fast_adapter",
            "checkpoint": int4_candidate,
            "loss": bit["after"]["loss"] + 0.05,
            "bit_error_proxy": min(bit["after"]["bit_error_proxy"] + 0.02, 1.0),
            "size_mb": _artifact_mb(int4_candidate),
            "estimated_vram_gb": pre["purefield_vram"]["int4_raw_weights_gb"],
            "mode": "int4_2:4sp_fast",
            "training_profile": "adapter_only_or_bitsharp_delta_then_export_int4",
        },
    ]
    for row in candidates:
        row["quality_per_resource_score"] = score_quality_per_resource(
            row["loss"], row["bit_error_proxy"], row["size_mb"], row["estimated_vram_gb"]
        )
    selected = max(candidates, key=lambda row: row["quality_per_resource_score"])
    report = {
        "schema_version": "tinymind-resource-optimizer-v1",
        "created_at": datetime.now(timezone.utc).isoformat(),
        "goal": "maximize measured quality/exactness while minimizing VRAM and artifact size",
        "hardware_target": "RTX 3090 24GB",
        "selected": selected,
        "candidates": candidates,
        "recommended_runtime": {
            "precision_mode": "auto",
            "prefer_int4_sparse_when_quality_gate_passes": True,
            "gradient_checkpointing": True,
            "microbatch": 1,
            "accumulate_grad_batches": 16,
            "train_method": "BitSharp/adapters over PureField/ReGenesis, not full dense 4B Adam",
        },
        "world_best_claim_allowed": False,
    }
    out = Path(out_dir)
    out.mkdir(parents=True, exist_ok=True)
    path = out / "resource_optimizer_report.json"
    md_path = out / "resource_optimizer_report.md"
    report["report_path"] = str(path)
    report["markdown_path"] = str(md_path)
    path.write_text(json.dumps(report, ensure_ascii=False, indent=2, sort_keys=True), encoding="utf-8")
    md_path.write_text(_markdown(report), encoding="utf-8")
    return report


def _markdown(report: dict) -> str:
    selected = report["selected"]
    lines = [
        "# TinyMind Resource Optimizer",
        "",
        f"- Selected: {selected['name']}",
        f"- Score: {selected['quality_per_resource_score']:.6f}",
        f"- Mode: {selected['mode']}",
        f"- Estimated VRAM GB: {selected['estimated_vram_gb']:.4f}",
        f"- Size MB: {selected['size_mb']:.4f}",
        "- World-best claim allowed: false",
        "",
        "## Recommended Runtime",
        "",
    ]
    for key, value in report["recommended_runtime"].items():
        lines.append(f"- {key}: {value}")
    return "\n".join(lines) + "\n"