Unmap tensors on CPU to reduce temp VRAM overhead while loading

Author: turboderp

exllamav2/exllamav2_ext/cpp/safetensors.cpp

@@ -453,4 +453,65 @@ void safetensors_read_fb(uintptr_t handle, size_t beg, size_t size, torch::Tenso
             remaining -= chunk;
         }
     }
-}
+}
+
+void tensor_remap
+(
+    torch::Tensor tensor,
+    torch::Tensor index
+)
+{
+    TORCH_CHECK_SHAPES(tensor, 1, index, 0, 1);
+    TORCH_CHECK_DTYPE(tensor, kInt);
+    TORCH_CHECK_DTYPE(index, kInt);
+
+    int rows = tensor.size(0);
+    int cols = tensor.size(1);
+    uint32_t* temp = (uint32_t*) calloc(cols, sizeof(int));
+    uint32_t* a = (uint32_t*) tensor.data_ptr();
+    uint32_t* idx = (uint32_t*) index.data_ptr();
+
+    for (int r = 0; r < rows; ++r)
+    {
+        memcpy(temp, a, sizeof(uint32_t) * cols);
+        for (int c = 0; c < cols; ++c)
+        {
+            *a++ = temp[idx[c]];
+        }
+    }
+    free(temp);
+}
+
+void tensor_remap_4bit
+(
+    torch::Tensor tensor,
+    torch::Tensor index
+)
+{
+    TORCH_CHECK_SHAPES(index, 0, tensor, 1, 8);
+    TORCH_CHECK_DTYPE(tensor, kInt);
+    TORCH_CHECK_DTYPE(index, kInt);
+
+    int rows = tensor.size(0);
+    int cols = index.size(0);
+    uint32_t* temp = (uint32_t*) calloc(cols / 8, sizeof(int));
+    uint32_t* a = (uint32_t*) tensor.data_ptr();
+    uint32_t* idx = (uint32_t*) index.data_ptr();
+
+    for (int r = 0; r < rows; ++r)
+    {
+        memcpy(temp, a, sizeof(uint32_t) * cols / 8);
+        for (int c = 0; c < cols;)
+        {
+            uint32_t rv = 0;
+            for (int b = 0; b < 8; ++b, ++c)
+            {
+                uint32_t i = idx[c];
+                uint32_t v = (temp[i / 8] >> ((i & 7) * 4) & 0x0f);
+                rv |= v << (b * 4);
+            }
+            *a++ = rv;
+        }
+    }
+    free(temp);
+}

exllamav2/exllamav2_ext/cpp/safetensors.h

@@ -47,4 +47,17 @@ uintptr_t safetensors_open_fb(const char* filename);
 void safetensors_close_fb(uintptr_t handle);
 void safetensors_read_fb(uintptr_t handle, size_t beg, size_t size, torch::Tensor target);
 
+void tensor_remap
+(
+    torch::Tensor tensor,
+    torch::Tensor index
+);
+
+void tensor_remap_4bit
+(
+    torch::Tensor tensor,
+    torch::Tensor index
+);
+
+
 #endif

exllamav2/exllamav2_ext/ext_bindings.cpp

@@ -55,6 +55,8 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m)
     m.def("safetensors_pinned_buffer", &safetensors_pinned_buffer, "safetensors_pinned_buffer");
     m.def("safetensors_free_pinned_buffer", &safetensors_free_pinned_buffer, "safetensors_free_pinned_buffer");
     m.def("safetensors_read_fb", &safetensors_read_fb, "safetensors_read_fb");
+    m.def("tensor_remap", &tensor_remap, "tensor_remap");
+    m.def("tensor_remap_4bit", &tensor_remap_4bit, "tensor_remap_4bit");
 
     // qmatrix
 

exllamav2/ext.py

@@ -173,9 +173,9 @@ def find_msvc():
     # gcc / cl.exe flags
 
     if windows:
-        extra_cflags = ["/Ox", "/openmp"]
+        extra_cflags = ["/Ox"]
     else:
-        extra_cflags = ["-Ofast", "-fopenmp"]
+        extra_cflags = ["-Ofast"]
 
     if ext_debug:
         extra_cflags += ["-ftime-report", "-DTORCH_USE_CUDA_DSA"]

exllamav2/linear.py

@@ -8,6 +8,7 @@
 from exllamav2.compat import safe_move_tensor
 from exllamav2.tensor_p import BROADCAST_VC
 from exllamav2.util import unpack_4bit, pack_4bit
