WIP: Compose TP + FSDP2

[S1ro1]

Reopening this: managed working TP + FSDP2 @ 2a13375. Run accelerate launch examples/fsdp2/fsdp2_tp.py --apply-tp --apply-fsdp on 8 GPUs (runs TP2 FSDP4) and can be compared with CUDA_VISIBLE_DEVICES=0,1,2,3 accelerate launch examples/fsdp2/fsdp2_tp.py --apply-fsdp with vanila FSDP2. Needs to be rechecked for some hanging memory though, as the memory seems to not fix.

Will further refactor to enable easier composition with FSDP2 and other parallelisms (namely refactor of #3604 to have a better api).

Will continuosly rely on this transformers pr.

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

[S1ro1]

Managed to make a minimal repro of composable tp+fsdp2 working, though it requires nightly torch 🚀

[github-actions]

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the contributing guidelines are likely to be ignored.

[kmehant]

pulse ping

[S1ro1]

@kmehant we have internally decided to have as much logic as possible in transformers, so this is postponed until that is ~resolved. You can watch that here

[kmehant]

I see @S1ro1 !

[github-actions]

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the contributing guidelines are likely to be ignored.

[S1ro1]

This is maybe gonna be a mess, I started to play and had success. At 2a13375 we have a working fsdp2+tp example, gonna try to clean this up a bit so we don't have CP related changes. This will also include the design for FSDP2 + {any}Parallelism that will then be used in CP PR. cc @SunMarc

[S1ro1]

Also superseeded by #3682