[BUG] Zero3 mode generates an error of "ValueError: Fan in and fan out can not be computed for tensor with fewer than 2"

model =AlignModel.from_pretrained(self.args.vm_path)
This is the code that produces the error, simply using the AlignModel in transformers to read the local pre-trained weights.

  File "/opt/licy/MyVLM/model/builder.py", line 152, in load_model
    model =AlignModel.from_pretrained(self.args.vm_path)
  File "/usr/local/lib/python3.8/dist-packages/transformers/modeling_utils.py", line 3307, in from_pretrained
    ) = cls._load_pretrained_model(
  File "/usr/local/lib/python3.8/dist-packages/transformers/modeling_utils.py", line 3559, in _load_pretrained_model
    model.apply(model._initialize_weights)
  File "/usr/local/lib/python3.8/dist-packages/torch/nn/modules/module.py", line 885, in apply
    fn(self)
  File "/usr/local/lib/python3.8/dist-packages/transformers/modeling_utils.py", line 1388, in _initialize_weights
    self._init_weights(module)
  File "/usr/local/lib/python3.8/dist-packages/transformers/models/align/modeling_align.py", line 1189, in _init_weights
    nn.init.xavier_uniform_(module.text_projection.weight)
  File "/usr/local/lib/python3.8/dist-packages/torch/nn/init.py", line 323, in xavier_uniform_
    fan_in, fan_out = _calculate_fan_in_and_fan_out(tensor)
  File "/usr/local/lib/python3.8/dist-packages/torch/nn/init.py", line 287, in _calculate_fan_in_and_fan_out
    raise ValueError("Fan in and fan out can not be computed for tensor with fewer than 2 dimensions")
ValueError: Fan in and fan out can not be computed for tensor with fewer than 2 dimensions

But when I train with zero3 mode it produces the following error (which doesn't happen in zero2) and it doesn't report an error when I use CLipVisionModel or ConvnextModel with zero3.Can you help me with this? Thank you.

My model consists of an LLM and a multimodal model (ALIGNMODEL), and I was wondering if I could specify that ALIGNMODEL should not be split at zero3? If I can how do I do that?

Start with model/builder.py at line 152 and reproduce AlignModel.from_pretrained under ZeRO-3, comparing the behavior with ZeRO-2 and the listed CLIPVisionModel or ConvnextModel cases. Inspect the reported transformers modeling_align.py initialization path and DeepSpeed ZeRO-3 configuration. Done should identify the compatibility cause and establish whether AlignModel can be excluded from ZeRO-3, with a documented configuration or confirmed fix.