File size: 5,892 Bytes
f48ad7a | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 | """
bitlinear.py — TriChronos-0.1B
Absmean ternary weight quantization (BitNet-style) with Straight-Through Estimator.
Per-token activation quantization applied before the linear projection.
Reference: "The Era of 1-bit LLMs" (Ma et al., 2024)
"""
import math
import torch
import torch.nn as nn
import torch.nn.functional as F
# ---------------------------------------------------------------------------
# Weight quantization
# ---------------------------------------------------------------------------
def absmean_quantize(W: torch.Tensor, eps: float = 1e-6) -> torch.Tensor:
"""
Absmean ternary quantization.
1. Compute the scale γ = mean(|W|) (a single scalar per weight matrix)
2. Divide W by γ, round to nearest integer, clip to {-1, 0, +1}
During the forward pass only the ternary version is used; gradients flow
through unmodified (Straight-Through Estimator applied at call site).
"""
gamma = W.abs().mean().clamp(min=eps)
W_scaled = W / gamma
W_ternary = W_scaled.round().clamp(-1, 1)
return W_ternary, gamma
# ---------------------------------------------------------------------------
# Activation quantization
# ---------------------------------------------------------------------------
def per_token_quant(x: torch.Tensor, bits: int = 8, eps: float = 1e-6) -> torch.Tensor:
"""
Symmetric per-token (per-row) quantization of activations to `bits` bits.
x : (..., d_model) — last dim is the feature dimension, all other dims
are treated as independent "tokens".
Each token is scaled independently so that its dynamic range fills the
integer grid [-2^(bits-1)+1, 2^(bits-1)-1]. We use absmax scaling
(same approach as LLM.int8()).
"""
qmax = 2 ** (bits - 1) - 1 # 127 for 8-bit
# absmax per token; keep dims for broadcasting
scale = x.abs().amax(dim=-1, keepdim=True).clamp(min=eps)
x_scaled = (x / scale * qmax).round().clamp(-qmax, qmax)
# dequantize back to float so the rest of the graph stays in FP16/BF16
x_dequant = x_scaled / qmax * scale
return x_dequant
# ---------------------------------------------------------------------------
# BitLinear layer
# ---------------------------------------------------------------------------
class BitLinear(nn.Linear):
"""
Drop-in replacement for nn.Linear that:
• Keeps FP16/BF16 master weights (for optimizer stability)
• Quantizes weights to ternary {-1, 0, +1} during the forward pass
• Applies per-token activation quantization to the input
• Uses STE: gradients pass through the quantization step unchanged
Usage is identical to nn.Linear:
layer = BitLinear(in_features, out_features, bias=False)
y = layer(x)
"""
def __init__(
self,
in_features: int,
out_features: int,
bias: bool = False, # bias=False is the BitNet default
act_bits: int = 8,
weight_eps: float = 1e-6,
):
super().__init__(in_features, out_features, bias=bias)
self.act_bits = act_bits
self.weight_eps = weight_eps
# LayerNorm before quantization stabilises activations (BitNet paper §3)
self.norm = nn.LayerNorm(in_features, elementwise_affine=True)
def forward(self, x: torch.Tensor) -> torch.Tensor:
# 1. Normalise activations
x = self.norm(x)
# 2. Quantize activations per-token (straight through — no rounding
# in the backward graph, we use the dequantized value directly)
x_q = per_token_quant(x, bits=self.act_bits, eps=self.weight_eps)
# 3. Quantize weights with STE
# Forward: use ternary weights scaled by γ
# Backward: gradient flows through as if weights were continuous
W_ternary, gamma = absmean_quantize(self.weight, eps=self.weight_eps)
# STE: detach the rounding error so gradients see a straight-through copy
W_ste = self.weight + (W_ternary - self.weight).detach()
W_ste_scaled = W_ste * gamma # restore scale for correct output magnitude
return F.linear(x_q, W_ste_scaled, self.bias)
# ------------------------------------------------------------------
# Convenience: export ternary weights for storage (packed int8)
# ------------------------------------------------------------------
@torch.no_grad()
def ternary_weights(self) -> torch.Tensor:
"""Return quantized weights as int8 tensor (values in {-1, 0, 1})."""
W_ternary, _ = absmean_quantize(self.weight, eps=self.weight_eps)
return W_ternary.to(torch.int8)
@torch.no_grad()
def weight_scale(self) -> torch.Tensor:
"""Return the absmean scale γ used during quantization."""
_, gamma = absmean_quantize(self.weight, eps=self.weight_eps)
return gamma
# ---------------------------------------------------------------------------
# Quick sanity check (run directly: python bitlinear.py)
# ---------------------------------------------------------------------------
if __name__ == "__main__":
torch.manual_seed(42)
layer = BitLinear(64, 128)
x = torch.randn(4, 16, 64) # (batch, seq_len, d_model)
y = layer(x)
print(f"Input shape : {x.shape}")
print(f"Output shape : {y.shape}")
assert y.shape == (4, 16, 128), "Unexpected output shape"
# Verify STE: backward pass should not raise
loss = y.sum()
loss.backward()
print(f"Weight grad : {layer.weight.grad.shape} (STE working)")
# Ternary check
W_t = layer.ternary_weights()
unique_vals = W_t.unique().tolist()
print(f"Ternary unique values: {unique_vals} (expected subset of [-1, 0, 1])")
assert all(v in (-1, 0, 1) for v in unique_vals)
print("bitlinear.py - all checks passed OK")
|