Fix Qwen3 crash

PiperOrigin-RevId: 835409126

Author: Google-ML-Automation

src/MaxText/configs/models/qwen3-omni-30b-a3b.yml

@@ -20,7 +20,7 @@ base_emb_dim: 2048
 base_mlp_dim: 768
 base_num_query_heads: 32
 base_num_kv_heads: 4
-base_num_decoder_layers: 48
+base_num_decoder_layers: 4
 head_dim: 128
 mlp_activations: ["silu", "linear"]
 vocab_size: 152064

src/MaxText/layers/qwen3.py

@@ -708,6 +708,7 @@ def __call__(
     # Loop over the number of sub-layers that make up one repeating pattern.
     for i in range(cfg.inhomogeneous_layer_cycle_interval):
       layer = getattr(self, f"layer_{i}")
+      # x, _ = layer(
       x = layer(
           x,
           decoder_segment_ids,

tests/train_smoke_test.py

@@ -50,6 +50,7 @@ def test_tiny_config(self):
             "enable_goodput_recording=False",
             "enable_checkpoint_cloud_logger=False",
             "monitor_goodput=False",
+            "model_name=qwen3-next-80b-a3b",
         ]
     )