[GGUF] Reduce peak RAM usage by casting dequantized tensors early during load

src/transformers/modeling_gguf_pytorch_utils.py

@@ -463,7 +463,7 @@ def get_gguf_hf_weights_map(
     return gguf_to_hf_name_map
 
 
-def load_gguf_checkpoint(gguf_checkpoint_path, return_tensors=False, model_to_load=None):
+def load_gguf_checkpoint(gguf_checkpoint_path, return_tensors=False, model_to_load=None, torch_dtype=None):
     """
     Load a GGUF file and return a dictionary of parsed parameters containing tensors, the parsed
     tokenizer and config attributes.
@@ -474,6 +474,12 @@ def load_gguf_checkpoint(gguf_checkpoint_path, return_tensors=False, model_to_lo
         return_tensors (`bool`, defaults to `False`):
             Whether to read the tensors from the file and return them. Not doing so is faster
             and only loads the metadata in memory.
+        model_to_load (`nn.Module`, *optional*):
+            The model to load the weights into. This is used to map GGUF tensor names to
+            Transformers parameter names.
+        torch_dtype (`torch.dtype`, *optional*):
+            The desired `torch.dtype` for the loaded tensors. If provided, tensors will be
+            converted to this dtype immediately after dequantization to save memory.
     """
     if is_gguf_available() and is_torch_available():
         from gguf import GGUFReader, dequantize
@@ -644,7 +650,10 @@ def load_gguf_checkpoint(gguf_checkpoint_path, return_tensors=False, model_to_lo
 
             name = tensor_key_mapping[name]
 
-            parsed_parameters["tensors"][name] = torch.from_numpy(np.copy(weights))
+            tensor = torch.from_numpy(np.copy(weights))
+            if torch_dtype is not None:
+                tensor = tensor.to(torch_dtype)
+            parsed_parameters["tensors"][name] = tensor
 
     if len(reader_keys) > 0:
         logger.info(f"Some keys of the GGUF file were not considered: {reader_keys}")

src/transformers/modeling_utils.py

@@ -782,14 +782,16 @@ def _get_dtype(
                         dtype = sharded_metadata["dtype"]
                     elif state_dict is not None:
                         dtype = get_state_dict_dtype(state_dict)
+                    elif checkpoint_files is not None and checkpoint_files[0].endswith(".gguf"):
+                        dtype = torch.float32
                     else:
                         state_dict = load_state_dict(
                             checkpoint_files[0], map_location="meta", weights_only=weights_only
                         )
                         dtype = get_state_dict_dtype(state_dict)
                     logger.info(
-                        "Since the `dtype` attribute can't be found in model's config object, "
-                        "will use dtype={dtype} as derived from model's weights"
+                        f"Since the `dtype` attribute can't be found in model's config object, "
+                        f"will use dtype={dtype} as derived from model's weights"
                     )
             elif hasattr(torch, dtype):
                 dtype = getattr(torch, dtype)
@@ -4078,21 +4080,22 @@ def from_pretrained(
 
         is_quantized = hf_quantizer is not None
 
+        # Find the correct dtype based on current state
+        config, dtype = _get_dtype(
+            dtype, checkpoint_files, config, sharded_metadata, state_dict, weights_only, hf_quantizer
+        )
+
         if gguf_file:
             from .modeling_gguf_pytorch_utils import load_gguf_checkpoint
 
             # we need a dummy model to get the state_dict - for this reason, we keep the state_dict as if it was
             # passed directly as a kwarg from now on
             with torch.device("meta"):
                 dummy_model = cls(config)
-            state_dict = load_gguf_checkpoint(checkpoint_files[0], return_tensors=True, model_to_load=dummy_model)[
-                "tensors"
-            ]
 
-        # Find the correct dtype based on current state
-        config, dtype = _get_dtype(
-            dtype, checkpoint_files, config, sharded_metadata, state_dict, weights_only, hf_quantizer
-        )
+            state_dict = load_gguf_checkpoint(
+                checkpoint_files[0], return_tensors=True, model_to_load=dummy_model, torch_dtype=dtype
+            )["tensors"]
 
         config.name_or_path = pretrained_model_name_or_path