Metadata-Version: 2.4
Name: zenbit
Version: 0.2.0
Summary: TPU sharding (data-parallel, tensor-parallel) and quantization (int8, NF4) utilities for PyTorch/XLA training
Author-email: Riki <alkenzyhaikal@gmail.com>
License: MIT
Project-URL: Homepage, https://github.com/RikZD/Zenbit
Project-URL: Repository, https://github.com/RikZD/Zenbit
Keywords: pytorch,tpu,xla,quantization,nf4,sharding,tensor-parallel,qlora
Classifier: Development Status :: 3 - Alpha
Classifier: Intended Audience :: Developers
Classifier: Intended Audience :: Science/Research
Classifier: License :: OSI Approved :: MIT License
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3.9
Classifier: Programming Language :: Python :: 3.10
Classifier: Programming Language :: Python :: 3.11
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Requires-Python: >=3.9
Description-Content-Type: text/markdown
Requires-Dist: torch>=2.0
Requires-Dist: zadapter>=0.1.0
Provides-Extra: tpu
Requires-Dist: torch_xla>=2.0; extra == "tpu"

# Zenbit

TPU sharding and quantization utilities for PyTorch/XLA training. Built as
a companion to [ZAdapter](https://pypi.org/project/zadapter/), but usable
standalone for any PyTorch/XLA training setup.

## Install

```bash
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

## Quick example

```python
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:

```python
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
