Skip fsdp2 import if running with deploy (meta-pytorch#2483)

Summary:
Pull Request resolved: meta-pytorch#2483

title, this breaks deploy models

Reviewed By: sayitmemory

Differential Revision: D64237929

fbshipit-source-id: 0cc908549c3cb7f6e66eb6d1ef0d0ccd3241d122

Author: s4ayub

torchrec/distributed/train_pipeline/utils.py

@@ -31,9 +31,17 @@
     Union,
 )
 
+import torch
 from torch import distributed as dist
 
-from torch.distributed._composable.fsdp.fully_shard import FSDPModule as FSDP2
+if not torch._running_with_deploy():
+    from torch.distributed._composable.fsdp.fully_shard import FSDPModule as FSDP2
+else:
+
+    class FSDP2:
+        pass
+
+
 from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
 from torch.fx.immutable_collections import (
     immutable_dict as fx_immutable_dict,