Fix ModelParallel layout map axis order for Phi3 QKV kernels

Phi3's query/key/value kernels are shaped (hidden, heads, head_dim),
with hidden as the contracting dim and heads as the output dim --
the opposite axis order from Gemma's (heads, hidden, head_dim)
kernels. The layout map had the data/model axes swapped for these
kernels (and for token_embedding/reverse_embeddings), which put the
small, GQA-sized key/value heads axis on the model-parallel dim
instead of hidden, and the large hidden axis on the batch axis
instead of the smaller heads axis for query -- both backwards from
the intended Megatron-style column-parallel sharding, and prone to
IndivisibleError whenever the model-mesh dimension didn't evenly
divide the (typically much smaller) key/value head count.

Corrected the axis order so heads land on the model-parallel dim
(matching attention_output.kernel's existing convention) and hidden
lands on the data-parallel dim, with key/value heads left fully
replicated on the model dim regardless of divisibility (mirrors
Gemma's kv convention).

Also expands phi3_backbone_test.py with a mesh-shape sweep across
this repo's capped device-mesh shapes, covering both the real
"mini" (1:1 query:kv head ratio) and a synthetic GQA-ratio width
class, plus a live-preset sweep that fetches real config.json files
from every registered Phi3 preset and validates the layout map
against their actual dims (skipping any build whose estimated
memory footprint would be unsafe to run locally).

Author: pctablet505

keras_hub/src/models/phi3/phi3_backbone.py

@@ -207,3 +207,92 @@ def get_config(self):
             }
         )
         return config
+
+    @staticmethod
+    def get_layout_map(
+        device_mesh,
+        model_parallel_dim_name="model",
+        data_parallel_dim_name="batch",
+    ):
+        """Get a `keras.distribution.LayoutMap` for model parallel distribution.
+
+        The returned `LayoutMap` contains the sharding spec for the Phi3
+        backbone weights.
+
+        Args:
+            device_mesh: keras.distribution.DeviceMesh. The device mesh
+                instance for distribution.
+            model_parallel_dim_name: str. The axis name of the device mesh,
+                where the weights should be partitioned on.
+            data_parallel_dim_name: str. The axis name of the device mesh,
+                where the data should be partitioned on.
+
+        Returns:
+            `keras.distribution.LayoutMap` that contains the sharding spec
+            for all the model weights.
+        """
+        if not isinstance(device_mesh, keras.distribution.DeviceMesh):
+            raise ValueError(
+                "Invalid device_mesh type. Expected "
+                f"`keras.distribution.DeviceMesh`, got {type(device_mesh)}"
+            )
+        if model_parallel_dim_name not in device_mesh.axis_names:
+            raise ValueError(
+                f"{model_parallel_dim_name} is not found in the "
+                f"device_mesh.axis_names. {device_mesh.axis_names=}"
+            )
+        if data_parallel_dim_name not in device_mesh.axis_names:
+            raise ValueError(
+                f"{data_parallel_dim_name} is not found in the "
+                f"device_mesh.axis_names. {device_mesh.axis_names=}"
+            )
+
+        data_dim = data_parallel_dim_name
+        model_dim = model_parallel_dim_name
+        layout_map = keras.distribution.LayoutMap(device_mesh)
+        layout_map["token_embedding/embeddings"] = (model_dim, data_dim)
+        layout_map["token_embedding/reverse_embeddings"] = (
+            data_dim,
+            model_dim,
+        )
+        # Query/key/value kernels are `(hidden, heads, head_dim)` (einsum
+        # `bqm,muh->bquh` / `bkm,mvh->bkvh`), i.e. hidden is the contracting
+        # dim and heads is the output dim -- the opposite axis order from
+        # Gemma's `(heads, hidden, head_dim)` kernels. Column-parallel
+        # (Megatron) sharding therefore puts the *heads* axis on the model
+        # dim and the *hidden* (contracting) axis on the data dim, matching
+        # `attention_output.kernel`'s already-heads-on-model convention
+        # below. Key/value kernels are sized by num_key_value_heads (GQA),
+        # which is independent of and typically much smaller than
+        # num_query_heads -- sharding that small heads axis on the model dim
+        # would raise an IndivisibleError whenever the model mesh dimension
+        # doesn't divide it, so key/value heads are left fully replicated on
+        # that axis, while hidden is still sharded on the data dim for
+        # memory savings (mirrors Gemma's kv convention).
+        layout_map["transformer_layer.*attention.*query.kernel"] = (
+            data_dim,
+            model_dim,
+            None,
+        )
+        layout_map["transformer_layer.*attention.*(key|value).kernel"] = (
+            data_dim,
+            None,
+            None,
+        )
+        layout_map["transformer_layer.*attention_output.kernel"] = (
+            model_dim,
+            None,
+            data_dim,
+        )
+        layout_map[
+            "transformer_layer.*feedforward_intermediate_dense.kernel"
+        ] = (data_dim, model_dim)
+        layout_map["transformer_layer.*feedforward_gate_dense.kernel"] = (
+            data_dim,
+            model_dim,
+        )
+        layout_map["transformer_layer.*feedforward_output_dense.kernel"] = (
+            model_dim,
+            data_dim,
+        )
+        return layout_map