| |
| """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() |
|
|
