[qgallouedec]

Resolves huggingface/transformers#48841.

Trainer.evaluate() before train() prepares the model alone, and under FSDP2 prepare raises without an optimizer. The optimizer is only required so its parameters can be swapped for the sharded ones after fully_shard, so this keeps the error for an optimizer passed without its model and lets a model through alone. It comes out sharded, as with an optimizer. When train() then prepares it again with its optimizer, activation checkpointing is not applied a second time.

• 2x B200, tiny Qwen3, Trainer + FSDP2: evaluate() on a fresh Trainer now runs, with the same eval loss as without FSDP (6.957570). With activation checkpointing, evaluate then train ends at the same eval loss as train alone (6.773301).

[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.

[mangeshraut712]

Relaxing the check to "optimizer requires its model" matches what _prepare_fsdp2 actually needs, and the new test covers the model-only prepare. My question is about what happens next in the Trainer flow from the linked issue: evaluate() (model alone), then train() (model + optimizer). The test doesn't exercise that sequence, and I think it may not be idempotent. This is reasoned from the code, not run.

[SunMarc]

Thanks ! cc @winglian I think you wanted something similar ? maybe relax the condition on passing an optimizer together with the model