Minor changes to deepseek_google_test.py

PiperOrigin-RevId: 817743737

Author: jesselu-google

src/MaxText/configs/models/deepseek3-tiny.yml

@@ -0,0 +1,50 @@
+# Copyright 2023–2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# Tiny version of DeepSeek V3 for testing.
+
+base_emb_dim: 64
+base_num_query_heads: 4
+base_num_kv_heads: 4
+base_mlp_dim: 64
+base_moe_mlp_dim: 64
+base_num_decoder_layers: 61
+first_num_dense_layers: 3
+mlp_activations: ["silu","linear"]
+vocab_size: 129280
+enable_dropout: False
+logits_via_embedding: False
+normalization_layer_epsilon: 1.0e-6
+num_experts: 16
+num_experts_per_tok: 8
+shared_experts: 1
+routed_scaling_factor: 2.5
+routed_score_func: "sigmoid"
+routed_bias: True
+decoder_block: "deepseek"
+# MLA
+attention_type: "mla"
+q_lora_rank: 32
+kv_lora_rank: 16
+qk_nope_head_dim: 128
+qk_rope_head_dim: 64
+v_head_dim: 128
+mscale: 1.0
+# RoPE
+rope_type: "yarn"
+rope_max_timescale: 10_000 # DeepSeek uses  "rope_theta": 10000
+max_position_embeddings: 163840
+original_max_position_embeddings: 4096
+rope_factor: 40
+beta_fast: 32

src/MaxText/layers/deepseek_batchsplit.py

@@ -16,7 +16,7 @@
 
 """Alternative DeepSeek model definition with batch-split schedule."""
 
-from flax import linen as nnx
+from flax import linen as nn
 import jax
 import jax.numpy as jnp
 from MaxText import common_types
@@ -29,7 +29,7 @@
 from MaxText.layers import quantizations
 
 
-class DeepSeekGenericLayer(nnx.Module):
+class DeepSeekGenericLayer(nn.Module):
   """Generic DeepSeek layer with Multi-Head Latent Attention.
 
   This is to be used as a base class for DeepSeek layers with dense/sparse MLPs.
@@ -101,7 +101,7 @@ def logical_axis_names(self):
       )
 
   def with_logical_constraint(self, x):
-    return nnx.with_logical_constraint(x, self.logical_axis_names)
+    return nn.with_logical_constraint(x, self.logical_axis_names)
 
   def rms_norm_layer(self, name):
     return normalizations.rms_norm(
@@ -187,7 +187,7 @@ def mlp(self, x, deterministic):
     raise NotImplementedError()
 
   def dropout_layer(self):
-    return nnx.Dropout(rate=self.config.dropout_rate, broadcast_dims=(-2,))
+    return nn.Dropout(rate=self.config.dropout_rate, broadcast_dims=(-2,))
 
   def dropout(self, x, deterministic):
     return self.with_logical_constraint(

src/MaxText/pyconfig.py

@@ -372,6 +372,7 @@ def validate_model_name(s: str) -> bool:
       "deepseek2-236b",
       "deepseek3-671b",
       "deepseek3-test",
+      "deepseek3-tiny",
       "kimi-k2-1t",
       "gemma-7b",
       "gemma-2b",