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

tests/unit/qwen3_5_layers_test.py only covers TestQwen3_5MoeVisionEncoderEndToEnd, so it never reaches this code and still passes. I didn't add a test. One test can check the parameter count of a dense config. Let me know and I'll add it here.

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.