Clear the Optimizer before wrapping with FSDP (#1732)

* Adding in performance registration

* clearing optimizer param groups before FSDP wrapping

* fixing comments

Co-authored-by: Brandon Cui <bcui@Brandons-MBP.hsd1.ca.comcast.net>

composer/trainer/dist_strategy.py

@@ -151,6 +151,17 @@ def prepare_fsdp_module(model: torch.nn.Module, optimizers: Optional[Union[torch
     from torch.distributed.fsdp import (BackwardPrefetch, CPUOffload, FullyShardedDataParallel, MixedPrecision,
                                         ShardingStrategy)
 
+    if optimizers:
+        optimizers_tuple = ensure_tuple(optimizers)
+        if len(optimizers_tuple) != 1:
+            raise NotImplementedError(f'Only one optimizer is supported; found {len(optimizers_tuple)} optimizers')
+
+        # clearing optimizer param groups and state
+        # that will be recreated at the end of prepare_fsdp_module
+        optim = optimizers_tuple[0]
+        optim.param_groups.clear()
+        optim.state.clear()
+
     sharding_map = {
         'NO_SHARD': ShardingStrategy.NO_SHARD,
         'SHARD_GRAD_OP': ShardingStrategy.SHARD_GRAD_OP,
@@ -281,8 +292,6 @@ def _check_fn(module: torch.nn.Module) -> bool:
     # Rebuild optimizer now that parameters are sharded
     if optimizers:
         optimizers_tuple = ensure_tuple(optimizers)
-        if len(optimizers_tuple) != 1:
-            raise NotImplementedError(f'Only one optimizer is supported; found {len(optimizers_tuple)} optimizers')
         optim = optimizers_tuple[0]
-        optim.param_groups = []
+        optim.param_groups.clear()
         optim.add_param_group({'params': list(model.parameters())})