diff --git a/src/maxtext/models/qwen3.py b/src/maxtext/models/qwen3.py index e46a41dda9..34d62393f6 100644 --- a/src/maxtext/models/qwen3.py +++ b/src/maxtext/models/qwen3.py @@ -1138,32 +1138,39 @@ def __init__(self, config: Config, mesh: Mesh, quant: None | Quant = None, *, rn rngs=rngs, ) - # 2. Instantiate and apply the shared expert. - self.shared_expert = MlpBlock( - config=cfg, - mesh=mesh, - in_features=cfg.emb_dim, - intermediate_dim=cfg.moe_mlp_dim, - activations=cfg.mlp_activations, - intermediate_dropout_rate=cfg.dropout_rate, - dtype=cfg.dtype, - weight_dtype=cfg.weight_dtype, - quant=self.quant, - model_mode=config.model_call_mode, - rngs=rngs, - ) + # 2. Instantiate and apply the shared expert, if the config asks for one. + # cfg.shared_experts defaults to 0. Building the expert every time puts an + # extra full-size MLP in every layer of a dense configuration. + self.use_shared_expert = cfg.shared_experts > 0 + if self.use_shared_expert: + self.shared_expert = MlpBlock( + config=cfg, + mesh=mesh, + in_features=cfg.emb_dim, + intermediate_dim=cfg.moe_mlp_dim, + activations=cfg.mlp_activations, + intermediate_dropout_rate=cfg.dropout_rate, + dtype=cfg.dtype, + weight_dtype=cfg.weight_dtype, + quant=self.quant, + model_mode=config.model_call_mode, + rngs=rngs, + ) - # 3. Instantiate and apply the gate for the shared expert. - self.shared_expert_gate = DenseGeneral( - in_features_shape=cfg.emb_dim, - out_features_shape=1, - use_bias=False, # Qwen3-Next shared_expert_gate does not have a bias - dtype=cfg.dtype, - kernel_init=max_initializers.nd_dense_init(cfg.dense_init_scale, "fan_in", "truncated_normal"), - kernel_axes=("embed", None), - matmul_precision=cfg.matmul_precision, - rngs=rngs, - ) + # 3. Instantiate and apply the gate for the shared expert. + self.shared_expert_gate = DenseGeneral( + in_features_shape=cfg.emb_dim, + out_features_shape=1, + use_bias=False, # Qwen3-Next shared_expert_gate does not have a bias + dtype=cfg.dtype, + kernel_init=max_initializers.nd_dense_init(cfg.dense_init_scale, "fan_in", "truncated_normal"), + kernel_axes=("embed", None), + matmul_precision=cfg.matmul_precision, + rngs=rngs, + ) + else: + self.shared_expert = None + self.shared_expert_gate = None def __call__(self, hidden_states: Array, deterministic: bool) -> tuple[Array, Array | None]: """ @@ -1181,6 +1188,9 @@ def __call__(self, hidden_states: Array, deterministic: bool) -> tuple[Array, Ar # 1. Apply the routed experts block. routed_output, load_balance_loss, _ = self.routed_experts(hidden_states) + if not self.use_shared_expert: + return routed_output, load_balance_loss + # 2. Apply the shared expert. shared_expert_output = self.shared_expert(hidden_states, deterministic=deterministic) diff --git a/tests/unit/qwen3_next_shared_expert_test.py b/tests/unit/qwen3_next_shared_expert_test.py new file mode 100644 index 0000000000..201a5e8148 --- /dev/null +++ b/tests/unit/qwen3_next_shared_expert_test.py @@ -0,0 +1,82 @@ +# Copyright 2026 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. +"""The Qwen3-Next sparse MoE block builds a shared expert only when asked.""" + +import unittest + +from flax import nnx +import jax +from jax.sharding import Mesh + +from maxtext.configs import pyconfig +from maxtext.models.qwen3 import Qwen3NextSparseMoeBlock +from maxtext.utils import maxtext_utils +from tests.utils.test_helpers import get_test_config_path + + +def _count_params(module) -> int: + return sum(x.size for x in jax.tree.leaves(nnx.state(module, nnx.Param))) + + +class Qwen3NextSharedExpertTest(unittest.TestCase): + """`shared_experts` defaults to 0, so building the expert unconditionally put a + second full-size MLP in every layer of a dense configuration.""" + + def _block(self, shared_experts: int) -> Qwen3NextSparseMoeBlock: + config = pyconfig.initialize( + [None, get_test_config_path()], + run_name="qwen3_next_shared_expert_test", + enable_checkpointing=False, + override_model_config=True, + model_name="qwen3-next-80b-a3b", + dtype="bfloat16", + weight_dtype="bfloat16", + megablox=False, + sparse_matmul=False, + max_target_length=8, + per_device_batch_size=1, + base_emb_dim=8, + base_moe_mlp_dim=16, + base_num_decoder_layers=1, + num_experts=2, + num_experts_per_tok=1, + shared_experts=shared_experts, + ) + mesh = Mesh(maxtext_utils.create_device_mesh(config), config.mesh_axes) + return Qwen3NextSparseMoeBlock(config=config, mesh=mesh, rngs=nnx.Rngs(params=jax.random.PRNGKey(0))) + + def test_no_shared_expert_when_zero(self): + block = self._block(shared_experts=0) + self.assertFalse(block.use_shared_expert) + self.assertIsNone(block.shared_expert) + self.assertIsNone(block.shared_expert_gate) + + def test_shared_expert_when_one(self): + block = self._block(shared_experts=1) + self.assertTrue(block.use_shared_expert) + self.assertIsNotNone(block.shared_expert) + self.assertIsNotNone(block.shared_expert_gate) + + def test_the_only_difference_is_the_shared_expert(self): + """The parameter count with the flag off is short by the shared expert and + its gate, and by nothing else.""" + with_expert = self._block(shared_experts=1) + without = self._block(shared_experts=0) + expert_params = _count_params(with_expert.shared_expert) + _count_params(with_expert.shared_expert_gate) + self.assertGreater(expert_params, 0) + self.assertEqual(_count_params(with_expert) - _count_params(without), expert_params) + + +if __name__ == "__main__": + unittest.main()