Quantization compresses neural network weights from their native floating-point representation (FP32/FP16) into lower-precision formats (Int8, Int4, NF4), reducing memory consumption by 2–8x with minimal impact on model quality. For large language models, quantization is often the difference between fitting a model in available memory or not. A 7-billion-parameter model requires ~14 GB in FP16 but only ~3.5 GB in 4-bit quantization—making it feasible to run on consumer GPUs or even CPUs with sufficient RAM.
A Go PEFT library’s quantization package implements NF4 (Normal Float 4-bit), Int4, and Int8 quantization schemes with per-group scaling, configurable group sizes, and support for multiple compute data types. The NF4 format is specifically designed for QLoRA: it uses a non-uniform quantization grid derived from the normal distribution, which better preserves the information content of pretrained weights compared to uniform integer quantization. Weights are stored as packed nibble tensors with separate metadata for scales and zero points.
In systems that serve embedding and classification models on CPU-only deployment targets, quantization is essential. The quantization pipeline converts pretrained model weights from safetensors format into quantized representations, pairs them with LoRA adapters for fine-tuning, and dequantizes on-the-fly during inference to the compute dtype (BF16 or FP32). This combination of quantized base weights plus full-precision LoRA adapters is the core of the QLoRA training methodology.
Quantization maps continuous floating-point values to a discrete set of levels. For b-bit quantization, there are 2^b possible levels. The mapping requires a scale (how much each quantization step represents) and a zero point (which quantized value corresponds to floating-point zero). Symmetric quantization sets the zero point to 0 and centers the range around zero. Asymmetric quantization allows the zero point to shift, better representing distributions that aren’t centered at zero.
NF4 is a 4-bit data type where the 16 quantization levels are spaced according to the quantiles of a standard normal distribution, rather than uniformly. Since pretrained neural network weights approximately follow a normal distribution, NF4 places more quantization levels near zero (where most weights cluster) and fewer levels in the tails. This information-theoretically optimal spacing preserves more precision where it matters most, yielding better model quality than uniform Int4 quantization at the same bit width.
Instead of using a single scale and zero point for an entire weight tensor, per-group quantization divides the tensor into groups of G elements (typically 32, 64, or 128) and computes separate scale/zero-point values for each group. This dramatically reduces quantization error because each group’s scale is calibrated to its local value range. A group size of 64 means that for every 64 weight values, there are 2 extra FP16 values (scale and zero point) stored as metadata—a 6.25% overhead that buys substantially better precision.
4-bit values don’t align to byte boundaries, so two 4-bit weights are packed into a single byte. The packing scheme stores the first value in the lower nibble (bits 0–3) and the second value in the upper nibble (bits 4–7). Unpacking requires bit masking and shifting: val1 = byte & 0x0F, val2 = byte >> 4. This halves the storage compared to using one byte per value, but means that random access to individual weights requires computing the byte offset and nibble position.
Quantized weights are stored in 4-bit or 8-bit format but all arithmetic happens in a higher-precision compute dtype (BF16 or FP32). During inference, each quantized weight group is dequantized on-the-fly: float_val = scale * (quant_val - zero_point). The dequantized values are used for the matrix multiplication and then discarded. This means quantization saves memory (weights are stored compressed) but doesn’t reduce compute (arithmetic still happens in full precision). The memory savings are what matter for fitting large models on limited hardware.
QLoRA introduces double quantization to further reduce the memory overhead of quantization metadata. The per-group scales (which are FP16 values) are themselves quantized to 8-bit, with a second level of grouping. For a weight tensor with group size 64, the scales are grouped in blocks of 256 and quantized to Int8 with their own FP32 scale. This reduces the metadata overhead from ~0.5 bits per parameter to ~0.37 bits per parameter—a small but meaningful saving when hosting models with billions of parameters.
QLoRA combines NF4-quantized base weights with full-precision (FP16/BF16) LoRA adapter weights. During the forward pass, base weights are dequantized to the compute dtype, matrix-multiplied with the input, and the LoRA contribution is added in full precision. During backpropagation, gradients flow through the dequantization operation (using a straight-through estimator) to the LoRA weights only—base weights are frozen and receive no updates. This architecture fine-tunes a 65B model on a single 48GB GPU, which would otherwise require >780 GB for full FP32 fine-tuning.
Define quantization schemes with per-group scaling and compute dtype selection:
package quantization
// QuantScheme identifies the quantization method.
type QuantScheme int
const (
SchemeNF4 QuantScheme = iota // Normal Float 4-bit
SchemeInt4 // Uniform Integer 4-bit
SchemeInt8 // Uniform Integer 8-bit
)
// QuantConfig specifies how to quantize a weight tensor.
type QuantConfig struct {
Scheme QuantScheme
GroupSize int // Elements per quantization group (32, 64, 128)
ComputeDtype DType // Dequantization target: BF16 or FP32
DoubleQuant bool // Apply double quantization to scales
}
// DefaultQLoRAConfig returns the standard QLoRA quantization config.
func DefaultQLoRAConfig() QuantConfig {
return QuantConfig{
Scheme: SchemeNF4,
GroupSize: 64,
ComputeDtype: BFloat16,
DoubleQuant: true,
}
}
// QuantizedTensor holds a quantized weight tensor with its metadata.
type QuantizedTensor struct {
Data []byte // Packed quantized values (nibbles for 4-bit)
Scales []float16 // Per-group scale factors
ZeroPoints []float16 // Per-group zero points (asymmetric only)
Shape []int // Original tensor shape
Config QuantConfig
NumGroups int // len(Scales) == NumGroups
}The NF4 lookup table maps 4-bit indices (0–15) to non-uniform float values:
// nf4Table contains the 16 NF4 quantization levels, computed as
// quantiles of a standard normal distribution normalized to [-1, 1].
var nf4Table = [16]float32{
-1.0, -0.6962, -0.5251, -0.3949,
-0.2844, -0.1848, -0.0911, 0.0,
0.0796, 0.1609, 0.2461, 0.3379,
0.4407, 0.5626, 0.7230, 1.0,
}
// QuantizeNF4Group quantizes a group of float32 values to NF4.
func QuantizeNF4Group(values []float32) ([]uint8, float32) {
// Find absmax for symmetric scaling
absmax := float32(0)
for _, v := range values {
if abs := absFloat32(v); abs > absmax {
absmax = abs
}
}
scale := absmax
if scale == 0 {
scale = 1.0 // avoid division by zero
}
invScale := 1.0 / scale
// Quantize each value to nearest NF4 level
quantized := make([]uint8, len(values))
for i, v := range values {
normalized := v * invScale // map to [-1, 1]
quantized[i] = findNearestNF4(normalized)
}
return quantized, scale
}
// findNearestNF4 returns the NF4 index closest to the normalized value.
func findNearestNF4(normalized float32) uint8 {
bestIdx := uint8(0)
bestDist := absFloat32(normalized - nf4Table[0])
for i := uint8(1); i < 16; i++ {
dist := absFloat32(normalized - nf4Table[i])
if dist < bestDist {
bestDist = dist
bestIdx = i
}
}
return bestIdx
}Quantize an entire weight tensor with per-group scaling and pack into nibbles:
// QuantizeTensor quantizes a float32 weight tensor to the specified scheme.
func QuantizeTensor(weights []float32, shape []int,
cfg QuantConfig) (*QuantizedTensor, error) {
totalElements := 1
for _, dim := range shape {
totalElements *= dim
}
if totalElements != len(weights) {
return nil, fmt.Errorf("shape mismatch: %d vs %d", totalElements, len(weights))
}
numGroups := (totalElements + cfg.GroupSize - 1) / cfg.GroupSize
scales := make([]float16, numGroups)
quantized := make([]uint8, totalElements)
// Quantize each group
for g := 0; g < numGroups; g++ {
start := g * cfg.GroupSize
end := start + cfg.GroupSize
if end > totalElements {
end = totalElements
}
group := weights[start:end]
var groupQuantized []uint8
var scale float32
switch cfg.Scheme {
case SchemeNF4:
groupQuantized, scale = QuantizeNF4Group(group)
case SchemeInt4:
groupQuantized, scale = QuantizeInt4Group(group)
default:
return nil, fmt.Errorf("unsupported scheme: %d", cfg.Scheme)
}
copy(quantized[start:end], groupQuantized)
scales[g] = float16(scale)
}
// Pack 4-bit values into nibbles (2 values per byte)
packed := packNibbles(quantized)
return &QuantizedTensor{
Data: packed,
Scales: scales,
Shape: shape,
Config: cfg,
NumGroups: numGroups,
}, nil
}
// packNibbles packs pairs of 4-bit values into bytes.
func packNibbles(values []uint8) []byte {
packed := make([]byte, (len(values)+1)/2)
for i := 0; i < len(values); i += 2 {
lo := values[i] & 0x0F
hi := uint8(0)
if i+1 < len(values) {
hi = values[i+1] & 0x0F
}
packed[i/2] = lo | (hi << 4)
}
return packed
}Quantization is a step function with zero gradient almost everywhere, which would block gradient flow during training. QLoRA uses a straight-through estimator (STE): during the forward pass, weights are quantized and dequantized normally; during the backward pass, the gradient passes through the quantization/dequantization as if it were an identity function. This allows gradients to reach the LoRA adapter weights despite the non-differentiable quantization step in the computation graph. The STE introduces a small gradient approximation error, but in practice it works well because only adapter weights are updated.
The primary application is the QLoRA fine-tuning pipeline in the library. The pipeline loads a pretrained model in NF4 format, attaches LoRA adapters to the attention and MLP layers, and trains on task-specific data. The quantized base model uses ~3.5 GB for a 7B model, the LoRA adapters add ~16 MB, and gradient buffers for the adapters add another ~32 MB. Total training memory is under 4 GB, compared to >56 GB for full FP32 fine-tuning.
Quantization makes it practical to serve embedding models on CPU-only infrastructure. In a typical CPU-only deployment, embedding models are quantized to Int4 or NF4 and loaded into system RAM. Each inference request dequantizes the required weight groups to FP32 for computation. While slower than GPU inference, the reduced memory footprint means models that would require a GPU can run on standard cloud instances, significantly reducing serving cost.
4-bit quantized models are 4x smaller than FP16 and 8x smaller than FP32, which matters for distribution. A model registry storing quantized model blobs can distribute them in seconds rather than minutes. Edge deployments pull quantized models over constrained network links, and the smaller file sizes reduce storage costs in the registry.
Not all layers need the same precision. The library supports mixed-precision quantization where critical layers (first and last transformer blocks, embedding layers) keep higher precision (Int8 or FP16) while middle layers use aggressive NF4 quantization. This per-layer configuration balances model quality against memory savings, typically recovering 80–90% of the quality lost by uniform 4-bit quantization.
Convert quantized tensors back to compute dtype on-the-fly:
// Dequantize converts a quantized tensor back to float32 for computation.
func (qt *QuantizedTensor) Dequantize() []float32 {
totalElements := 1
for _, dim := range qt.Shape {
totalElements *= dim
}
result := make([]float32, totalElements)
// Unpack nibbles first
unpacked := unpackNibbles(qt.Data, totalElements)
// Dequantize each group
for g := 0; g < qt.NumGroups; g++ {
start := g * qt.Config.GroupSize
end := start + qt.Config.GroupSize
if end > totalElements {
end = totalElements
}
scale := float32(qt.Scales[g])
for i := start; i < end; i++ {
switch qt.Config.Scheme {
case SchemeNF4:
// Look up the NF4 float value and apply scale
result[i] = nf4Table[unpacked[i]] * scale
case SchemeInt4:
// Linear dequantization: val = (quant - 8) * scale / 7
result[i] = (float32(unpacked[i]) - 8.0) * scale / 7.0
}
}
}
return result
}
// unpackNibbles extracts 4-bit values from packed bytes.
func unpackNibbles(packed []byte, count int) []uint8 {
result := make([]uint8, count)
for i := 0; i < count; i++ {
byteIdx := i / 2
if i%2 == 0 {
result[i] = packed[byteIdx] & 0x0F
} else {
result[i] = packed[byteIdx] >> 4
}
}
return result
}Wire quantized weights into GoMLX’s computation graph with on-the-fly dequantization:
// QuantizedLinear implements a linear layer with quantized weights.
type QuantizedLinear struct {
Weight *QuantizedTensor
Bias []float32 // Bias stays in full precision
InFeatures int
OutFeatures int
}
// Forward computes output = input @ dequantize(weight)^T + bias
// The dequantization happens per-group during the matmul to minimize
// peak memory: only one group of weights is in FP32 at a time.
func (ql *QuantizedLinear) Forward(ctx *context.Context,
input *Node) *Node {
// Dequantize weights to compute dtype
// In a real implementation this integrates with GoMLX's graph
// to fuse dequantization with the matmul kernel
weightFloat := ql.Weight.Dequantize()
weightNode := ConstantFromSlice(ctx, weightFloat, ql.Weight.Shape)
output := MatMul(input, Transpose(weightNode))
if ql.Bias != nil {
biasNode := ConstantFromSlice(ctx, ql.Bias, []int{ql.OutFeatures})
output = Add(output, biasNode)
}
return output
}
// MemoryFootprint returns the memory used by quantized weights in bytes.
func (ql *QuantizedLinear) MemoryFootprint() int64 {
dataBytes := int64(len(ql.Weight.Data))
scaleBytes := int64(ql.Weight.NumGroups) * 2 // FP16 scales
biasBytes := int64(len(ql.Bias)) * 4 // FP32 bias
return dataBytes + scaleBytes + biasBytes
}Smaller group sizes (32) give finer-grained scaling and less quantization error, but increase metadata overhead (more scale values to store). Larger group sizes (128) minimize overhead but introduce more error because diverse value ranges share a single scale. The sweet spot for NF4 quantization is a group size of 64: it provides good precision with only ~6% metadata overhead. For Int8 quantization, group sizes of 128 work well because the higher bit width tolerates coarser scaling.
lo | (hi << 4) but the dequantizer unpacks as hi | (lo << 4), every other weight pair is swapped. This produces subtly wrong outputs that may not be immediately obvious—test with known values and verify round-trip accuracy.absmax(group) < 65504 (FP16 max) or fall back to FP32 scales for outlier groups.