bitnet-qat-train

Fused quantization-aware training for BitNet b1.58 (W1.58 A8) on NVIDIA Ampere and newer, loadable through kernels. Training counterpart of bitnet-tc (CUDA inference) and bitnet-cpu (CPU inference); all three share the same 2-bit packing. The reference baseline is the eager fake-quant recipe, tracked to |Δloss| ≤ 0.0019 over 300 identical steps.

Quantization-aware training for ternary models is normally simulated: the quantized weight is materialized as a full bf16 tensor every step, the matmul runs in bf16, and autograd saves the full-precision operands. This kernel runs the real thing: ternarization, scaling, and 2-bit packing are one fused pass, the forward runs on INT8 tensor cores, and autograd saves INT8 activations and 2-bit weights, so the training loop is faster and lighter while following the same recipe to the same loss.

The fused kernel and the eager fake-quant recipe trained on identical data and seeds: maximum loss gap 0.0078 over 240 steps, while the fused step runs 4.8x faster (7 vs 32 ms) at 2.2x less peak memory (660 vs 1,475 MB) on a 2560-to-6912 stack.

Usage

import torch
from kernels import get_kernel

qat = get_kernel("phanerozoic/bitnet-qat-train", version=1, trust_remote_code=True)

layer = qat.QATBitLinear(2560, 6912, device="cuda")   # master weights, bf16
x = torch.randn(8, 512, 2560, dtype=torch.bfloat16, device="cuda")
y = layer(x)                                          # [8, 512, 6912] bf16
y.sum().backward()                                    # STE gradients on layer.weight

layer = qat.QATBitLinear.from_linear(dense_linear)
w_packed, scale_wt = layer.export_inference()         # -> bitnet-tc / bitnet-cpu

version selects the release branch; trust_remote_code is required by kernels for publishers without the trusted-publisher mark.

API

Method

Forward computes quant(x) @ ternarize(W)^T * s_x * gamma with INT32 accumulation, which is exact integer arithmetic; the eager recipe's bf16 matmul of the same quantized values is the less precise side. Backward is the straight-through estimator: dX = dY @ W_hat, dW = dY^T @ x_hat, with the operands reconstructed transiently from the saved 2-bit and INT8 state, so saved weight state is 8x smaller and saved activation state 2x smaller than the eager recipe. gamma's dependence on W is treated as constant, matching the standard recipe.

Measured

On an L4 (sm_89), torch 2.12, against the eager fake-quant recipe:

• Gradients: with references built from the kernel's own quantized operands, dX and dW match bitwise across 12 shape/dtype cases; the forward residual is bf16 output rounding (< 7e-3 maximum relative).
• Training: a 4-layer transformer LM trained 300 steps on identical batches and seeds tracks the eager recipe to |Δloss| ≤ 0.0019 over the first 50 steps and 0.0006 at termination, at 1.43x the eager wall-clock.
• Memory: peak allocation through an 8-layer stack (2560 ↔ 6912, batch 4x512, forward+backward) is 1.64x lower (1887 → 1151 MB).
• Step time: forward+backward on (2560 → 6912), batch 8x512, is 1.76x faster (15.48 → 8.80 ms median of 20).
• Ternarization is bitwise deterministic across repeated invocations.
• Exported weights produce bitwise-identical outputs through the phanerozoic/bitnet-tc inference kernel loaded in the same process.

Requirements and limits

• NVIDIA GPU with compute capability 8.0+ (Ampere, Ada, Hopper).
• K divisible by 32; bf16 activations; bf16 or f32 master weights.

References

Ma et al., "The Era of 1-bit LLMs" (BitNet b1.58, 2024); straight-through estimation (Bengio et al., 2013); the shared 2-bit packing of phanerozoic/bitnet-tc and Microsoft's bitnet.cpp.

License

Apache-2.0.