From b82be559dcf48a99a6f42e5d3a9d1e36c54b225a Mon Sep 17 00:00:00 2001
From: Christian Medina <37550954+cmedinasoriano@users.noreply.github.com>
Date: Fri, 3 Jul 2026 01:38:54 -0400
Subject: [PATCH] fix: use Linear4bit to preserve weight shapes during
 quantization

---
 quantize_streaming.py | 24 +++++++++++++++---------
 1 file changed, 15 insertions(+), 9 deletions(-)

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