Sharded weights support #2218
Sharded weights support #2218
Conversation
james77777778
force-pushed
the
sharded_weights_support
branch
from
9c92ba4 to
bf9966a
Compare
|
@james77777778 thanks will take a look! We don't need to be backwards compatible here, the error message you have which an action the user can take is as good as we can do here I think. |
keras_hub/src/utils/preset_utils.py
Outdated
| dtype = keras.backend.standardize_dtype(dtype) | ||
| dtype_size = int( | ||
| ( | ||
| dtype.replace("bfloat", "") |
There was a problem hiding this comment.
can you explain what's going on here? maybe flip this to a dtype_size function (just so you can add a quick docstring?)
There was a problem hiding this comment.
I have updated the code, and it should be more explicit by using regex.
| "use_post_attention_norm": True, | ||
| "use_sliding_window_attention": True, | ||
| } | ||
| backbone = GemmaBackbone(**init_kwargs) # ~4.4MB |
There was a problem hiding this comment.
Can we make this even smaller? Feel free to use bert or something simple if its easier. Try to make this run as fast as possible while testing the business logic.
There was a problem hiding this comment.
I have changed the config to make the backbone smaller (422KB). It now takes only 2.5 seconds from the start to the end of the test.
james77777778
force-pushed
the
sharded_weights_support
branch
from
bf9966a to
2b70a6f
Compare
Please see the colab for an example using Gemma2 2B:
https://colab.research.google.com/drive/1iF_Psb6aEV2pkajT-q9ZBjpoO4RX4-Qa?usp=sharing
This PR adds support for sharded weights in
KerasPresetSaverandKerasPresetLoader.The default
max_shard_sizeis set to 10GB.Kindly ping @divyashreepathihalli @mattdangerw
Note: This feature requires the latest Keras (
git+https://github.com/keras-team/keras.git). It is difficult to ensure the backward compatibility.Related to #2084