Repository: huggingface/transformers
Pull request: Fix checkpoint loading with DeepSpeed ZeRO3

# What does this PR do?

With the [latest revision](https://github.com/huggingface/transformers/commit/6f6095e0cf509f7384d3ce0c1804013ef6cafd5f), loading weights fails with an error when DeepSpeed ZeRO3 is enabled.
This PR resolves the issue.

Here is the stack trace of the error.
```
  File "/home/runner/work/DeepSpeed/DeepSpeed/tests/unit/runtime/zero/test_zero_nesting_init.py", line 70, in test_nested_parallel_init
    model = VisionEncoderDecoderModel.from_pretrained(
            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/runner/work/DeepSpeed/DeepSpeed/unit-test-venv/lib/python3.12/site-packages/transformers/modeling_utils.py", line 270, in _wrapper
    return func(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/runner/work/DeepSpeed/DeepSpeed/unit-test-venv/lib/python3.12/site-packages/transformers/modeling_utils.py", line 4122, in from_pretrained
    model, missing_keys, unexpected_keys, mismatched_keys, offload_index, error_msgs = cls._load_pretrained_model(
                                                                                       ^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/runner/work/DeepSpeed/DeepSpeed/unit-test-venv/lib/python3.12/site-packages/transformers/modeling_utils.py", line 4236, in _load_pretrained_model
    error_msgs += _load_state_dict_into_zero3_model(model, state_dict)
                  ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/runner/work/DeepSpeed/DeepSpeed/unit-test-venv/lib/python3.12/site-packages/transformers/integrations/deepspeed.py", line 302, in _load_state_dict_into_zero3_model
    state_dict = state_dict.copy()
                 ^^^^^^^^^^^^^^^
AttributeError: 'NoneType' object has no attribute 'copy'
```

This is the minimal repro. You can run this with `torchrun`.

```
def main() -> None:
    from transformers import VisionEncoderDecoderModel, __version__ as transformers_version
    from transformers.integrations.deepspeed import HfDeepSpeedConfig

    import deepspeed  # noqa: F401  # ensure DeepSpeed is available inside the process

    ds_config = {
        "train_batch_size": 1,
        "zero_optimization": {
            "stage": 3,
        },
    }
    # Keep the config object alive so transformers detects ZeRO-3.
    dschf = HfDeepSpeedConfig(ds_config)

    print("Attempting to load model with transformers version:", transformers_version)
    VisionEncoderDecoderModel.from_pretrained(
        "hf-internal-testing/tiny-random-VisionEncoderDecoderModel-vit-gpt2")
    # Keep the config alive until the end of the function.
    assert dschf is not None


if __name__ == "__main__":
    main()
```

## Before submitting
- [ ] This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
- [x] Did you read the [contributor guideline](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md#create-a-pull-request),
      Pull Request section?
- [ ] Was this discussed/approved via a Github issue or the [forum](https://discuss.huggingface.co/)? Please add a link
      to it if that's the case.
- [ ] Did you make sure to update the documentation with your changes? Here are the
      [documentation guidelines](https://github.com/huggingface/transformers/tree/main/docs), and
      [here are tips on formatting docstrings](https://github.com/huggingface/transformers/tree/main/docs#writing-source-documentation).
- [ ] Did you write any new necessary tests?


## Who can review?

Let me mention a few people who contributed to the related commits.

- model loading (from pretrained, etc): @CyrilVallez
- distributed: @3outeille @ArthurZucker
- Big Model Inference: @SunMarc

Changed files:
filename: src/transformers/modeling_utils.py
patch: @@ -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()