Add ModelParallel layout map for Qwen3

keras_hub/src/models/qwen3/qwen3_backbone.py

@@ -187,3 +187,95 @@ 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 Qwen3
+        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)
+        # The reverse (output projection) embedding is `(hidden, vocab)`, the
+        # transpose of the input embedding above -- vocab-parallel output
+        # projection wants vocab sharded on the model axis, so this is
+        # `(data_dim, model_dim)`, not a copy of the input embedding's tuple.
+        layout_map["token_embedding/reverse_embeddings"] = (
+            data_dim,
+            model_dim,
+        )
+        # QKV kernels are `(hidden, num_heads, head_dim)` (see the
+        # `bqm,muh->bquh` / `bkm,mvh->bkvh` einsum equations in
+        # `Qwen3Attention.build`), so `hidden` is the contracting dim and
+        # must be sharded on the data axis, with heads sharded on the model
+        # axis -- this is Megatron column-parallel and matches this model's
+        # own `attention_output` rule below, which already shards heads on
+        # `model_dim`. 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 axis on the model-parallel
+        # 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 still shards on data for
+        # memory savings.
+        layout_map["transformer_layer.*self_attention.*query.kernel"] = (
+            data_dim,
+            model_dim,
+            None,
+        )
+        layout_map["transformer_layer.*self_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