| """Throughput and memory of the upstream training step at candidate scales. |
| |
| Drives the same per-chunk path as `train.py read`: peek/want experts, one |
| FileReader.step (which re-forwards the whole window), backward, clip, AdamW over |
| the trunk/pool split. Data is random bytes: cost does not depend on content. |
| """ |
|
|
| import argparse |
| import copy |
| import json |
| import os |
| import tempfile |
| import time |
|
|
| import numpy as np |
| import torch |
| import yaml |
|
|
| from upstream import UPSTREAM, use_config |
|
|
| BASE = yaml.safe_load(open(os.path.join(UPSTREAM, "config.yaml"))) |
|
|
| SCALES = { |
| |
| "xs": (dict(d_model=256, n_head=4, d_ff=704, max_steps=8, train_steps_mean=4.0, |
| bptt_window=8, halt_prior=0.2), |
| dict(experts=32, width=512, top_k=4, resident=16), 256, 512), |
| "s": (dict(d_model=256, n_head=4, d_ff=704, max_steps=8, train_steps_mean=4.0, |
| bptt_window=8, halt_prior=0.2), |
| dict(experts=32, width=512, top_k=4, resident=16), 512, 1024), |
| "m": (dict(d_model=384, n_head=6, d_ff=1024, max_steps=12, train_steps_mean=6.4, |
| bptt_window=12, halt_prior=0.144), |
| dict(experts=48, width=1024, top_k=6, resident=24), 512, 1024), |
| "full": ({}, {}, 2048, 2048), |
| } |
|
|
|
|
| def make_config(scale): |
| mo, po, chunk, ctx = SCALES[scale] |
| c = copy.deepcopy(BASE) |
| c["model"].update(mo) |
| c["model"]["context_start"] = ctx |
| c["model"]["context_end"] = ctx |
| c["pool"].update(po) |
| c["training"]["chunk"] = chunk |
| return c, chunk, ctx |
|
|
|
|
| def build(scale, device): |
| from minagi.create import create |
| from minagi.recur import RecurCoder |
| import train as T |
|
|
| c, chunk, ctx = make_config(scale) |
| tmp = tempfile.mkdtemp(prefix=f"bench_{scale}_") |
| cpath = os.path.join(tmp, "config.yaml") |
| yaml.safe_dump(c, open(cpath, "w")) |
| use_config(cpath) |
| wdir = os.path.join(tmp, "weights") |
| create(wdir, force=True, verbose=False) |
| model, cfg, pool, man = T.build_paged(wdir, device) |
| m = c["model"] |
| cfg.train_steps_mean = float(m["train_steps_mean"]) |
| cfg.min_steps = int(m["min_steps"]) |
| cfg.bptt_window = min(int(m["bptt_window"]), cfg.max_steps) |
| cfg.halt_prior = float(m["halt_prior"]) |
| cfg.halt_thresh = float(m["halt_thresh"]) |
| cfg.ponder_beta = float(m["ponder_beta"]) |
| trunk, pool_ps = T._split_trunk_pool(model) |
| lr = c["training"]["lr"] |
| opt = torch.optim.AdamW( |
| [{"params": trunk, "lr": lr * 0.1, "weight_decay": 0.1}, |
| {"params": pool_ps, "lr": lr, "weight_decay": 0.1}], |
| betas=(0.9, 0.95), fused=True) |
| pool.attach_optimiser(opt) |
| return model, cfg, pool, opt, chunk, ctx, trunk, pool_ps |
|
|
|
|
| def run(scale, seconds): |
| from minagi.precision import set_compute_dtype |
| from minagi.stream import FileReader |
| device = torch.device("cuda") |
| set_compute_dtype("bf16") |
| torch.manual_seed(0) |
| model, cfg, pool, opt, chunk, ctx, trunk, pool_ps = build(scale, device) |
| model.train() |
| rng = np.random.default_rng(0) |
| data = rng.integers(32, 127, size=4_000_000).astype(np.uint16) |
| r = FileReader(model, data, "bench", chunk, ctx, device) |
| torch.cuda.reset_peak_memory_stats() |
| chars, t0, steps = 0, None, 0 |
| while True: |
| nxt = r.peek() |
| model.want_experts(nxt) |
| opt.zero_grad(set_to_none=True) |
| loss = r.step(learn=True, aux_weight=cfg.pool_aux) |
| loss.backward() |
| torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) |
| opt.step() |
| steps += 1 |
| if steps == 5: |
| torch.cuda.synchronize() |
| t0 = time.time() |
| elif steps > 5: |
| chars += chunk |
| if time.time() - t0 > seconds: |
| break |
| torch.cuda.synchronize() |
| dt = time.time() - t0 |
| out = {"scale": scale, "chunk": chunk, "context": ctx, |
| "params_total_M": round(pool.n_params() / 1e6 + sum(p.numel() for p in trunk) / 1e6, 2), |
| "trunk_M": round(sum(p.numel() for p in trunk) / 1e6, 2), |
| "vram_M": round(pool.vram_params() / 1e6, 2), |
| "char_per_s": round(chars / dt), "peak_GB": round(torch.cuda.max_memory_allocated() / 1e9, 2), |
| "loss": round(float(loss), 3)} |
| print(json.dumps(out), flush=True) |
|
|
|
|
| if __name__ == "__main__": |
| ap = argparse.ArgumentParser() |
| ap.add_argument("scale", choices=list(SCALES)) |
| ap.add_argument("--seconds", type=float, default=30) |
| a = ap.parse_args() |
| run(a.scale, a.seconds) |
|
|
