Fix checkpoint loading with DeepSpeed ZeRO3 (#42201)

fix checkpoint loading with DeepSpeed ZeRO3

Signed-off-by: Masahiro Tanaka <[email protected]>
Co-authored-by: Ferdinand Mom <[email protected]>

Author: tohtana

src/transformers/modeling_utils.py

@@ -4233,6 +4233,19 @@ def _load_pretrained_model(
         error_msgs = []
 
         if is_deepspeed_zero3_enabled() and not is_quantized:
+            if state_dict is None:
+                if checkpoint_files is None:
+                    raise ValueError(
+                        "DeepSpeed ZeRO-3 initialization requires a state_dict or checkpoint files to load from."
+                    )
+                merged_state_dict = {}
+                for ckpt_file in checkpoint_files:
+                    merged_state_dict.update(
+                        load_state_dict(
+                            ckpt_file, is_quantized=is_quantized, map_location="cpu", weights_only=weights_only
+                        )
+                    )
+                state_dict = merged_state_dict
             error_msgs += _load_state_dict_into_zero3_model(model, state_dict)
             # This is not true but for now we assume only best-case scenario with deepspeed, i.e. perfectly matching checkpoints
             missing_keys, unexpected_keys, mismatched_keys, misc = set(), set(), set(), set()