Common Attention Utilities

Shared utilities for attention modules.

Provides quantization helpers (BitNet b1.58 ternary weights, 8-bit activations) and the BitLinear layer.

Normalization layers (DynamicTanhNorm, Derf) and the get_norm factory have been moved to src/model/norm/.

References

Ma et al. (2024), “The Era of 1-bit LLMs: All Large Language Models are in 1.58 Bits”, arXiv:2402.17764.

class src.model.attention.common.BitConv1d(*args: Any, **kwargs: Any)[source]

Bases: Conv1d

BitNet b1.58 1D convolution with ternary weight quantization.

A drop-in replacement for nn.Conv1d that applies the same ternary weight ({-1, 0, 1}) + int8 activation quantization as BitLinear, via the straight-through estimator (STE). Inputs are LayerNorm-normalized over the channel dimension before activation quantization, mirroring the BitLinear formulation.

Used by FactorizedEmbedding so the embedding Conv1d pre-projection also becomes ternary when use_bitnet is enabled.

Parameters:

  • in_channels – Number of input channels.

  • out_channels – Number of output channels.

  • kernel_size – Size of the convolving kernel.

  • stride – Stride of the convolution. Default: 1.

  • padding – Padding added to both sides. Default: 0.

  • groups – Number of blocked connections. Default: 1.

  • bias – If True, adds a learnable bias. Defaults to False.

Reference:

Ma et al. (2024), “The Era of 1-bit LLMs”, arXiv:2402.17764.

bake_ternary_weights()

Replace the master weight with its ternary-quantized value.

Idempotent: a no-op if the weight is already ternary. See BitLinear.bake_ternary_weights() for details.

forward(x)[source]

Forward pass with activation and weight quantization.

Parameters:

x – Input tensor of shape (B, C_in, L).

Returns:

Output tensor of shape (B, C_out, L_out).

class src.model.attention.common.BitLinear(*args: Any, **kwargs: Any)[source]

Bases: Linear

BitNet b1.58 linear layer with ternary weight quantization.

A drop-in replacement for nn.Linear that reduces VRAM usage by 3-4x through ternary weight quantization and 8-bit activation quantization. Applies LayerNorm to inputs before quantization for stability, as prescribed by the BitNet b1.58 formulation.

Reference:

Ma et al. (2024), “The Era of 1-bit LLMs”, arXiv:2402.17764.

Parameters:

  • in_features – Size of each input sample.

  • out_features – Size of each output sample.

  • bias – If True, adds a learnable bias to the output. Defaults to False.

bake_ternary_weights()

Replace the master weight with its ternary-quantized value.

Applies weight_quant() once and stores the result ({-1, 0, 1} * scale) as the learnable parameter, removing the straight-through estimator (STE). After baking, the weight holds the faithful ternary values and is no longer a full-precision master.

Idempotent: if the weight is already ternary (all nonzero elements equal +/-max(|w|)), the call is a no-op. This prevents the per-tensor absmean scale from drifting on repeated baking.

Used at export/deployment time to produce compact, self-describing ternary checkpoints. Training should not continue after baking (the STE gradient path is gone).

The bias (if any) is left untouched.

forward(x)[source]

Forward pass with activation and weight quantization.

Parameters:

x – Input tensor of shape (..., in_features).

Returns:

Output tensor of shape (..., out_features).

src.model.attention.common.activation_quant(x)[source]

Quantize activations to 8-bit via straight-through estimator (STE).

Scales activations to the range [-128, 127], rounds to integer, then rescales back. The STE passes gradients through the rounding operation unchanged, enabling end-to-end training with quantized activations.

Parameters:

x – Input tensor of any shape. The last dimension is used for per-token scaling.

Returns:

Tensor of same shape as x, with values quantized to 8-bit precision. Gradients flow through the quantization via STE.

src.model.attention.common.is_bitlinear_module(module: torch.nn.Module) → bool[source]

Return True if module is a BitNet-quantized layer.

Covers both BitLinear and BitConv1d (the ternary Conv1d used by the factorized embedding pre-projection).

Parameters:

module – Any nn.Module.

Returns:

Whether the module performs BitNet b1.58 ternary quantization.

src.model.attention.common.weight_quant(w)[source]

Quantize weights to ternary values {-1, 0, 1} via straight-through estimator.

Scales weights by the inverse of their mean absolute value, rounds to the nearest integer in [-1, 1], then rescales. The STE passes gradients through the rounding operation unchanged.

Parameters:

w – Weight tensor of any shape.

Returns:

Tensor of same shape as w, with values quantized to ternary precision. Gradients flow through the quantization via STE.