diff --git a/src/transformers/modeling_gguf_pytorch_utils.py b/src/transformers/modeling_gguf_pytorch_utils.py index 66306b6f71f6..4e6c21d8517f 100644 --- a/src/transformers/modeling_gguf_pytorch_utils.py +++ b/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}") diff --git a/src/transformers/modeling_utils.py b/src/transformers/modeling_utils.py index c6d694f68218..02c6b6c150d9 100644 --- a/src/transformers/modeling_utils.py +++ b/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,6 +4080,11 @@ 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 @@ -4085,14 +4092,10 @@ def from_pretrained( # 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