fix converting checkpoint and tp

Author: yanboliang

mixtral-moe/scripts/convert_hf_checkpoint.py

@@ -79,7 +79,7 @@ def convert_hf_checkpoint(
         elif "w1" in key or "w3" in key:
             final_result[key] = final_result[key].reshape(config.num_experts, config.intermediate_size, config.dim).contiguous()
         elif "w2" in key:
-            final_result[key] = final_result[key].reshape(config.num_experts, config.dim, config.intermediate_size).contiguous()
+            final_result[key] = final_result[key].reshape(config.num_experts, config.intermediate_size, config.dim).permute(0, 2, 1).contiguous()
         elif "gate" in key:
             final_result[key] = final_result[key].contiguous()
 

mixtral-moe/tp.py

@@ -104,7 +104,7 @@ def _apply_tp_moe_ffn(mlp: MOEFeedForward) -> None:
     if hasattr(mlp.cond_ffn, "scales1"):
         mlp.cond_ffn.scales1 = nn.Parameter(shard(mlp.cond_ffn.scales1, 1), requires_grad=False)
         mlp.cond_ffn.scales3 = nn.Parameter(shard(mlp.cond_ffn.scales3, 1), requires_grad=False)
-        mlp.cond_ffn.scales2 = nn.Parameter(shard(mlp.cond_ffn.scales2, 1), requires_grad=False)
+        mlp.cond_ffn.scales2 = nn.Parameter(mlp.cond_ffn.scales2, requires_grad=False)
 
     world_size = _get_world_size()
     mlp.cond_ffn.register_forward_hook(lambda _module, _input, output: funcol.all_reduce(