enable xla fsdp (#32048)

* enable xla fsdp * add acceleration version check for xla fsdp

enable xla fsdp (#32048)
* enable xla fsdp * add acceleration version check for xla fsdp
46d09af4 · append-only · GitHub · 7ad784ae · 46d09af4
Unverified Commit 46d09af4 authored Aug 07, 2024 by append-only Committed by GitHub Aug 07, 2024
Hide whitespace changes
Inline Side-by-side

Showing with 1 addition and 1 deletion

src/transformers/training_args.py src/transformers/training_args.py +1 -1

No files found.
--- a/src/transformers/training_args.py
+++ b/src/transformers/training_args.py
@@ -1902,7 +1902,7 @@ class TrainingArguments:
                warnings.warn("`--xla_fsdp_grad_ckpt` is useful only when `--xla` is set to true.")

        # accelerate integration for FSDP
-        if len(self.fsdp) > 0 and not self.fsdp_config["xla"]:
+        if len(self.fsdp) > 0 and is_accelerate_available("0.28.0"):
            os.environ["ACCELERATE_USE_FSDP"] = "true"
            from accelerate.utils.constants import (
                FSDP_AUTO_WRAP_POLICY,