Zenbit

TPU sharding and quantization utilities for PyTorch/XLA training. Built as a companion to ZAdapter, but usable standalone for any PyTorch/XLA training setup.

Install

pip install zenbit[tpu]

The [tpu] extra pulls in torch_xla. Without it, quantization utilities (int8, NF4) still work on CPU/GPU for offline testing — only the sharding/tensor-parallel pieces require an actual TPU runtime.

What's inside

• zenbit.int8 — uniform int8 weight quantization (~50% memory savings vs fp16)
• zenbit.nf4 — NF4 blockwise quantization with double quantization, QLoRA-style (~75-89% memory savings vs fp32)
• zenbit.sharding — data-parallel training loop utilities: device detection, dataset sharding, gradient sync, checkpointing, logging
• zenbit.tensor_parallel — SPMD tensor parallel sharding (Megatron-style column/row parallel linear layers) for models too large to replicate on a single TPU core
• zenbit.pallas_nf4 (optional, pip install zenbit[pallas]) — genuinely fused NF4 dequant+matmul as a custom Pallas TPU kernel (JAX backend). Weights are streamed compressed into VMEM and dequantized on-chip, never materializing the full-precision weight in HBM — closer to how bitsandbytes fuses dequant+matmul on GPU. Includes a fused backward pass (custom VJP, dL/dX computed without materializing the dequantized weight — dL/dW is not computed, matching a frozen quantized base weight), double quantization for the scale factors, tiling-size helpers (suggest_tile_sizes), and an optional Flax nn.Module drop-in layer (zenbit.pallas_nf4.flax_layer.NF4DenseFused, pip install zenbit[pallas-flax]). Verified on Kaggle TPU (v5e), forward and backward both bit-exact (max_diff=0.0) vs reference; not currently verified on Google Colab TPU due to a Mosaic module version mismatch in Colab's pinned jax/jaxlib/libtpu combination (unrelated to kernel correctness).

Quick example

from transformers import AutoModelForCausalLM
from zadapter import inject_adapter
from zenbit import quantize_model_nf4, NF4Config, train_loop, TrainConfig

model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen2.5-Coder-7B")
model = quantize_model_nf4(model, NF4Config(double_quant=True))
model = inject_adapter(model, r=64)

train_loop(model, train_dataset, loss_fn, TrainConfig(lr=1e-4, num_epochs=3))

For models too large for a single TPU core (30B+), combine with tensor parallel:

from zenbit import setup_spmd_mesh, build_tp_zadapter_model, TPConfig

mesh = setup_spmd_mesh(TPConfig(num_shards=8))
model = build_tp_zadapter_model(model, mesh)

Design notes

• Why data-parallel by default: adapter parameters are small (~0.5-3% of the model), so gradient sync overhead is minimal. Data parallel keeps communication to one all-reduce per step, avoiding the per-layer communication overhead of full model parallelism.
• Why tensor parallel is opt-in: it's necessary once the base model no longer fits on a single TPU core's HBM (roughly 30B+ params even with NF4 quantization), but comes with real communication overhead per layer. Use it only when the model size requires it.

License

MIT