aws / sagemaker-huggingface-inference-toolkit

Apache License 2.0
240 stars 60 forks source link

Sharded checkpoint support #93

Closed JimAllanson closed 6 months ago

JimAllanson commented 1 year ago

Description of changes:

In version 4.18.0 of Transformers, support for sharded checkpoints was added: huggingface/transformers#16343 Models using this sharded format now exist in the wild, for example in Blip2 (PyTorch) or ESM-2 (TensorFlow+PyTorch)

These models currently fail to load in the AWS Deep Learning Containers with huggingface support, as the model files are filtered out during the initial cache population stage.

This change adds support for matching model files with filenames in the sharded format, as well as the sharded model index files that list the parts.


By submitting this pull request, I confirm that you can use, modify, copy, and redistribute this contribution, under the terms of your choice.

philschmid commented 1 year ago

Hey @BaiqingL, this PR https://github.com/aws/sagemaker-huggingface-inference-toolkit/pull/104 should have fixed that.