Adding support for FSDP2

[rithwik-db]

Moving all of the FSDP.summon_full_params to support both FSDP1 and FSDP2 (fully_shard). Unfortunately, functionality doesn't exist by default so I created a workaround to handle the args that we use within the codebase.

Tests:

• FSDP1: compose-rl-grpo-test-WHoytM
• FSDP2: compose-rl-grpo-test-fsdp2-q0Gy6Q <- tested with the changes from Composer's PR

They seem to fail due to unrelated issues that occur regardless of the changes (saving the checkpoints to UC Volumes), but the graphs look reasonable (link)

[bowenyang008]

After a read (w/o tests) I think the cognitive difficulty is too high for me, e.g., there are recurse=True/False, FSDP and non-FSDP module, DTensor vs non-DTensor etc and their combination makes it hard to reason its correctness. Here is my suggestion:

1. implements the summon_full_params in Composer, as it is a quite generic util
2. only supports either FSDP1 or DTensor (not just FSDP2), not a hybrid of them
3. if it is FSDP1, use FSDP1 summon; if there is DTensor, replace it with a full tensor, cache the FQN and tying for context swap as you have done otherwise do nothing
4. use this context regardless of if a module is FSDP module or not

let me know if this will work or if I missed anything

[rithwik-db]

@bowenyang008 re: this comment yeah this should work, moving it to Composer makes sense as I can just call it a summon_full_params_fsdp fn, although I will say that within compose-rl, if we do want to continue supporting FSDP1 alongside FSDP2, the changes in vllm_utils.py might have to stay the same since FSDP.summon_full_params() might require the module to be FSDP wrapped. Not 100% sure though, will check tomorrow.

[bowenyang008]

@bowenyang008 re: this comment yeah this should work, moving it to Composer makes sense as I can just call it a summon_full_params_fsdp fn, although I will say that within compose-rl, if we do want to continue supporting FSDP1 alongside FSDP2, the changes in vllm_utils.py might have to stay the same since FSDP.summon_full_params() might require the module to be FSDP wrapped. Not 100% sure though, will check tomorrow.

I think it is covered by my 3rd point above, summon_full_params can work on FSDP1 module or module w/ or w/o DTensor, the last case it is basically a null context

[bowenyang008]

also LGTM, before merge can you run one more e2e test to make sure it produces the same result?

[gupta-abhay]

lgtm - would prefer merging the composer PR in first and not having to push arbitrary branches

[rithwik-db]

also LGTM, before merge can you run one more e2e test to make sure it produces the same result?

Will merge after the tests pass once the composer PR goes in + E2E test (but will only merge both once I get GPUs 😭 )