#!/usr/bin/env python3 """Train a tiny native-ternary BitNet b1.58-style causal LM on CPU. Weights are ternary {-1,0,+1} in the forward pass (absmean, per tensor) with straight-through gradients; activations are 8-bit absmax per token. Master weights are kept in FP32 for the optimizer. The exported GGUF stores the ternary weights (TQ1_0/TQ2_0). """ import argparse, json, math, os, random, time from pathlib import Path import numpy as np import sentencepiece as spm import torch import torch.nn as nn import torch.nn.functional as F from transformers import BitNetConfig, BitNetForCausalLM torch.set_num_threads(int(os.environ.get('TORCH_THREADS', '2'))) try: torch.set_num_interop_threads(1) except RuntimeError: pass torch.manual_seed(42) random.seed(42) np.random.seed(42) def weight_quant(w): scale = w.abs().mean().clamp(min=1e-5) q = (w / scale).round().clamp(-1, 1) * scale return w + (q - w).detach() def act_quant(x): scale = x.abs().amax(dim=-1, keepdim=True).clamp(min=1e-5) / 127.0 q = (x / scale).round().clamp(-127, 127) * scale return x + (q - x).detach() class BitLinear(nn.Module): def __init__(self, lin: nn.Linear): super().__init__() self.weight = nn.Parameter(lin.weight.detach().clone()) self.bias = None if lin.bias is None else nn.Parameter(lin.bias.detach().clone()) def forward(self, x): return F.linear(act_quant(x), weight_quant(self.weight), self.bias) def bitnetize(model): for layer in model.model.layers: for attn in [layer.self_attn]: for name in ['q_proj', 'k_proj', 'v_proj', 'o_proj']: setattr(attn, name, BitLinear(getattr(attn, name))) for name in ['gate_proj', 'up_proj', 'down_proj']: setattr(layer.mlp, name, BitLinear(getattr(layer.mlp, name))) return model def load_tokens(path, sp): import array bin_path = str(path) + '.u16.bin' if not os.path.exists(bin_path): arr = array.array('H') with open(path, encoding='utf-8') as f: for line in f: line = line.strip() if not line: continue arr.append(sp.bos_id()) arr.extend(sp.encode(line, out_type=int)) arr.append(sp.eos_id()) with open(bin_path, 'wb') as f: arr.tofile(f) del arr return np.memmap(bin_path, dtype=np.uint16, mode='r') def get_batch(tokens, bs, seq, device): ix = np.random.randint(0, len(tokens) - seq - 1, size=(bs,)) idx = ix[:, None] + np.arange(seq + 1)[None, :] batch = torch.from_numpy(tokens[idx].astype(np.int64)).to(device) return batch[:, :-1], batch[:, 1:] @torch.no_grad() def evaluate(model, tokens, bs, seq, device, n=16): model.eval() losses = [] for _ in range(n): x, y = get_batch(tokens, bs, seq, device) logits = model(x).logits losses.append(F.cross_entropy(logits.reshape(-1, logits.size(-1)), y.reshape(-1)).item()) model.train() return float(np.mean(losses)) @torch.no_grad() def sample(model, sp, device, prompt='User: Hello Assistant:', max_new_tokens=60): model.eval() ids = [sp.bos_id()] + sp.encode(prompt, out_type=int) x = torch.tensor([ids], dtype=torch.long, device=device) out = model.generate(x, max_new_tokens=max_new_tokens, do_sample=True, temperature=0.8, top_p=0.9, eos_token_id=sp.eos_id(), pad_token_id=sp.pad_id()) model.train() return sp.decode(out[0].tolist()) def main(): ap = argparse.ArgumentParser() ap.add_argument('--data', required=True) ap.add_argument('--out', required=True) ap.add_argument('--steps', type=int, default=6000) ap.add_argument('--batch-size', type=int, default=16) ap.add_argument('--seq', type=int, default=128) ap.add_argument('--lr', type=float, default=1e-3) ap.add_argument('--warmup', type=int, default=200) ap.add_argument('--hidden', type=int, default=256) ap.add_argument('--layers', type=int, default=4) ap.add_argument('--intermediate', type=int, default=512) ap.add_argument('--heads', type=int, default=4) ap.add_argument('--kv-heads', type=int, default=2) ap.add_argument('--eval-every', type=int, default=250) ap.add_argument('--log-every', type=int, default=25) args = ap.parse_args() data = Path(args.data); out = Path(args.out); out.mkdir(parents=True, exist_ok=True) device = torch.device('cpu') sp = spm.SentencePieceProcessor(model_file=str(data / 'tokenizer.model')) train_tokens = load_tokens(data / 'formatted_train.txt', sp) val_tokens = load_tokens(data / 'formatted_val.txt', sp) print('train tokens', len(train_tokens), 'val tokens', len(val_tokens), flush=True) cfg = BitNetConfig( vocab_size=sp.get_piece_size(), hidden_size=args.hidden, intermediate_size=args.intermediate, num_hidden_layers=args.layers, num_attention_heads=args.heads, num_key_value_heads=args.kv_heads, max_position_embeddings=args.seq, hidden_act='relu2', rms_norm_eps=1e-5, rope_theta=10000.0, tie_word_embeddings=True, attention_bias=False, attention_dropout=0.0, bos_token_id=sp.bos_id(), eos_token_id=sp.eos_id(), pad_token_id=sp.pad_id(), initializer_range=0.02, ) model = bitnetize(BitNetForCausalLM(cfg)).to(device) n_params = sum(p.numel() for p in model.parameters()) print('parameters', n_params, flush=True) decay, no_decay = [], [] for name, p in model.named_parameters(): if p.ndim < 2 or 'embed' in name or 'norm' in name: no_decay.append(p) else: decay.append(p) opt = torch.optim.AdamW([ {'params': decay, 'weight_decay': 0.01}, {'params': no_decay, 'weight_decay': 0.0}, ], lr=args.lr, betas=(0.9, 0.95)) def lr_at(step): if step < args.warmup: return args.lr * (step + 1) / args.warmup t = (step - args.warmup) / max(1, args.steps - args.warmup) return args.lr * (0.1 + 0.9 * 0.5 * (1 + math.cos(math.pi * t))) best = float('inf') history = [] t0 = time.time() model.train() for step in range(args.steps): lr = lr_at(step) for g in opt.param_groups: g['lr'] = lr x, y = get_batch(train_tokens, args.batch_size, args.seq, device) logits = model(x).logits loss = F.cross_entropy(logits.reshape(-1, logits.size(-1)), y.reshape(-1)) opt.zero_grad(set_to_none=True) loss.backward() torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) opt.step() if step % args.log_every == 0 or step == args.steps - 1: dt = time.time() - t0 tok_s = (step + 1) * args.batch_size * args.seq / max(dt, 1e-9) print(f'step {step:5d} loss {loss.item():.4f} lr {lr:.2e} tok/s {tok_s:.0f}', flush=True) if step % args.eval_every == 0 or step == args.steps - 1: val = evaluate(model, val_tokens, args.batch_size, args.seq, device) txt = sample(model, sp, device) rec = {'step': step, 'train_loss': loss.item(), 'val_loss': val, 'sample': txt, 'elapsed_s': time.time() - t0} history.append(rec) print('EVAL', json.dumps(rec, ensure_ascii=False), flush=True) if val < best: best = val ref = BitNetForCausalLM(cfg) ref.load_state_dict(model.state_dict(), strict=True) ref.save_pretrained(out / 'best_model') for f in ['tokenizer.model', 'tokenizer.vocab', 'tokenizer_config.json', 'special_tokens_map.json', 'chat_template.jinja']: src = data / f if src.exists(): (out / 'best_model' / f).write_bytes(src.read_bytes()) (out / 'best.json').write_text(json.dumps(rec, indent=2, ensure_ascii=False)) (out / 'history.json').write_text(json.dumps(history, indent=2, ensure_ascii=False)) print('DONE best_val', best, flush=True) if __name__ == '__main__': main()