IndexError: index 0 is out of bounds for dimension 0 with size 0 with deepspeed zero3 init and BertModel
Author: MichelDucartier
Pull request: no
Labels: name: bug

### System Info

- `deepspeed` version: 0.18.4
- `transformers` version: 5.0.0
- Platform: Linux-5.14.21-150500.55.65_13.0.74-cray_shasta_c_64k-aarch64-with-glibc2.39
- Python version: 3.12.3
- Huggingface_hub version: 1.3.4
- Safetensors version: 0.5.3
- Accelerate version: 1.12.0
- Accelerate config: 	not found
- DeepSpeed version: 0.18.4
- PyTorch version (accelerator?): 2.8.0a0+5228986c39.nv25.05 (CUDA)
- GPU type: NVIDIA GH200 120GB

### Who can help?

_No response_

### Information

- [ ] The official example scripts
- [x] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [x] My own task or dataset (give details below)

### Reproduction

Create a `script.py`:
```py
import deepspeed
from transformers import AutoModel

with deepspeed.zero.Init():
    model = AutoModel.from_pretrained("google-bert/bert-base-uncased")

print(model)
```
And run
```
torchrun --nproc-per-node 1 script.py
```

I get the following error:
```
[rank0]: Traceback (most recent call last):
[rank0]:   File "/users/mzhang/meditron/multimodal/biomed/test.py", line 28, in <module>
[rank0]:     model = BertModel.from_pretrained("google-bert/bert-base-uncased")
[rank0]:             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/usr/local/lib/python3.12/dist-packages/transformers/modeling_utils.py", line 4072, in from_pretrained
[rank0]:     model = cls(config, *model_args, **model_kwargs)
[rank0]:             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/usr/local/lib/python3.12/dist-packages/deepspeed/runtime/zero/partition_parameters.py", line 529, in wrapper
[rank0]:     f(module, *args, **kwargs)
[rank0]:   File "/usr/local/lib/python3.12/dist-packages/transformers/models/bert/modeling_bert.py", line 629, in __init__
[rank0]:     self.post_init()
[rank0]:   File "/usr/local/lib/python3.12/dist-packages/transformers/modeling_utils.py", line 1362, in post_init
[rank0]:     self.init_weights()
[rank0]:   File "/usr/local/lib/python3.12/dist-packages/transformers/modeling_utils.py", line 3080, in init_weights
[rank0]:     self.initialize_weights()
[rank0]:   File "/usr/local/lib/python3.12/dist-packages/torch/utils/_contextlib.py", line 116, in decorate_context
[rank0]:     return func(*args, **kwargs)
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/usr/lib/python3.12/contextlib.py", line 81, in inner
[rank0]:     return func(*args, **kwds)
[rank0]:            ^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/usr/local/lib/python3.12/dist-packages/transformers/modeling_utils.py", line 2373, in initialize_weights
[rank0]:     self.smart_apply(self._initialize_weights)
[rank0]:   File "/usr/local/lib/python3.12/dist-packages/transformers/modeling_utils.py", line 2366, in smart_apply
[rank0]:     module.smart_apply(fn)
[rank0]:   File "/usr/local/lib/python3.12/dist-packages/transformers/modeling_utils.py", line 2366, in smart_apply
[rank0]:     module.smart_apply(fn)
[rank0]:   File "/usr/local/lib/python3.12/dist-packages/transformers/modeling_utils.py", line 2367, in smart_apply
[rank0]:     fn(self)
[rank0]:   File "/usr/local/lib/python3.12/dist-packages/transformers/modeling_utils.py", line 2344, in _initialize_weights
[rank0]:     self._init_weights(module)
[rank0]:   File "/usr/local/lib/python3.12/dist-packages/torch/utils/_contextlib.py", line 116, in decorate_context
[rank0]:     return func(*args, **kwargs)
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/usr/local/lib/python3.12/dist-packages/transformers/models/bert/modeling_bert.py", line 566, in _init_weights
[rank0]:     super()._init_weights(module)
[rank0]:   File "/usr/local/lib/python3.12/dist-packages/torch/utils/_contextlib.py", line 116, in decorate_context
[rank0]:     return func(*args, **kwargs)
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/usr/local/lib/python3.12/dist-packages/transformers/modeling_utils.py", line 2305, in _init_weights
[rank0]:     init.zeros_(module.weight[module.padding_idx])
[rank0]:                 ~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/usr/local/lib/python3.12/dist-packages/torch/utils/_device.py", line 104, in __torch_function__
[rank0]:     return func(*args, **kwargs)
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^
[rank0]: IndexError: index 0 is out of bounds for dimension 0 with size 0
```

### Expected behavior

The model should load the model properly using deepspeed zero 3. This works in transformers pre v5
I think the bug is linked to the fact that deepspeed shards the parameters before transformers model construction.

Comments:
cc @sunmarc maybe?

I realize this example might not be the recommended way of using Deepspeed. I provided a better example in https://github.com/huggingface/transformers/issues/43638

Closing this issue in favor of the other one