bitnet-qat-train

Fused quantization-aware training for BitNet b1.58 (W1.58 A8) on NVIDIA Ampere and newer (sm_80 / sm_86 / sm_89 / sm_90). Training counterpart of phanerozoic/bitnet-tc (CUDA inference) and phanerozoic/bitnet-cpu (CPU inference); all three share the same 2-bit packing.

Implements the BitNet b1.58 recipe: per-tensor absmean weight scale, RoundClip ternarization to {-1, 0, +1}, per-token absmax INT8 activation quantization, straight-through estimator backward. Differences from the eager fake-quant implementation:

• Weight ternarization, scaling, and 2-bit packing are one fused pass; the bf16 fake-quant weight tensor is never materialized.
• The forward runs on INT8 tensor cores (the bitnet-tc GEMM is compiled into this extension) instead of simulating quantization in bf16.
• Autograd saves INT8 activations and 2-bit weights; backward reconstructs operands transiently. Saved weight state is 8x smaller, saved activation state 2x smaller, than the eager recipe.

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

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

Replace existing layers and export for inference:

layer = qat.QATBitLinear.from_linear(dense_linear)
# ... train ...
w_packed, scale_wt = layer.export_inference()
# drop into phanerozoic/bitnet-tc (CUDA) or phanerozoic/bitnet-cpu BitLinear

API

Semantics

Forward computes quant(x) @ ternarize(W)^T * s_x * gamma with INT32 accumulation (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 W_hat/x_hat reconstructed from the saved 2-bit/INT8 state. gamma's dependence on W is treated as constant, matching the standard recipe.

Measured behavior

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

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

License

Apache-2.0.