| |
| |
| |
| import argparse, json, math, os, time, sys |
| from pathlib import Path |
| import numpy as np, torch, torch.nn.functional as F |
| sys.path.insert(0, '/root/dna') |
| from model_dna import DnaChat, counts |
| from model_dna_tern import ternarize_, controller_byte_budget |
|
|
| ap = argparse.ArgumentParser() |
| ap.add_argument('--base', default='/root/dna/ckpt/base.pt') |
| ap.add_argument('--data-dir', default='/root/dna/data') |
| ap.add_argument('--out', default='/root/dna/ckpt-tern') |
| ap.add_argument('--batch', type=int, default=16) |
| ap.add_argument('--seq', type=int, default=512) |
| ap.add_argument('--steps', type=int, default=1500) |
| ap.add_argument('--lr', type=float, default=1e-4) |
| ap.add_argument('--warmup', type=int, default=100) |
| ap.add_argument('--valid-iters', type=int, default=20) |
| ap.add_argument('--smoke', action='store_true') |
| a = ap.parse_args() |
| dev = 'cuda' if torch.cuda.is_available() else 'cpu' |
| Path(a.out).mkdir(parents=True, exist_ok=True); torch.manual_seed(0) |
|
|
| shards = sorted(Path(a.data_dir).glob('shard_*.u16')) |
| _cur = {'i': -1, 'arr': None} |
| def batch(): |
| i = np.random.randint(0, len(shards)) |
| if i != _cur['i']: _cur['arr'] = np.memmap(shards[i], np.uint16, 'r'); _cur['i'] = i |
| arr = _cur['arr']; n = len(arr) - a.seq - 1 |
| ix = np.random.randint(0, n, size=a.batch) |
| return torch.from_numpy(np.stack([np.asarray(arr[j:j+a.seq+1], np.int64) for j in ix])).to(dev) |
|
|
| @torch.no_grad() |
| def valid_ce(m, iters): |
| m.eval(); tot = 0.0 |
| for _ in range(iters): |
| ids = batch() |
| with torch.autocast(dev, dtype=torch.bfloat16, enabled=(dev=='cuda')): |
| feat, _, _ = m.features(ids) |
| ce = F.cross_entropy(F.linear(feat[:, :-1].reshape(-1, m.d), m.embed.weight), |
| ids[:, 1:].reshape(-1)) |
| tot += ce.item() |
| m.train(); return tot / iters |
|
|
| ck = torch.load(a.base, map_location=dev, weights_only=False) |
| m = DnaChat(**ck['config']).to(dev) |
| m.load_state_dict({k.replace('_orig_mod.', ''): v for k, v in ck['model'].items()}, strict=True) |
| ctrl, _ = counts(m) |
| if a.smoke: |
| a.steps = 3; a.valid_iters = 2 |
|
|
| ce_fp = valid_ce(m, a.valid_iters) |
| n_swapped = ternarize_(m) |
| bud = controller_byte_budget(m) |
| ce_q0 = valid_ce(m, a.valid_iters) |
| print('TERN_INIT', json.dumps({ |
| 'ctrl_params': ctrl, 'ternarized_linears': n_swapped, |
| 'valid_ce_fp16': round(ce_fp, 4), 'valid_ce_ternary_raw': round(ce_q0, 4), |
| 'bytes_per_token_bf16': round(bud['bytes_per_token_bf16']/1e6, 1), |
| 'bytes_per_token_tern': round(bud['bytes_per_token_tern']/1e6, 1), |
| 'tern_params': bud['tern_params'], 'fp_linear_params': bud['fp_linear_params'], |
| }), flush=True) |
|
|
| opt = torch.optim.AdamW(m.parameters(), lr=a.lr, betas=(.9, .95), weight_decay=0.0) |
| def lrf(s): return s/a.warmup if s < a.warmup else max(0.1, 0.5*(1+math.cos(math.pi*(s-a.warmup)/max(1, a.steps-a.warmup)))) |
| t0 = time.time() |
| for step in range(1, a.steps + 1): |
| for g in opt.param_groups: g['lr'] = a.lr * lrf(step) |
| ids = batch(); opt.zero_grad(set_to_none=True) |
| with torch.autocast(dev, dtype=torch.bfloat16, enabled=(dev=='cuda')): |
| feat, bal, _ = m.features(ids) |
| ce = F.cross_entropy(F.linear(feat[:, :-1].reshape(-1, m.d), m.embed.weight), ids[:, 1:].reshape(-1)) |
| loss = ce + 1e-4 * bal |
| loss.backward(); torch.nn.utils.clip_grad_norm_(m.parameters(), 1.0); opt.step() |
| if step % 100 == 0 or step == a.steps: |
| print(f'qat step={step}/{a.steps} ce={ce.item():.4f} ppl={math.exp(min(20,ce.item())):.1f} ' |
| f'tok_s={step*a.batch*a.seq/(time.time()-t0):,.0f}', flush=True) |
|
|
| ce_qat = valid_ce(m, a.valid_iters) |
| torch.save({'model': m.state_dict(), 'config': m.config(), 'ternary': True}, Path(a.out)/'tern.pt') |
| print('TERN_DONE', json.dumps({ |
| 'valid_ce_fp16': round(ce_fp, 4), 'valid_ce_ternary_raw': round(ce_q0, 4), |
| 'valid_ce_ternary_qat': round(ce_qat, 4), |
| 'ppl_fp16': round(math.exp(min(20, ce_fp)), 1), |
| 'ppl_ternary_qat': round(math.exp(min(20, ce_qat)), 1), |
| 'bytes_per_token_bf16': round(bud['bytes_per_token_bf16']/1e6, 1), |
| 'bytes_per_token_tern': round(bud['bytes_per_token_tern']/1e6, 1), |
| }), flush=True) |
|
|
