Add shared ModelParallel distribution-test helper; refactor Gemma onto it

Adds TestCase.run_distribution_test, a shared helper used by every
model's test_distribution: sharding-spec + full-weight-coverage
assertions, a real forward/backward/optimizer step (not just build()),
and a numerical-parity + short-training-convergence check against an
undistributed twin model.

Adds a CI-safe mesh-shape sweep (test_layout_map_mesh_shapes) across
representative 2D/3D device-mesh shapes sized to mirror real
accelerator-pod topologies, using frozen realistic dimensions with a
memory-conscious scale (full production-scale dims are exercised by
the separate live-preset sweep, gated behind a per-width-class memory
estimate so it never attempts an unsafe local build).

Adds a manual live-preset sweep (test_layout_map_live_presets) that
fetches every real preset's config from the registry and validates
sharding across the same mesh-shape matrix.

conftest.py now exposes 8 virtual CPU devices to JAX so distribution
tests run instead of silently skipping on a single-device host, and
registers the multi_device pytest marker.

Gemma's existing test_distribution and test_distribution_with_lora
are refactored onto the new helper; Gemma's layout map itself is
unchanged (already correct).

Author: pctablet505

conftest.py

@@ -1,7 +1,11 @@
 import os
 
-import keras
-import pytest
+if os.environ.get("KERAS_BACKEND") == "jax" and "XLA_FLAGS" not in os.environ:
+    # Expose virtual CPU devices so keras.distribution tests run in CI.
+    os.environ["XLA_FLAGS"] = "--xla_force_host_platform_device_count=8"
+
+import keras  # noqa: E402
+import pytest  # noqa: E402
 
 # OpenVINO supported test paths
 OPENVINO_SUPPORTED_PATHS = [
@@ -95,6 +99,10 @@ def pytest_configure(config):
         "markers",
         "kaggle_key_required: mark test needing a kaggle key",
     )
+    config.addinivalue_line(
+        "markers",
+        "multi_device: mark test as needing multiple devices",
+    )
 
 
 def pytest_collection_modifyitems(config, items):

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,