Quantization

BitNet b1.58 quantization and compression utilities for Frankenstein.

Provides faithful ternary weight packing ({-1, 0, 1} at ~1.58 bits per weight), INT8 activation quantization, model size estimation, and fast serialization/deserialization for quantized checkpoints.

Design contract

Only BitLinear modules (instances of src.model.attention.common.BitLinear) are packed to ternary. All other parameters — full-precision nn.Linear weights (e.g. routing and scoring projections when bitnet_routers is False), biases, embeddings, LayerNorm/normalization parameters — are stored verbatim as float. This correctly honours the use_bitnet / bitnet_routers schema flags because the model itself decides at construction time which layers become BitLinear.

Reference:

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

class src.deploy.quantization.ActivationQuantizer

Bases: object

Runtime activation quantization for inference optimization.

Provides static methods for quantizing activations to INT8 and dequantizing back to float, enabling efficient integer-arithmetic computation during inference.

static dequantize_activation_int8(x_q: torch.Tensor, scale: float) → torch.Tensor

Dequantize INT8 activations back to float.

Parameters:

• x_q – Quantized INT8 tensor.

• scale – The scale factor from quantize_activation_int8().

Returns:

Dequantized float tensor.

static quantize_activation_int8(x: torch.Tensor) → Tuple[torch.Tensor, float]

Quantize activations to INT8 range [-128, 127].

Parameters:

x – Activation tensor of any shape.

Returns:

Tuple of (quantized_tensor, scale) where scale is the factor used for quantization (for later dequantization).

class src.deploy.quantization.BitNetQuantizer

Bases: object

Quantization manager for BitNet ternary weight models.

Handles packing/unpacking of ternary weights ({-1, 0, 1}) into 2-bit-per-element byte arrays (4 weights per byte) and dequantization back to float tensors. Only BitLinear modules are quantized; every other parameter is preserved at full precision.

quantization_config

Dictionary describing the quantization scheme (weight_bits = 1.58 ternary, activation_bits = 8).

__init__()

Initialize the quantizer with default BitNet b1.58 config.

dequantize_model_weights(quantized_state: Dict[str, Any], model: torch.nn.Module) → None

Load quantized weights back into a model.

Ternary-packed tensors are dequantized; full-precision tensors are loaded verbatim. Missing keys are ignored (strict=False).

Parameters:

• quantized_state – Dictionary produced by quantize_model_weights().

• model – PyTorch model to load weights into.

static dequantize_ternary_weights(packed: numpy.ndarray, scale: float, original_shape: Tuple[int, ...]) → torch.Tensor

Dequantize packed ternary weights back to a float tensor.

Parameters:

• packed – Packed ternary weights (uint8).

• scale – The absmean scaling factor.

• original_shape – Original tensor shape.

Returns:

Dequantized float tensor of original_shape.

quantize_model_weights(model: torch.nn.Module) → Dict[str, Any]

Quantize all BitLinear weights in the model to ternary format.

Only BitLinear modules are packed; every other parameter is stored verbatim as a float numpy array. This honours use_bitnet and bitnet_routers because the model itself decided which layers are BitLinear at construction time.

Parameters:

model – PyTorch model to quantize.

Returns:

Dictionary with weights, scales, shapes, quantized_tensors (names of packed tensors), and config.

static quantize_ternary_weights(weight: torch.Tensor) → Tuple[numpy.ndarray, float]

Quantize weights to ternary {-1, 0, 1} and pack for storage.

Uses the BitNet b1.58 per-tensor absmean scale, unless the weight is already ternary (baked), in which case the existing scale (max(|w|)) is preserved so the representation is reproduced bit-exactly.

Parameters:

weight – Float tensor to quantize (master or already-baked).

Returns:

Tuple of (packed_weights, scale) where packed_weights is a uint8 array (2 bits per weight, 4 weights per byte) and scale is the scaling factor for dequantization.

src.deploy.quantization.bake_bitnet_weights(model: torch.nn.Module) → int

Bake ternary weights into every BitNet module in model.

Applies bake_ternary_weights() once to each BitNet-quantized layer (BitLinear and BitConv1d) so the stored master weight becomes the faithful {-1, 0, 1} * scale value (no STE). Used at export/deploy time to produce compact, self-describing ternary checkpoints.

Parameters:

model – Root model whose BitNet layers should be baked.

Returns:

Number of BitNet modules that were baked.

src.deploy.quantization.estimate_model_size(model: torch.nn.Module) → Dict[str, float]

Estimate model size in different formats.

Parameters:

model – PyTorch model.

Returns:

Dictionary with size estimates in MB for fp32, fp16, and BitNet 1.58.

src.deploy.quantization.load_quantized_checkpoint(load_path: str, model: torch.nn.Module) → Dict[str, Any]

Load a quantized checkpoint into a model.

Parameters:

• load_path – Path to a quantized checkpoint.

• model – Model to load weights into.

Returns:

Dictionary with additional metadata from the checkpoint (everything except quantized_weights).

src.deploy.quantization.save_quantized_checkpoint(model: torch.nn.Module, save_path: str, additional_data: Dict[str, Any] = None) → None

Save a model as a quantized (ternary-packed) checkpoint.

Only BitLinear weights are packed; other parameters are stored as float. Additional metadata (config, tokenizer info, etc.) is merged in.

Parameters:

• model – PyTorch model to save.

• save_path – Path to save the checkpoint.

• additional_data – Additional metadata to merge into the checkpoint.