fix: CPU RAM efficient loading for nd or HSDP parallelisms

[kmehant]

What does this PR do?

This fixes cpu ram efficient loading bug for nd parallel or even HSDP (replicate + shard) which needs broadcasting the tensor to all world ranks from global rank 0. Since from_pretrained from transformers is designed to load on global rank 0 in ram efficient loading case.

Fixes

    dist.broadcast(full_tensor, src=0, group=device_mesh.get_group())
                                             ^^^^^^^^^^^^^^^^^^^^^^^
  File "/workspace/lib/python3.12/site-packages/torch/distributed/device_mesh.py", line 752, in get_group
    raise RuntimeError: ('Found the DeviceMesh have 2 dimensions', 'Optional kwarg `mesh_dim` needs to be specified when device_mesh.ndim > 1.', 'If you want to get the list of all the ProcessGroups in the DeviceMesh,please use `get_all_groups()` instead.')

Before submitting

• This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
• Did you read the contributor guideline,
  Pull Request section?
• Was this discussed/approved via a Github issue or the forum? 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, and
  here are tips on formatting docstrings.
• Did you write any new necessary tests?

Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

@S1ro1 @SunMarc @zach-huggingface

[kmehant]

@S1ro1 @SunMarc can we attend to this pressing bug please? Thank you.

[S1ro1]

LGTM!

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

[kmehant]

@S1ro1 Thank you, can we merge this?

[S1ro1]

Yep, sorry forgot to come back to this after checks passed haha!

[kmehant]

Hi @S1ro1 @SunMarc When are we making a release on accelerate and including these patches? Would be helpful having these released.

[SunMarc]

Hi @S1ro1 @SunMarc When are we making a release on accelerate and including these patches? Would be helpful having these released.

I will try to do a patch end of the week if I manage to fix some minor bugs !

[kmehant]

Thank you @SunMarc