Remove torch._running_with_deploy() from fbcode, Fix exception handling for torch.ops.load_libraries

Differential Revision: D78583233

Author: nipung90

torchrec/distributed/comm_ops.py

@@ -38,7 +38,7 @@
     torch.ops.load_library(
         "//deeplearning/fbgemm/fbgemm_gpu:merge_pooled_embeddings_cpu"
     )
-except OSError:
+except (OSError, RuntimeError):
     pass
 
 # OSS

torchrec/distributed/dist_data.py

@@ -105,8 +105,8 @@
 try:
     torch.ops.load_library("//deeplearning/fbgemm/fbgemm_gpu:sparse_ops")
     torch.ops.load_library("//deeplearning/fbgemm/fbgemm_gpu:sparse_ops_cpu")
-except OSError:
-    pass
+except (OSError, RuntimeError):
+    from fbgemm_gpu import sparse_ops, sparse_ops_cpu  # noqa: F401, E402
 
 
 logger: logging.Logger = logging.getLogger(__name__)

torchrec/distributed/embedding.py

@@ -113,8 +113,8 @@
     torch.ops.load_library("//deeplearning/fbgemm/fbgemm_gpu:sparse_ops")
     torch.ops.load_library("//deeplearning/fbgemm/fbgemm_gpu:sparse_ops_cpu")
     torch.ops.load_library("//deeplearning/fbgemm/fbgemm_gpu/codegen:index_select_ops")
-except OSError:
-    pass
+except (OSError, RuntimeError):
+    from fbgemm_gpu import sparse_ops, sparse_ops_cpu  # noqa: F401, E402
 
 
 def _pin_and_move(tensor: torch.Tensor, device: torch.device) -> torch.Tensor:

torchrec/distributed/embeddingbag.py

@@ -58,8 +58,8 @@
 try:
     torch.ops.load_library("//deeplearning/fbgemm/fbgemm_gpu:sparse_ops")
     torch.ops.load_library("//deeplearning/fbgemm/fbgemm_gpu:sparse_ops_cpu")
-except OSError:
-    pass
+except (OSError, RuntimeError):
+    from fbgemm_gpu import sparse_ops, sparse_ops_cpu  # noqa: F401, E402
 
 
 _DDP_STATE_DICT_PREFIX = "module."

torchrec/distributed/model_parallel.py

@@ -100,8 +100,8 @@
 try:
     torch.ops.load_library("//deeplearning/fbgemm/fbgemm_gpu:sparse_ops")
     torch.ops.load_library("//deeplearning/fbgemm/fbgemm_gpu:sparse_ops_cpu")
-except OSError:
-    pass
+except (OSError, RuntimeError):
+    from fbgemm_gpu import sparse_ops, sparse_ops_cpu  # noqa: F401, E402
 
 
 logger: logging.Logger = logging.getLogger(__name__)

torchrec/distributed/quant_embedding.py

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

torchrec/distributed/train_pipeline/tracing.py

@@ -80,8 +80,7 @@
 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.

torchrec/distributed/train_pipeline/train_pipelines.py

@@ -30,7 +30,7 @@
         "//deeplearning/fbgemm/fbgemm_gpu:intraining_embedding_pruning_gpu"
     )
 except OSError:
-    pass
+    from fbgemm_gpu import intraining_embedding_pruning_gpu  # noqa: F401, E402
 
 logger: logging.Logger = logging.getLogger(__name__)
 

torchrec/modules/itep_modules.py

@@ -79,8 +79,8 @@
 try:
     torch.ops.load_library("//deeplearning/fbgemm/fbgemm_gpu:sparse_ops")
     torch.ops.load_library("//deeplearning/fbgemm/fbgemm_gpu:sparse_ops_cpu")
-except OSError:
-    pass
+except (OSError, RuntimeError):
+    from fbgemm_gpu import sparse_ops, sparse_ops_cpu  # noqa: F401, E402
 
 # OSS
 try: