Add safetensors support

[jihnenglin]

In response to DeepSpeed-MII issue #272

Introduce compatibility for HF repositories containing solely .safetensors checkpoints. In case of coexistence, prioritize .safetensors checkpoints over .bin ones.

[jihnenglin]

@microsoft-github-policy-service agree company="Gigabyte Technology"

[jeffra]

This is great, thank you @jihnenglin! Can you also create a PR on the MII side to update https://github.com/microsoft/DeepSpeed-MII/blob/main/requirements/requirements.txt to add safetensors as a dependency?

[jihnenglin]

This is great, thank you @jihnenglin! Can you also create a PR on the MII side to update https://github.com/microsoft/DeepSpeed-MII/blob/main/requirements/requirements.txt to add safetensors as a dependency?

Sure, will do!

Edit: DeepSpeed-MII PR #296

[jihnenglin]

I've made it automatically fetch *.safentensors checkpoints by default, and only fallback to download *.bin checkpoints if there are no *.safentensors.

[mrwyattii]

I've made it automatically fetch *.safentensors checkpoints by default, and only fallback to download *.bin checkpoints if there are no *.safentensors.

Thanks @jihnenglin! I cleaned up your implementation a little bit so we could avoid some of the deeper nested if-else branching and make it a bit easier to read the code. LGTM

[loadams]

The only a6000 test failing is a known failure, so we could bypass this if we need.