[BUG] - <title>When I use fsdp, Because the flattened parameters, I always meet some question

[sqzhang-lazy]

Add Link

When I use fsdp, Because the flattened parameters, I always meet some question.
for examples:
RuntimeError: mat2 must be a matrix, got 1-D tensor
and
RuntimeError: weight should have at least three dimensions
It always occurred in some flattened model weights, sucn as conv, linear etc.
How can I solve this problem?

Describe the bug

When I use fsdp, Because the flattened parameters, I always meet some question
for examples:
RuntimeError: mat2 must be a matrix, got 1-D tensor
and
RuntimeError: weight should have at least three dimensions
It always occurred in some flattened model weights, sucn as conv, linear etc.
How can I solve this problem?

Describe your environment

Pytorch 2.1.0

cc @osalpekar @H-Huang @kwen2501

[yiliu30]

Hi @648652443, could please provide some code snippets to reproduce the bug?