#!/usr/bin/env python3 """ Ternary (1.58-bit) SlothLM-E trainer — a drop-in replacement for `reproduce/train_slothlm_e.py` that adds: * BitLinear (ternary {-1,0,+1} weights, absMEDIAN scale per output channel, int8 activations, straight-through estimator) in place of nn.Linear inside every transformer block. * Mixed-precision islands: the syllable embedding, the char head, and the first/last `--fp-boundary` blocks stay full precision (int8 at deploy). * Progressive fp->ternary annealing (`--anneal-frac`). * "Extra RMSNorm before every linear" (SubLN) for ternary stability. * Knowledge distillation from a full-precision SlothLM-E teacher (`--teacher`). Module names MIRROR reproduce/train_slothlm_e.py exactly, so: - a full-precision teacher checkpoint saved by the original script loads into an all-fp SlothE_T for distillation, and - the saved ternary checkpoint carries a `config` dict that the companion gate/export scripts use to rebuild the model. Design rationale + citations: see docs/GUIDE_TERNARY_SLOTHLM_E.md in this repo. Example (on the RTX 5090): python3 train_slothe_ternary.py \ --data train_e_g2pw.bin --vocab syl_vocab.json --tokenizer tokenizer \ --out slothe_t_20m \ --dim 320 --depth 16 --heads 8 --kv-heads 2 --ffn 880 --embed-norm \ --teacher slothe_32m --distill-alpha 0.7 --distill-temp 2.0 \ --anneal-frac 0.15 --weight-quant median --pre-norm \ --batch 384 --epochs 8 --lr 2.5e-3 """ import argparse import json import math import os import numpy as np import torch import torch.distributed as dist import torch.nn as nn import torch.nn.functional as F from torch.utils.data import DataLoader, Dataset # ====================================================================== # Quantization primitives (ternary weights + int8 activations, STE) # ====================================================================== def activation_quant(x, eps=1e-5): """Per-token int8 (absmax) fake-quant; returns a dequantized tensor.""" scale = 127.0 / x.abs().amax(dim=-1, keepdim=True).clamp_(min=eps) return (x * scale).round().clamp_(-128, 127) / scale def weight_quant_ternary(w, mode="median", eps=1e-5): """Ternarize to {-1,0,+1} * scale, per OUTPUT channel. mode="median" uses the absmedian scale (BitNet b1.58 Reloaded, arXiv:2407.09527 — more robust than absmean for SMALL models). mode="mean" is vanilla b1.58. """ if mode == "median": scale = w.abs().median(dim=1, keepdim=True).values.clamp_(min=eps) else: scale = w.abs().mean(dim=1, keepdim=True).clamp_(min=eps) return (w / scale).round().clamp_(-1, 1) * scale def _ste(real, quant, alpha): """Straight-through estimator with anneal blend (alpha in [0,1]).""" blended = real + alpha * (quant - real) return real + (blended - real).detach() class RMSNorm(nn.Module): # param name `.w` matches reproduce/train_slothlm_e.py so teacher weights load def __init__(self, d, eps=1e-6): super().__init__() self.w = nn.Parameter(torch.ones(d)) self.eps = eps def forward(self, x): x = x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps) return x * self.w class BitLinear(nn.Linear): """nn.Linear drop-in. weight_bits='ternary' -> W1.58A8 QAT; 'fp' -> plain fp. pre_norm=True inserts an RMSNorm before the projection (the "extra RMSNorm" stabilizer, arXiv:2505.08823). For 'fp' islands we keep it off so the state dict matches the original SlothE exactly (teacher loading). """ def __init__(self, in_f, out_f, bias=False, weight_bits="ternary", act_bits=8, weight_quant="median", pre_norm=True): super().__init__(in_f, out_f, bias=bias) self.weight_bits = weight_bits self.act_bits = act_bits self.wq_mode = weight_quant self.pre = RMSNorm(in_f) if (pre_norm and weight_bits == "ternary") else None self.register_buffer("quant_alpha", torch.tensor(1.0), persistent=True) def set_quant_alpha(self, a): self.quant_alpha.fill_(float(a)) def forward(self, x): if self.pre is not None: x = self.pre(x) if self.weight_bits == "fp": return F.linear(x, self.weight, self.bias) a = self.quant_alpha # tensor keeps torch.compile stable across anneal if self.act_bits: x = _ste(x, activation_quant(x), a) w = _ste(self.weight, weight_quant_ternary(self.weight, self.wq_mode), a) return F.linear(x, w, self.bias) # ====================================================================== # Model — mirrors reproduce/train_slothlm_e.py, Linear -> BitLinear # ====================================================================== def rope(x, pos, dim): half = dim // 2 freq = 1.0 / (10000 ** (torch.arange(0, half, device=x.device) / half)) ang = pos[:, None].float() * freq[None, :] cos = torch.cat([ang.cos(), ang.cos()], -1)[None, None] sin = torch.cat([ang.sin(), ang.sin()], -1)[None, None] x1, x2 = x[..., :half], x[..., half:] rot = torch.cat([-x2, x1], -1) return x * cos + rot * sin def _lin(in_f, out_f, fp, wq, pre_norm): return BitLinear(in_f, out_f, bias=False, weight_bits="fp" if fp else "ternary", weight_quant=wq, pre_norm=pre_norm) class Attn(nn.Module): def __init__(self, dim, heads, kv, fp, wq, pre_norm): super().__init__() self.h, self.kv, self.dh = heads, kv, dim // heads self.q = _lin(dim, heads * self.dh, fp, wq, pre_norm) self.k = _lin(dim, kv * self.dh, fp, wq, pre_norm) self.v = _lin(dim, kv * self.dh, fp, wq, pre_norm) self.o = _lin(heads * self.dh, dim, fp, wq, pre_norm) self.qn = RMSNorm(self.dh) self.kn = RMSNorm(self.dh) def forward(self, x, pos, amask): B, T, _ = x.shape q = self.q(x).view(B, T, self.h, self.dh).transpose(1, 2) k = self.k(x).view(B, T, self.kv, self.dh).transpose(1, 2) v = self.v(x).view(B, T, self.kv, self.dh).transpose(1, 2) q, k = self.qn(q), self.kn(k) q, k = rope(q, pos, self.dh), rope(k, pos, self.dh) rep = self.h // self.kv k = k.repeat_interleave(rep, 1) v = v.repeat_interleave(rep, 1) o = F.scaled_dot_product_attention(q, k, v, attn_mask=amask[:, None, None, :]) o = o.transpose(1, 2).reshape(B, T, -1) return self.o(o) class SwiGLU(nn.Module): def __init__(self, dim, hidden, fp, wq, pre_norm): super().__init__() self.w1 = _lin(dim, hidden, fp, wq, pre_norm) self.w3 = _lin(dim, hidden, fp, wq, pre_norm) self.w2 = _lin(hidden, dim, fp, wq, pre_norm) def forward(self, x): return self.w2(F.silu(self.w1(x)) * self.w3(x)) class Block(nn.Module): def __init__(self, dim, heads, kv, ffn, fp, wq, pre_norm): super().__init__() self.n1 = RMSNorm(dim) self.attn = Attn(dim, heads, kv, fp, wq, pre_norm) self.n2 = RMSNorm(dim) self.ffn = SwiGLU(dim, ffn, fp, wq, pre_norm) def forward(self, x, pos, amask): x = x + self.attn(self.n1(x), pos, amask) x = x + self.ffn(self.n2(x)) return x class SlothE_T(nn.Module): """Ternary SlothLM-E. embed + head + first/last `fp_boundary` blocks stay fp.""" def __init__(self, n_syl, n_char, dim=320, depth=16, heads=8, kv=2, ffn=880, embed_norm=True, weight_bits="ternary", weight_quant="median", fp_boundary=1, pre_norm=True, act_bits=8, char_hints=False, tie_hints=False): super().__init__() self.embed = nn.Embedding(n_syl, dim) # fp island self.hint_tied = char_hints and tie_hints self.hint_embed = nn.Embedding(n_char + 1, dim, padding_idx=0) if (char_hints and not tie_hints) else None self.hint_none = nn.Parameter(torch.zeros(dim)) if self.hint_tied else None self.embed_norm = RMSNorm(dim) if embed_norm else None blocks = [] for i in range(depth): is_fp = (weight_bits == "fp") or i < fp_boundary or i >= depth - fp_boundary blocks.append(Block(dim, heads, kv, ffn, is_fp, weight_quant, pre_norm)) self.blocks = nn.ModuleList(blocks) self.norm = RMSNorm(dim) self.head = nn.Linear(dim, n_char, bias=False) # fp island (int8 at deploy) self.apply(self._init) def _init(self, m): if isinstance(m, nn.Linear): nn.init.normal_(m.weight, std=0.02) elif isinstance(m, nn.Embedding): nn.init.normal_(m.weight, std=0.02) def set_quant_alpha(self, a): for m in self.modules(): if isinstance(m, BitLinear): m.set_quant_alpha(a) def forward(self, syl, amask, hints=None): pos = torch.arange(syl.shape[1], device=syl.device) x = self.embed(syl) if hints is not None and self.hint_tied: has = (hints > 0).unsqueeze(-1).to(x.dtype) he = F.embedding((hints - 1).clamp(min=0), self.head.weight) x = x + he * has + self.hint_none * (1.0 - has) elif self.hint_embed is not None and hints is not None: x = x + self.hint_embed(hints) if self.embed_norm is not None: x = self.embed_norm(x) for b in self.blocks: x = b(x, pos, amask) return self.head(self.norm(x)) # ====================================================================== # Data (reuses the AlignedBin format from the original pipeline) # ====================================================================== class AlignedBin(Dataset): CTX_MAX = 12 def __init__(self, path, ctx_p=0.0, typo=None, typo_p=0.0, teacher_path=None): self.data = np.fromfile(path, dtype=np.uint16) self.tdata = np.fromfile(teacher_path, dtype=np.uint16) if teacher_path else None self.ctx_p, self.typo, self.typo_p = ctx_p, typo, typo_p self.idx = [] i, d = 0, self.data while i < len(d): n = int(d[i]); self.idx.append((i + 1, n)); i += 1 + 2 * n print(f"{len(self.idx)} aligned pairs") def __len__(self): return len(self.idx) def __getitem__(self, k): s, n = self.idx[k] syl = self.data[s:s + n].astype(np.int64) tsyl = (self.tdata[s:s + n].astype(np.int64) if self.tdata is not None else syl.copy()) chr_ = self.data[s + n:s + 2 * n].astype(np.int64) chr_[chr_ == 65535] = -100 if self.typo is not None and self.typo_p > 0 and np.random.rand() < self.typo_p: zh = np.nonzero(chr_ >= 0)[0] if len(zh): i = int(np.random.choice(zh)); nb = self.typo.get(int(syl[i])) if nb: syl[i] = int(nb[np.random.randint(len(nb))]) forced = np.zeros(n, dtype=np.int64) if self.ctx_p > 0 and k >= 2 and np.random.rand() < self.ctx_p: ps, pn = self.idx[k - 2] pch = self.data[ps + pn:ps + 2 * pn].astype(np.int64) pch = pch[pch != 65535][-self.CTX_MAX:] if len(pch): L = len(pch) syl = np.concatenate([np.zeros(L, dtype=np.int64), syl]) tsyl = np.concatenate([np.zeros(L, dtype=np.int64), tsyl]) chr_ = np.concatenate([np.full(L, -100, dtype=np.int64), chr_]) forced = np.concatenate([pch + 1, forced]) return (torch.from_numpy(syl), torch.from_numpy(tsyl), torch.from_numpy(chr_), torch.from_numpy(forced)) def collate(batch, pad=0): n = max(len(s) for s, _, _, _ in batch); B = len(batch) syl = torch.zeros(B, n, dtype=torch.long) tsyl = torch.zeros(B, n, dtype=torch.long) chr_ = torch.full((B, n), -100, dtype=torch.long) mask = torch.zeros(B, n, dtype=torch.bool) forced = torch.zeros(B, n, dtype=torch.long) for i, (sq, tq, c, f) in enumerate(batch): syl[i, :len(sq)] = sq; tsyl[i, :len(tq)] = tq chr_[i, :len(c)] = c; mask[i, :len(sq)] = True; forced[i, :len(f)] = f return syl, tsyl, chr_, mask, forced def masked_ce(logits, target, syl, legal_mask, eps): """CE + label-smoothing restricted to the per-position legal char set (mirrors serve-time masking): the ternary net optimizes only the ~1-50 legal chars eval scores on, not all 8342. Positions whose gold char is not legal (train/serve skew) are dropped so the loss never goes -inf. logits (B,T,V); target (B,T,-100 pad); syl (B,T) input syllable ids; legal_mask (n_syl,V) bool.""" lm = legal_mask[syl] logits = logits.masked_fill(~lm, float("-inf")) logp = F.log_softmax(logits, dim=-1) tgt = target.clamp(min=0) gold_legal = lm.gather(-1, tgt.unsqueeze(-1)).squeeze(-1) valid = (target != -100) & gold_legal ce_hard = -logp.gather(-1, tgt.unsqueeze(-1)).squeeze(-1) nlegal = lm.sum(-1).clamp(min=1) ls_term = -(logp.masked_fill(~lm, 0.0)).sum(-1) / nlegal loss = (1.0 - eps) * ce_hard + eps * ls_term return loss[valid].mean() def load_teacher(path, dev): """Load a frozen all-fp SlothE_T teacher (KD). Mirrors gate build().""" ck = torch.load(os.path.join(path, "slothe.pt"), map_location="cpu") c = ck["config"] m = SlothE_T(c["n_syl"], c["n_char"], c["dim"], c["depth"], c["heads"], c["kv"], c["ffn"], embed_norm=c.get("embed_norm", False), weight_bits=c.get("weight_bits") or "fp", weight_quant=c.get("weight_quant", "median"), fp_boundary=c.get("fp_boundary", 1), pre_norm=c.get("pre_norm", True), act_bits=c.get("act_bits", 8), char_hints=c.get("char_hints", False), tie_hints=c.get("tie_hints", False)) m.load_state_dict(ck["model"], strict=False) # teacher predates quant_alpha buffers (unused for fp) m.set_quant_alpha(1.0) m.eval().to(dev) for prm in m.parameters(): prm.requires_grad_(False) return m def distill_loss(logits, t_logits, target, alpha, temp): """(1-alpha)*CE + alpha*KL(student||teacher) at temperature `temp`.""" V = logits.shape[-1] flat = target.reshape(-1) ce = F.cross_entropy(logits.reshape(-1, V), flat, ignore_index=-100) keep = flat != -100 s = logits.reshape(-1, V)[keep] t = t_logits.reshape(-1, V)[keep] kl = F.kl_div(F.log_softmax(s / temp, -1), F.softmax(t / temp, -1), reduction="batchmean") * (temp * temp) return (1.0 - alpha) * ce + alpha * kl, ce.detach(), kl.detach() def _dachen_key_adjacency(): """Bopomofo symbol -> set of symbols on physically 8-adjacent keys, per the TAAI-2024 keyboard 8-adjacency error model (Li/Yeh/Chang, 大城市 = 284t/6g4): a real mis-key lands on one of the 8 neighbours in the QWERTY 3x3 block ('h' -> t/y/u/g/j/b/n/m). Used to make typo-noise reflect finger slips instead of arbitrary edit-distance-1 substitutions.""" rows = ["1234567890-=", "qwertyuiop[]", "asdfghjkl;'", "zxcvbnm,./"] pos = {k: (r, c) for r, row in enumerate(rows) for c, k in enumerate(row)} key_nb = {} for k, (r, c) in pos.items(): key_nb[k] = {k2 for k2, (r2, c2) in pos.items() if abs(r2 - r) <= 1 and abs(c2 - c) <= 1 and k2 != k} DACHEN = {"1": "\u3105", "q": "\u3106", "a": "\u3107", "z": "\u3108", "2": "\u3109", "w": "\u310a", "s": "\u310b", "x": "\u310c", "e": "\u310d", "d": "\u310e", "c": "\u310f", "r": "\u3110", "f": "\u3111", "v": "\u3112", "5": "\u3113", "t": "\u3114", "g": "\u3115", "b": "\u3116", "y": "\u3117", "h": "\u3118", "n": "\u3119", "u": "\u3127", "j": "\u3128", "m": "\u3129", "8": "\u311a", "i": "\u311b", "k": "\u311c", ",": "\u311d", "9": "\u311e", "o": "\u311f", "l": "\u3120", ".": "\u3121", "0": "\u3122", "p": "\u3123", ";": "\u3124", "/": "\u3125", "-": "\u3126"} sym_key = {v: k for k, v in DACHEN.items()} return {sym: {DACHEN[k2] for k2 in key_nb.get(k, ()) if k2 in DACHEN} for sym, k in sym_key.items()} def build_typo_neighbors(syl_vocab): TONES = set("ˊˇˋ˙"); by_key = {} sym_adj = _dachen_key_adjacency() # TAAI 8-adjacency slip model for sy, i in syl_vocab.items(): if sy.startswith("<"): continue tone = "".join(c for c in sy if c in TONES); base = tuple(c for c in sy if c not in TONES) by_key.setdefault((tone, len(base)), []).append((base, i)) out = {} for sy, i in syl_vocab.items(): if sy.startswith("<"): continue tone = "".join(c for c in sy if c in TONES); base = tuple(c for c in sy if c not in TONES); nb = [] for dl in (-1, 0, 1): for other, j in by_key.get((tone, len(base) + dl), []): if j == i: continue a, b = base, other if len(a) == len(b): diffs = [(x, y) for x, y in zip(a, b) if x != y] if len(diffs) == 1 and diffs[0][1] in sym_adj.get(diffs[0][0], ()): nb.append(j) # only physically-adjacent-key slips else: s1, s2 = (a, b) if len(a) < len(b) else (b, a) for pp in range(len(s2)): if s2[:pp] + s2[pp+1:] == s1: nb.append(j); break if nb: out[i] = nb return out def main(): ap = argparse.ArgumentParser() ap.add_argument("--data", required=True) ap.add_argument("--vocab", default="syl_vocab.json") ap.add_argument("--tokenizer", default="tokenizer") ap.add_argument("--out", default="slothe_t") # architecture (recommended ~20M ternary: dim320 depth16 heads8 kv2 ffn880) ap.add_argument("--dim", type=int, default=320) ap.add_argument("--depth", type=int, default=16) ap.add_argument("--heads", type=int, default=8) ap.add_argument("--kv-heads", type=int, default=2) ap.add_argument("--ffn", type=int, default=880) ap.add_argument("--embed-norm", action="store_true") # quantization ap.add_argument("--quant", choices=["ternary", "fp"], default="ternary") ap.add_argument("--weight-quant", choices=["median", "mean"], default="median") ap.add_argument("--fp-boundary", type=int, default=1) ap.add_argument("--pre-norm", action="store_true", default=True) ap.add_argument("--no-pre-norm", dest="pre_norm", action="store_false") ap.add_argument("--act-bits", type=int, default=8) ap.add_argument("--anneal-frac", type=float, default=0.15, help="ramp quant_alpha 0->1 over this fraction of steps") # distillation ap.add_argument("--teacher", default="", help="dir with a fp teacher slothe.pt") ap.add_argument("--teacher-data", default="", help="exact-syllable .bin for cross-input distillation (teacher input)") ap.add_argument("--distill-alpha", type=float, default=0.7) ap.add_argument("--distill-temp", type=float, default=2.0) # optimization (retuned for small ternary: higher LR ok, see 2407.09527) ap.add_argument("--batch", type=int, default=384) ap.add_argument("--epochs", type=float, default=8.0) ap.add_argument("--lr", type=float, default=2.5e-3) ap.add_argument("--weight-decay", type=float, default=0.1) ap.add_argument("--steps", type=int, default=0) ap.add_argument("--resume", default="") ap.add_argument("--char-hints", action="store_true") ap.add_argument("--tie-hints", action="store_true") ap.add_argument("--typo-noise", type=float, default=0.0) ap.add_argument("--compile", action="store_true", help="torch.compile the model (compile BEFORE DDP wrap)") ap.add_argument("--compile-mode", default="reduce-overhead", help="reduce-overhead | cudagraphs | default") ap.add_argument("--legal-mask", default="", help="npz (n_syl,V) legal-char mask -> mask CE + LS to the legal set") ap.add_argument("--label-smoothing", type=float, default=0.0, help="teacher-free soft targets for direct (pure-CE) training") ap.add_argument("--context", type=float, default=0.0) ap.add_argument("--save-every", type=int, default=0, help="snapshot {out}_ep{N} every N epochs (for accuracy-vs-epoch curve)") args = ap.parse_args() from transformers import AutoTokenizer tok = AutoTokenizer.from_pretrained(args.tokenizer) syl_vocab = json.load(open(args.vocab, encoding="utf-8")) ddp = int(os.environ.get("WORLD_SIZE", 1)) > 1 if ddp: dist.init_process_group("nccl") local_rank = int(os.environ["LOCAL_RANK"]); rank = dist.get_rank(); world = dist.get_world_size() torch.cuda.set_device(local_rank); dev = f"cuda:{local_rank}" else: rank, local_rank, world = 0, 0, 1 dev = "cuda" if torch.cuda.is_available() else "cpu" is_main = (rank == 0) if dev.startswith("cuda"): torch.backends.cuda.matmul.allow_tf32 = True torch.backends.cudnn.allow_tf32 = True torch.backends.cudnn.benchmark = True typo_map = build_typo_neighbors(syl_vocab) if args.typo_noise > 0 else None ds = AlignedBin(args.data, ctx_p=args.context, typo=typo_map, typo_p=args.typo_noise, teacher_path=(args.teacher_data or None)) sampler = torch.utils.data.DistributedSampler(ds, shuffle=True) if ddp else None dl = DataLoader(ds, batch_size=args.batch, shuffle=(sampler is None), sampler=sampler, num_workers=8, collate_fn=collate, pin_memory=True, drop_last=True) model = SlothE_T(len(syl_vocab), len(tok), args.dim, args.depth, args.heads, args.kv_heads, args.ffn, embed_norm=args.embed_norm, weight_bits=args.quant, weight_quant=args.weight_quant, fp_boundary=args.fp_boundary, pre_norm=args.pre_norm, act_bits=args.act_bits, char_hints=args.char_hints, tie_hints=args.tie_hints).to(dev) if args.resume: model.load_state_dict(torch.load(args.resume, map_location=dev)["model"]) if int(os.environ.get("RANK","0"))==0: print("resumed from "+args.resume, flush=True) np_ = sum(p.numel() for p in model.parameters()) print(f"SlothLM-E-T {np_/1e6:.1f}M params | quant={args.quant} " f"wq={args.weight_quant} fp_boundary={args.fp_boundary} pre_norm={args.pre_norm}") teacher = load_teacher(args.teacher, dev) if args.teacher else None legal_mask = (torch.from_numpy(np.load(args.legal_mask)["mask"]).to(dev) if args.legal_mask else None) raw_model = model if args.compile: model = (torch.compile(model, backend="cudagraphs") if args.compile_mode == "cudagraphs" else torch.compile(model, mode=args.compile_mode)) if ddp: model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[local_rank], broadcast_buffers=False) opt = torch.optim.AdamW(model.parameters(), lr=args.lr, betas=(0.9, 0.95), weight_decay=args.weight_decay) total = args.steps or int(len(dl) * args.epochs) sched = torch.optim.lr_scheduler.OneCycleLR(opt, args.lr, total_steps=total, pct_start=0.03) anneal_steps = max(1, int(args.anneal_frac * total)) model.train() step = 0 for ep in range(math.ceil(args.epochs)): if ddp: sampler.set_epoch(ep) for syl, tsyl, chr_, mask, forced in dl: syl, tsyl, chr_, mask, forced = (syl.to(dev), tsyl.to(dev), chr_.to(dev), mask.to(dev), forced.to(dev)) hints = None if args.char_hints: reveal_p = torch.rand(syl.shape[0], 1, device=dev) * 0.3 reveal = (torch.rand(syl.shape, device=dev) < reveal_p) & (chr_ >= 0) hints = torch.where(reveal, chr_ + 1, torch.zeros_like(chr_)) hints = torch.where(forced > 0, forced, hints) raw_model.set_quant_alpha(1.0 if args.resume else min(1.0, step / anneal_steps)) with torch.autocast(dev, dtype=torch.bfloat16, enabled=dev == "cuda"): logits = model(syl, mask, hints) if teacher is not None: with torch.no_grad(): t_logits = teacher(tsyl, mask, hints) # exact syllables loss, ce, kl = distill_loss(logits, t_logits, chr_, args.distill_alpha, args.distill_temp) elif legal_mask is not None: loss = masked_ce(logits, chr_, syl, legal_mask, args.label_smoothing) ce = kl = loss.detach() else: loss = F.cross_entropy(logits.reshape(-1, logits.shape[-1]), chr_.reshape(-1), ignore_index=-100, label_smoothing=args.label_smoothing) ce = kl = loss.detach() opt.zero_grad(set_to_none=True) loss.backward() torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) opt.step(); sched.step(); step += 1 if step % 50 == 0 and is_main: a = min(1.0, step / anneal_steps) print(f"step {step}/{total} loss {loss.item():.3f} " f"ce {ce.item():.3f} kl {kl.item():.3f} alpha {a:.2f}", flush=True) if step >= total: break if args.save_every and (ep + 1) % args.save_every == 0: if is_main: _snap = f"{args.out}_ep{ep + 1}" os.makedirs(_snap, exist_ok=True) torch.save({"model": raw_model.state_dict(), "config": {"n_syl": len(syl_vocab), "n_char": len(tok), "dim": args.dim, "depth": args.depth, "heads": args.heads, "kv": args.kv_heads, "ffn": args.ffn, "embed_norm": args.embed_norm, "weight_bits": args.quant, "weight_quant": args.weight_quant, "fp_boundary": args.fp_boundary, "pre_norm": args.pre_norm, "act_bits": args.act_bits, "char_hints": args.char_hints, "tie_hints": args.tie_hints}}, os.path.join(_snap, "slothe.pt")) json.dump(syl_vocab, open(os.path.join(_snap, "syl_vocab.json"), "w", encoding="utf-8"), ensure_ascii=False) print(f"snapshot saved to {_snap}", flush=True) if ddp: dist.barrier() if step >= total: break raw_model.set_quant_alpha(1.0) if ddp: dist.barrier() if not is_main: if ddp: dist.destroy_process_group() return os.makedirs(args.out, exist_ok=True) torch.save({"model": raw_model.state_dict(), "config": {"n_syl": len(syl_vocab), "n_char": len(tok), "dim": args.dim, "depth": args.depth, "heads": args.heads, "kv": args.kv_heads, "ffn": args.ffn, "embed_norm": args.embed_norm, "weight_bits": args.quant, "weight_quant": args.weight_quant, "fp_boundary": args.fp_boundary, "pre_norm": args.pre_norm, "act_bits": args.act_bits, "char_hints": args.char_hints, "tie_hints": args.tie_hints}}, os.path.join(args.out, "slothe.pt")) _sd = raw_model.state_dict() json.dump(syl_vocab, open(os.path.join(args.out, "syl_vocab.json"), "w", encoding="utf-8"), ensure_ascii=False) print(f"saved to {args.out}") if __name__ == "__main__": main()