Added FSDP distributed strategy (#3026)

ludwig-ai · Jan 30, 2023 · b5139c7 · b5139c7
1 parent 35fde29
commit b5139c7
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 1 deletion.
diff --git a/ludwig/distributed/__init__.py b/ludwig/distributed/__init__.py
@@ -10,6 +10,12 @@ def load_ddp():
     return DDPStrategy
 
 
+def load_fsdp():
+    from ludwig.distributed.fsdp import FSDPStrategy
+
+    return FSDPStrategy
+
+
 def load_horovod():
     from ludwig.distributed.horovod import HorovodStrategy
 
@@ -20,7 +26,7 @@ def load_local():
     return LocalStrategy
 
 
-STRATEGIES = {"ddp": load_ddp, "horovod": load_horovod, "local": load_local}
+STRATEGIES = {"ddp": load_ddp, "fsdp": load_fsdp, "horovod": load_horovod, "local": load_local}
 
 
 def get_current_dist_strategy(allow_local=True) -> Type[DistributedStrategy]:

diff --git a/ludwig/distributed/ddp.py b/ludwig/distributed/ddp.py
@@ -16,6 +16,9 @@
 class DDPStrategy(DistributedStrategy):
     def __init__(self):
         self._local_rank, self._local_size = local_rank_and_size()
+        self._log_on_init()
+
+    def _log_on_init(self):
         logging.info("Using DDP strategy")
 
     def wrap_model(self, model: nn.Module) -> nn.Module: