| """Text-only pure-ternary training. |
| |
| For the maintained CLI entrypoint prefer ``python -m arbitor.train`` or the |
| ``arbs-train`` console script after ``pip install -e .``. |
| """ |
| import os, sys, time, math, json, torch |
| sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..")) |
| from torch.utils.tensorboard import SummaryWriter |
| from arbitor import ARBModel, VOCAB, CTX |
| from arbitor.kernel.ternary_scale import TScaleType |
| from arbitor.kernel.ternary_audit import freeze_float_parameters, audit_model, format_audit |
|
|
|
|
| def load_text_data(source, ctx=CTX): |
| """Load byte dataset from .txt or .pt file, or HuggingFace dataset name.""" |
| if source.endswith('.pt'): |
| data = torch.load(source, weights_only=True) |
| elif os.path.isfile(source): |
| data = torch.tensor(list(open(source, 'rb').read()), dtype=torch.long) |
| else: |
| raise ValueError(f"Dataset not found: {source}") |
| n = int(0.9 * len(data)) |
| return data[:n].cuda(), data[n:].cuda() |
|
|
|
|
| if __name__ == "__main__": |
| import argparse |
| parser = argparse.ArgumentParser(description="ARB text training") |
| parser.add_argument("--steps", type=int, default=5000) |
| parser.add_argument("--batch", type=int, default=8) |
| parser.add_argument("--ctx", type=int, default=min(CTX, 128)) |
| parser.add_argument("--data", type=str, default="training/data/tinyshakespeare.txt") |
| parser.add_argument("--run", type=str, default="text") |
| parser.add_argument("--eval-interval", type=int, default=250) |
| parser.add_argument("--log-interval", type=int, default=10) |
| parser.add_argument("--backend", choices=("triton", "torch", "auto", "tilelang"), default="triton") |
| parser.add_argument("--full-text-stack", action="store_true", |
| help="Enable text VQ + KG/MoEGraph + KV attention. Default is byte-text only.") |
| parser.add_argument("--with-vq", action="store_true") |
| parser.add_argument("--with-graph", action="store_true") |
| parser.add_argument("--with-moe", action="store_true") |
| parser.add_argument("--with-attention", action="store_true") |
| args = parser.parse_args() |
| os.environ["ARB_TERNARY_BACKEND"] = args.backend |
| if args.backend == "tilelang" and os.environ.get("ARB_TILELANG_TRAINING", "0").lower() not in {"1", "true", "yes"}: |
| raise ValueError("TileLang BigInt training is unfinished. Use --backend triton for training.") |
|
|
| enable_vq = args.full_text_stack or args.with_vq |
| enable_graph = (args.full_text_stack or args.with_graph) and enable_vq |
| enable_moe = args.full_text_stack or args.with_moe |
| enable_attention = (args.full_text_stack or args.with_attention) and enable_graph |
|
|
| model = ARBModel(enable_image=False, enable_audio=False, |
| enable_vq=enable_vq, enable_graph=enable_graph, |
| enable_memory_modules=False, enable_moe=enable_moe, |
| max_moe_iters=4, |
| enable_attention=enable_attention, |
| enable_output_router=False, |
| enable_video_output=False, |
| enable_talker_output=False).cuda() |
| freeze_float_parameters(model) |
|
|
| print(format_audit(audit_model(model))) |
| print( |
| "Text trainer config: " |
| f"backend={args.backend}, vq={enable_vq}, graph={enable_graph}, " |
| f"moe={enable_moe}, attention={enable_attention}" |
| ) |
| train_params = [p for p in model.parameters() if p.requires_grad] |
| if train_params: |
| raise RuntimeError("Text trainer is pure ternary; use training/finetuning for float LoRA adapters.") |
| print("Pure ternary update path is active") |
|
|
| train_data, val_data = load_text_data(args.data, args.ctx) |
| run_dir = f"models/checkpoints/{args.run}" |
| os.makedirs(run_dir, exist_ok=True) |
| writer = SummaryWriter(run_dir) |
| step = 0 |
| best = float('inf') |
|
|
| while step < args.steps: |
| ix = torch.randint(0, len(train_data) - args.ctx - 1, (args.batch,)) |
| x = torch.stack([train_data[i:i+args.ctx] for i in ix]) |
| t = x[:, 3:] |
|
|
| _, losses, _, _ = model(x, targets=t) |
| model.prepare_ternary_backward(losses.total.detach(), update_scales=True) |
| losses.total.backward() |
|
|
| model._ternary_update_memory(accum_threshold=3, update_scales=True, loss_signal=losses.total) |
| model.zero_grad(set_to_none=True) |
|
|
| step += 1 |
| if args.log_interval and step % args.log_interval == 0: |
| writer.add_scalar("loss/train", losses.total.item(), step) |
| print(f"step {step:>5d} train={losses.total.item():.3f}") |
|
|
| if args.eval_interval and step % args.eval_interval == 0: |
| with torch.no_grad(): |
| ix_v = torch.randint(0, len(val_data) - args.ctx - 1, (args.batch,)) |
| xv = torch.stack([val_data[i:i+args.ctx] for i in ix_v]) |
| tv = xv[:, 3:] |
| _, lv, _, _ = model(xv, targets=tv) |
| val = lv.total.item() |
| writer.add_scalar("loss/train", losses.total.item(), step) |
| writer.add_scalar("loss/eval", val, step) |
| if val < best: |
| best = val |
| torch.save(model.state_dict(), f"{run_dir}/best.pt") |
| print(f"step {step:>5d} train={losses.total.item():.3f} eval={val:.3f} best={best:.3f}") |
|
|
