Add Linen SPMD logical partition annotations for Qwen3 model sharding

Author: amepas

src/maxdiffusion/models/qwen3_flax.py

@@ -84,18 +84,21 @@ def __call__(self, x):
     gate_proj = nn.Dense(
         self.config.intermediate_size,
         use_bias=False,
+        kernel_init=nn.with_logical_partitioning(nn.initializers.lecun_normal(), ("embed", "mlp")),
         dtype=self.config.dtype,
         name="gate_proj",
     )
     up_proj = nn.Dense(
         self.config.intermediate_size,
         use_bias=False,
+        kernel_init=nn.with_logical_partitioning(nn.initializers.lecun_normal(), ("embed", "mlp")),
         dtype=self.config.dtype,
         name="up_proj",
     )
     down_proj = nn.Dense(
         self.config.hidden_size,
         use_bias=False,
+        kernel_init=nn.with_logical_partitioning(nn.initializers.lecun_normal(), ("mlp", "embed")),
         dtype=self.config.dtype,
         name="down_proj",
     )
@@ -172,24 +175,28 @@ def __call__(
     q_proj = nn.Dense(
         self.config.num_attention_heads * self.config.head_dim,
         use_bias=False,
+        kernel_init=nn.with_logical_partitioning(nn.initializers.lecun_normal(), ("embed", "heads")),
         dtype=self.config.dtype,
         name="q_proj",
     )
     k_proj = nn.Dense(
         self.config.num_key_value_heads * self.config.head_dim,
         use_bias=False,
+        kernel_init=nn.with_logical_partitioning(nn.initializers.lecun_normal(), ("embed", "heads")),
         dtype=self.config.dtype,
         name="k_proj",
     )
     v_proj = nn.Dense(
         self.config.num_key_value_heads * self.config.head_dim,
         use_bias=False,
+        kernel_init=nn.with_logical_partitioning(nn.initializers.lecun_normal(), ("embed", "heads")),
         dtype=self.config.dtype,
         name="v_proj",
     )
     o_proj = nn.Dense(
         self.config.hidden_size,
         use_bias=False,
+        kernel_init=nn.with_logical_partitioning(nn.initializers.lecun_normal(), ("heads", "embed")),
         dtype=self.config.dtype,
         name="o_proj",
     )
@@ -347,7 +354,9 @@ def __call__(
     embed_tokens = nn.Embed(
         num_embeddings=self.config.vocab_size,
         features=self.config.hidden_size,
-        embedding_init=nn.initializers.normal(stddev=self.config.hidden_size**-0.5),
+        embedding_init=nn.with_logical_partitioning(
+            nn.initializers.normal(stddev=self.config.hidden_size**-0.5), ("vocab", "embed")
+        ),
         dtype=self.config.dtype,
         name="embed_tokens",
     )