Fix excessive CPU memory usage with FSDP and cpu_ram_efficient_loading
Author: matthewdouglas
Pull request: yes
Labels: name: PyTorch FSDP

# What does this PR do?

This PR fixes an issue with FSDP + CPU_RAM_EFFICIENT_LOADING where a copy of the parameters are loaded into CPU memory for each rank. The change offloads to CPU only for rank 0, and the rest on the meta device. On a typical 8x node this will dramatically decrease the system RAM overhead required to load a large model.

This is split from a previously reverted PR #32276 originally contributed by @winglian. The revert was due to issues we had with validating the change that have since been resolved.

The issue we encountered was specific to our cluster environment on AWS. With the [AWS EFI plugin for NCCL](https://github.com/aws/aws-ofi-nccl), we encountered consistent hangs. If we upgrade NCCL from the version bundled with PyTorch (2.20.5) to NCCL 2.22.3 via `pip install nvidia-nccl-cu12==2.22.3`, this issue is resolved. ([Internal discussion](https://huggingface.slack.com/archives/C043JTYE1MJ/p1724260357959149))

<!--
Congratulations! You've made it this far! You're not quite done yet though.

Once merged, your PR is going to appear in the release notes with the title you set, so make sure it's a great title that fully reflects the extent of your awesome contribution.

Then, please replace this with a description of the change and which issue is fixed (if applicable). Please also include relevant motivation and context. List any dependencies (if any) that are required for this change.

Once you're done, someone will review your PR shortly (see the section "Who can review?" below to tag some potential reviewers). They may suggest changes to make the code even better. If no one reviewed your PR after a week has passed, don't hesitate to post a new comment @-mentioning the same persons---sometimes notifications get lost.
-->

<!-- Remove if not applicable -->

Fixes #31721, #31577

## Before submitting
- [ ] This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
- [x] Did you read the [contributor guideline](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md#create-a-pull-request),
      Pull Request section?
- [ ] Was this discussed/approved via a Github issue or the [forum](https://discuss.huggingface.co/)? Please add a link
      to it if that's the case.
- [ ] Did you make sure to update the documentation with your changes? Here are the
      [documentation guidelines](https://github.com/huggingface/transformers/tree/main/docs), and
      [here are tips on formatting docstrings](https://github.com/huggingface/transformers/tree/main/docs#writing-source-documentation).
- [ ] Did you write any new necessary tests?


## Who can review?
@ArthurZucker @LysandreJik


Comments:
The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/transformers/pr_33154). All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

Ping me when this is ready for review! 

@ArthurZucker Ready!

thanks @matthewdouglas !

@ArthurZucker I've added more background to the description.

> The issue we encountered was specific to our cluster environment on AWS. With the [AWS EFI plugin for NCCL](https://github.com/aws/aws-ofi-nccl), we encountered consistent hangs. If we upgrade NCCL from the version bundled with PyTorch (2.20.5) to NCCL 2.22.3 via pip install nvidia-nccl-cu12==2.22.3, this issue is resolved. ([Internal discussion](https://huggingface.slack.com/archives/C043JTYE1MJ/p1724260357959149))

@ArthurZucker @matthewdouglas I tried this fix but im having similar NCCL issues as what you had. Unfortunately your suggestion to upgrade to latest is not working. I understand you have some internal debugging discussions on this topic. Is it possible to share NCCL env settings and other package versions, that might shed  light on the root cause?

**Update**: found the root cause and it was not an NCCL issue. Have submitted a fix to TRL for it