fix: use Linear4bit to preserve weight shapes during quantization

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):