Fix fsdp weight tying

[bcui19]

What does this PR do?

When initializing FSDP with device='meta' it undoes weight tying. This is a known issue in PyTorch with deferred initialization. Additionally, in order to address this, all weight tied modules have to be in the same FSDP module, as a result we try our best to force the FSDP parameters into the same module.

What issue(s) does this change relate to?

CO-1511

Before submitting

• Have you read the contributor guidelines?
• Is this change a documentation change or typo fix? If so, skip the rest of this checklist.
• Was this change discussed/approved in a GitHub issue first? It is much more likely to be merged if so.
• Did you update any related docs and document your change?
• Did you update any related tests and add any new tests related to your change? (see testing)
• Did you run the tests locally to make sure they pass?
• Did you run pre-commit on your change? (see the pre-commit section of prerequisites)

[dakinggg]

LGTM as far as I can tell, will let abhi or other approve. Also, have you tested with the examples repo (without Vitaliy's recent fix)? I'd like to know two things 1) Does it properly respect the tied weights there? and 2) Does it change the memory/throughput?

[bcui19]

1. Does it properly respect the tied weights there? and 2) Does it change the memory/throughput?

1. Yes so long as the tied modules end up in the same FSDP module, which can be subject to 'min_params'
2. I checked the base memory usage of each model (for 125m params), without the fix it uses 2.69GB, with the fix it uses 2.23GB. Throughput is mildly affected https://wandb.ai/mosaic-ml/meta-tensors-python?workspace=user-bcui (< 1% change)

[vchiley]

Left a comment, but it generally looks good.

Also whats the point of file composer/scratch?