Build the Qwen3-Next shared expert only when shared_experts > 0

[WandLZhang]

Qwen3NextSparseMoeBlock builds a shared expert and its gate in every layer. It doesn't read cfg.shared_experts, which defaults to 0 in base.yml.

A dense configuration therefore carries one unrequested MlpBlock per layer. The parameter count rises and the extra weights receive gradients.

Effect

Measured on a Qwen 3.5 dense configuration, 32 layers, emb_dim 4096, moe_mlp_dim 12288:

	Parameters
Reported by MaxText	11.756B
Published for this architecture	6.924B

The 4.8B difference is one unrequested MLP per layer. At 11.756B the model doesn't fit a v6e-8: HLO temporaries reach 40.50 GB against 31.24 GB of HBM. With this change the same configuration reports 6.924B and trains on 8 chips.

Scope

The block now builds the shared expert only when cfg.shared_experts > 0, and skips it in the forward pass on the same condition. A configuration that sets the field keeps its current behaviour, so shipped Qwen3-Next MoE configurations are unaffected.

Tests

Adds tests/unit/qwen3_next_shared_expert_test.py, three cases, CPU only. With shared_experts: 0 the expert and its gate are both None. With 1, both are built. And the parameter-count difference between the two equals the shared expert plus its gate, and nothing else.

On main the third case fails with 0 != 392 and the other two error on the missing use_shared_expert.

tests/unit/qwen3_5_layers_test.py only covers TestQwen3_5MoeVisionEncoderEndToEnd, so it never reaches this code and still passes.

cc @mmcsa

[gemini-code-assist]

Code Review

This pull request conditionally instantiates the shared expert and its gate in the Qwen3 model based on the configuration, preventing unnecessary MLP blocks from being added. The review feedback recommends explicitly initializing these attributes to None when they are not used to prevent potential AttributeError issues and static analysis warnings.

[gemini-code-assist]

Conditionally defining instance attributes can lead to AttributeErrors if they are accessed elsewhere, and it often triggers warnings or errors in static type checkers like Pytype or MyPy. It is a best practice to always initialize all instance attributes in __init__.

Consider explicitly initializing self.shared_expert and self.shared_expert_gate to None when self.use_shared_expert is False.

    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,
      )
    else:
      self.shared_expert = None
      self.shared_expert_gate = None

[WandLZhang]

Both attributes are now set to None in an else branch.