Deepspeed + Trainer issue with zero shapes (Medusa training)
Author: Narsil
Pull request: no

### System Info

- `transformers` version: 4.38.1
- Platform: Linux-5.15.0-1048-aws-x86_64-with-glibc2.31
- Python version: 3.11.6
- Huggingface_hub version: 0.20.3
- Safetensors version: 0.4.2
- Accelerate version: 0.27.2
- Accelerate config:    not found
- PyTorch version (GPU?): 2.3.0.dev20240218+cu121 (True)
- Tensorflow version (GPU?): not installed (NA)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Using GPU in script?: Yes H100 
- Using distributed or parallel set-up in script?: Yes DS + Trainer





### Who can help?

@pacman100 

### 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


```python


model = AutoModel.from_pretrained(...)
medusa = MedusaHeads(model, ...)

...

def __init__(self, ...):
      self.head = nn.Linear(*model.lm_head.shape)
      
```

This doesn't work with Deepspeed because while using DeepSpeed, the model is loaded onto `meta` device, therefore the model has no shape.

On a similar note, I tried creating the shape from the config, but there are also issue when settting the weigths from the original model.


```python

def __init__(self, ...):
      self.head = nn.Linear(*model.lm_head.shape)
      self.head.weight.data[:] = model.lm_head.weight.data[:]
```

Fails with 
```
RuntimeError: The expanded size of the tensor (4096) must match the existing size (0) at non-singleton dimension 1.  Target sizes: [32000, 4096].  Tensor sizes: [0]
```

It's the same reason, because the weight.data is not loaded anymore.

How can I actually create the heads starting from the lm_head data while using DeepSpeed for TP ?

### Expected behavior

Getting the actual weigths from somewhere so the model heads can be started from somewhere

Comments:
Hello,

Thank you for raising the issue. This happens because the zero_init is enabled by default when using ZeRO-3 which shards the model as it is loaded in `AutoModel.from_pretrained()` call. Now, as the model is sharded the shapes of the weights are 0 and if the shapes of existing layers are used to initialize the new Medusa heads it will lead to heads with no elements which in return results in below error:
```
[rank0]:     assert param.ds_status == ZeroParamStatus.AVAILABLE, param.ds_summary()
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]: AssertionError: {'id': 995, 'status': 'NOT_AVAILABLE', 'numel': 0, 'ds_numel': 0, 'shape': (0,), 'ds_shape'
: (0, 0), 'requires_grad': True, 'grad_shape': None, 'persist': True, 'active_sub_modules': {1674}, 'ds_tensor.shape
': torch.Size([0])}
```

Now, to overcome this you have used the values from the model config which avoids the zero element heads. However, when you are initializing the medusa heads with existing lm head weights, as the weights of lm head are sharded it results in the error:
```
RuntimeError: The expanded size of the tensor (4096) must match the existing size (0) at non-singleton dimension 1.  Target sizes: [32000, 4096].  Tensor sizes: [0]
```

To overcome this, you need to use `deepspeed.zero.GatheredParameters(lm_head.weight)` context manager which gathers the weights of the specified layers so that each rank has all the params for that layer. After this, you should be able to copy the params of lm head to the new medusa heads.  Please refer the usage example of this at https://github.com/huggingface/transformers/blob/89c64817ce4172bc8bb58c675c445a63f16d0e38/src/transformers/modeling_utils.py#L1928-L1929

Finally, to resolve issue to enable only saving of the medusa weights instead of the whole model via `trainer.save_model`/`trainer.save_pretrained()`, you need to do do the following:

```
full_state_dict = trainer.accelerator.get_state_dict(self.deepspeed)
#choose the keys with medusa heads
medusa_sd ={k: v for k,v in full_state_dict.items() if "medusa" in k}
# save using save pretrained
if trainer.accelerator.is_main_process:
    model.save_pretrained(output_dir, state_dict=medusa_sd)
```

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md) are likely to be ignored.

Closing as this was resolved!