Refactor Gemma distribution tests onto the shared helper

Depends on keras-team#2809 (adds TestCase.run_distribution_test) for tests to
pass -- this PR's own diff is independent (based directly on master,
not stacked) so it stays small and reviewable on its own.

Refactors Gemma's existing test_distribution/test_distribution_with_lora
onto the shared helper. Adds test_layout_map_mesh_shapes (CI-safe sweep
across representative 2D/3D mesh shapes) and test_layout_map_live_presets
(manual, Kaggle-gated sweep across the full preset registry). Gemma's
get_layout_map itself is unchanged -- already correct.

Author: pctablet505

keras_hub/src/models/gemma/gemma_backbone.py

@@ -296,11 +296,22 @@ def get_layout_map(
         # See https://arxiv.org/abs/2403.08295
         layout_map = keras.distribution.LayoutMap(device_mesh)
         layout_map["token_embedding/embeddings"] = (model_dim, data_dim)
-        layout_map["decoder_block.*attention.*(query|key|value).kernel"] = (
+        # Key/value kernels are sized by num_key_value_heads (GQA/MQA),
+        # which is independent of and typically much smaller than
+        # num_query_heads -- sharding that small axis on the model-parallel
+        # dim would raise an IndivisibleError whenever the device count
+        # doesn't divide it, so they are left replicated on that axis while
+        # query stays sharded.
+        layout_map["decoder_block.*attention.*query.kernel"] = (
             model_dim,
             data_dim,
             None,
         )
+        layout_map["decoder_block.*attention.*(key|value).kernel"] = (
+            None,
+            data_dim,
+            None,
+        )
         layout_map["decoder_block.*attention_output.kernel"] = (
             model_dim,
             None,