enable xla fsdp (#32048)

* enable xla fsdp * add acceleration version check for xla fsdp
huggingface · Aug 7, 2024 · 46d09af · 46d09af
1 parent 7ad784a
commit 46d09af
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/src/transformers/training_args.py b/src/transformers/training_args.py
@@ -1902,7 +1902,7 @@ def __post_init__(self):
                 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,