Bantam (bantamtrain)

Train small language models and weather models faster and cheaper on limited compute, without throwing away data. Made by Nathan, Jamaica.

Bantam is the library behind a set of measured results on real data (TinyStories from HuggingFace, ERA5 from WeatherBench2), all obtained in a 4-core CPU sandbox; the same code runs on a GPU with mixed precision and torch.compile.

pip install bantamtrain

What is inside

Quick start

import torch
from bantamtrain.models import ZoomGPT
from bantamtrain.optim import build_optimizer
from bantamtrain.zoom import ZoomSchedule, lr_at

model = ZoomGPT(vocab=4096, d=256, n_layer=6, n_head=4, zoom_ks=[2])
opt = build_optimizer(model, "muon", lr=1.6e-2, weight_decay=0.1, state_bits=8)
sched = ZoomSchedule.parse("2:0.5,1:0.5")      # k=2 for the first half of the budget

x = torch.randint(0, 4096, (16, 258))
for frac in (0.1, 0.6):                         # fraction of the compute budget spent
    k = sched.k_at(frac)
    for g in opt.param_groups:
        g["lr"] = lr_at(frac, 1.6e-2, sched="wsd")
    loss = model(x[:, :256], x[:, k:256 + k], k=k)
    opt.zero_grad(); loss.backward(); opt.step()

Store a trained model small:

from bantamtrain import qweights, weightpack
weightpack.save("model.zwpk", model.state_dict(), bits=8)    # int8 checkpoint file
qweights.quantize_model(model, bits=4)                        # int4 weights in RAM, float compute

Size a run on one GPU:

python -m bantamtrain.memory --params 5e9 --gpu_gb 96 --tokens 1e11

Honest notes

• Muon, 8-bit optimizer state, stochastic-rounding bf16 and blockwise weight quantization are published techniques re-implemented here; Zoom Training and the storage formats are this project's own.
• All quality numbers above were measured on a 5.8M-parameter model; gains at larger scale are expected to be smaller for the optimizer and are untested for Zoom. Nothing has been run on a GPU yet.
• Source, experiments, tests and the full write-up: https://github.com/Wisk-inc/AboxRebuildcode