[zero3] nn.utils.weight_norm gotchas

@samyam and I have sorted out nn.utils.weight_norm gotchas in wav2vec2 HF code, and are looking at some better solutions but I will document this for posterity:

So here self.conv.weight Param is created and then instantly replaced by self.conv.weight_v and self.conv.weight_g which deepspeed's zero.Init currently misses since they happen outside of Conv1D init.

So below is the solution to the first part:

class Wav2Vec2PositionalConvEmbedding(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.conv = nn.Conv1d(
            config.hidden_size,
            config.hidden_size,
            kernel_size=config.num_conv_pos_embeddings,
            padding=config.num_conv_pos_embeddings // 2,
            groups=config.num_conv_pos_embedding_groups,
        )

        from transformers.integrations import is_deepspeed_zero3_enabled
        if is_deepspeed_zero3_enabled():
            import deepspeed
            with deepspeed.zero.GatheredParameters(self.conv.weight, modifier_rank=0):
                self.conv = nn.utils.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)
        else:
            self.conv = nn.utils.weight_norm(self.conv, name="weight", dim=2)

        self.padding = Wav2Vec2SamePadLayer(config.num_conv_pos_embeddings)
        self.activation = ACT2FN[config.feat_extract_activation]

The problem at the moment is that register_external_parameter doesn't partition the weights, so we have to continue manually gather things - and it will only kick in during the first forward call.

Therefore the workaround doesn't stop here, there is another place to fix:

class Wav2Vec2PreTrainedModel(PreTrainedModel):
[...]
    config_class = Wav2Vec2Config
    base_model_prefix = "wav2vec2"
    _keys_to_ignore_on_load_missing = [r"position_ids"]

    def _init_weights(self, module):
        """Initialize the weights"""
        if isinstance(module, nn.Linear):
            # Slightly different from the TF version which uses truncated_normal for initialization
            # cf https://github.com/pytorch/pytorch/pull/5617
            module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
        elif isinstance(module, (nn.LayerNorm, nn.GroupNorm)):
            module.bias.data.zero_()
            module.weight.data.fill_(1.0)
        elif isinstance(module, nn.Conv1d):
            from transformers.integrations import is_deepspeed_zero3_enabled
            if is_deepspeed_zero3_enabled():
                import deepspeed
                if hasattr(module, "weight_v") and hasattr(module, "weight_g"):
                    with deepspeed.zero.GatheredParameters([module.weight_v, module.weight_g], modifier_rank=0):
                        torch.nn.init.kaiming_normal_(module.weight.data)
                else:
                    with deepspeed.zero.GatheredParameters(module.weight, modifier_rank=0):
                        torch.nn.init.kaiming_normal_(module.weight.data)
            else:
                torch.nn.init.kaiming_normal_(module.weight.data)

sometimes Conv1D is used with nn.utils.weight_norm and at other times it's not - leading to this code...

This is far from "ease-of-use" so we are looking at some better options, but at the very least it now works! Yay!

Not asking for anything here, just documenting for posterity.

No target file, test, or acceptance criteria is named; the issue explicitly documents a workaround rather than requesting a change. If this is revisited, locate Wav2Vec2PositionalConvEmbedding and Wav2Vec2PreTrainedModel, then define the intended follow-up and how the DeepSpeed Zero3 behavior should be verified.