Guard thunder/distributed/transforms/ by torch.distributed.is_availab…

…le (#212)
Lightning-AI · Apr 19, 2024 · 91ac188 · 91ac188
1 parent bb80c5a
commit 91ac188
Showing 1 changed file with 7 additions and 2 deletions.
diff --git a/thunder/distributed/transforms/__init__.py b/thunder/distributed/transforms/__init__.py
@@ -1,6 +1,11 @@
-from .ddp import optimize_allreduce_in_ddp_backward
-from .fsdp import FSDPCommBucketing
+import torch
 
+if torch.distributed.is_available():
+    from .ddp import optimize_allreduce_in_ddp_backward
+    from .fsdp import FSDPCommBucketing
+else:
+    optimize_allreduce_in_ddp_backward = None
+    FSDPCommBucketing = None
 
 __all__ = [
     "optimize_allreduce_in_ddp_backward",