Model parameters don't update with deepspeed integration
Author: avivbrokman
Pull request: no

### System Info

- `transformers` version: 4.31.0.dev0
- Platform: Linux-5.10.173-154.642.amzn2.x86_64-x86_64-with-glibc2.26
- Python version: 3.10.10
- Huggingface_hub version: 0.14.1
- Safetensors version: 0.3.1
- PyTorch version (GPU?): 2.0.0 (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
- Using distributed or parallel set-up in script?: yes

### Who can help?

@pacman100

### Information

- [X] The official example scripts
- [X] My own modified scripts

### Tasks

- [X] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [ ] My own task or dataset (give details below)

### Reproduction

(1) I noticed that performance wasn't improving over training epochs, so I added in the following code to `examples/pytorch/translation/run_translate.py` at line 250 in order to monitor parameter value changes:

```
from transformers import TrainerCallback
class MonitorParameterCallback(TrainerCallback):
    def on_train_begin(self, args, state, control, model, **kwargs):
        self.original_value = model.encoder.block[0].layer[0].SelfAttention.q.weight.clone()

    def on_epoch_end(self, args, state, control, model, **kwargs):
        new_value = model.encoder.block[0].layer[0].SelfAttention.q.weight.clone()
        change = new_value - self.original_value
        change_norm = float(change.square().sum())
        print('change = ', change_norm)
```

However, any other method of confirming parameter changes would suffice.  

(2) The following code (no deepspeed) prints out positive values for training loss and change in parameter values:

```
python examples/pytorch/translation/run_translation.py \
    --model_name_or_path t5-small \
    --do_train \
    --source_lang en \
    --target_lang ro \
    --dataset_name wmt16 \
    --dataset_config_name ro-en \
    --output_dir /tmp/tst-translation \
    --per_device_train_batch_size=4 \
    --per_device_eval_batch_size=4 \
    --overwrite_output_dir \
    --predict_with_generate \
    --max_train_samples 16
```
    
(3) The following code (with deepspeed) prints out positive values for training loss but 0 for change in parameter values:

```
deepspeed examples/pytorch/translation/run_translation.py \
    --model_name_or_path t5-small \
    --do_train \
    --source_lang en \
    --target_lang ro \
    --dataset_name wmt16 \
    --dataset_config_name ro-en \
    --output_dir /tmp/tst-translation \
    --per_device_train_batch_size=4 \
    --per_device_eval_batch_size=4 \
    --overwrite_output_dir \
    --predict_with_generate \
    --max_train_samples 16 \
    --deepspeed tests/deepspeed/ds_config_zero3.json
```


### Expected behavior

When training with deepspeed, parameter values of the model should update.

Comments:
I apologize if this is actually a deepspeed issue rather than a deepspeed integration issue—I'm having trouble parsing which option is the case.

I think it is expected behaviour, because Zero3 partitions the model weights. You should use `deepspeed.zero.GatheredParameters` context manager or you can check the partitioned parameters stored in the `param.ds_tensor` attribute. To prove it, you can check your `model.encoder.block[0].layer[0].SelfAttention.q.weight.data.shape`, it should be empty.

@1ytic, you're right—it was empty. It's not clear to me how to use the suggestions you gave me—can you provide a little more detail?


Try to use this tensor `model.encoder.block[0].layer[0].SelfAttention.q.weight.ds_tensor.clone()` in your `MonitorParameterCallback`.

Thanks!  It turns out, the issue was due to NaN loss from fp16.  

Hi @avivbrokman,  did you fix the problem? Do you mind sharing how you fixed it? Thanks 

Hi @avivbrokman, I met a similar problem when using zero3, the model parameters that are supposed to update does not update when I print them in forward()