Class BitLinear

utils.py:88–146

Drop-in replacement for nn.Linear implementing BitNet b1.58 training. Forward pass (training): 1. Normalise input with internal parameter-free RMSNorm. 2. Quantise normalised input activations to 8-bit per token (STE). 3. Quantise weight matrix to ternary {-1, 0, +1} via

# ──────────────────────────────────────────────────────────────────────────────

class BitLinear(nn.Linear):
    """
    Drop-in replacement for nn.Linear implementing BitNet b1.58 training.

    Forward pass (training):
      1. Normalise input with internal parameter-free RMSNorm.
      2. Quantise normalised input activations to 8-bit per token (STE).
      3. Quantise weight matrix to ternary {-1, 0, +1} via absmean (STE).
      4. Compute F.linear(x_q, w_q) using the regular fp matmul kernel
         (specialised ternary kernels only needed at inference time).
      5. Rescale output: multiply by (weight_scale × activation_scale).

    Bias is always disabled as per the paper.
    """

    def __init__(
        self,
        in_features: int,
        out_features: int,
        bias: bool = False,          # paper: no bias
        dtype: torch.dtype | None = None,
    ) -> None:
        super().__init__(in_features, out_features, bias=bias, dtype=dtype)
        self.norm = _RMSNorm(in_features)

    # ------------------------------------------------------------------
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # 1. Internal RMSNorm
        x_normed = self.norm(x)

        # 2. Quantise activations (STE, keeps gradients)
        x_q, act_scale = _quantize_activations_int8(x_normed)

        # 3. Quantise weights (STE, keeps gradients through round)
        w_q, w_scale = _quantize_weights_ternary(self.weight)

        # 4. Linear using the (still fp) quantised values
        out = F.linear(x_q, w_q, self.bias)

        # 5. Rescale: broadcast act_scale [..., 1] and scalar w_scale
        out = out * act_scale * w_scale

        return out

    # ------------------------------------------------------------------
    @classmethod
    def from_linear(cls, linear: nn.Linear) -> "BitLinear":
        """Construct a BitLinear from an existing nn.Linear, copying its weight."""
        bl = cls(
            in_features=linear.in_features,
            out_features=linear.out_features,
            bias=linear.bias is not None,
            dtype=linear.weight.dtype,
        )
        with torch.no_grad():
            bl.weight.copy_(linear.weight)
            if linear.bias is not None and bl.bias is not None:
                bl.bias.copy_(linear.bias)