Fix fsdp checkpoint strategy

composer/trainer/dist_strategy.py

@@ -149,6 +149,7 @@ def prepare_fsdp_module(model: torch.nn.Module, optimizers: Optional[Union[torch
         raise RuntimeError('To use FSDP with Composer, you must use torch>=1.12.0.')
     from torch.distributed.fsdp import (BackwardPrefetch, CPUOffload, FullyShardedDataParallel, MixedPrecision,
                                         ShardingStrategy)
+    from torch.distributed.fsdp.flatten_params_wrapper import FlattenParamsWrapper
 
     if optimizers:
         optimizers_tuple = ensure_tuple(optimizers)
@@ -261,6 +262,8 @@ def _auto_wrap_policy(module: torch.nn.Module, recurse: bool, unwrapped_params:
                 # If module has attribute `module._activation_checkpointing = ...`, always respect it
                 # Otherwise checkpoint if root object `obj.activation_checkpointing_fn(module)` is true
                 def _check_fn(module: torch.nn.Module) -> bool:
+                    if isinstance(module, (FullyShardedDataParallel, FlattenParamsWrapper)):
+                        return False
                     if hasattr(module, '_activation_checkpointing'):
                         return bool(module._activation_checkpointing)
                     if hasattr(obj, 'activation_checkpointing_fn') and isinstance(obj.activation_checkpointing_fn,