Wav2vec2 model has unknown attributes weight_g/weight_v when DeepSpeed ZeRO-3 is enabled
Author: jonnyli1125
Pull request: no
Labels: name: DeepSpeed

name: Audio

### System Info

transformers==4.40.1
torch==2.3.0
Python 3.8


### Who can help?

@sanchit-gandhi  @pacman100 

### Information

- [ ] The official example scripts
- [X] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [ ] My own task or dataset (give details below)

### Reproduction

Instantiate the model with deepspeed ZeRO-3 enabled.
```py
AutoModel.from_pretrained("facebook/wav2vec2-base-960h")
```

Error:
```
  File "/tmp/ray/session_2024-05-16_06-57-20_885527_8/runtime_resources/py_modules_files/_ray_pkg_bb8e0dd54cc12418/eclipse/train.py", line 301, in __init__
    self.audio_model = AutoModel.from_pretrained(config.audio_model_id)
  File "/tmp/ray/session_2024-05-16_06-57-20_885527_8/runtime_resources/pip/188f865d478e54893fae918526a104ea43ba8d05/virtualenv/lib/python3.8/site-packages/transformers/models/auto/auto_factory.py", line 563, in from_pretrained
    return model_class.from_pretrained(
  File "/tmp/ray/session_2024-05-16_06-57-20_885527_8/runtime_resources/pip/188f865d478e54893fae918526a104ea43ba8d05/virtualenv/lib/python3.8/site-packages/transformers/modeling_utils.py", line 3550, in from_pretrained
    model = cls(config, *model_args, **model_kwargs)
  File "/tmp/ray/session_2024-05-16_06-57-20_885527_8/runtime_resources/pip/188f865d478e54893fae918526a104ea43ba8d05/virtualenv/lib/python3.8/site-packages/deepspeed/runtime/zero/partition_parameters.py", line 503, in wrapper
    f(module, *args, **kwargs)
  File "/tmp/ray/session_2024-05-16_06-57-20_885527_8/runtime_resources/pip/188f865d478e54893fae918526a104ea43ba8d05/virtualenv/lib/python3.8/site-packages/transformers/models/wav2vec2/modeling_wav2vec2.py", line 1453, in __init__
    self.encoder = Wav2Vec2Encoder(config)
  File "/tmp/ray/session_2024-05-16_06-57-20_885527_8/runtime_resources/pip/188f865d478e54893fae918526a104ea43ba8d05/virtualenv/lib/python3.8/site-packages/deepspeed/runtime/zero/partition_parameters.py", line 503, in wrapper
    f(module, *args, **kwargs)
  File "/tmp/ray/session_2024-05-16_06-57-20_885527_8/runtime_resources/pip/188f865d478e54893fae918526a104ea43ba8d05/virtualenv/lib/python3.8/site-packages/transformers/models/wav2vec2/modeling_wav2vec2.py", line 752, in __init__
    self.pos_conv_embed = Wav2Vec2PositionalConvEmbedding(config)
  File "/tmp/ray/session_2024-05-16_06-57-20_885527_8/runtime_resources/pip/188f865d478e54893fae918526a104ea43ba8d05/virtualenv/lib/python3.8/site-packages/deepspeed/runtime/zero/partition_parameters.py", line 503, in wrapper
    f(module, *args, **kwargs)
  File "/tmp/ray/session_2024-05-16_06-57-20_885527_8/runtime_resources/pip/188f865d478e54893fae918526a104ea43ba8d05/virtualenv/lib/python3.8/site-packages/transformers/models/wav2vec2/modeling_wav2vec2.py", line 384, in __init__
    deepspeed.zero.register_external_parameter(self, self.conv.weight_v)
  File "/tmp/ray/session_2024-05-16_06-57-20_885527_8/runtime_resources/pip/188f865d478e54893fae918526a104ea43ba8d05/virtualenv/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1709, in __getattr__
    raise AttributeError(f"'{type(self).__name__}' object has no attribute '{name}'")
AttributeError: 'ParametrizedConv1d' object has no attribute 'weight_v'
```

### Expected behavior

PyTorch has deprecated the `weight_g` and `weight_v` attributes on modules wrapped with weight_norm, as per their official [documentation](https://pytorch.org/docs/2.3/generated/torch.nn.utils.weight_norm.html). The new attribute names should be `original0` and `original1`, respectively.

This is still being used with the wav2vec2 model definition, specifically when deepspeed ZeRO-3 is enabled:
https://github.com/huggingface/transformers/blob/b6eb708bf1903f40499558b64e6d9392c1b889e1/src/transformers/models/wav2vec2/modeling_wav2vec2.py#L396-L402

Relevant issues:
- https://github.com/pytorch/pytorch/issues/102999
- https://github.com/huggingface/transformers/issues/26796
- https://github.com/huggingface/transformers/pull/24030

Comments:
Hey @jonnyli1125, thanks for opening this issue!
This is indeed a deprecation issue! Would you like to open a PR to correct this behavior? We could probably add an if condition to register the right weights. 

Let me know if that works for you, otherwise I can do this!

@ylacombe Sure, I can do this.