From cad495446dd484b8d85a0604df6a8e9572f241cc Mon Sep 17 00:00:00 2001 From: Marc Sun Date: Tue, 20 May 2025 18:39:46 +0200 Subject: [PATCH] quick fix --- src/diffusers/models/modeling_utils.py | 13 ++++++++++--- 1 file changed, 10 insertions(+), 3 deletions(-) diff --git a/src/diffusers/models/modeling_utils.py b/src/diffusers/models/modeling_utils.py index a6b00e8d11..3cbd09a3de 100644 --- a/src/diffusers/models/modeling_utils.py +++ b/src/diffusers/models/modeling_utils.py @@ -1185,8 +1185,16 @@ class ModelMixin(torch.nn.Module, PushToHubMixin): state_dict = None if not is_sharded: + map_location = "cpu" + if ( + device_map is not None + and hf_quantizer is not None + and hf_quantizer.quantization_config.quant_method == QuantizationMethod.TORCHAO + and hf_quantizer.quantization_config.quant_type in ["int4_weight_only", "autoquant"] + ): + map_location = torch.device([d for d in device_map.values() if d not in ["cpu", "disk"]][0]) # Time to load the checkpoint - state_dict = load_state_dict(resolved_model_file[0], disable_mmap=disable_mmap, dduf_entries=dduf_entries) + state_dict = load_state_dict(resolved_model_file[0], disable_mmap=disable_mmap, dduf_entries=dduf_entries, map_location=map_location) # We only fix it for non sharded checkpoints as we don't need it yet for sharded one. model._fix_state_dict_keys_on_load(state_dict) @@ -1443,10 +1451,9 @@ class ModelMixin(torch.nn.Module, PushToHubMixin): device_map is not None and hf_quantizer is not None and hf_quantizer.quantization_config.quant_method == QuantizationMethod.TORCHAO - and hf_quantizer.quantfization_config.quant_type in ["int4_weight_only", "autoquant"] + and hf_quantizer.quantization_config.quant_type in ["int4_weight_only", "autoquant"] ): map_location = torch.device([d for d in device_map.values() if d not in ["cpu", "disk"]][0]) - for shard_file in resolved_model_file: state_dict = load_state_dict(shard_file, dduf_entries=dduf_entries, map_location=map_location)