Add handle for empty fsdp_config[xla_fsdp_v2] (#1706)

huggingface · Feb 23, 2024 · 990c203 · 990c203
1 parent cd81dbd
commit 990c203
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/optimum/onnxruntime/training_args.py b/optimum/onnxruntime/training_args.py
@@ -397,6 +397,7 @@ def __post_init__(self):
         ):
             raise ValueError("`min_num_params` and `transformer_layer_cls_to_wrap` are mutually exclusive.")
         self.fsdp_config["xla"] = self.fsdp_config.get("xla", False)
+        self.fsdp_config["xla_fsdp_v2"] = self.fsdp_config.get("xla_fsdp_v2", False)
         self.fsdp_config["xla_fsdp_grad_ckpt"] = self.fsdp_config.get("xla_fsdp_grad_ckpt", False)
         if self.fsdp_config["xla"]:
             if len(self.fsdp) > 0: