You signed in with another tab or window. Reload to refresh your session.You signed out in another tab or window. Reload to refresh your session.You switched accounts on another tab or window. Reload to refresh your session.Dismiss alert
We want model parallelism to be easy to use across the library.
We should add a get_layout_map() implementation to all backbones. This should be mostly copy paste from the Gemma version, since all transformers are pretty much the same weight structure.
We should also keep the docstring for the method on the Backbone base class. And factor out all the error checking somehow. That way the per model code here could be really minimal.
We want model parallelism to be easy to use across the library.
We should add a
get_layout_map()
implementation to all backbones. This should be mostly copy paste from the Gemma version, since all transformers are pretty much the same weight structure.keras-nlp/keras_nlp/src/models/gemma/gemma_backbone.py
Lines 227 to 315 in a00efc2
The text was updated successfully, but these errors were encountered: