layout_map patch for gemma-2b-it-keras

models.py

@@ -40,11 +40,27 @@ def get_default_layout_map(preset_name, device_mesh):
         or "vicuna" in preset_name
     ):
         layout_map = keras_hub.models.Llama3Backbone.get_layout_map(device_mesh)
+        # Default layout map patch:
         # This line is missing for some Llama models (TODO: fix this in keras_hub)
         layout_map["token_embedding/reverse_embeddings"] = ("batch", "model")
         return layout_map
+
     elif "gemma" in preset_name:
-        return keras_hub.models.GemmaBackbone.get_layout_map(device_mesh)
+        layout_map = keras_hub.models.GemmaBackbone.get_layout_map(device_mesh)
+
+        if "gemma-2b-" in preset_name:
+            # Default layout map patch:
+            # Gemma QKV weigts are shaped [NB_HEADS, EMBED_DIM, INNER_DIM]
+            # Llama QKV weights are shaped [EMBED_DIM, NB_HEADS, INNER_DIM]
+            # However:
+            # The default layout map for KQV weights on Gemma is: (model_dim,data_dim,None)
+            # Which means sharding NB_HEADS on the "model" dimension.
+            # But gemma-2b-it-keras has only 1 head so this won't work: must patch it
+            # TODO: fix this in the Gemma layout map in Keras hub.
+            patch_key = "decoder_block.*attention.*(query|key|value).kernel"
+            layout_map.pop(patch_key)
+            layout_map[patch_key] = (None, "model", "batch")
+            return layout_map
 
 
 def log_applied_layout_map(model):