[amanzoni1]

What does this PR do?

Follow-up to #42521, as suggested when it was closed.
TrainingArguments._process_fsdp_args copies a fixed set of keys from fsdp_config into the kwargs of Accelerate's FullyShardedDataParallelPlugin. Three valid plugin options are not in that set, so they are silently dropped when set through fsdp_config, while the same options work through an Accelerate config file:

• mixed_precision_policy
• ignored_modules
• activation_checkpointing_offload (FSDP2, accelerate>=1.15.0)

This PR forwards them and documents them in the fsdp_config docstring. A mixed_precision_policy dict has its dtype strings converted to torch.dtype on a copy, so fsdp_config stays JSON serializable. activation_checkpointing_offload raises a clear error on Accelerate versions that don't have the option.

Repro

from transformers import TrainingArguments

args = TrainingArguments(
    output_dir="out",
    fsdp=True,
    fsdp_config={
        "mixed_precision_policy": {"param_dtype": "bf16", "reduce_dtype": "fp32", "output_dtype": "bf16"},
        "ignored_modules": "lm_head",
        "activation_checkpointing": True,
        "activation_checkpointing_offload": True,
    },
)
for key in ("mixed_precision_policy", "ignored_modules", "activation_checkpointing_offload"):
    print(key, "->", args.fsdp_plugin_args.get(key))

Before:

mixed_precision_policy -> None
ignored_modules -> None
activation_checkpointing_offload -> None

With this PR:

mixed_precision_policy -> {'param_dtype': torch.bfloat16, 'reduce_dtype': torch.float32, 'output_dtype': torch.bfloat16}
ignored_modules -> lm_head
activation_checkpointing_offload -> True

Who can review?

@SunMarc @ArthurZucker

[github-actions]

Thank you for your contribution 🤗!

CI Security Gate — waiting for maintainer approval

CI is approved automatically only for authors who already have a contribution merged in this repository. A maintainer will review this PR and approve CI to run on it.

[ArthurZucker]

I don't think we need the unitest here

[amanzoni1]

Removed the unit tests, thanks.