Convert DenseGeneral to NNX

Co-authored-by: Branden Vandermoon <[email protected]>
Co-authored-by: Cristian Garcia <[email protected]>

Author: bvandermoon

MaxText/layers/attentions.py

@@ -64,7 +64,7 @@ class AttentionType(enum.Enum):
 DType = common_types.DType
 Mesh = common_types.Mesh
 PRNGKey = common_types.PRNGKey
-DenseGeneral = linears.DenseGeneral
+dense_general = linears.dense_general
 RMSNorm = linears.RMSNorm
 RotaryEmbedding = embeddings.RotaryEmbedding
 YarnRotaryEmbedding = embeddings.YarnRotaryEmbedding
@@ -1330,7 +1330,8 @@ def query_init(*args):
     kernel_axes = (
         (None, None, None) if self.config.ici_context_autoregressive_parallelism > 1 else ("embed", "q_heads", "kv")
     )
-    query_proj = DenseGeneral(
+    query_proj = dense_general(
+        inputs_shape=inputs_q.shape,
         features=(self.num_query_heads, self.head_dim),
         axis=-1,
         kernel_init=query_init,
@@ -1366,7 +1367,8 @@ def kv_projection(self, inputs_kv: Array, proj_name: str) -> Array:
         else ("embed", "kv_heads", "kv_head_dim")
     )
 
-    kv_proj = DenseGeneral(
+    kv_proj = dense_general(
+        inputs_shape=inputs_kv.shape,
         features=(self.num_kv_heads, self.head_dim),
         axis=-1,
         kernel_init=self.kernel_init,
@@ -1382,7 +1384,8 @@ def kv_projection(self, inputs_kv: Array, proj_name: str) -> Array:
   def qkv_projection(self, inputs: Array, proj_name: str):
     """Fused QKV projection"""
 
-    qkv_proj = DenseGeneral(
+    qkv_proj = dense_general(
+        inputs_shape=inputs.shape,
         features=(3, self.num_query_heads, self.head_dim),
         axis=-1,
         kernel_init=self.kernel_init,
@@ -1402,7 +1405,8 @@ def out_projection(self, output_dim: int, out: Array) -> Array:
     out_kernel_axis = (
         (None, None, None) if self.config.ici_context_autoregressive_parallelism > 1 else ("heads", "kv", "embed")
     )
-    out_proj = DenseGeneral(
+    out_proj = dense_general(
+        inputs_shape=out.shape,
         features=output_dim,
         axis=(-2, -1),
         kernel_init=self.kernel_init,
@@ -1660,7 +1664,8 @@ def setup(self):
 
     if self.q_lora_rank == 0:
       # Standard Q projection (without LoRA).
-      self.query_proj = DenseGeneral(
+      self.query_proj = dense_general(
+          in_features=self.config.emb_dim,
           features=(self.num_query_heads, self.qk_head_dim),
           axis=-1,
           kernel_init=self.kernel_init,
@@ -1673,7 +1678,8 @@ def setup(self):
       )
     else:
       # LoRA path for Q.
-      self.wq_a = DenseGeneral(
+      self.wq_a = dense_general(
+          in_features=self.config.emb_dim,
           features=self.q_lora_rank,
           axis=-1,
           kernel_init=self.kernel_init,
@@ -1691,7 +1697,8 @@ def setup(self):
           epsilon=self.config.normalization_layer_epsilon,
           kernel_axes=("norm",),
       )
-      self.wq_b = DenseGeneral(
+      self.wq_b = dense_general(
+          in_features=self.q_lora_rank,
           features=(self.num_query_heads, self.qk_head_dim),
           axis=-1,
           kernel_init=self.kernel_init,
@@ -1704,7 +1711,8 @@ def setup(self):
       )
 
     # KV LoRA path.
-    self.wkv_a = DenseGeneral(
+    self.wkv_a = dense_general(
+        in_features=self.config.emb_dim,
         features=self.kv_lora_rank + self.qk_rope_head_dim,
         axis=-1,
         kernel_init=self.kernel_init,
@@ -1722,8 +1730,12 @@ def setup(self):
         epsilon=self.config.normalization_layer_epsilon,
         kernel_axes=("norm",),
     )
-    self.wkv_b = DenseGeneral(
-        features=(self.num_query_heads, (self.qk_nope_head_dim + self.v_head_dim)),
+    self.wkv_b = dense_general(
+        in_features=self.kv_lora_rank,
+        features=(
+            self.num_query_heads,
+            (self.qk_nope_head_dim + self.v_head_dim),
+        ),
         axis=-1,
         kernel_init=self.kernel_init,
         kernel_axes=("kv_lora", "kv_heads", "kv_head_dim"),
@@ -1933,3 +1945,4 @@ def __hash__(self):
             self.q_sequence.tobytes() if self.q_sequence is not None else None,
         )
     )
+  

MaxText/layers/gpt3.py

@@ -49,7 +49,7 @@
 D_KV = common_types.D_KV
 EMBED = common_types.EMBED
 
-DenseGeneral = linears.DenseGeneral
+dense_general = linears.dense_general
 NdInitializer = initializers.NdInitializer
 Initializer = initializers.Initializer
 nd_dense_init = initializers.nd_dense_init
@@ -163,7 +163,8 @@ class Gpt3MultiHeadAttention(nn.Module):
   def qkv_projection(self, inputs: Array, proj_name: str):
     """Fused QKV projection"""
 
-    qkv_proj = DenseGeneral(
+    qkv_proj = dense_general(
+        inputs_shape=inputs.shape,
         features=(3, self.num_heads, self.head_dim),
         axis=-1,
         kernel_init=self.kernel_init,
@@ -181,7 +182,8 @@ def qkv_projection(self, inputs: Array, proj_name: str):
 
   def projection(self, inputs: Array, proj_name: str) -> Array:
     """individual projection for one of q, k and v."""
-    proj = DenseGeneral(
+    proj = dense_general(
+        inputs_shape=inputs.shape,
         features=(self.num_heads, self.head_dim),
         axis=-1,
         kernel_init=self.kernel_init,
@@ -197,7 +199,8 @@ def projection(self, inputs: Array, proj_name: str) -> Array:
 
   def out_projection(self, output_dim: int, out: Array) -> Array:
     """output projection"""
-    out_proj = DenseGeneral(
+    out_proj = dense_general(
+        inputs_shape=out.shape,
         features=output_dim,
         axis=(-2, -1),
         kernel_init=self.kernel_init,