Merge pull request bitsandbytes-foundation#1228 from jiqing-feng/4bit

fix blocksize
Xia-Weiwen · May 30, 2024 · dba8376 · dba8376
2 parents 2dbf876 + 36fe1a0
commit dba8376
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/bitsandbytes/backends/cpu.py b/bitsandbytes/backends/cpu.py
@@ -135,6 +135,8 @@ def quantize_4bit(
         quant_type: Literal["fp4", "nf4"] = "fp4",
         quant_storage=torch.uint8,
     ) -> Tuple[torch.Tensor, QuantState]:
+        if blocksize is None:
+            blocksize = 64
         assert_on_cpu([A, absmax, out])
         assert quant_storage == torch.uint8, "CPU backend only supports uint8 quant_storage"
         return quantize_4bit_impl(A, absmax, out, blocksize, compress_statistics, quant_type)
@@ -148,6 +150,8 @@ def dequantize_4bit(
         blocksize: int = 64,
         quant_type: Literal["fp4", "nf4"] = "fp4",
     ) -> torch.Tensor:
+        if blocksize is None:
+            blocksize = 64
         assert_on_cpu([A, absmax, out])
         return dequantize_4bit_impl(A, quant_state, absmax, out, blocksize, quant_type)