File size: 4,232 Bytes
e9c8366 | 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 | """Double quantization — quantize the per-group scales themselves to int8.
QLoRA paper (Dettmers et al., 2023) observes that the per-group fp16/fp32 scales
of NF4 have their own non-trivial storage overhead (~32 bit per group). Double
quantization compresses the scales: each block of `block_size` consecutive
scales (within one output row) gets one fp32 `block_scale`
(= max(abs(scales_in_block)) / 127), and the scales inside the block are
stored as int8.
Storage reduction (per group):
without double quant: scale = fp32 = 32 bit/group
with double quant: scale_int8 = 8 bit/group + block_scale fp32 per block
= 8 + 32/block_size bit/group
for block_size=256: ~8.125 bit/group (3.9x reduction on scales)
The block_scale itself could be quantized further (2nd level), but QLoRA stops
at one level — diminishing returns + added complexity. We follow that.
Quantization is performed PER ROW (each output channel's scales quantized
independently), so dequant restores [out, num_groups] without cross-row mixing.
The last block in each row may be shorter than `block_size` (no padding) —
we store the true num_groups per row and trim on dequant. This avoids the
padding overhead that would otherwise dominate for small layers.
"""
import torch
def double_quantize_scales_2d(
scales: torch.Tensor,
block_size: int = 256,
) -> tuple[torch.Tensor, torch.Tensor]:
"""Quantize per-row scales [out, num_groups] to int8 + per-block fp32 scale.
Each output row is quantized independently: scales in row `o` are split
into blocks of up to `block_size` consecutive elements (the last block may
be shorter — no padding is added), each block gets one fp32 block_scale,
the scales inside the block are int8.
Args:
scales: fp32/fp16 tensor [out, num_groups] of positive per-group scales.
block_size: max number of consecutive scales per double-quant block.
Returns:
(scales_int8, block_scale) where:
scales_int8: int8 [out, num_groups] (no padding)
block_scale: fp32 [out, num_blocks_per_row] (ceil(num_groups / block_size))
"""
s = scales.detach().to(torch.float32)
out_features, num_groups = s.shape
# Number of blocks per row (last block may be short)
num_blocks = (num_groups + block_size - 1) // block_size
block_scale = torch.empty(out_features, num_blocks, dtype=torch.float32)
scales_int8 = torch.empty(out_features, num_groups, dtype=torch.int8)
for b in range(num_blocks):
start = b * block_size
end = min(start + block_size, num_groups)
block = s[:, start:end] # [out, block_len]
block_max = block.abs().amax(dim=1).clamp(min=1e-12) # [out]
bs = block_max / 127.0
block_scale[:, b] = bs
q = torch.clamp(torch.round(block / bs.unsqueeze(1)), min=-127, max=127).to(torch.int8)
scales_int8[:, start:end] = q
return scales_int8, block_scale
def dequantize_scales_2d(
scales_int8: torch.Tensor,
block_scale: torch.Tensor,
block_size: int = 256,
num_groups: int | None = None,
) -> torch.Tensor:
"""Reconstruct per-row scales [out, num_groups] from int8 + per-block scale.
Args:
scales_int8: int8 [out, num_groups] (from double_quantize_scales_2d).
block_scale: fp32 [out, num_blocks_per_row].
block_size: max block size used in quantization.
num_groups: original num_groups per row. If None, uses scales_int8.shape[1].
Returns:
fp32 tensor [out, num_groups].
"""
out_features = scales_int8.shape[0]
ng = scales_int8.shape[1] if num_groups is None else num_groups
num_blocks = block_scale.shape[1]
deq = torch.empty(out_features, scales_int8.shape[1], dtype=torch.float32)
for b in range(num_blocks):
start = b * block_size
end = min(start + block_size, scales_int8.shape[1])
block_i = scales_int8[:, start:end].to(torch.float32)
bs = block_scale[:, b].to(torch.float32).unsqueeze(1)
deq[:, start:end] = block_i * bs
if num_groups is not None and num_groups < deq.shape[1]:
deq = deq[:, :num_groups]
return deq |