From 5ada9a07188e62359d275157cdcbfe432cda0c89 Mon Sep 17 00:00:00 2001 From: Travis Addair Date: Mon, 30 Jan 2023 11:39:04 -0800 Subject: [PATCH] Added FSDP distributed strategy --- ludwig/distributed/__init__.py | 8 +++++++- ludwig/distributed/ddp.py | 3 +++ 2 files changed, 10 insertions(+), 1 deletion(-) diff --git a/ludwig/distributed/__init__.py b/ludwig/distributed/__init__.py index 4da807c3c6d..b354cd8ef9b 100644 --- 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 index 31ad262836f..29c5810e430 100644 --- 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: