Clear the Optimizer before wrapping with FSDP

[bcui19]

What does this PR do?

If a model is created on device 'cpu' instead of 'meta,' the optimizer holds references to unsharded modules on GPU, thus causing the GPU to OOM when a model could fit onto memory.

We now clear the optimizer param groups and states before wrapping the model with FSDP.

Testing
Link to the gist where we now OOM on the forwards pass but not in the model wrapping:

https://gist.github.com/bcui19/c972aa2743e2149baf631b5f71a9fae8

What issue(s) does this change relate to?

Fixes CO-1409

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)

[abhi-mosaic]

Looks good to me, thank you!

[bandish-shah]

👍