Repository: huggingface/transformers
Pull request: Fix tests fsdp

# What does this PR do?

This PR fixes all fsdp tests as there were some regressions:
- fsdp typing changed so hf parser had issues
- `fused_adamw_torch` (new default for users with  torch >=2.8) + fp16 + FSDP don't work cc @winglian @qgallouedec did you face this issue ? 
- Fix cmd by using `fsdp_config`, this was broken for way too long 

## Run tests with 2 gpus

`RUN_SLOW=True pytest tests/fsdp/test_fsdp.py`

Changed files:
filename: src/transformers/training_args.py
patch: @@ -473,7 +473,7 @@ class TrainingArguments:
             When resuming training, whether or not to skip the epochs and batches to get the data loading at the same
             stage as in the previous training. If set to `True`, the training will begin faster (as that skipping step
             can take a long time) but will not yield the same results as the interrupted training would have.
-        fsdp (`bool`, `str` or list of [`~trainer_utils.FSDPOption`], *optional*, defaults to `[]`):
+        fsdp (`bool`, `str` or list of [`~trainer_utils.FSDPOption`], *optional*, defaults to `None`):
             Use PyTorch Distributed Parallel Training (in distributed training only).
 
             A list of options along the following:
@@ -1146,8 +1146,8 @@ class TrainingArguments:
             )
         },
     )
