|
|
|
|
|
|
|
|
|
|
| from __future__ import annotations
|
| import argparse, json, math, pathlib, random, time, os
|
| from contextlib import nullcontext
|
| from typing import Dict, Any, List, Optional, Tuple
|
| from collections import deque
|
|
|
| import torch
|
| import torch.nn as nn
|
| import torch.nn.functional as F
|
| from datasets import load_dataset
|
| from transformers import AutoTokenizer, logging as hf_log
|
| from tqdm.auto import tqdm
|
|
|
|
|
| hf_log.set_verbosity_error()
|
| DEV = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
| torch.backends.cuda.matmul.allow_tf32 = True
|
| try:
|
| torch.set_float32_matmul_precision("high")
|
| except Exception:
|
| pass
|
|
|
|
|
| TOKENIZER_ID = os.environ.get(
|
| "TOKENIZER_ID",
|
| "Qwen/Qwen3-235B-A22B-Thinking-2507"
|
| )
|
|
|
|
|
| tok = AutoTokenizer.from_pretrained(TOKENIZER_ID, use_fast=True, trust_remote_code=True)
|
| if tok.pad_token is None:
|
| tok.add_special_tokens({"pad_token": "[PAD]"})
|
| VOCAB, BLANK, EOS = (
|
| max(tok.get_vocab().values()) + 1,
|
| tok.pad_token_id,
|
| tok.eos_token_id if tok.eos_token_id is not None else tok.sep_token_id
|
| )
|
|
|
| PRESETS: Dict[str, Dict[str, int]] = {
|
| "small": dict(d=512, layers=8, heads=16, rank=64),
|
| "smallx2": dict(d=512, layers=16, heads=16, rank=64),
|
| "base": dict(d=768, layers=12, heads=24, rank=96),
|
| }
|
|
|
|
|
| DEFAULT_BLOCK = 576
|
| SAT_BLOCK = 2
|
| LR_CORE, LR_HEAD = 5e-5, 2e-4
|
| EMIT_LAMBDA = 0.1
|
| DEFAULT_SAVE_SEC = 8 * 24 * 3600
|
| CKDIR = pathlib.Path("ckpts_joint")
|
|
|
|
|
|
|
| def rng_state():
|
| if DEV.type == "cuda":
|
| try:
|
| return torch.cuda.get_rng_state(DEV)
|
| except TypeError:
|
| return torch.cuda.get_rng_state()
|
| return torch.get_rng_state()
|
|
|
|
|
| def _is_probably_ckpt(path: pathlib.Path) -> bool:
|
| try:
|
| return path.is_file() and path.suffix == ".pt" and not path.name.endswith(".pt.tmp") and path.stat().st_size > (1<<20)
|
| except Exception:
|
| return False
|
|
|
|
|
| def _resolve_ckpt(path: pathlib.Path) -> pathlib.Path | None:
|
| """
|
| Return a solid .pt (never .tmp). If 'path' is dir, pick newest *.pt.
|
| If not usable, return None.
|
| """
|
| try:
|
| if path.is_dir():
|
| cands = sorted([p for p in path.glob("*.pt") if _is_probably_ckpt(p)],
|
| key=lambda p: p.stat().st_mtime, reverse=True)
|
| return cands[0] if cands else None
|
| if path.suffix == ".tmp":
|
| solid = path.with_suffix("")
|
| return solid if _is_probably_ckpt(solid) else _resolve_ckpt(path.parent)
|
| return path if _is_probably_ckpt(path) else _resolve_ckpt(path.parent)
|
| except Exception:
|
| return None
|
|
|
|
|
| def _try_load(path: pathlib.Path, map_location="cpu"):
|
| try:
|
|
|
| return torch.load(path, map_location=map_location)
|
| except Exception as e:
|
| print(f"[ckpt-skip] {path} not usable: {e}")
|
| return None
|
|
|
|
|
|
|
| try:
|
| from torch.amp import autocast as _ac, GradScaler
|
| except ImportError:
|
| from torch.cuda.amp import autocast as _ac, GradScaler
|
|
|
| def _auto_amp_dtype():
|
| if DEV.type == "cuda":
|
| try:
|
|
|
| if torch.cuda.is_bf16_supported():
|
| return torch.bfloat16
|
| return torch.float16
|
| except Exception:
|
| return torch.float16
|
| return torch.float32
|
|
|
| def amp(enabled):
|
| return nullcontext() if not enabled else _ac(device_type="cuda", dtype=_auto_amp_dtype())
|
|
|
|
|
|
|
| def token_stream(ds_name: str, target: int, seed: int = 42):
|
| ds = load_dataset(ds_name, split="train", streaming=True)
|
| ds = ds.shuffle(buffer_size=10_000, seed=seed)
|
| emitted = 0
|
| for ex in ds:
|
|
|
| enc = tok.encode(ex["text"])
|
| if EOS is not None and (len(enc) == 0 or enc[-1] != EOS):
|
| enc = enc + [EOS]
|
| for t in enc:
|
| yield t
|
| emitted += 1
|
| if emitted >= target:
|
| return
|
|
|
|
|
|
|
| def _alibi_slopes(n_heads: int):
|
| import math
|
| def pow2slopes(n):
|
| start = 2 ** (-2 ** -(math.log2(n) - 3))
|
| ratio = start
|
| return [start * (ratio ** i) for i in range(n)]
|
| if math.log2(n_heads).is_integer():
|
| vals = pow2slopes(n_heads)
|
| else:
|
| closest = 2 ** math.floor(math.log2(n_heads))
|
| vals = pow2slopes(closest)
|
| extra = pow2slopes(2 * closest)
|
| vals += extra[0::2][: n_heads - closest]
|
| return torch.tensor(vals, device=DEV).view(1, n_heads, 1, 1)
|
|
|
| def alibi_bias(n_heads: int, n_tokens: int):
|
| i = torch.arange(n_tokens, device=DEV).view(1, 1, n_tokens, 1)
|
| j = torch.arange(n_tokens, device=DEV).view(1, 1, 1, n_tokens)
|
| dist = (j - i).clamp_min(0)
|
| slopes = _alibi_slopes(n_heads)
|
| return -slopes * dist
|
|
|
|
|
|
|
| class LowRankMHA(nn.Module):
|
| """
|
| Cache-aware MHA with low-rank projections; supports kv caching for decode.
|
| """
|
| def __init__(self, d: int, h: int, r: int, use_relpos: bool = True):
|
| super().__init__()
|
| assert d % h == 0, "d must be divisible by number of heads"
|
| self.h, self.dk = h, d // h
|
| self.use_relpos = use_relpos
|
| self.q = nn.Linear(d, d, bias=False)
|
| self.k = nn.Linear(d, d, bias=False)
|
| self.v = nn.Linear(d, d, bias=False)
|
| self.U = nn.Parameter(torch.randn(self.dk, r))
|
| nn.init.orthogonal_(self.U)
|
| self.proj = nn.Linear(h * r, d, bias=False)
|
| self.drop = nn.Dropout(0.1)
|
|
|
| def _proj(self, x):
|
| B, N, _ = x.shape
|
| return (x.view(B, N, self.h, self.dk).transpose(1, 2) @ self.U)
|
|
|
| def forward(
|
| self,
|
| x: torch.Tensor,
|
| mask: Optional[torch.Tensor] = None,
|
| rel_bias_tokens: Optional[int] = None,
|
| kv_cache: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
|
| use_cache: bool = False,
|
| ):
|
| q = self._proj(self.q(x))
|
| k_new = self._proj(self.k(x))
|
| v_new = self._proj(self.v(x))
|
|
|
| if kv_cache is None:
|
| k, v = k_new, v_new
|
| else:
|
| k, v = kv_cache
|
| if use_cache:
|
| k = torch.cat([k, k_new], dim=2)
|
| v = torch.cat([v, v_new], dim=2)
|
|
|
| att = (q @ k.transpose(-1, -2)) / math.sqrt(self.dk)
|
|
|
| if q.size(2) == k.size(2):
|
| if self.use_relpos and rel_bias_tokens is not None:
|
| att = att + alibi_bias(self.h, rel_bias_tokens)
|
| if mask is not None:
|
| att = att + mask
|
|
|
| z = (att.softmax(-1) @ v).transpose(1, 2)
|
| z = z.reshape(x.size(0), x.size(1), -1)
|
| out = self.drop(self.proj(z))
|
| return (out, (k, v)) if use_cache else out
|
|
|
|
|
| class Block(nn.Module):
|
| def __init__(self, d: int, h: int, r: int):
|
| super().__init__()
|
| self.ln1, self.ln2 = nn.LayerNorm(d), nn.LayerNorm(d)
|
| self.mha = LowRankMHA(d, h, r, use_relpos=True)
|
| self.ff = nn.Sequential(nn.Linear(d, 4 * d), nn.ReLU(), nn.Linear(4 * d, d))
|
|
|
| def forward(
|
| self,
|
| x: torch.Tensor,
|
| mask: Optional[torch.Tensor],
|
| kv: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
|
| use_cache: bool = False
|
| ):
|
| n = x.size(1)
|
| if use_cache:
|
| y, new_kv = self.mha(self.ln1(x), mask, rel_bias_tokens=n if mask is not None else None, kv_cache=kv, use_cache=True)
|
| x = x + y
|
| x = x + self.ff(self.ln2(x))
|
| return x, new_kv
|
| else:
|
| x = x + self.mha(self.ln1(x), mask, rel_bias_tokens=n)
|
| return x + self.ff(self.ln2(x))
|
|
|
|
|
| class Encoder(nn.Module):
|
| """
|
| Transformer encoder with optional kv caching (for AR/SAT decode).
|
| """
|
| def __init__(self, cfg: Dict[str, int]):
|
| super().__init__()
|
| d, l, h, r = cfg["d"], cfg["layers"], cfg["heads"], cfg["rank"]
|
| self.emb = nn.Embedding(VOCAB, d)
|
| self.blocks = nn.ModuleList(Block(d, h, r) for _ in range(l))
|
| self.ln = nn.LayerNorm(d)
|
|
|
| def forward(
|
| self,
|
| ids: torch.Tensor,
|
| mask: Optional[torch.Tensor],
|
| kv_caches: Optional[List[Optional[Tuple[torch.Tensor, torch.Tensor]]]] = None,
|
| use_cache: bool = False
|
| ):
|
| x = self.emb(ids)
|
| if not use_cache:
|
| for blk in self.blocks:
|
| x = blk(x, mask)
|
| return self.ln(x)
|
|
|
| new_kvs: List[Tuple[torch.Tensor, torch.Tensor]] = []
|
| for i, blk in enumerate(self.blocks):
|
| kv = kv_caches[i] if (kv_caches is not None) else None
|
| x, kv_out = blk(x, mask, kv, use_cache=True)
|
| new_kvs.append(kv_out)
|
| return self.ln(x), new_kvs
|
|
|
|
|
| class ARHead(nn.Module):
|
| def __init__(self, d):
|
| super().__init__()
|
| self.proj = nn.Linear(d, VOCAB)
|
| def forward(self, h): return self.proj(h)
|
|
|
|
|
| class NATHead(nn.Module):
|
| def __init__(self, d):
|
| super().__init__()
|
| self.proj = nn.Linear(d, VOCAB)
|
| def forward(self, h): return self.proj(h)
|
|
|
|
|
| class SATHead(nn.Module):
|
| def __init__(self, d, mode="var"):
|
| super().__init__()
|
| self.proj = nn.Linear(d, VOCAB)
|
| self.mode = mode
|
| self.gate = nn.Linear(d, 2) if mode == "var" else None
|
| def forward(self, h_last):
|
| logits = self.proj(h_last)
|
| gate = self.gate(h_last[:, 0]) if self.gate is not None else None
|
| return logits, gate
|
|
|
|
|
|
|
| def causal_mask(n):
|
| m = torch.full((1, 1, n, n), float("-inf"), device=DEV)
|
| return torch.triu(m, 1)
|
|
|
| def sat_mask(n, block=SAT_BLOCK):
|
| idx = torch.arange(n, device=DEV)
|
| grp = idx.unsqueeze(0) // block
|
| allow = (grp.T == grp) | (grp.T > grp)
|
| return torch.where(allow, 0.0, float("-inf")).unsqueeze(0).unsqueeze(0)
|
|
|
|
|
|
|
| def save_ckpt(
|
| path: pathlib.Path,
|
| core: nn.Module,
|
| ar_h: nn.Module,
|
| nat_h: nn.Module,
|
| sat_h: nn.Module,
|
| opt: torch.optim.Optimizer,
|
| scaler: GradScaler,
|
| meta: Dict[str, Any],
|
| ):
|
| path.parent.mkdir(exist_ok=True, parents=True)
|
| tmp = path.with_suffix(path.suffix + ".tmp")
|
| state = {
|
| "core": core.state_dict(),
|
| "ar": ar_h.state_dict(),
|
| "nat": nat_h.state_dict(),
|
| "sat": sat_h.state_dict(),
|
| "opt": opt.state_dict(),
|
| "scaler": scaler.state_dict(),
|
| "cfg": meta.get("cfg"),
|
| "tokenizer_id": TOKENIZER_ID,
|
| **{k: v for k, v in meta.items() if k not in {"cfg"}},
|
| }
|
| torch.save(state, tmp, _use_new_zipfile_serialization=False)
|
| tmp.replace(path)
|
| (path.parent / "latest.json").write_text(json.dumps({"path": str(path), "step": meta["step"]}))
|
| print(f"\n✓ saved checkpoint {path.name}")
|
|
|
|
|
| def load_ckpt(
|
| path: pathlib.Path,
|
| core: nn.Module,
|
| ar_h: nn.Module,
|
| nat_h: nn.Module,
|
| sat_h: nn.Module,
|
| opt: torch.optim.Optimizer,
|
| scaler: GradScaler,
|
| ):
|
| p = _resolve_ckpt(path) or path
|
| ck = _try_load(p, map_location=DEV)
|
| if ck is None:
|
| raise FileNotFoundError(f"No valid checkpoint at {p}")
|
| core.load_state_dict(ck["core"])
|
| ar_h.load_state_dict(ck["ar"])
|
| nat_h.load_state_dict(ck["nat"])
|
| sat_h.load_state_dict(ck["sat"])
|
| opt.load_state_dict(ck["opt"])
|
| scaler.load_state_dict(ck["scaler"])
|
| return ck.get("step", 0), ck.get("seen_tok", 0), ck.get("wall_time", time.time())
|
|
|
|
|
| def _safe_load_any(path: pathlib.Path, tgt: nn.Module, key: str | None = None, rename: str | None = None):
|
| p = _resolve_ckpt(path) or path
|
| if not p.exists(): return 0
|
| ck = _try_load(p, map_location=DEV)
|
| if ck is None: return 0
|
| sd = ck.get(key, ck) if key else ck
|
| if isinstance(sd, dict) and "state_dict" in sd:
|
| sd = sd["state_dict"]
|
| if rename:
|
| sd = {k.replace(rename, "proj."): v for k, v in sd.items() if rename in k}
|
| tgt_sd = tgt.state_dict()
|
| filt = {k: v for k, v in sd.items() if k in tgt_sd and v.shape == tgt_sd[k].shape}
|
| if filt:
|
| tgt.load_state_dict(filt, strict=False)
|
| return len(filt)
|
|
|
|
|
| def infer_cfg_from_ckpt(path: pathlib.Path):
|
| p = _resolve_ckpt(path) or path
|
| if not p.exists(): return None
|
| sd = _try_load(p, map_location="cpu")
|
| if sd is None: return None
|
| if isinstance(sd, dict) and "cfg" in sd and isinstance(sd["cfg"], dict):
|
| return dict(sd["cfg"])
|
| core = sd.get("core")
|
| if core is None: return None
|
| emb_w = core.get("emb.weight")
|
| if emb_w is None: return None
|
| d = emb_w.shape[1]
|
| layer_ids = []
|
| for k in core.keys():
|
| if k.startswith("blocks."):
|
| parts = k.split(".")
|
| if len(parts) > 2 and parts[1].isdigit():
|
| layer_ids.append(int(parts[1]))
|
| layers = (max(layer_ids) + 1) if layer_ids else None
|
| U = core.get("blocks.0.mha.U")
|
| heads = rank = None
|
| if U is not None:
|
| dk, r = U.shape
|
| rank = r
|
| heads = d // dk if dk > 0 else None
|
| out = {"d": d}
|
| if layers is not None: out["layers"] = layers
|
| if heads is not None: out["heads"] = heads
|
| if rank is not None: out["rank"] = rank
|
| return out
|
|
|
|
|
|
|
| def _parse_grow_plan(s: str) -> List[int]:
|
| steps = []
|
| for part in s.split(","):
|
| part = part.strip()
|
| if part:
|
| v = int(part)
|
| if v >= 128:
|
| steps.append(v)
|
| return sorted(set(steps))
|
|
|
|
|
| def train(args):
|
| cfg = PRESETS[args.preset].copy()
|
|
|
|
|
| if not args.fresh:
|
| src_probe = pathlib.Path(args.warmstart_from) if args.warmstart_from else pathlib.Path(args.save_dir) / "final.pt"
|
| prev_cfg = infer_cfg_from_ckpt(src_probe)
|
| else:
|
| prev_cfg = None
|
|
|
| if prev_cfg:
|
| cfg["d"] = prev_cfg.get("d", cfg["d"])
|
| if prev_cfg.get("heads"):
|
| cfg["heads"] = prev_cfg["heads"]
|
| if args.rank is None and prev_cfg.get("rank"):
|
| cfg["rank"] = prev_cfg["rank"]
|
|
|
| if prev_cfg.get("layers"):
|
| cfg["layers"] = prev_cfg["layers"]
|
|
|
| if args.x2 and prev_cfg.get("layers"):
|
| cfg["layers"] = max(cfg["layers"], prev_cfg["layers"] * 2)
|
| if args.rank:
|
| cfg["rank"] = args.rank
|
| if args.x2 and not prev_cfg:
|
| cfg["layers"] *= 2
|
|
|
| BLOCK = args.block or DEFAULT_BLOCK
|
|
|
| core = Encoder(cfg).to(DEV)
|
| ar_h, nat_h = ARHead(cfg["d"]).to(DEV), NATHead(cfg["d"]).to(DEV)
|
| sat_h = SATHead(cfg["d"], mode="var").to(DEV)
|
|
|
|
|
| loaded = 0
|
| if not args.fresh:
|
| src = pathlib.Path(args.warmstart_from) if args.warmstart_from else pathlib.Path(args.save_dir) / "final.pt"
|
| src = _resolve_ckpt(src)
|
| if src:
|
| loaded += _safe_load_any(src, core, key="core")
|
| loaded += _safe_load_any(src, ar_h, key="ar")
|
| loaded += _safe_load_any(src, nat_h, key="nat")
|
| loaded += _safe_load_any(src, sat_h, key="sat")
|
| if loaded:
|
| print(f"Warm-start: loaded {loaded} matching tensors from {src}")
|
|
|
| opt = torch.optim.AdamW(
|
| [
|
| {"params": core.parameters(), "lr": LR_CORE},
|
| {"params": ar_h.parameters(), "lr": LR_HEAD},
|
| {"params": nat_h.parameters(), "lr": LR_HEAD},
|
| {"params": sat_h.parameters(), "lr": LR_HEAD},
|
| ]
|
| )
|
| scaler = GradScaler(enabled=args.amp and DEV.type == "cuda")
|
|
|
| ce_tok = nn.CrossEntropyLoss(label_smoothing=0.1)
|
| ctc = nn.CTCLoss(blank=BLANK, zero_infinity=True)
|
| ce_gate = nn.CrossEntropyLoss()
|
|
|
|
|
| start_step, seen_tok = 0, 0
|
| last_save_time = time.time()
|
| if args.resume and not args.fresh:
|
| start_step, seen_tok, last_save_time = load_ckpt(
|
| pathlib.Path(args.resume), core, ar_h, nat_h, sat_h, opt, scaler
|
| )
|
| print(f"✓ resumed from step {start_step:,}, seen_tokens={seen_tok:,}")
|
|
|
|
|
| if args.target_tokens:
|
| target_tokens = args.target_tokens
|
| else:
|
| param_count = sum(p.numel() for p in core.parameters())
|
| target_tokens = int(25 * param_count)
|
|
|
| new_tokens_needed = target_tokens - seen_tok
|
| if new_tokens_needed <= 0:
|
| print("Target already reached – nothing to train.")
|
| return
|
| new_steps = new_tokens_needed // BLOCK
|
| if args.steps:
|
| new_steps = min(new_steps, args.steps)
|
| new_tokens_needed = new_steps * BLOCK
|
|
|
| total_tokens_needed = seen_tok + new_tokens_needed
|
| print(f"[auto-steps] {new_steps:,} training steps (@ {BLOCK} tokens/step)")
|
|
|
|
|
| grow_plan = _parse_grow_plan(args.grow_plan) if args.auto_grow else []
|
| if args.auto_grow:
|
| if BLOCK not in grow_plan:
|
| grow_plan = sorted(set(grow_plan + [BLOCK]))
|
| print(f"[auto-grow] plan: {grow_plan} every {args.grow_every_steps} steps")
|
|
|
| stream = token_stream(args.source, target_tokens, seed=42)
|
| buf: list[int] = []
|
| pbar = tqdm(total=total_tokens_needed, initial=seen_tok, unit="tok")
|
| step = start_step
|
| steps_since_last_grow = 0
|
|
|
| while seen_tok < total_tokens_needed:
|
|
|
| try:
|
| while len(buf) < BLOCK:
|
| buf.append(next(stream))
|
| except StopIteration:
|
| break
|
| ids = torch.tensor(buf[:BLOCK], device=DEV).unsqueeze(0)
|
| buf = buf[BLOCK:]
|
|
|
| tgt_ar = ids.clone()
|
| ids_nat = torch.repeat_interleave(ids, 2, 1)
|
|
|
| try:
|
| with amp(args.amp):
|
|
|
| h_ar = core(ids, causal_mask(ids.size(1)))
|
| logits_ar = ar_h(h_ar)[:, :-1]
|
| loss_ar = ce_tok(logits_ar.reshape(-1, VOCAB), tgt_ar[:, 1:].reshape(-1))
|
|
|
|
|
| h_nat = core(ids_nat, None)
|
| log_nat = nat_h(h_nat).log_softmax(-1).transpose(0, 1)
|
| ilen = tlen = torch.tensor([ids_nat.size(1) // 2], device=DEV)
|
| loss_nat = ctc(log_nat, tgt_ar, ilen, tlen)
|
|
|
|
|
| h_sat = core(ids, sat_mask(ids.size(1)))
|
| logits_sat, gate = sat_h(h_sat[:, -SAT_BLOCK:])
|
| tgt_sat = ids[:, 1:SAT_BLOCK+1]
|
| loss_sat = ce_tok(logits_sat.reshape(-1, VOCAB), tgt_sat.reshape(-1))
|
| if gate is not None:
|
| loss_sat += EMIT_LAMBDA * ce_gate(gate, torch.ones(ids.size(0), device=DEV, dtype=torch.long))
|
|
|
| loss = loss_ar + loss_nat + loss_sat
|
|
|
|
|
| scaler.scale(loss).backward()
|
| scaler.unscale_(opt)
|
| nn.utils.clip_grad_norm_(core.parameters(), 1.0)
|
| scaler.step(opt)
|
| scaler.update()
|
| opt.zero_grad(set_to_none=True)
|
|
|
| except RuntimeError as e:
|
| msg = str(e).lower()
|
| if "out of memory" in msg or "cuda error" in msg:
|
| new_block = max(128, BLOCK // 2)
|
| if new_block < BLOCK:
|
| print(f"\n[OOM] reducing block from {BLOCK} -> {new_block}")
|
| BLOCK = new_block
|
| if DEV.type == "cuda":
|
| torch.cuda.empty_cache()
|
| buf = ids[0].tolist() + buf
|
| steps_since_last_grow = 0
|
| continue
|
| raise
|
|
|
|
|
| step += 1
|
| seen_tok += BLOCK
|
| pbar.update(BLOCK)
|
| pbar.set_postfix(loss=f"{loss.item():.3f}", block=BLOCK)
|
|
|
|
|
| now = time.time()
|
| time_due = (now - last_save_time) >= args.save_every_sec > 0
|
| step_due = args.save_every_steps > 0 and step % args.save_every_steps == 0
|
| if time_due or step_due:
|
| ck_name = f"step{step:08d}.pt"
|
| save_ckpt(
|
| pathlib.Path(args.save_dir) / ck_name,
|
| core, ar_h, nat_h, sat_h, opt, scaler,
|
| meta={
|
| "cfg": cfg,
|
| "step": step,
|
| "seen_tok": seen_tok,
|
| "wall_time": now,
|
| "py_state": random.getstate(),
|
| "torch_state": rng_state(),
|
| },
|
| )
|
| last_save_time = now
|
|
|
|
|
| if args.auto_grow:
|
| steps_since_last_grow += 1
|
| if steps_since_last_grow >= args.grow_every_steps:
|
| steps_since_last_grow = 0
|
| try:
|
| idx = grow_plan.index(BLOCK)
|
| if idx + 1 < len(grow_plan):
|
| candidate = grow_plan[idx + 1]
|
| print(f"[auto-grow] attempting BLOCK {BLOCK} -> {candidate}")
|
| BLOCK = candidate
|
| if DEV.type == "cuda":
|
| torch.cuda.empty_cache()
|
| else:
|
| print("[auto-grow] at max planned block; no further growth.")
|
| except ValueError:
|
| grow_plan = sorted(set(grow_plan + [BLOCK]))
|
| idx = grow_plan.index(BLOCK)
|
| if idx + 1 < len(grow_plan):
|
| candidate = grow_plan[idx + 1]
|
| print(f"[auto-grow] moving to planned BLOCK {candidate}")
|
| BLOCK = candidate
|
| if DEV.type == "cuda":
|
| torch.cuda.empty_cache()
|
|
|
| pbar.close()
|
|
|
|
|
| save_ckpt(
|
| pathlib.Path(args.save_dir) / "final.pt",
|
| core, ar_h, nat_h, sat_h, opt, scaler,
|
| meta={
|
| "cfg": cfg,
|
| "step": step,
|
| "seen_tok": seen_tok,
|
| "wall_time": time.time(),
|
| "py_state": random.getstate(),
|
| "torch_state": rng_state(),
|
| },
|
| )
|
| print("🎉 training complete")
|
|
|
|
|
|
|
| @torch.no_grad()
|
| def apply_repetition_penalty_(logits: torch.Tensor,
|
| recent_ids: deque | List[int],
|
| penalty: float,
|
| window: int) -> None:
|
| """
|
| In-place adjustment of logits using HF-style repetition penalty.
|
| Penalizes tokens that appeared in the last `window` ids.
|
| Works for shapes (V,) or (1,V). We assume batch=1 for decode.
|
| """
|
| if penalty is None or penalty <= 1.0 or window <= 0 or not recent_ids:
|
| return
|
|
|
| lview = logits
|
| while lview.dim() > 1:
|
| lview = lview[0]
|
| tail = list(recent_ids)[-window:]
|
| if not tail:
|
| return
|
| u, cnt = torch.unique(torch.tensor(tail, device=lview.device, dtype=torch.long), return_counts=True)
|
| powv = (torch.ones_like(cnt, dtype=lview.dtype) * penalty).pow(cnt.to(lview.dtype))
|
| sel = lview.index_select(0, u)
|
| sel = torch.where(sel > 0, sel / powv, sel * powv)
|
| lview.index_copy_(0, u, sel)
|
|
|
|
|
|
|
| def _sync_tokenizer_for_checkpoint(sd: dict):
|
| global tok, TOKENIZER_ID, VOCAB, BLANK, EOS
|
| ck_tok = sd.get("tokenizer_id")
|
| if isinstance(ck_tok, str) and ck_tok and ck_tok != TOKENIZER_ID:
|
| TOKENIZER_ID = ck_tok
|
| tok = AutoTokenizer.from_pretrained(TOKENIZER_ID, use_fast=True, trust_remote_code=True)
|
| if tok.pad_token is None:
|
| tok.add_special_tokens({"pad_token": "[PAD]"})
|
| VOCAB = max(tok.get_vocab().values()) + 1
|
| BLANK = tok.pad_token_id
|
| EOS = tok.eos_token_id if tok.eos_token_id is not None else tok.sep_token_id
|
|
|
|
|
| def load_joint(ckpt: str, preset: str):
|
| path = _resolve_ckpt(pathlib.Path(ckpt)) or pathlib.Path(ckpt)
|
|
|
| sd = _try_load(path, map_location="cpu")
|
| if sd is None:
|
| raise FileNotFoundError(f"No valid checkpoint at {path}")
|
| _sync_tokenizer_for_checkpoint(sd)
|
| cfg = sd["cfg"] if "cfg" in sd and isinstance(sd["cfg"], dict) else PRESETS[preset]
|
| core = Encoder(cfg).to(DEV)
|
| ar_h, nat_h = ARHead(cfg["d"]).to(DEV), NATHead(cfg["d"]).to(DEV)
|
| sat_h = SATHead(cfg["d"]).to(DEV)
|
| core.load_state_dict(sd["core"])
|
| ar_h.load_state_dict(sd["ar"])
|
| nat_h.load_state_dict(sd["nat"])
|
| sat_h.load_state_dict(sd["sat"])
|
| return core, ar_h, nat_h, sat_h
|
|
|
|
|
| @torch.no_grad()
|
| def ar_decode(core, ar_h, prompt: str, max_new: int, T: float,
|
| repetition_penalty: float = 1.0, rep_window: int = 256):
|
| ids = torch.tensor([tok.encode(prompt)], device=DEV)
|
| if ids.size(1) == 0:
|
| ids = torch.tensor([[EOS] if EOS is not None else [0]], device=DEV)
|
|
|
| h_full, kvs = core(ids, causal_mask(ids.size(1)), use_cache=True)
|
| logits = ar_h(h_full)[:, -1]
|
| recent = deque(ids[0].tolist(), maxlen=max(1, rep_window))
|
| apply_repetition_penalty_(logits, recent, repetition_penalty, rep_window)
|
|
|
| if T <= 1e-6:
|
| nxt = torch.argmax(logits, dim=-1, keepdim=True)
|
| else:
|
| nxt = torch.softmax(logits / T, -1).multinomial(1)
|
|
|
| ids = torch.cat([ids, nxt], 1)
|
| recent.append(nxt.item())
|
| start = time.time()
|
| for _ in range(max(0, max_new - 1)):
|
| x = ids[:, -1:]
|
| h_step, kvs = core(x, None, kv_caches=kvs, use_cache=True)
|
| logits = ar_h(h_step)[:, -1]
|
| apply_repetition_penalty_(logits, recent, repetition_penalty, rep_window)
|
|
|
| if T <= 1e-6:
|
| nxt = torch.argmax(logits, dim=-1, keepdim=True)
|
| else:
|
| nxt = torch.softmax(logits / T, -1).multinomial(1)
|
|
|
| ids = torch.cat([ids, nxt], 1)
|
| recent.append(nxt.item())
|
| print(tok.decode(ids[0].tolist(), skip_special_tokens=True))
|
| print(f"[{max_new} tok in {time.time() - start:.2f}s]")
|
|
|
|
|
| @torch.no_grad()
|
| def sat_decode(core, sat_h, prompt, max_new, T, var,
|
| repetition_penalty: float = 1.0, rep_window: int = 256):
|
| ids = torch.tensor([tok.encode(prompt)], device=DEV)
|
| recent = deque(ids[0].tolist(), maxlen=max(1, rep_window))
|
| added, t0 = 0, time.time()
|
| while added < max_new:
|
| h = core(ids, sat_mask(ids.size(1)))
|
| logits, gate = sat_h(h[:, -SAT_BLOCK:])
|
|
|
| if not var or gate is None:
|
| stride = SAT_BLOCK
|
| else:
|
| gprob = torch.softmax(gate, -1)
|
|
|
| stride = int(torch.multinomial(gprob, 1).item() + 1)
|
|
|
|
|
| logits = logits[:, :stride, :]
|
| for s in range(logits.size(1)):
|
| apply_repetition_penalty_(logits[:, s, :], recent, repetition_penalty, rep_window)
|
|
|
| if T <= 1e-6:
|
| nxt = torch.argmax(logits, dim=-1)
|
| else:
|
| probs = torch.softmax(logits / T, -1)
|
| flat = probs.view(-1, probs.size(-1))
|
| picks = torch.multinomial(flat, 1).view(1, -1)
|
| nxt = picks
|
|
|
| tok_ids = nxt.squeeze(0).tolist()
|
| for tid in tok_ids:
|
| ids = torch.cat([ids, torch.tensor([[tid]], device=ids.device)], 1)
|
| recent.append(tid)
|
| added += 1
|
| if added >= max_new:
|
| break
|
| print(tok.decode(ids[0].tolist(), skip_special_tokens=True))
|
| print(f"[{added} tok in {time.time() - t0:.2f}s]")
|
|
|
|
|
| @torch.no_grad()
|
| def nat_decode(core, nat_h, prompt, max_new, passes, streams):
|
| ids = torch.tensor([tok.encode(prompt) + [BLANK] * (max_new * 2)], device=DEV)
|
| t0 = time.time()
|
| for _ in range(passes):
|
| h = core(ids, None)
|
| logits = nat_h(h)
|
| logits[..., BLANK] = -1e9
|
| cand = logits.topk(streams, -1).indices.permute(2, 0, 1)
|
| best = (cand != BLANK).float().mean(-1).argmax(0)
|
| ids = cand[best, torch.arange(ids.size(0), device=DEV)][:, ::2]
|
| out = [t for t in ids[0].tolist() if t != BLANK]
|
| print(tok.decode(out, skip_special_tokens=True))
|
| print(f"[{len(out)} output tokens in {time.time() - t0:.2f}s]")
|
|
|
|
|
|
|
| def main():
|
| ap = argparse.ArgumentParser()
|
| sub = ap.add_subparsers(dest="cmd", required=True)
|
|
|
| tr = sub.add_parser("train")
|
| tr.add_argument("--preset", choices=PRESETS, default="small")
|
| tr.add_argument("--rank", type=int)
|
| tr.add_argument("--block", type=int, default=DEFAULT_BLOCK)
|
| tr.add_argument("--source", default="cerebras/SlimPajama-627B")
|
| tr.add_argument("--target_tokens", type=int)
|
| tr.add_argument("--steps", type=int)
|
| tr.add_argument("--amp", action="store_true")
|
| tr.add_argument("--save_every_sec", type=int, default=DEFAULT_SAVE_SEC)
|
| tr.add_argument("--save_every_steps", type=int, default=0)
|
| tr.add_argument("--save_dir", default=str(CKDIR))
|
| tr.add_argument("--resume", type=str)
|
| tr.add_argument("--x2", action="store_true", help="~2x params by doubling layers")
|
| tr.add_argument("--warmstart_from", type=str, default=None, help="Path to previous final.pt for shape-safe warm start")
|
| tr.add_argument("--fresh", action="store_true", help="Start from scratch: do not probe or load any checkpoints")
|
|
|
|
|
| tr.add_argument("--auto_grow", action="store_true", help="Automatically grow block size over time")
|
| tr.add_argument("--grow_plan", type=str, default="576,640,768,896,1024", help="Comma list of block sizes to try in order")
|
| tr.add_argument("--grow_every_steps", type=int, default=50000, help="Steps between growth attempts")
|
|
|
| inf = sub.add_parser("infer")
|
| inf.add_argument("--mode", choices=["ar", "nat", "sat"], required=True)
|
| inf.add_argument("--ckpt", required=True)
|
| inf.add_argument("--preset", default="small")
|
| inf.add_argument("--prompt", required=True)
|
| inf.add_argument("--max_new", type=int, default=120)
|
| inf.add_argument("--temperature", type=float, default=1.0)
|
| inf.add_argument("--var", action="store_true")
|
| inf.add_argument("--passes", type=int, default=1)
|
| inf.add_argument("--streams", type=int, default=5)
|
|
|
| inf.add_argument("--repetition_penalty", type=float, default=1.0,
|
| help=">1.0 discourages repeating recently emitted tokens (HF-style; default off)")
|
| inf.add_argument("--rep_window", type=int, default=256,
|
| help="Number of most-recent tokens to penalize (default 256)")
|
|
|
| args = ap.parse_args()
|
| if args.cmd == "train":
|
| train(args)
|
| else:
|
| core, ar_h, nat_h, sat_h = load_joint(args.ckpt, args.preset)
|
| if args.mode == "ar":
|
| ar_decode(core, ar_h, args.prompt, args.max_new, args.temperature,
|
| repetition_penalty=args.repetition_penalty, rep_window=args.rep_window)
|
| elif args.mode == "sat":
|
| sat_decode(core, sat_h, args.prompt, args.max_new, args.temperature, args.var,
|
| repetition_penalty=args.repetition_penalty, rep_window=args.rep_window)
|
| else:
|
| nat_decode(core, nat_h, args.prompt, args.max_new, args.passes, args.streams)
|
|
|
|
|
| if __name__ == "__main__":
|
| main()
|
|
|
