Spaces:
Paused
Paused
File size: 8,429 Bytes
6f75f79 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 | """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)
@torch.no_grad()
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()
|