Fix tests fsdp

[SunMarc]

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

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[cyyever]

What issues are in hf parser?

[SunMarc]

What issues are in hf parser?

It has trouble dealing with union of different types and there were some strange cases involving optional. As you can see in this pr, --fsdp arg was incorrectly parsed. We need to fix the parser at some point ;)

[cyyever]

@SunMarc The current behaviour was made by me. Now if not specified fsdp_config is []. What tests are now broken so I have a chance to look and fix? We also have to fix union in parser because Optional[A] can be represented as A | None in newer typing.

[SunMarc]

If you are interested in fixing the parser, you can use play with the following script and use for example --fsdp arg to see if the trainings_args are correctly parsed ! Without this PR for example, you still see that
python script.py --fsdp "full_shard auto_wrap offload" will fail.

from transformers import TrainingArguments, HfArgumentParser

parser = HfArgumentParser((TrainingArguments,))
training_args, = parser.parse_args_into_dataclasses()

The tests should be fixed by this PR but this is not a proper fix tbh

[cyyever]

@SunMarc Yes, give me time to reproduce and fix.

[Cyrilvallez]

Thanks! Trusting you here for fsdp!