Make env variables optional for FSDP (#2998)

* Bookmark * Tests pass! * Fix imports * Try with raw dict * Make diff easier * Add defaults to all relevent areas * Rest of refactor * Fix all of benjamin's nits * Adjust logic based on Benjamin's feedback * Adjust for new logic
huggingface · Aug 12, 2024 · 3bde615 · 3bde615
1 parent dc3b5ad
commit 3bde615
Show file tree

Hide file tree

Showing 3 changed files with 322 additions and 153 deletions.
diff --git a/src/accelerate/state.py b/src/accelerate/state.py
@@ -894,7 +894,7 @@ def __init__(
                 DistributedType.MULTI_NPU,
                 DistributedType.MULTI_XPU,
             ]:
-                if os.environ.get("ACCELERATE_USE_FSDP", "false") == "true":
+                if os.environ.get("ACCELERATE_USE_FSDP", "false") == "true" or fsdp_plugin is not None:
                     self.distributed_type = DistributedType.FSDP
                     if self._mixed_precision != "no":
                         fsdp_plugin.set_mixed_precision(self._mixed_precision)