| |
| |
| import argparse, json, math, os, time |
| from pathlib import Path |
| import numpy as np, torch, torch.nn.functional as F |
| import torch.utils.checkpoint as cp |
| import sys; sys.path.insert(0, '/root/dna') |
| from model_dna_peer import DnaPeer, peer_counts |
|
|
| ap = argparse.ArgumentParser() |
| ap.add_argument('--base', default='/root/dna/ckpt-peer/base.pt') |
| ap.add_argument('--sft', default='/root/dna/sft'); ap.add_argument('--out', default='/root/dna/ckpt-peer-sft') |
| ap.add_argument('--batch', type=int, default=16); ap.add_argument('--maxlen', type=int, default=768) |
| ap.add_argument('--steps', type=int, default=6000); ap.add_argument('--lr', type=float, default=2e-5); ap.add_argument('--warmup', type=int, default=150) |
| ap.add_argument('--hf-repo', default='jaivial/dna-diskchat-2b-peer-v1') |
| a = ap.parse_args(); dev = 'cuda'; Path(a.out).mkdir(parents=True, exist_ok=True); torch.manual_seed(3) |
| torch.backends.cuda.matmul.allow_tf32 = True |
|
|
| ck = torch.load(a.base, map_location=dev, weights_only=False) |
| m = DnaPeer(**ck['config']).to(dev) |
| m.load_state_dict({k.replace('_orig_mod.', ''): v for k, v in ck['model'].items()}, strict=True); m.train() |
| for i in range(len(m.blocks)): m.blocks[i] = torch.compile(m.blocks[i]) |
| flat = np.memmap(f'{a.sft}/sft_tokens.u16', np.uint16, 'r'); mflat = np.memmap(f'{a.sft}/sft_mask.u8', np.uint8, 'r') |
| lens = np.fromfile(f'{a.sft}/sft_lens.i32', np.int32); off = np.zeros(len(lens) + 1, np.int64); off[1:] = np.cumsum(lens) |
| print('SFT_CONFIG', json.dumps({'records': len(lens), 'ctrl_params': peer_counts(m)[0]}), flush=True) |
| opt = torch.optim.AdamW(m.parameters(), lr=a.lr, betas=(.9, .95), weight_decay=0.0) |
|
|
| def batch(): |
| idx = np.random.randint(0, len(lens), size=a.batch); L = min(a.maxlen, max(int(lens[i]) for i in idx)) |
| ids = np.zeros((a.batch, L), np.int64); msk = np.zeros((a.batch, L), np.float32) |
| for j, i in enumerate(idx): |
| n = min(int(lens[i]), L); s = off[i]; ids[j, :n] = flat[s:s+n]; msk[j, :n] = mflat[s:s+n] |
| return torch.from_numpy(ids).to(dev), torch.from_numpy(msk).to(dev) |
| def fwd(ids): |
| B, T = ids.shape; x = m.embed(ids) |
| for blk in m.blocks: x = cp.checkpoint(blk, x, m.chunk, use_reentrant=False) |
| mem, bal, _ = m.route(x.reshape(B * T, m.d)); return m.norm(x + mem.reshape(B, T, m.d)), bal |
| 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, msk = batch(); opt.zero_grad(set_to_none=True) |
| with torch.autocast('cuda', dtype=torch.bfloat16): |
| feat, bal = fwd(ids); lo = F.linear(feat, m.embed.weight) |
| tgt = ids[:, 1:]; mm = msk[:, 1:] |
| ce = F.cross_entropy(lo[:, :-1].reshape(-1, m.vocab), tgt.reshape(-1), reduction='none').view_as(tgt) |
| loss = (ce * mm).sum() / mm.sum().clamp_min(1) + 1e-4 * bal |
| loss.backward(); torch.nn.utils.clip_grad_norm_(m.parameters(), 1.0); opt.step() |
| if step % 50 == 0: |
| print(f'sft step={step}/{a.steps} loss={loss.item():.4f} ppl={math.exp(min(20, loss.item())):.1f} rec_s={step*a.batch/(time.time()-t0):,.1f}', flush=True) |
| if step % 1000 == 0: torch.save({'model': m.state_dict(), 'config': m.config(), 'step': step}, Path(a.out)/'sft.pt') |
| torch.save({'model': m.state_dict(), 'config': m.config(), 'step': a.steps}, Path(a.out)/'sft.pt') |
| print('PEER_SFT_DONE', flush=True) |
| try: |
| from huggingface_hub import HfApi |
| api = HfApi(token=os.environ.get('HUGGING_FACE')) |
| api.upload_file(path_or_fileobj=str(Path(a.out)/'sft.pt'), path_in_repo='sft/sft.pt', repo_id=a.hf_repo) |
| print('PEER_SFT_HF_PUSH ok', flush=True) |
| except Exception as e: print('PEER_SFT_HF_PUSH_FAIL', str(e)[:160], flush=True) |
|
|
