Fix checkpoint loading with DeepSpeed ZeRO3
Author: tohtana
Pull request: yes

# 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

Comments:
The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/transformers/pr_42201). All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

tried the repro and indeed there was a bug. LGTM, thanks for the PR 

Thanks! We expected it to break and were gonna fix it but thanks a lot for doing it ahead of us @tohtana ! 
Actually we'd want to have deepseep use our parallelism when possible (because we shard in this func) but @3outeille and I will have a look at that once stuff are stabilized!

Thank you @3outeille and @ArthurZucker for merging this so quickly!

> Actually we'd want to have deepseep use our parallelism when possible (because we shard in this func) but @3outeille and I will have a look at that once stuff are stabilized!

That sounds great! Could you elaborate a bit more on what kind of integration you have in mind?
From the DeepSpeed side, we’re definitely open to extending our APIs to enable better interoperability with HF’s parallelism.

I need to have a look at how you implement TP today to give a better answer, but mostly if it can leverage the normal `from_pretrained` path for parallelism it would simplify things! 🤗 

You mean [AutoTP](https://github.com/deepspeedai/DeepSpeed/blob/master/blogs/huggingface-tp/README.md) in DeepSpeed? It's a very important feature for us. Please let me know if there is anything we could do on DeepSpeed side (e.g. adding APIs /tests).

Not sure if anyone is already aware of this issue - with this PR ZeRO3 is still not loading model weights properly, no error thrown but training loss looks abnormal (around 12, normally would be 1-2). @tohtana @ArthurZucker 

@zhangwj618 would you have a small reproducer? cc @3outeille 

@ArthurZucker I reran the test with the latest code and found out this issue has been resolved by #42626 .

Cool ty for checking!