Misnamed mapping

microsoft · Dec 16, 2023 · c1e90a3 · c1e90a3
1 parent d28c327
commit c1e90a3
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/deepspeed/inference/v2/model_implementations/mixtral/container.py b/deepspeed/inference/v2/model_implementations/mixtral/container.py
@@ -26,7 +26,7 @@ class MixtralTransformerContainer(LayerContainer):
         "self_attn.k_proj.weight": "qkv_w.k_params",
         "self_attn.v_proj.weight": "qkv_w.v_params",
         "self_attn.o_proj.weight": "attn_out_w.params",
-        "moe_gate.weight": "moe_gate.params",
+        "block_sparse_moe.gate.weight": "moe_gate.params",
         "block_sparse_moe.experts.*.w1.weight": "moe_mlp_1.gating_experts",
         "block_sparse_moe.experts.*.w3.weight": "moe_mlp_1.up_experts",
         "block_sparse_moe.experts.*.w2.weight": "moe_mlp_2.experts",