Remove usage of _running_with_deploy in torchrec

Differential Revision: D78667510

Author: nipung90

torchrec/distributed/comm_ops.py

@@ -54,10 +54,10 @@ def get_gradient_division() -> bool:
 
 
 def set_use_sync_collectives(val: bool) -> None:
-    if val and torch._running_with_deploy():
-        raise RuntimeError(
-            "TorchRec sync_collectives are not supported in torch.deploy."
-        )
+    # if val and torch._running_with_deploy():
+    #     raise RuntimeError(
+    #         "TorchRec sync_collectives are not supported in torch.deploy."
+    #     )
 
     global USE_SYNC_COLLECTIVES
     USE_SYNC_COLLECTIVES = val
@@ -2356,7 +2356,7 @@ def backward(ctx, grad_output: Tensor) -> Tuple[None, None, Tensor]:
         return (None, None, myreq.dummy_tensor)
 
 
-if not torch._running_with_deploy():  # noqa C901
+if True:  # not torch._running_with_deploy():  # noqa C901
     # Torch Library op def can not be used in Deploy
     class AllToAllSingle(torch.autograd.Function):
         @staticmethod

torchrec/distributed/train_pipeline/tracing.py

@@ -13,13 +13,7 @@
 
 import torch
 
-if not torch._running_with_deploy():
-    from torch.distributed._composable.fsdp.fully_shard import FSDPModule as FSDP2
-else:
-
-    class FSDP2:
-        pass
-
+from torch.distributed._composable.fsdp.fully_shard import FSDPModule as FSDP2
 
 from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
 from torch.fx.immutable_collections import (

torchrec/distributed/train_pipeline/train_pipelines.py

@@ -80,8 +80,8 @@
 except ImportError:
     logger.warning("torchrec_use_sync_collectives is not available")
 
-if not torch._running_with_deploy():
-    torch.ops.import_module("fbgemm_gpu.sparse_ops")
+
+torch.ops.import_module("fbgemm_gpu.sparse_ops")
 
 
 # Note: doesn't make much sense but better than throwing.