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

Summary:

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

@@ -25,7 +25,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/dist_data.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/embedding.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/embeddingbag.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/model_parallel.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/quant_embedding.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/train_pipeline/tracing.py

@@ -12,7 +12,9 @@
 from typing import Any, Dict, List, Optional, Set, Tuple, Union
 
 import torch
-
+from torch.distributed._composable.fsdp.fully_shard import FSDPModule as FSDP2
+||||||| base:   7fac7a14884d - noreply+581177899995665: [version bump]: horiz...
+if not torch._running_with_deploy():
 from torch.distributed._composable.fsdp.fully_shard import FSDPModule as FSDP2
 
 from torch.distributed.fsdp import FullyShardedDataParallel as FSDP

torchrec/quant/embedding_modules.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/sparse/jagged_tensor.py

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