[maxtext] avoid creating kernel param in serve mode

self.kernel is not used in serve model anyway. In serve mode, the checkpoint is quantized. Quantized weights will be read by aqt, so we don't need to create this param.

PiperOrigin-RevId: 772594747

Author: maxtext authors

MaxText/layers/linears.py

@@ -131,16 +131,17 @@ def __init__(
         len(self.axis), len(self.axis) + len(self.out_features)
     )
 
-    self.kernel = nnx.Param(
-        self.kernel_init(
-            rngs.params(),
-            kernel_shape,
-            self.weight_dtype,
-            kernel_in_axis,
-            kernel_out_axis,
-        ),
-        sharding=self.kernel_axes,
-    )
+    if not quantizations.in_serve_mode(self.quant):
+      self.kernel = nnx.Param(
+          self.kernel_init(
+              rngs.params(),
+              kernel_shape,
+              self.weight_dtype,
+              kernel_in_axis,
+              kernel_out_axis,
+          ),
+          sharding=self.kernel_axes,
+      )
 
     if self.use_bias:
       bias_axes = self.kernel_axes[-len(self.out_features) :]