int4 weight only quantized model using torchao

2025-05-09 20:27:25 +08:00
parent 3f1813d0d5
commit 054f599ebd
4 changed files with 288 additions and 37 deletions
@@ -10,11 +10,13 @@ class CpuOffloader:
        self.original_dtype = model.dtype
    
    def __enter__(self):
-        self.model.to(self.original_device, dtype=self.original_dtype)
+        if not hasattr(self.model,"torchao_quantized"):
+            self.model.to(self.original_device, dtype=self.original_dtype)
        return self.model
    
    def __exit__(self, *args):
-        self.model.to("cpu")
+        if not hasattr(self.model,"torchao_quantized"):
+            self.model.to("cpu")
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
            torch.cuda.synchronize()