+import gc
 
 from typing import TYPE_CHECKING
 
@@ -118,7 +119,7 @@ def load(self,
         cfg = self.model.config
 
         if self.f_key: w = self.load_weight_fused(self.f_key, self.f_beg, self.f_end, self.in_features, self.out_features, self.altpack_qkv)
-        if w is None: w = self.load_weight()
+        if w is None: w = self.load_weight(cpu = output_map is not None)
 
         # Load quantized linear layer from dictionary
 
@@ -137,7 +138,7 @@ def load(self,
             self.q_tensors = w
 
             if unmap and "q_perm" in w:
-                perm = w["q_perm"]
+                perm = w["q_perm"].cpu()
                 del w["q_perm"]
                 del w["q_invperm"]
                 # w["q_perm"] = torch.arange(0, w["q_perm"].shape[-1], dtype = w["q_perm"].dtype, device = w["q_perm"].device)
@@ -146,8 +147,10 @@ def load(self,
                 perm = None
 
             if output_map is not None:
-                w["q_weight"] = w["q_weight"][:, output_map]
-                w["q_scale"] = pack_4bit(unpack_4bit(w["q_scale"])[:, output_map])
+                ext_c.tensor_remap(w["q_weight"], output_map)
+                ext_c.tensor_remap_4bit(w["q_scale"], output_map)
+                for k in w.keys():
+                    w[k] = safe_move_tensor(w[k], self.device())
 
             self.q_handle = ext.make_q_matrix(w,
                                               self.temp_dq,

exllamav2/module.py

@@ -60,7 +60,8 @@ def device(self) -> str:
     def load_multi(self,
                    key: str,
                    keys: list[str],
-                   measure: bool = False) -> int | dict[str: torch.Tensor]:
+                   measure: bool = False,
+                   cpu: bool = False) -> int | dict[str: torch.Tensor]:
 
         tensors = {}
         submap = {}
@@ -85,13 +86,14 @@ def load_multi(self,
                 if measure:
                     size += stfile.measure(key + "." + k)
                 else:
-                    tensors[k] = stfile.get_tensor(key + "." + k, device = self.device())
+                    tensors[k] = stfile.get_tensor(key + "." + k, device = self.device() if not cpu else "cpu")
 
         return size if measure else tensors
 
 
     def load_weight(self,
-                    override_key: str | None = None):
+                    override_key: str | None = None,
+                    cpu: bool = False):
 
         if override_key is not None:
             keys = [override_key]
@@ -105,14 +107,14 @@ def load_weight(self,
             # EXL2
 
             if key + ".q_weight" in self.model.config.tensor_file_map:
-                qtensors = self.load_multi(key, ["q_weight", "q_invperm", "q_scale", "q_scale_max", "q_groups", "q_perm", "bias"])
+                qtensors = self.load_multi(key, ["q_weight", "q_invperm", "q_scale", "q_scale_max", "q_groups", "q_perm", "bias"], cpu = cpu)
                 qtensors["q_perm"] = torch.argsort(qtensors["q_invperm"]).to(torch.int)
                 return qtensors
 
             # GPTQ
 
             if key + ".qweight" in self.model.config.tensor_file_map:
-                qtensors = self.load_multi(key, ["qweight", "qzeros", "scales", "g_idx", "bias"])
+                qtensors = self.load_multi(key, ["qweight", "qzeros", "scales", "g_idx", "bias"], cpu = cpu)
                 if "bias" in qtensors and torch.all(qtensors["bias"].eq(0)):
                     del qtensors["bias"]
                 qtensors["scales"] = qtensors["scales"].half()
@@ -122,14 +124,14 @@ def load_weight(self,
 
             if key + ".weight" in self.model.config.tensor_file_map:
                 if key + ".bias" in self.model.config.tensor_file_map:
-                    tensors = self.load_multi(key, ["weight", "bias"])
+                    tensors = self.load_multi(key, ["weight", "bias"], cpu = cpu)
                     tensor = tensors["weight"].half()
                     bias = tensors["bias"].half()
                     if self.model.config.arch.orig_weights_transposed and len(tensor.shape) == 2:
                         tensor = tensor.T
                     return nn.Parameter(tensor, requires_grad = False), nn.Parameter(bias, requires_grad = False)
                 else:
-                    tensors = self.load_multi(key, ["weight"])
+                    tensors = self.load_multi(key, ["weight"], cpu = cpu)
                     tensor = tensors["weight"].half()
                     # if self.model.config.arch.orig_weights_transposed:
                     #     tensor = tensor.T