Add NNXFlaxQwen3Model and NNX Qwen3 sub-modules

Author: amepas

src/maxdiffusion/models/qwen3_flax.py

@@ -16,6 +16,7 @@
 
 import math
 from typing import Any, List, Optional, Tuple
+from flax import nnx
 import flax.linen as nn
 import jax
 import jax.numpy as jnp
@@ -400,6 +401,224 @@ def __call__(
     return hidden_states, all_hidden_states
 
 
+# -----------------------------------------------------------------------------
+# NNX Transformer Model Implementations for Qwen3
+# -----------------------------------------------------------------------------
+
+
+class NNXFlaxQwen3RMSNorm(nnx.Module):
+
+  def __init__(self, rngs: nnx.Rngs, dim: int, eps: float = 1e-6, dtype: jnp.dtype = jnp.float32):
+    self.eps = eps
+    self.dtype = dtype
+    self.weight = nnx.Param(jnp.ones((dim,), dtype=dtype))
+
+  def __call__(self, x: jnp.ndarray) -> jnp.ndarray:
+    x_float = x.astype(jnp.float32)
+    variance = jnp.mean(jnp.square(x_float), axis=-1, keepdims=True)
+    normed = x_float * jax.lax.rsqrt(variance + self.eps)
+    return (normed.astype(self.dtype)) * self.weight.value
+
+
+class NNXFlaxQwen3MLP(nnx.Module):
+
+  def __init__(self, rngs: nnx.Rngs, config: FlaxQwen3Config):
+    self.config = config
+    self.gate_proj = nnx.Linear(
+        in_features=config.hidden_size,
+        out_features=config.intermediate_size,
+        use_bias=False,
+        kernel_init=nnx.with_partitioning(nnx.initializers.lecun_normal(), ("embed", "mlp")),
+        dtype=config.dtype,
+        rngs=rngs,
+    )
+    self.up_proj = nnx.Linear(
+        in_features=config.hidden_size,
+        out_features=config.intermediate_size,
+        use_bias=False,
+        kernel_init=nnx.with_partitioning(nnx.initializers.lecun_normal(), ("embed", "mlp")),
+        dtype=config.dtype,
+        rngs=rngs,
+    )
+    self.down_proj = nnx.Linear(
+        in_features=config.intermediate_size,
+        out_features=config.hidden_size,
+        use_bias=False,
+        kernel_init=nnx.with_partitioning(nnx.initializers.lecun_normal(), ("mlp", "embed")),
+        dtype=config.dtype,
+        rngs=rngs,
+    )
+
+  def __call__(self, x: jnp.ndarray) -> jnp.ndarray:
+    return self.down_proj(jax.nn.silu(self.gate_proj(x)) * self.up_proj(x))
+
+
+class NNXFlaxQwen3Attention(nnx.Module):
+
+  def __init__(self, rngs: nnx.Rngs, config: FlaxQwen3Config):
+    self.config = config
+    self.num_heads = config.num_attention_heads
+    self.num_kv_heads = config.num_key_value_heads
+    self.head_dim = config.head_dim
+
+    self.q_proj = nnx.Linear(
+        in_features=config.hidden_size,
+        out_features=config.num_attention_heads * config.head_dim,
+        use_bias=False,
+        kernel_init=nnx.with_partitioning(nnx.initializers.lecun_normal(), ("embed", "heads")),
+        dtype=config.dtype,
+        rngs=rngs,
+    )
+    self.k_proj = nnx.Linear(
+        in_features=config.hidden_size,
+        out_features=config.num_key_value_heads * config.head_dim,
+        use_bias=False,
+        kernel_init=nnx.with_partitioning(nnx.initializers.lecun_normal(), ("embed", "heads")),
+        dtype=config.dtype,
+        rngs=rngs,
+    )
+    self.v_proj = nnx.Linear(
+        in_features=config.hidden_size,
+        out_features=config.num_key_value_heads * config.head_dim,
+        use_bias=False,
+        kernel_init=nnx.with_partitioning(nnx.initializers.lecun_normal(), ("embed", "heads")),
+        dtype=config.dtype,
+        rngs=rngs,
+    )
+    self.o_proj = nnx.Linear(
+        in_features=config.num_attention_heads * config.head_dim,
+        out_features=config.hidden_size,
+        use_bias=False,
+        kernel_init=nnx.with_partitioning(nnx.initializers.lecun_normal(), ("heads", "embed")),
+        dtype=config.dtype,
+        rngs=rngs,
+    )
+
+    self.q_norm = NNXFlaxQwen3RMSNorm(rngs=rngs, dim=config.head_dim, eps=config.rms_norm_eps, dtype=config.dtype)
+    self.k_norm = NNXFlaxQwen3RMSNorm(rngs=rngs, dim=config.head_dim, eps=config.rms_norm_eps, dtype=config.dtype)
+
+  def __call__(
+      self,
+      x: jnp.ndarray,
+      attention_mask: Optional[jnp.ndarray] = None,
+      cos_table: Optional[jnp.ndarray] = None,
+      sin_table: Optional[jnp.ndarray] = None,
+  ) -> jnp.ndarray:
+    batch_size, seq_len, _ = x.shape
+
+    q = self.q_proj(x)
+    k = self.k_proj(x)
+    v = self.v_proj(x)
+
+    q = q.reshape(batch_size, seq_len, self.num_heads, self.head_dim)
+    k = k.reshape(batch_size, seq_len, self.num_kv_heads, self.head_dim)
+    v = v.reshape(batch_size, seq_len, self.num_kv_heads, self.head_dim)
+
+    q = self.q_norm(q)
+    k = self.k_norm(k)
+
+    if cos_table is not None and sin_table is not None:
+      q, k = apply_qwen3_rotary_pos_emb(q, k, cos_table, sin_table)
+
+    if self.num_kv_heads != self.num_heads:
+      num_repeats = self.num_heads // self.num_kv_heads
+      k = jnp.repeat(k, num_repeats, axis=2)
+      v = jnp.repeat(v, num_repeats, axis=2)
+
+    q = jnp.transpose(q, (0, 2, 1, 3))
+    k = jnp.transpose(k, (0, 2, 1, 3))
+    v = jnp.transpose(v, (0, 2, 1, 3))
+
+    scale = 1.0 / math.sqrt(self.head_dim)
+    scores = jnp.matmul(q, jnp.transpose(k, (0, 1, 3, 2))) * scale
+
+    if attention_mask is not None:
+      scores = scores + attention_mask
+
+    attn_probs = jax.nn.softmax(scores, axis=-1)
+    output = jnp.matmul(attn_probs, v)
+    output = jnp.transpose(output, (0, 2, 1, 3))
+    output = output.reshape(batch_size, seq_len, -1)
+    output = self.o_proj(output)
+    return output
+
+
+class NNXFlaxQwen3DecoderLayer(nnx.Module):
+
+  def __init__(self, rngs: nnx.Rngs, config: FlaxQwen3Config):
+    self.config = config
+    self.input_layernorm = NNXFlaxQwen3RMSNorm(
+        rngs=rngs, dim=config.hidden_size, eps=config.rms_norm_eps, dtype=config.dtype
+    )
+    self.self_attn = NNXFlaxQwen3Attention(rngs=rngs, config=config)
+    self.post_attention_layernorm = NNXFlaxQwen3RMSNorm(
+        rngs=rngs, dim=config.hidden_size, eps=config.rms_norm_eps, dtype=config.dtype
+    )
+    self.mlp = NNXFlaxQwen3MLP(rngs=rngs, config=config)
+
+  def __call__(
+      self,
+      x: jnp.ndarray,
+      attention_mask: Optional[jnp.ndarray] = None,
+      cos_table: Optional[jnp.ndarray] = None,
+      sin_table: Optional[jnp.ndarray] = None,
+  ) -> jnp.ndarray:
+    residual = x
+    normed_x = self.input_layernorm(x)
+    attn_out = self.self_attn(
+        normed_x,
+        attention_mask=attention_mask,
+        cos_table=cos_table,
+        sin_table=sin_table,
+    )
+    x = residual + attn_out
+
+    mlp_out = self.mlp(self.post_attention_layernorm(x))
+    x = x + mlp_out
+    return x
+
+
+class NNXFlaxQwen3Model(nnx.Module):
+
+  def __init__(self, rngs: nnx.Rngs, config: FlaxQwen3Config):
+    self.config = config
+    self.embed_tokens = nnx.Embed(
+        num_embeddings=config.vocab_size,
+        features=config.hidden_size,
+        embedding_init=nnx.with_partitioning(nnx.initializers.normal(stddev=config.hidden_size**-0.5), ("vocab", "embed")),
+        dtype=config.dtype,
+        rngs=rngs,
+    )
+    self.layers = [NNXFlaxQwen3DecoderLayer(rngs=rngs, config=config) for _ in range(config.num_hidden_layers)]
+    self.norm = NNXFlaxQwen3RMSNorm(rngs=rngs, dim=config.hidden_size, eps=config.rms_norm_eps, dtype=config.dtype)
+
+  def __call__(
+      self,
+      input_ids: jnp.ndarray,
+      attention_mask: Optional[jnp.ndarray] = None,
+  ) -> Tuple[jnp.ndarray, List[jnp.ndarray]]:
+    hidden_states = self.embed_tokens(input_ids)
+    all_hidden_states = [hidden_states]
+
+    cos_table, sin_table = precompute_qwen3_freqs_cis(
+        head_dim=self.config.head_dim,
+        max_seq_len=self.config.max_position_embeddings,
+        theta=self.config.rope_theta,
+    )
+
+    for layer in self.layers:
+      hidden_states = layer(
+          hidden_states,
+          attention_mask=attention_mask,
+          cos_table=cos_table,
+          sin_table=sin_table,
+      )
+      all_hidden_states.append(hidden_states)
+
+    hidden_states = self.norm(hidden_states)
+    return hidden_states, all_hidden_states
+
+
 # -----------------------------------------------------------------------------
 # Weight Mapping & Conversion Utilities
 # -----------------------------------------------------------------------------