Merge pull request #1295 from matthewdouglas/fsdp-load-prequantized

[FSDP] Enable loading prequantized weights with bf16/fp16/fp32 quant_storage
bitsandbytes-foundation · Jul 29, 2024 · 3a6911f · 3a6911f
2 parents 5cc91b1 + a96d2f0
commit 3a6911f
Showing 1 changed file with 5 additions and 0 deletions.
diff --git a/bitsandbytes/nn/modules.py b/bitsandbytes/nn/modules.py
@@ -273,6 +273,7 @@ def from_prequantized(
         quantized_stats: Dict[str, Any],
         requires_grad: bool = False,
         device="cuda",
+        module: Optional["Linear4bit"] = None,
         **kwargs,
     ) -> "Params4bit":
         self = torch.Tensor._make_subclass(cls, data.to(device))
@@ -284,6 +285,10 @@ def from_prequantized(
         self.bnb_quantized = True
 
         self.quant_storage = data.dtype
+        self.module = module
+
+        if self.module is not None:
+            self.module.quant_state = self.quant_state
 
         return self