Generalize Optimizer validation to accommodate both FSDP 1.x and 2.0 #16732

Description & Motivation

Generalize Optimizer validation to accommodate both FSDP 1.x and 2.x

In preparing the downstream package I maintain (finetuning-scheduler) to support PyTorch 2.0's version of FSDP, I noticed that Lightning has yet to be adapted to accommodate FSDP 2.x's useful use_orig_params params mode.

Currently, false-positive "The optimizer does not seem to reference any FSDP parameters" ValueErrors will be triggered when the use_orig_params FSDP configuration is used because such a configuration does not require FlatParameter instances be present in the Optimizer.

https://github.com/Lightning-AI/lightning/blob/c35fee92f4ffdcfe5137c762f614b44ca7f8b7d0/src/lightning/fabric/strategies/fsdp.py#L344-L347

https://github.com/Lightning-AI/lightning/blob/c35fee92f4ffdcfe5137c762f614b44ca7f8b7d0/src/lightning/pytorch/strategies/fsdp.py#L282-L287

FSDP versions 1.13 and 2.0 both ensure FSDP flattened parameters have the _fsdp_flattened attribute set, allowing for a single check to work across both versions (this is the check used internally by PyTorch FSDP). For BC, we need to keep the existing FlatParameter-based check for versions below 1.13, but we should be able to remove that version guard once 1.12 is deprecated.

def _optimizer_has_flat_params(optimizer: Optimizer) -> bool:
    if _TORCH_GREATER_EQUAL_1_13:
        return any(getattr(param, _FSDP_FLATTENED, False) for param in optimizer.param_groups[0]["params"])
    else:
        from torch.distributed.fsdp import FlatParameter

        return any(isinstance(param, FlatParameter) for param in optimizer.param_groups[0]["params"])

I'm sure support for this (along with other FSDP 2.x features) will be arriving as the team's priorities dictate but I'm hoping this issue and the PR I'll be opening momentarily will allow PyTorch 2.x development/testing to continue for use_orig_params in a BC-compatible manner while also adding some useful 1.x and 2.x FSDP tests.

Pitch

The PR I'll be opening for this feature makes 4 minor changes:

1. Generalizes _optimizer_has_flat_params to accommodate FSDP 1.x and 2.0 usage.
   • To avoid directly depending upon PyTorch's protected _is_fsdp_flattened method and protected constant (which have both changed locations between versions), I've included the relevant constant and attribute check in fabric/strategies/fsdp.py directly rather than relying on a protected attribute whose location is version dependent.
2. Adds an FSDP 2.0 compatible auto_wrap_policy test
3. Adds an FSDP 2.0 _FSDPPolicy-based test that uses use_orig_params to validate the multi-version Optimizer validation while also testing the preferred _FSDPPolicy-based auto-wrapping approach moving forward
4. Guards the 1.x compatible auto_wrap_policy behind the necessary max_torch mark

Alternatives

One could use an FSDP instance-based check but that seems both more complicated and brittle. The _fsdp_flattened attribute is used extensively internally and was the best candidate attribute for this check IMHO.

Additional context

No response

cc @Borda @awaelchli @carmocca