Wraps sharded model for proper access to it state_dict in FSDP strategy

[SpirinEgor]

What does this PR do?

Fixes #16526 by following previously deleted DDPFullyShardedStrategy

Does your PR introduce any breaking changes?

No, it doesn't.

[carmocca]

Thanks for looking into this!

This is only an issue in master, not 1.9, correct? In that case, this doesn't need a CHANGELOG entry.

Can you write a test?

[SpirinEgor]

This is only an issue in master, not 1.9, correct?

No, I faced this in 1.9 in the fsdp_native strategy. The master renames it to fsdp without any fixes from ddp_fully_sharded that are removed.

[SpirinEgor]

@carmocca I wrote tests for this. And I notice that the problem is actually occurred only for layers that is not wrapped, e.g. small layers. So, I parametrized tests with different wrapping policies and slightly changed BoringModel for correct asserts.

My build has one failed test, but it seems not my fault. What should I do?

And do you have any suggestions about modifying _LightningModuleWrapperBase?

[SpirinEgor]

@awaelchli I implemented what we discussed. For now, FSDP always aggregate full state dict on zero rank. Cpu offload depends on CPUOffload from initialized strategy. I'm not sure if this is okay, but this was enough to pass tests 😅 Setting cpu_offload=True in FullStateDictConfig lead to errors on CI, but not in my environment (details).

[SpirinEgor]

I rethought this logic with offloading to the CPU. It's not good to reuse this variable as it's intended for a completely different purpose. We need to figure out why this doesn't work on CI. Because it works on my setup (2xA100).

[awaelchli]

I'm looking into it!