Add MoE to Gemma4 TP plan (huggingface#45219)

reduce memory for gemma4 moe model in tp

Signed-off-by: Wang, Yi <yi.a.wang@intel.com>
Co-authored-by: Cyril Vallez <cyril.vallez@huggingface.co>

Author: 2 people

src/transformers/models/gemma4/configuration_gemma4.py

@@ -132,6 +132,9 @@ class Gemma4TextConfig(PreTrainedConfig):
         "layers.*.mlp.gate_proj": "colwise",
         "layers.*.mlp.up_proj": "colwise",
         "layers.*.mlp.down_proj": "rowwise",
+        "layers.*.experts.gate_up_proj": "packed_colwise",
+        "layers.*.experts.down_proj": "rowwise",
+        "layers.*.experts": "moe_tp_experts",
     }
     base_model_pp_plan = {
         "embed_tokens": (["input_ids"], ["inputs_embeds"]),