Spaces:
Paused
Paused
| """Console entry points: a char-LM parity gate and a training-peak memory gate. | |
| memory-native-charlm --kinds dense,qat,counter,counter_rms [--data-path FILE] ... | |
| memory-native-memgate --config tiny [--device cuda] | |
| Both run on stock PyTorch (CPU/CUDA), no engine, and fall back to a synthetic corpus offline. | |
| """ | |
| from __future__ import annotations | |
| import argparse | |
| import time | |
| import torch | |
| from .data import get_batch, load_corpus | |
| from .memory import fmt_bytes, memory_report, peak_training_memory | |
| from .models import CONFIGS, GPT | |
| from .optimizers import available_optimizers, build_optimizer | |
| def _device(name: str) -> torch.device: | |
| if name == "cuda" and not torch.cuda.is_available(): | |
| print("cuda requested but not available; using cpu") | |
| return torch.device("cpu") | |
| return torch.device(name) | |
| def _train_one(kind: str, args, vocab, train_data, val_data, device) -> dict: | |
| cfg = CONFIGS[args.config] | |
| counter_kw = dict(lr=args.fs_lr if kind == "counter" else args.fs_lr_rms, | |
| lr_scale=args.fs_lr_scale, C=args.C, tile_rows=args.tile_rows) | |
| kw = counter_kw if kind in ("counter", "counter_rms") else {} | |
| model = GPT(kind, cfg, **kw).to(device) | |
| model.train() | |
| lr = args.base_lr if kind in ("dense", "qat") else args.fp_lr | |
| opt = build_optimizer(args.optimizer, model.trainable_parameters(), lr) | |
| g = torch.Generator().manual_seed(args.seed) | |
| def evaluate(): | |
| model.eval() | |
| tot = 0.0 | |
| for _ in range(args.eval_iters): | |
| x, y = get_batch(val_data, cfg.block_size, args.batch, device, g) | |
| _, loss = model(x, y) | |
| tot += loss.item() | |
| model.train() | |
| return tot / args.eval_iters | |
| if device.type == "cuda": | |
| torch.cuda.synchronize(device) | |
| torch.cuda.reset_peak_memory_stats(device) | |
| started = time.time() | |
| for step in range(args.steps + 1): | |
| x, y = get_batch(train_data, cfg.block_size, args.batch, device, g) | |
| _, loss = model(x, y) | |
| opt.zero_grad(set_to_none=True) | |
| loss.backward() # counter layers self-update during backward | |
| opt.step() | |
| if step % max(1, args.steps // 6) == 0: | |
| print(f" [{kind:12s}] step {step:5d} train {loss.item():.4f}") | |
| if device.type == "cuda": | |
| torch.cuda.synchronize(device) | |
| elapsed = time.time() - started | |
| final_val = evaluate() | |
| rep = memory_report(model) | |
| tokens = (args.steps + 1) * args.batch * cfg.block_size | |
| tok_s = tokens / max(elapsed, 1e-9) | |
| peak = int(torch.cuda.max_memory_allocated(device)) if device.type == "cuda" else 0 | |
| peak_str = f" peak {fmt_bytes(peak)}" if peak else "" | |
| print(f" [{kind:12s}] final val {final_val:.4f} persistent {fmt_bytes(rep['persistent_bytes'])}" | |
| f" {tok_s:,.0f} tok/s{peak_str} ({elapsed:.0f}s)") | |
| return {"final_val": final_val, "report": rep, "tok_s": tok_s, "peak_bytes": peak} | |
| def charlm_main(argv=None) -> None: | |
| ap = argparse.ArgumentParser(description="char-LM parity gate (counter vs AdamW baselines)") | |
| ap.add_argument("--config", choices=list(CONFIGS), default="tiny") | |
| ap.add_argument("--kinds", default="dense,qat,counter,counter_rms") | |
| ap.add_argument("--steps", type=int, default=600) | |
| ap.add_argument("--batch", type=int, default=32) | |
| ap.add_argument("--eval-iters", type=int, default=20) | |
| ap.add_argument("--C", type=int, default=8) | |
| ap.add_argument("--fs-lr", type=float, default=0.04) | |
| ap.add_argument("--fs-lr-rms", type=float, default=3e-3) | |
| ap.add_argument("--fs-lr-scale", type=float, default=2e-4) | |
| ap.add_argument("--fp-lr", type=float, default=2e-3) | |
| ap.add_argument("--base-lr", type=float, default=3e-3) | |
| ap.add_argument("--tile-rows", type=int, default=0, | |
| help="0=untiled update (fast, default); >0 tiles grad_w to never " | |
| "materialize the full [out,in] gradient (strict, ~3x slower)") | |
| ap.add_argument("--seed", type=int, default=0) | |
| ap.add_argument("--device", default="cpu") | |
| ap.add_argument("--data-path", default=None) | |
| ap.add_argument("--no-download", action="store_true", | |
| help="never hit the network: use --data-path / cache, else synthetic corpus") | |
| ap.add_argument("--optimizer", default="adamw", choices=available_optimizers(), | |
| help="optimizer for the trainable (non-counter) params: " | |
| "adamw | bnb8 | galore | lomo") | |
| args = ap.parse_args(argv) | |
| device = _device(args.device) | |
| train_data, val_data, vocab = load_corpus(device, args.data_path, download=not args.no_download) | |
| # vocab from the corpus overrides the config's placeholder vocab_size | |
| cfg = CONFIGS[args.config] | |
| CONFIGS[args.config] = cfg.__class__(vocab, cfg.block_size, cfg.n_layer, cfg.n_head, cfg.n_embd) | |
| print(f"corpus vocab={vocab} config={args.config} device={device}") | |
| print("=" * 64) | |
| kinds = [k.strip() for k in args.kinds.split(",") if k.strip()] | |
| res = {k: _train_one(k, args, vocab, train_data, val_data, device) for k in kinds} | |
| print("=" * 64) | |
| ref_key = "dense" if "dense" in res else next(iter(res)) | |
| ref = res[ref_key]["final_val"] | |
| print("FINAL val loss:") | |
| for k in kinds: | |
| v = res[k]["final_val"] | |
| print(f" {k:14s} {v:.4f} gap {v-ref:+.4f} ({(v-ref)/ref*100:+.1f}% vs {ref_key})") | |
| if "qat" in res: | |
| qv = res["qat"]["final_val"] | |
| for k in ("counter", "counter_rms"): | |
| if k in res: | |
| iso = res[k]["final_val"] - qv | |
| print(f" ISOLATION {k} - qat = {iso:+.4f} ({iso/qv*100:+.1f}%) = counter-optimizer cost") | |
| def memgate_main(argv=None) -> None: | |
| ap = argparse.ArgumentParser( | |
| description="training-peak memory gate: counter_rms vs dense across memory-efficient optimizers") | |
| ap.add_argument("--config", choices=list(CONFIGS), default="tiny") | |
| ap.add_argument("--batch", type=int, default=16) | |
| ap.add_argument("--device", default="cpu") | |
| ap.add_argument("--optimizers", default="adamw,galore,lomo", | |
| help="comma list of dense baselines to compare against: " | |
| f"{','.join(available_optimizers())}") | |
| args = ap.parse_args(argv) | |
| device = _device(args.device) | |
| cfg = CONFIGS[args.config] | |
| import torch.nn.functional as F | |
| x = torch.randint(0, cfg.vocab_size, (args.batch, cfg.block_size), device=device) | |
| y = torch.randint(0, cfg.vocab_size, (args.batch, cfg.block_size), device=device) | |
| def step(model, opt=None): | |
| if opt is not None: | |
| opt.zero_grad(set_to_none=True) | |
| logits, _ = model(x) | |
| loss = F.cross_entropy(logits.reshape(-1, logits.size(-1)), y.reshape(-1)) | |
| loss.backward() | |
| if opt is not None: | |
| opt.step() | |
| # counter_rms: self-updating, no optimizer over its ternary weights. | |
| counter = GPT("counter_rms", cfg).to(device).train() | |
| cpeak = peak_training_memory(lambda: step(counter), device) | |
| cr = memory_report(counter) | |
| print(f"device={device} config={args.config} batch={args.batch}") | |
| print(f"persistent state: counter_rms={fmt_bytes(cr['persistent_bytes'])} " | |
| f"(counter weights packed-to-6bit would be {fmt_bytes(cr['counter_packed_6bit_bytes'])})") | |
| rows = [] | |
| for oname in [o.strip() for o in args.optimizers.split(",") if o.strip()]: | |
| dense = GPT("dense", cfg).to(device).train() | |
| try: | |
| dopt = build_optimizer(oname, dense.trainable_parameters(), lr=3e-3) | |
| except RuntimeError as exc: | |
| print(f" dense+{oname:7s}: SKIP ({exc})") | |
| continue | |
| dpeak = peak_training_memory(lambda: step(dense, dopt), device) | |
| rows.append((oname, dpeak, memory_report(dense))) | |
| if device.type == "cuda": | |
| print(f"TRAINING PEAK (max_memory_allocated): counter_rms = {fmt_bytes(cpeak)}") | |
| for oname, dpeak, _ in rows: | |
| print(f" dense+{oname:7s} = {fmt_bytes(dpeak):>12s} " | |
| f"{dpeak/max(cpeak,1):.2f}x of counter_rms") | |
| else: | |
| print("training peak: CPU has no allocator-peak API; run with --device cuda for the " | |
| "real peak. Persistent-state comparison:") | |
| for oname, _, dr in rows: | |
| print(f" dense+{oname:7s} persistent {fmt_bytes(dr['persistent_bytes'])}") | |
| if __name__ == "__main__": | |
| charlm_main() | |