Skip to content

Commit

Permalink
Model loading speed optimization (huggingface#5635)
Browse files Browse the repository at this point in the history
Move unchanging operation out of loop for speed benefit.
  • Loading branch information
RyanJDick authored Nov 3, 2023
1 parent 9c299f1 commit 2c764e9
Showing 1 changed file with 2 additions and 1 deletion.
3 changes: 2 additions & 1 deletion models/modeling_utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -134,6 +134,8 @@ def load_model_dict_into_meta(model, state_dict, device=None, dtype=None, model_
device = device or torch.device("cpu")
dtype = dtype or torch.float32

accepts_dtype = "dtype" in set(inspect.signature(set_module_tensor_to_device).parameters.keys())

unexpected_keys = []
empty_state_dict = model.state_dict()
for param_name, param in state_dict.items():
Expand All @@ -147,7 +149,6 @@ def load_model_dict_into_meta(model, state_dict, device=None, dtype=None, model_
f"Cannot load {model_name_or_path_str}because {param_name} expected shape {empty_state_dict[param_name]}, but got {param.shape}. If you want to instead overwrite randomly initialized weights, please make sure to pass both `low_cpu_mem_usage=False` and `ignore_mismatched_sizes=True`. For more information, see also: https://github.com/huggingface/diffusers/issues/1619#issuecomment-1345604389 as an example."
)

accepts_dtype = "dtype" in set(inspect.signature(set_module_tensor_to_device).parameters.keys())
if accepts_dtype:
set_module_tensor_to_device(model, param_name, device, value=param, dtype=dtype)
else:
Expand Down

0 comments on commit 2c764e9

Please sign in to comment.