Remove torch._running_with_deploy() from fbcode and add RuntimeError handling for load_library (#3207)

Summary:
Pull Request resolved: #3207

As per https://fb.workplace.com/groups/pytorch.dev/permalink/1828123831099422 we can now safely remove “torch.is_deploy_running”. This commit does this!

Additionally we are handling RuntimeErrors while using load_library in torchrec

Differential Revision: D78525065

Author: PaliC

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,7 +105,7 @@
 try:
     torch.ops.load_library("//deeplearning/fbgemm/fbgemm_gpu:sparse_ops")
     torch.ops.load_library("//deeplearning/fbgemm/fbgemm_gpu:sparse_ops_cpu")
-except OSError:
+except (OSError, RuntimeError):
     pass
 
 

torchrec/distributed/embedding.py

@@ -113,7 +113,7 @@
     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:
+except (OSError, RuntimeError):
     pass
 
 

torchrec/distributed/embeddingbag.py

@@ -58,7 +58,7 @@
 try:
     torch.ops.load_library("//deeplearning/fbgemm/fbgemm_gpu:sparse_ops")
     torch.ops.load_library("//deeplearning/fbgemm/fbgemm_gpu:sparse_ops_cpu")
-except OSError:
+except (OSError, RuntimeError):
     pass
 
 

torchrec/distributed/model_parallel.py

@@ -100,7 +100,7 @@
 try:
     torch.ops.load_library("//deeplearning/fbgemm/fbgemm_gpu:sparse_ops")
     torch.ops.load_library("//deeplearning/fbgemm/fbgemm_gpu:sparse_ops_cpu")
-except OSError:
+except (OSError, RuntimeError):
     pass
 
 

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

@@ -79,7 +79,7 @@
 try:
     torch.ops.load_library("//deeplearning/fbgemm/fbgemm_gpu:sparse_ops")
     torch.ops.load_library("//deeplearning/fbgemm/fbgemm_gpu:sparse_ops_cpu")
-except OSError:
+except (OSError, RuntimeError):
     pass
 
 # OSS

torchrec/quant/embedding_modules.py

@@ -44,7 +44,7 @@
     torch.ops.load_library(
         "//deeplearning/fbgemm/fbgemm_gpu:permute_multi_embedding_ops_gpu"
     )
-except OSError:
+except (OSError, RuntimeError):
     pass