Add SimpleFSDP to leaf node

Summary: Torchrec should not trace into SimpleFSDP which will cause failures. We thus added it to the leaf node

Differential Revision: D71450848

Author: Microve

torchrec/distributed/train_pipeline/utils.py

@@ -39,6 +39,7 @@
 
 if not torch._running_with_deploy():
     from torch.distributed._composable.fsdp.fully_shard import FSDPModule as FSDP2
+    from torch.distributed.fb.simple_fsdp import SimpleFSDPModule as SimpleFSDP
 else:
 
     class FSDP2:
@@ -743,6 +744,7 @@ def is_leaf_module(self, m: torch.nn.Module, module_qualified_name: str) -> bool
             or module_qualified_name in self._leaf_modules
             or isinstance(m, FSDP)
             or isinstance(m, FSDP2)
+            or isinstance(m, SimpleFSDP)
         ):
             return True
         return super().is_leaf_module(m, module_qualified_name)