Partial checkpoint support for SMP

[cavdard]

What does this PR do?

• Adds 3 new training args(smp_save_partial and smp_load_partial) to support partial checkpointing with SMP. smp_tensor_parallel_full_model to apply tensor parallelism to whole model.
• Uses the right ranks for partial checkpoint saving in should_save.
• Uses local_state_dict() with partial checkpoint saving.
• Uses smp.save instead of torch.save when partial checkpoint saving is enabled.
• Uses smp.load instead of torch.load when partial checkpoint loading is enabled. Reorders partial checkpoint loading to happen after wrapping of model, since smp.load can only load to a smp model.
• Updated checks for the existence of checkpoint files since smp partial checkpoints contain postfixes in addition to filename(example: filename_0_0 or filename_0_0_0).
• Skip checkpoint sharding when smp is enabled.
• smp_gather is causing increased memory usage on GPU0 when tensor parallelism is enabled. Switches to distributed_concat for ddp.
• adds load_best_model_at_end support for SMP.

Fixes # (issue)

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.

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint.

[philschmid]

@cavdard could you please run make style to apply the correct coding formatting?

[cavdard]

@cavdard could you please run make style to apply the correct coding formatting?

Update: Resolved by running pip install -e .[quality]

@philschmid Having this error. Am I missing a step?

make style
black examples tests src utils
make: black: No such file or directory

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