Fix bug when subset of model parameters is passed into optimizer with FSDP

[sashaDoubov]

What does this PR do?

This PR addresses #3493, where we were adding all model parameters into the optimizer state after FSDP wrapping. This change uses the previous flow for param_groups > 1, where we:

1. Get the mapping between params in the optimizer group and their names by using the pointers (found with id) of the tensors
2. Store the extra param group info, such as LR, wd, etc. in a separate dict
3. Iterate over the fsdp wrapped parameters, and populates the param_groups based on the names of the fsdp wrapped tensor and the mapping from before.

This change relies on use_orig_params=True, so an error is thrown if a subset of parameters is passed into the param_groups.

NOTE: at the moment, tensor parallelism breaks with multiple param_groups (as the ptrs of TP'd model.parameters() != ptrs of elements in optimizer param_groups), so a special case was added to handle this.

What issue(s) does this change relate to?

Fixes #3493

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)

[mvpatel2000]

LGTM!

why are gc collect necessary? will approve after that

[sashaDoubov]

@mvpatel2000 I've now removed the gc.collect() calls , including the old one in test_fsdp_param_groups