This repository is a fork of Microsoft BitNet. Custom NPU128 bare-metal kernels, Triton-fused QAT, hardened inference pipeline by Huntwter.
- What is bitone?
- How 1.58-bit Works
- Features
- Supported Pre-trained Models
- Installation
- Quick Start
- Architecture
- Security
- Credits
- License
bitone is an engine for running and training 1.58-bit large language models. Every weight in the neural network is one of exactly three values: -1, 0, or +1.
This means:
- No multiplication needed — multiplying by -1 is flipping a sign, by 0 is skipping, by +1 is nothing. The entire model runs on addition and subtraction.
- Tiny memory footprint — a 3B parameter model fits in 600 MB instead of 6 GB.
- Runs on weak hardware — no need for expensive GPUs or modern instruction sets.
In a normal neural network, every weight is a 16-bit or 32-bit float. In BitNet 1.58-bit, every weight is constrained to only three values:
Standard: -0.4821, 0.1337, 1.0042, -0.0003, ...
BitNet: -1, 0, +1, 0, ...
Each weight needs only ~1.58 bits (log₂(3) = 1.585). The model must be trained from scratch with these constraints — you cannot convert a normal model to 1.58-bit.
import bitone
import numpy as np
model = bitone.BitNetModel(model_path="models/bitnet_b1_58-3B.gguf")
output = model.generate(np.array([1, 15043, 29871], dtype=np.int32), max_new_tokens=128)Under the hood:
- Input passed as zero-copy NumPy buffer directly to C++ — no data copying.
- GIL released during all C++ compute — Python doesn't slow down inference.
- Accepts NumPy arrays, Python lists, and PyTorch tensors.
- Hardened input validation: ndim, dtype, contiguity, bounds (
MAX_SEQ_LEN=131K), and null pointer checks prevent buffer overflow attacks.
for token in model.generate(input_ids, max_new_tokens=128, stream=True):
print(tokenizer.decode([token]), end="", flush=True)- Returns a
TokenStreamiterator backed by a stateful C++InferenceSession. - KV-cache persists between tokens — no recomputation.
- Implements Python's
__iter__/__next__protocol with automaticStopIteration.
Concept: A large, smart Teacher model (e.g., Llama 3 70B) transfers its knowledge to a small BitNet Student.
Step 1: Give Teacher a prompt → "The capital of France is"
Step 2: Teacher produces probability distribution over ALL words:
"Paris"=85%, "Lyon"=3%, "Berlin"=0.1% ...
These soft probabilities contain "dark knowledge"
Step 3: Student tries to mimic Teacher's distribution
Step 4: Loss = 70% × KL(teacher ∥ student) + 30% × CrossEntropy(correct answer)
Step 5: Update Student's ternary weights. Repeat.
Temperature scaling (T=4) softens the Teacher's sharp probabilities to reveal richer relationships between words.
import bitone
config = bitone.BitNetStudentConfig(hidden_dim=1024, n_layers=12, n_heads=8)
distiller = bitone.BitNetDistiller(
teacher_api_url="https://your-vllm-server:8000",
student_config=config,
temperature=4.0,
alpha=0.7, # 70% soft loss, 30% hard loss
enforce_tls=True, # mandatory TLS verification
)
prompts = bitone.PromptDataset(["The meaning of life is", ...])
model = distiller.distill(prompts, epochs=3, batch_size=16, tokenizer=tok)Every linear layer uses BitLinear1_58b with custom Triton GPU kernels:
Forward pass (fused in GPU SRAM — no intermediate tensors in HBM):
- Compute
γ = mean(|W|)in a separate 1D reduction kernel (single W read) - GEMM kernel receives precomputed γ, quantizes
W_q = round(clip(W/γ, -1, 1))per tile in SRAM - Compute
Y = X @ W_q^T × γ - W is read exactly once — fixed the double-read performance bug
Why this matters: The previous autograd-based approach materialized a full W_q tensor in GPU memory. For a 3B model, that's ~6GB of intermediate storage. The Triton kernel eliminates this entirely.
from bitone.qat import BitLinear1_58b
layer = BitLinear1_58b(1024, 1024)
output = layer(torch.randn(1, 10, 1024)) # ternary QAT forward
layer.freeze_weights() # lock to {-1, 0, +1}Problem with hard STE: Weights outside the [-γ, +γ] clipping range get zero gradient and are permanently frozen. This affects 5-15% of weights early in training.
Solution: Replace the hard mask with a smooth sech²(β·w) function:
Hard STE: grad_scale = 1 if |w_scaled| ≤ 1, else 0 ← binary cutoff
Soft STE: grad_scale = 1 - tanh²(β·w_scaled) ← smooth decay
At w_scaled=0: grad_scale = 1.00 (full gradient)
At w_scaled=1: grad_scale ≈ 0.07 (attenuated, NOT zero)
At w_scaled=3: grad_scale ≈ 0.0001 (effectively zero, but recoverable)
No weight is ever permanently frozen. β=2 by default (configurable via SOFT_STE_BETA).
Problem: Full teacher logits for vocab=128K are 16.4 GB per batch over the network.
Solution: Teacher API returns only Top-10 logits per position (2.5 MB — 6500× smaller). The KL divergence uses a closed-form residual for the remaining vocab:
KL_total = KL_topk + KL_residual
KL_residual = t_res × (log(t_res) - log(s_res))
where t_res = 1 - Σ(top-k teacher probs) ← teacher's "tail" mass
s_res = 1 - Σ(student probs at top-k) ← student's "tail" mass
No dense teacher distribution is ever materialized.
Problem: Running asyncio inside the PyTorch training loop causes GIL contention and CUDA deadlocks.
Solution: Two completely separate processes:
Process A (Fetcher): Process B (Trainer):
asyncio + aiohttp Pure PyTorch, zero asyncio
TLS-verified API calls Reads from mp.Queue
Pydantic schema validation non_blocking H2D transfers
Writes pinned CPU tensors AMP + gradient clipping
↓ ↑
mp.Queue(maxsize=4) ───────────────────────→
- Backpressure: If GPU is faster, fetcher blocks on
queue.put(). If network is faster, trainer blocks onqueue.get(). - Pinned memory: All tensors are
pin_memory()for async CUDA H2D transfer vianon_blocking=True.
Custom C++ GEMM kernel for 128-bit in-order NPUs with no transparent cache:
- LUT-based weight unpacking —
PSHUFB/VTBLdoes 16 parallel lookups in 1 cycle (vs 48 shift-mask operations) - Double-buffered DMA — Tile N+1 loads asynchronously while Tile N computes (30:1 compute-to-DMA ratio)
- Software-pipelined inner loop — 12-cycle schedule, 12/16 registers used
- MPU-protected TCM — LUT region locked read-only after initialization
Following a comprehensive Red Team audit, bitone includes:
| Attack Vector | Mitigation |
|---|---|
| RCE via model files | GGUF magic validation, dimension bounds, SHA-256 integrity check |
| TCM arbitrary R/W | MPU region table with per-region permissions, DMA validation before every transfer |
| pybind11 buffer overflow | ndim, itemsize, contiguity, bounds, null checks on all inputs |
| MITM on teacher API | TLS 1.2+ enforced, ssl.create_default_context(), certificate verification |
| OOM via API response | 10 MB body cap, Pydantic schema validation, array length limits |
| Malformed JSON | Pydantic CompletionResponse schema with strict parsing, graceful batch dropping |
⚠️ bitone only supports models trained natively with 1.58-bit ternary weights. You CANNOT convert a standard model (Llama, GPT, Mistral) to BitNet.
| Model | Parameters | HuggingFace |
|---|---|---|
| BitNet b1.58 Large | ~700M | 1bitLLM/bitnet_b1_58-large |
| BitNet b1.58 3B | 3B | 1bitLLM/bitnet_b1_58-3B |
| Llama3 8B 1.58-bit | 8B | HF1BitLLM/Llama3-8B-1.58-100B-tokens |
| Model | FP16 | GPTQ 4-bit | bitone (1.58-bit) |
|---|---|---|---|
| 700M | 1.4 GB | 0.35 GB | 0.14 GB |
| 3B | 6 GB | 1.5 GB | 0.6 GB |
| 8B | 16 GB | 4 GB | 1.6 GB |
pip install bitone
# With training dependencies (PyTorch, Triton, aiohttp, Pydantic)
pip install bitone[train]git clone https://github.com/Huntwter/bitone.git
cd bitone/python
pip install -e ".[train]"cmake -B build -DBITNET_NPU128=ON
cmake --build build --config Releaseimport bitone
import numpy as np
model = bitone.BitNetModel("models/bitnet_b1_58-3B.gguf")
# Batch generation
output = model.generate(np.array([1, 2, 3], dtype=np.int32), max_new_tokens=50)
# Streaming
for tok in model.generate([1, 2, 3], max_new_tokens=50, stream=True):
print(tok, end=" ")import bitone
config = bitone.BitNetStudentConfig(hidden_dim=1024, n_layers=12, n_heads=8)
distiller = bitone.BitNetDistiller(
teacher_api_url="https://your-server:8000",
student_config=config,
enforce_tls=True,
)
prompts = bitone.PromptDataset(["Once upon a time", "The speed of light is", ...])
model = distiller.distill(prompts, epochs=3, batch_size=16, tokenizer=my_tokenizer)from bitone.qat import BitLinear1_58b, compute_gamma
import torch
layer = BitLinear1_58b(512, 512)
out = layer(torch.randn(1, 10, 512)) # Triton-fused ternary QAT forward
layer.freeze_weights() # permanently lock to {-1, 0, +1}┌──────────────────────────────────────────────────────┐
│ bitone Python Package │
│ │
│ model.py distiller.py qat.py │
│ Inference API Process-isolated Triton-fused │
│ HF-compatible KD orchestrator BitLinear │
│ Streaming TLS + Pydantic Soft-STE │
│ │ │ Top-K KD │
│ ▼ ▼ │
│ wrapper.cpp mp.Queue │
│ pybind11 pinned tensors │
│ Zero-copy non_blocking H2D │
│ Validated I/O │
└───────┬──────────────────────────────────────────────┘
│
┌───────┴──────────────────────────────────────────────┐
│ C++ Inference Engine │
│ │
│ model_loader.h ggml-bitnet-npu128.cpp │
│ GGUF magic + SHA-256 LUT unpack + DMA GEMM │
│ Bounds validation Ping-pong pipeline │
│ │
│ npu128_hal.h │
│ MPU-protected TCM v128 ISA (SSE/NEON/Scalar) │
│ Validated DMA Architecture-aware yields │
└──────────────────────────────────────────────────────┘
bitone/
├── include/
│ ├── ggml-bitnet.h # API declarations
│ ├── gemm-config.h # Tile/block sizing
│ ├── npu128_hal.h # HAL: MPU, DMA, v128 ISA
│ └── model_loader.h # Hardened GGUF loader + SHA-256
├── src/
│ ├── ggml-bitnet-npu128.cpp # NPU128 GEMM kernel
│ ├── ggml-bitnet-mad.cpp # Original AVX2/NEON kernels
│ └── ggml-bitnet-lut.cpp # Original LUT kernels
├── python/
│ ├── bitone/ # pip package
│ │ ├── __init__.py # Lazy imports
│ │ ├── model.py # BitNetModel + TokenStream
│ │ ├── qat.py # Triton kernels + Soft-STE + TopKKDLoss
│ │ ├── distiller.py # Process-isolated KD orchestrator
│ │ └── _ext/ # C++ extension
│ ├── wrapper.cpp # pybind11 bridge
│ ├── pyproject.toml # PyPI metadata
│ └── setup.py # Build config
├── CMakeLists.txt
├── LICENSE
└── README.md
See Feature 9: Security Hardening for the full list of mitigations.
To report a security vulnerability, please email the maintainer directly. Do not open a public issue.
This project is a fork of Microsoft BitNet, originally developed by Microsoft Research.
- Original: Microsoft Research — BitNet: Scaling 1-bit Transformers
- Fork by: Huntwter
- NPU128 kernel, bitone framework, Triton QAT, security hardening: Huntwter
MIT License — see LICENSE.
Original copyright © Microsoft Corporation. Fork modifications © Huntwter.