ofou's picture
download
raw
3.16 kB
#!/usr/bin/env python3
"""Claim 2 proxy: solution density vs width/depth after mixed pretrain (tinygrad)."""
from __future__ import annotations
import json
from pathlib import Path
import numpy as np
from datasets import load_data
from models import eval_model, eval_next_step
from train import Args, create_model, pretrain_base_model, set_seed
OUT = Path(__file__).resolve().parent / "outputs" / "claim2"
def density_for(width: int, depth: int, N: int = 200, sigma: float = 0.002, seed: int = 0) -> dict:
args = Args(
ctx_sz=10,
fut_sz=40,
width=width,
depth=depth,
N=N,
K=10,
pretrain_iters=600,
posttrain_dataset_sz=128,
sigma=sigma,
global_seed=seed,
)
set_seed(seed)
base = create_model(args)
base = pretrain_base_model(base, "mixed", args)
snap = base.snapshot_weights()
work = base.clone()
_, ctx_y, _, fut_y = load_data(args.posttrain_dataset_sz, "one_line", args)
base_mse = eval_next_step(base, ctx_y, fut_y)
deltas = []
for i in range(N):
work.perturb_from_snapshot(snap, i, sigma)
deltas.append(base_mse - eval_next_step(work, ctx_y, fut_y))
arr = np.asarray(deltas)
useful = float(np.mean(arr > 0.05 * max(base_mse, 1e-8)))
order = np.argsort(-arr)
work.perturb_from_snapshot(snap, int(order[0]), sigma)
top1_ar = eval_model(work, ctx_y, fut_y, args.fut_sz)
base_ar = eval_model(base, ctx_y, fut_y, args.fut_sz)
return {
"width": width,
"depth": depth,
"n_params": base.n_params(),
"base_mse_ns": base_mse,
"base_mse_ar": base_ar,
"top1_mse_ar": top1_ar,
"solution_density": float(np.mean(arr > 0)),
"useful_density_rel5pct": useful,
"mean_delta": float(arr.mean()),
"p90_delta": float(np.quantile(arr, 0.9)),
"best_delta": float(arr.max()),
"N": N,
"sigma": sigma,
"framework": "tinygrad",
"density_metric": "next_step_mse",
}
def main() -> None:
OUT.mkdir(parents=True, exist_ok=True)
rows = []
for w in [16, 32, 64, 128]:
row = density_for(w, depth=5)
rows.append(row)
print(row, flush=True)
for d in [3, 5, 7]:
row = density_for(64, depth=d)
row["sweep"] = "depth"
rows.append(row)
print(row, flush=True)
width_rows = [r for r in rows if r.get("sweep") != "depth"]
useful = [r["useful_density_rel5pct"] for r in width_rows]
best = [r["best_delta"] for r in width_rows]
claim = best[-1] >= best[0] and (useful[-1] >= useful[0] or best[-1] > best[0])
summary = {
"framework": "tinygrad",
"scope": "toy",
"maps_to": "Claim 2 Figure 3 (Qwen scale) via 1D width/depth proxy",
"blocker": "HF Jobs 402 insufficient prepaid credits; no Qwen2.5 multi-scale in tinygrad",
"claim2_proxy_larger_width_better_experts": claim,
"rows": rows,
}
(OUT / "results.json").write_text(json.dumps(summary, indent=2) + "\n")
print(json.dumps(summary, indent=2), flush=True)
if __name__ == "__main__":
main()

Xet Storage Details

Size:
3.16 kB
·
Xet hash:
dd9fc36014a30d4ce662e3456c51bf1c4d1539d881a283d69314e611357a332a

Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.