FSDP init can crash with shared parameters

[rohan-varma]

🐛 Describe the bug

FSDP initialization can crash when modules with shared params are wrapped separately. For example, if wrap https://github.com/facebookresearch/multimodal/blob/679f3596e4c44b483c68d4023b24e3c7f77292b3/torchmultimodal/modules/losses/flava.py#L138 linear (decoder) separately from the main module and then wrap the main module with device_id argument, this will raise an error due to bias param being shared. The bias param would have already been moved to GPU by the linear wrapped FSDP unit, but then the higher-level wrapper would still expect it to be on CPU, resulting in this error:

pytorch/torch/distributed/fsdp/fully_sharded_data_parallel.py

Line 814 in 9e65e93

f"FSDP only supports single device modules, but got params on {module_devices}"

Versions

main

cc @ezyang @gchanan @zou3519 @pietern @mrshenli @pritamdamania87 @zhaojuanmao @satgera @rohan-varma @gqchen @aazzolini @osalpekar @jiayisuse @SciPioneer @H-Huang @kwen2501

[awgu]

What is the desired behavior here instead?

[rohan-varma]

@awgu Not sure - I understand we don't support shared params wrapped in different FSDP units, but we should probably have a different error than this which makes it very hard to root cause what the issue is and provides a poor UX.