diff --git a/quantize_streaming.py b/quantize_streaming.py index b85546f..6a2e6f1 100644 --- a/quantize_streaming.py +++ b/quantize_streaming.py @@ -7,22 +7,28 @@ import torch import concurrent.futures from pathlib import Path from safetensors.torch import load_file, save_file -from bitsandbytes.functional import quantize_nf4 +from bitsandbytes.nn import Linear4bit +from torch import nn from transformers import AutoConfig def quantize_weight_nf4(weight: torch.Tensor): - """Quantize a single weight tensor to NF4 using bitsandbytes functional API.""" + """Quantize a single weight tensor to NF4 using Linear4bit.""" if weight.dim() != 2: return weight, None - # quantize_nf4 returns (quantized_tensor, quant_state) - qweight, quant_state = quantize_nf4( - weight, - blocksize=64, - compress_statistics=True, - ) - return qweight, quant_state + # Create Linear4bit and quantize + in_features = weight.size(1) + out_features = weight.size(0) + linear = Linear4bit(in_features, out_features, bias=False, compute_dtype=torch.float16, quant_type="nf4") + + with torch.no_grad(): + linear.weight = nn.Parameter(weight.clone()) + # Force quantization + _ = linear.weight.quant_state + + # Return quantized weight + return linear.weight, None def streaming_quantize(model_path: str, output_path: str):