Checkpoint sharding (#16343)

* Sharded checkpoint support

* Handle distant sharded checkpoints

* Add tests

* TODO is done

* Apply suggestions from code review

Co-authored-by: Stas Bekman <stas00@users.noreply.github.com>

* Fix docstring

* Add example and format

* Address review comments

* More review comments

* End of merge

* Revert unintentional change

* VsCode what did you do?

* Style

* Changes

* Address final comments

* Quality

* Moar tests

* Move import beneath is_pt_available

Co-authored-by: Stas Bekman <stas00@users.noreply.github.com>

Author: sgugger

src/transformers/file_utils.py

@@ -52,6 +52,7 @@
     USE_JAX,
     USE_TF,
     USE_TORCH,
+    WEIGHTS_INDEX_NAME,
     WEIGHTS_NAME,
     ContextManagers,
     DummyObject,