tinymind-native-colab-handoff/bundle/train/native_scaling_ladder.py

from __future__ import annotations

from datetime import datetime, timezone
import json
from pathlib import Path
import time
from typing import Any

from train.native_micro_train import run_native_micro_train


def _write_colab_handoff(out: Path, dataset: str | Path, pending_stages: list[dict[str, int]], config: dict[str, Any]) -> dict[str, str]:
    handoff = out / "colab_handoff"
    handoff.mkdir(parents=True, exist_ok=True)
    pending_layers = sorted({stage["layers"] for stage in pending_stages})
    pending_dims = sorted({stage["dim"] for stage in pending_stages})
    manifest = {
        "schema": "tinymind.native_scaling_colab_handoff.v1",
        "created_at": datetime.now(timezone.utc).isoformat(),
        "dataset": str(dataset),
        "pending_stages": pending_stages,
        "pending_layers": pending_layers,
        "pending_dims": pending_dims,
        "config": config,
        "reason": "Local run guard decided these stages should be run on Colab or a larger GPU runtime.",
    }
    manifest_path = handoff / "colab_handoff_manifest.json"
    manifest_path.write_text(json.dumps(manifest, ensure_ascii=False, indent=2, sort_keys=True) + "\n", encoding="utf-8")
    notebook_path = handoff / "native_scaling_ladder_colab.ipynb"
    notebook = {
        "cells": [
            {
                "cell_type": "markdown",
                "metadata": {},
                "source": ["# TinyMind Native Scaling Ladder Colab Handoff\n", "Upload the repo or mount Drive, then run the staged native micro-train command.\n"],
            },
            {
                "cell_type": "code",
                "execution_count": None,
                "metadata": {},
                "outputs": [],
                "source": [
                    "!python -m train.cli native-scaling-ladder "
                    f"--dataset {dataset} --layers {','.join(str(x) for x in pending_layers)} "
                    f"--dims {','.join(str(x) for x in pending_dims)} "
                    f"--seq-len {config['seq_len']} --max-steps {config['max_steps']} "
                    f"--eval-records {config['eval_records']} --out-dir reports/native_scaling_ladder_colab\n"
                ],
            },
        ],
        "metadata": {"kernelspec": {"display_name": "Python 3", "language": "python", "name": "python3"}},
        "nbformat": 4,
        "nbformat_minor": 5,
    }
    notebook_path.write_text(json.dumps(notebook, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")
    return {"manifest": str(manifest_path), "notebook": str(notebook_path)}


def run_native_scaling_ladder(
    out_dir: str | Path,
    *,
    dataset: str | Path,
    layers: list[int],
    dims: list[int] | None = None,
    dim: int = 128,
    seq_len: int = 128,
    max_steps: int = 8,
    eval_records: int = 16,
    vocab_size: int = 512,
    learning_rate: float = 3e-4,
    local_layer_limit: int = 24,
    local_dim_limit: int = 256,
) -> dict[str, Any]:
    if not layers:
        raise ValueError("layers must not be empty")
    dim_values = dims if dims is not None else [dim]
    if not dim_values:
        raise ValueError("dims must not be empty")
    out = Path(out_dir)
    out.mkdir(parents=True, exist_ok=True)
    stages: list[dict[str, Any]] = []
    pending_colab: list[dict[str, int]] = []
    started = time.time()
    for dim_value in dim_values:
        for layer_count in layers:
            stage_identity = {"dim": dim_value, "layers": layer_count}
            if dim_value > local_dim_limit:
                pending_colab.append(stage_identity)
                stages.append(stage_identity | {"status": "colab_handoff", "reason": "above_local_dim_limit"})
                continue
            if layer_count > local_layer_limit:
                pending_colab.append(stage_identity)
                stages.append(stage_identity | {"status": "colab_handoff", "reason": "above_local_layer_limit"})
                continue
            stage_dir = out / f"dim_{dim_value}_layers_{layer_count}"
            try:
                stage_started = time.time()
                stage = run_native_micro_train(
                    stage_dir,
                    dataset=dataset,
                    max_steps=max_steps,
                    eval_records=eval_records,
                    dim=dim_value,
                    layers=layer_count,
                    seq_len=seq_len,
                    vocab_size=vocab_size,
                    learning_rate=learning_rate,
                )
                stages.append(
                    stage_identity
                    | {
                        "status": "completed",
                        "elapsed_s": time.time() - stage_started,
                        "report_path": stage["json_path"],
                        "summary": stage["summary"],
                        "metrics": stage["metrics"],
                    }
                )
            except RuntimeError as exc:
                pending_colab.append(stage_identity)
                stages.append(stage_identity | {"status": "colab_handoff", "reason": f"RuntimeError: {exc}"})

    config = {
        "dim": dim_values[0],
        "dims": dim_values,
        "seq_len": seq_len,
        "max_steps": max_steps,
        "eval_records": eval_records,
        "vocab_size": vocab_size,
        "learning_rate": learning_rate,
    }
    handoff = _write_colab_handoff(out, dataset, pending_colab, config) if pending_colab else None
    completed = [stage for stage in stages if stage["status"] == "completed"]
    report = {
        "schema": "tinymind.native_scaling_ladder.v1",
        "created_at": datetime.now(timezone.utc).isoformat(),
        "dataset": str(dataset),
        "summary": {
            "stage_count": len(stages),
            "completed_count": len(completed),
            "colab_handoff_count": len(pending_colab),
            "layers_requested": layers,
            "dims_requested": dim_values,
            "elapsed_s": time.time() - started,
        },
        "config": config | {"local_layer_limit": local_layer_limit, "local_dim_limit": local_dim_limit},
        "stages": stages,
        "colab_handoff": handoff,
        "claim_gate": {
            "native_scaling_evidence_ready": bool(completed),
            "scaling_point_claim_allowed": False,
            "tier0_claim_allowed": False,
            "world_best_claim_allowed": False,
            "reason": "Scaling ladder measures local smoke learning only; it does not establish broad capability.",
        },
    }
    path = out / "native_scaling_ladder_report.json"
    report["json_path"] = str(path)
    path.write_text(json.dumps(report, ensure_ascii=False, indent=2, sort_keys=True) + "\n", encoding="utf-8")
    return report