Weight file sharding for large (e.g. 10GB+) models

[divyashreepathihalli]

No description provided.

[phoeenniixx]

Hi @divyashreepathihalli, I would like to work on this issue. I am new to keras-hub so can you please help me through the process?
As much as I could understand, models can be loaded using from_preset so we need to update the preset-utils and the from_preset functions in the base classes? (please correct me if I am wrong)
We need to update the load_task and load_preprocessor functions of KerasPresetLoader to add the sharding? We can leave the option of sharding to the user by adding one parameter (bool) to the from_preset which if true does the sharding or we can just do it ourselves by checking if the file size is more than 10GB.
(from_preset will change imo only if we keep the sharding option to the user)

[phoeenniixx]

Hi @divyashreepathihalli, if you liked my approach can i raise a PR for this?

[phoeenniixx]

Hi @sachinprasadhs, @divyashreepathihalli, I tried to implement a class for weight sharding that supports torch and tensorflow. It is based on these functions from transformers. Please look at it and tell me if I can raise a PR for this?
Notebook: https://colab.research.google.com/drive/1oB2vYzeqhZWvH-tPIXSl6nb1ppmc2ZEG?usp=sharing

[divyashreepathihalli]

Hi @phoeenniixx sorry about the late response. Thank you so much for your interest in contributing. This work is underway and @james77777778 is working on this feature.

[phoeenniixx]

No problem @divyashreepathihalli, is there any other issue I could start working on?