Parallelize the model download with the weight loading

[LysandreJik]

Feature request

It could be cool to load each individual shard as it finishes being downloaded instead of waiting for all the shards to be downloaded before loading them.

Right now, with the dynamic weight loading, we see a significant improvement in model loading time; this could help push this further.

cc @Cyrilvallez and @ArthurZucker probably

Motivation

Speed!

Your contribution

I can test it out 😉

[Aznix07]

Hi @LysandreJik, I'd like to take a shot at implementing this. A few questions before I start:

1. Which file/module handles the shard downloading and loading logic currently?
2. How does the existing "dynamic weight loading" work? Any specific files I should study first?
3. Are there any edge cases to consider (.g., failed downloads, memory constraints, dependency between shards) ?
4. Would you prefer an async-based approach or threading for parallel download + load?

Thanx!

[Cyrilvallez]

Hey @Aznix07! For this kind of core features, we are usually extremely picky and prefer to do it ourselves as it requires deep knowledge and understanding of the whole Transformers loading mechanisms... This is not intended to discourage you, but rather a friendly warning so that you do not waste your time in trying to work on this! It's definitely not a good one to start working on if you do not have any prior knowledge of the loading pipeline, as it's too involved between many different aspects of the library!