diff --git a/src/maxtext/models/qwen3.py b/src/maxtext/models/qwen3.py index 22f3335ba9..290847a44c 100644 --- a/src/maxtext/models/qwen3.py +++ b/src/maxtext/models/qwen3.py @@ -1305,6 +1305,7 @@ def __init__(self, config: Config, mesh: Mesh, quant: None | Quant = None, *, rn out_features_shape=1, use_bias=False, # Qwen3-Next shared_expert_gate does not have a bias dtype=cfg.dtype, + weight_dtype=cfg.weight_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,