Wav2vec2 model has unknown attributes weight_g/weight_v when DeepSpeed ZeRO-3 is enabled

[jonnyli1125]

System Info

transformers==4.40.1
torch==2.3.0
Python 3.8

Who can help?

@sanchit-gandhi @pacman100

Information

• The official example scripts
• 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.

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. 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:

transformers/src/transformers/models/wav2vec2/modeling_wav2vec2.py

Lines 396 to 402 in b6eb708

	if is_deepspeed_zero3_enabled():
	import deepspeed
	with deepspeed.zero.GatheredParameters(self.conv.weight, modifier_rank=0):
	self.conv = weight_norm(self.conv, name="weight", dim=2)
	deepspeed.zero.register_external_parameter(self, self.conv.weight_v)
	deepspeed.zero.register_external_parameter(self, self.conv.weight_g)

Relevant issues:

• Preserve weight_g/weight_v accessors on new weight_norm pytorch/pytorch#102999
• (False?) warning about weight_g/weight_v missing on WeightNorm on PyTorch #26796
• Use new parametrization based weight norm if available #24030

[ylacombe]

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!

[jonnyli1125]

@ylacombe Sure, I can do this.