Models With Tied Weights Need Re-Tieing After FSDP Param Init

[fabianlim]

What does this PR do?

Currently in FullyShardedDataParallelPlugin, the param_init_fn is set when sync_module_states=True. This is required by FSDP to initialize the shards (i.e. rank > 0) params in a variety of situations, including the important one where the parameter was on torch.device("meta") because low_cpu_mem_mode was used.

However FullyShardedDataParallelPlugin.param_init_fn is now set to

self.param_init_fn = lambda x: x.to_empty(device=device, recurse=False)

which causes problems when there are tied weights. Consider the following scenario

1. there are two modules A and B that share a tied weight A.weight = B.weight
2. A.to_empty() is called, then A.weight will be reassigned to a new tensor.
3. Similarly B.empty() is called, and then now A.weight != B.weight.

This is observed to cause problems in the low_cpu_mem_mode=True case, because now when getting the managed_params in FSDP, see here,

• in rank = 0, param_init_fn is not called, because low_cpu_mem_mode will load weights in this shard. So managed_params will not have duplicates if weights are tied.
• In other ranks > 0, param_init_fn is called and managed_params may have duplicates.

Then when FSDP calls _sync_module_params_and_buffers, the torch.distributed._broadcast_coalesced will be trying to communicate different number of tensors. This is not as intended and causes unexpected behaviors.

Fixes inconsistency in current logic in general, in particular is required to fix huggingface/trl#2089 to completion.

Before submitting

• This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
• Did you read the contributor guideline,
  Pull Request section?
• Was this discussed/approved via a Github issue or the forum? Please add a link
  to it if that's the case.
• Did you make sure to update the documentation with your changes? Here are the
  documentation guidelines, and
  here are tips on formatting docstrings.
• Did you write any new necessary tests?

Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

@muellerzr

[muellerzr]

Thanks for the fix, just one question on when this should actually be being set