Repository: huggingface/transformers
Pull request: Fix deepspeed with quantization

# What does this PR do?

Alright, this should fix all deepspeed issues (finally). See here https://github.com/huggingface/transformers/pull/36963/files#diff-6b72b98c4c2dcfc6cc606843917733f5d858374fbc22a735ff483bbc0c1e63eaL4248-L4252, `low_cpu_mem_usage` was forbidden with deepspeed, but then still overriden here in the quantized case https://github.com/huggingface/transformers/pull/36963/files#diff-6b72b98c4c2dcfc6cc606843917733f5d858374fbc22a735ff483bbc0c1e63eaL4358-L4361

cc @stas00, @winglian 

@stas00, next steps will be making sure tests are robust, and once this is done we'll try to see how we can truly harmonize the paths, so that everything is taking the same one! But this will be done after the tests, to avoid potential similar issues as what happened these past days 🙃

Changed files:
filename: src/transformers/modeling_utils.py
patch: @@ -3719,19 +3719,14 @@ def float(self, *args):
 
     @classmethod
     def get_init_context(cls, is_quantized: bool, _is_ds_init_called: bool):
-        # With deepspeed, we cannot initialize the model on meta device
         if is_deepspeed_zero3_enabled():
             init_contexts = [no_init_weights()]
+            # We cannot initialize the model on meta device with deepspeed when not quantized
             if not is_quantized and not _is_ds_init_called:
                 logger.info("Detected DeepSpeed ZeRO-3: activating zero.init() for this model")
-                init_contexts.extend(
-                    [
-                        deepspeed.zero.Init(config_dict_or_path=deepspeed_config()),
-                        set_zero3_state(),
-                    ]
-                )
+                init_contexts.extend([deepspeed.zero.Init(config_dict_or_path=deepspeed_config()), set_zero3_state()])
             elif is_quantized:
-                init_contexts.append(set_quantized_state())
+                init_contexts.extend([init_empty_weights(), set_quantized_state()])
         else:
             init_contexts = [no_init_weights(), init_empty_weights()]
 
@@ -4800,7 +4795,11 @@ def _load_pretrained_model(
                 continue
 
             map_location = "cpu"
-            if shard_file.endswith(".safetensors") and not is_hqq_or_bnb and not is_deepspeed_zero3_enabled():
+            if (
+                shard_file.endswith(".safetensors")
+                and not is_hqq_or_bnb
+                and not (is_deepspeed_zero3_enabled() and not is_quantized)
+            ):
                 map_location = "meta"
             elif (
                 device_map is not None
@@ -4822,7 +4821,7 @@ def _load_pretrained_model(
             # Fix the key names
             state_dict = {key_renaming_mapping[k]: v for k, v in state_dict.items() if k in key_renaming_mapping}
 
-            if is_deepspeed_zero3_enabled():
+            if is_deepspeed_zero3_enabled() and not is_quantized:
                 error_msgs += _load_state_dict_into_zero3_model(model_to_load, state_dict)
             # Skip it with fsdp on ranks other than 0
             elif not (is_fsdp_enabled() and not is_local_dist_rank_0() and not is_quantized):