| |
| """Does the learn phase actually scale with torch threads? |
| |
| main_v5 spends 390 s/iter in learn vs 45 s in self-play with --threads 48. |
| If the update step doesn't scale past a few threads, those cores are wasted and |
| the run should instead use more parallel ARMS (or a smaller batch). |
| """ |
| import argparse, json, os, sys, time |
| sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) |
| import torch |
| from tinychess.checkpoint import load_model_only |
| from tinychess.config import REGISTRY |
| from tinychess.model import TinyChess |
| from tinychess.selfplay import SelfPlayConfig, play_games |
| from tinychess.replay import ReplayBuffer |
| import train_selfplay as T |
|
|
|
|
| def main(): |
| ap = argparse.ArgumentParser() |
| ap.add_argument("--ckpt", default="") |
| ap.add_argument("--batch", type=int, default=1024) |
| ap.add_argument("--games", type=int, default=24) |
| ap.add_argument("--reps", type=int, default=6) |
| ap.add_argument("--threads", default="1,2,4,8,16,32,48") |
| ap.add_argument("--out", default="runs/bench_learn.json") |
| a = ap.parse_args() |
|
|
| if a.ckpt and os.path.exists(a.ckpt): |
| m, _ = load_model_only(a.ckpt) |
| else: |
| m = TinyChess(REGISTRY["baseline"]()) |
| cfg = SelfPlayConfig(n_games=a.games, max_plies=160, steps=4, |
| adjudicate="material", seed=0) |
| recs, _ = play_games(m, None, cfg) |
| buf = ReplayBuffer() |
| for r in recs: |
| buf.add_game(r, gen=0) |
| print(f"buffer {len(buf)} samples", flush=True) |
|
|
| class A: pass |
| a2 = A(); a2.steps = 4; a2.norm_adv = 1; a2.ppo_clip = 0.2 |
| a2.c_entropy = 0.01; a2.c_value = 1.0; a2.clip = 1.0; a2.c_compute = 0.0 |
| a2.terminal_weight = 0.0; a2.terminal_tau = 10.0 |
| opt = torch.optim.Adam(m.parameters(), lr=1e-6) |
|
|
| res = {"cpu_count": os.cpu_count(), "batch": a.batch, "runs": []} |
| for nt in [int(x) for x in a.threads.split(",")]: |
| torch.set_num_threads(nt) |
| T.train_step(m, opt, buf.sample(a.batch), a2) |
| t = time.time() |
| for _ in range(a.reps): |
| T.train_step(m, opt, buf.sample(a.batch), a2) |
| ms = (time.time() - t) / a.reps * 1000 |
| res["runs"].append({"threads": nt, "ms_per_update": round(ms, 1)}) |
| print(f"threads={nt:3d}: {ms:8.1f} ms/update", flush=True) |
| json.dump(res, open(a.out, "w"), indent=2) |
| best = min(res["runs"], key=lambda r: r["ms_per_update"]) |
| res["best"] = best |
| json.dump(res, open(a.out, "w"), indent=2) |
| print(f"\nBEST: {best['threads']} threads = {best['ms_per_update']} ms/update") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
