Blame for scripts/train_bitnet.py

#!/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()