-    fsdp: Union[list[FSDPOption], str, bool] = field(
-        default_factory=list,
+    fsdp: Optional[Union[list[FSDPOption], str]] = field(
+        default=None,
         metadata={
             "help": (
                 "Whether or not to use PyTorch Fully Sharded Data Parallel (FSDP) training (in distributed training"
@@ -1734,10 +1734,13 @@ def __post_init__(self):
         if not isinstance(self.warmup_steps, int) or self.warmup_steps < 0:
             raise ValueError("warmup_steps must be of type int and must be 0 or a positive integer.")
 
-        if isinstance(self.fsdp, bool):
-            self.fsdp = [FSDPOption.FULL_SHARD] if self.fsdp else ""
-        if isinstance(self.fsdp, str):
+        if self.fsdp is None:
+            self.fsdp = []
+        elif self.fsdp is True:
+            self.fsdp = [FSDPOption.FULL_SHARD]
+        elif isinstance(self.fsdp, str):
             self.fsdp = [FSDPOption(s) for s in self.fsdp.split()]
+
         if self.fsdp == [FSDPOption.OFFLOAD]:
             raise ValueError(
                 "`--fsdp offload` can't work on its own. It needs to be added to `--fsdp full_shard` or "

filename: tests/fsdp/test_fsdp.py
patch: @@ -258,8 +258,12 @@ def test_fsdp_config_transformers_auto_wrap(self, sharding_strategy, dtype):
     def test_basic_run(self, sharding_strategy, dtype):
         launcher = get_launcher(distributed=True, use_accelerate=False)
         output_dir = self.get_auto_remove_tmp_dir()
+        fsdp_config = '{"fsdp_transformer_layer_cls_to_wrap": "BertLayer"}'
         args = self.get_base_args(output_dir, 1, 50).split() + [f"--{dtype}"]
-        fsdp_args = ["--fsdp", f"{sharding_strategy} auto_wrap", "--fsdp_transformer_layer_cls_to_wrap", "BertLayer"]
+        fsdp_args = ["--fsdp", f"{sharding_strategy} auto_wrap", "--fsdp_config", f"{fsdp_config}"]
+        if dtype == "fp16":
+            # fp16 + fsdp + fused adamw torch breaks so we switch optimizers
+            fsdp_args += ["--optim", "adamw_torch"]
         script = [f"{self.examples_dir_str}/pytorch/text-classification/run_glue.py"]
         cmd = launcher + script + args + fsdp_args
         execute_subprocess_async(cmd, env=self.get_env())
@@ -271,8 +275,12 @@ def test_basic_run(self, sharding_strategy, dtype):
     def test_basic_run_with_gradient_accumulation(self, sharding_strategy, dtype):
         launcher = get_launcher(distributed=True, use_accelerate=False)
         output_dir = self.get_auto_remove_tmp_dir()
+        fsdp_config = '{"fsdp_transformer_layer_cls_to_wrap": "BertLayer"}'
         args = self.get_base_args(output_dir, 1, 50).split() + [f"--{dtype}", "--gradient_accumulation_steps", "2"]
-        fsdp_args = ["--fsdp", f"{sharding_strategy} auto_wrap", "--fsdp_transformer_layer_cls_to_wrap", "BertLayer"]
+        fsdp_args = ["--fsdp", f"{sharding_strategy} auto_wrap", "--fsdp_config", f"{fsdp_config}"]
+        if dtype == "fp16":
+            # fp16 + fsdp + fused adamw torch breaks so we switch optimizers
+            fsdp_args += ["--optim", "adamw_torch"]
         script = [f"{self.examples_dir_str}/pytorch/text-classification/run_glue.py"]
         cmd = launcher + script + args + fsdp_args
         execute_subprocess_async(cmd, env=self.get_env())
@@ -285,7 +293,11 @@ def test_basic_run_with_cpu_offload(self, dtype):
         launcher = get_launcher(distributed=True, use_accelerate=False)
         output_dir = self.get_auto_remove_tmp_dir()
         args = self.get_base_args(output_dir, 1, 50).split() + [f"--{dtype}", "--max_steps", "10"]
-        fsdp_args = ["--fsdp", "full_shard auto_wrap offload", "--fsdp_transformer_layer_cls_to_wrap", "BertLayer"]
+        fsdp_config = '{"fsdp_transformer_layer_cls_to_wrap": "BertLayer"}'
+        fsdp_args = ["--fsdp", "full_shard auto_wrap offload", "--fsdp_config", f"{fsdp_config}"]
+        if dtype == "fp16":
+            # fp16 + fsdp + fused adamw torch breaks so we switch optimizers
+            fsdp_args += ["--optim", "adamw_torch"]
         script = [f"{self.examples_dir_str}/pytorch/text-classification/run_glue.py"]
         cmd = launcher + script + args + fsdp_args
         execute_subprocess_async(cmd, env=self.get_env())
@@ -295,7 +307,7 @@ def test_basic_run_with_cpu_offload(self, dtype):
     @run_first
     @slow
     def test_training_and_can_resume_normally(self, state_dict_type):
-        output_dir = self.get_auto_remove_tmp_dir("./xxx", after=False)
+        output_dir = self.get_auto_remove_tmp_dir()
 
         sharding_strategy = "full_shard"
         use_accelerate = state_dict_type == "SHARDED_STATE_DICT"
@@ -351,7 +363,7 @@ def test_fsdp_cpu_offloading(self):
     @require_fsdp_v2_version
     @require_accelerate_fsdp2
     def test_accelerate_fsdp2_integration(self):
-        output_dir = self.get_auto_remove_tmp_dir("./xxx", after=False)
+        output_dir = self.get_auto_remove_tmp_dir()
         sharding_strategy = "full_shard"
         use_accelerate = True
 
@@ -415,12 +427,8 @@ def test_fsdp2_cpu_offloading(self):
 
     def run_cmd_and_get_logs(self, use_accelerate, sharding_strategy, launcher, script, args, output_dir):
         if not use_accelerate:
-            fsdp_args = [
-                "--fsdp",
-                f"{sharding_strategy} auto_wrap",
-                "--fsdp_transformer_layer_cls_to_wrap",
-                "BertLayer",
-            ]
+            fsdp_config = '{"fsdp_transformer_layer_cls_to_wrap": "BertLayer"}'
+            fsdp_args = ["--fsdp", f"{sharding_strategy} auto_wrap", "--fsdp_config", f"{fsdp_config}"]
             cmd = launcher + script + args + fsdp_args
         else:
             fsdp_config = f"""