已合并
[pytorch][bugfix]fix tp-extend-ep in ckpt #2855
温一盏创建于 2025年6月19日
[pytorch][bugfix]fix tp-extend-ep in ckpt #2855
已合并
从refs/pull/2855/head合入到2.0.0
共 1 个文件变更+6-2
| @@ -649,8 +649,12 @@ class HuggingfaceModel(ModelBase): | |||
| 649 | fc1_weight = self.get_layers_mlp_experts_linear_fc1_weight(**kwargs) | 649 | fc1_weight = self.get_layers_mlp_experts_linear_fc1_weight(**kwargs) |
| 650 | if getattr(args, "swiglu", None): | 650 | if getattr(args, "swiglu", None): |
| 651 | gate_w, up_w = torch.chunk(fc1_weight, 2, dim=0) | 651 | gate_w, up_w = torch.chunk(fc1_weight, 2, dim=0) |
| 652 | - gate_w_list = torch.chunk(gate_w, getattr(self.args_cmd, 'target_tensor_parallel_size', 1), dim=0) | 652 | + if args.moe_tp_extend_ep: |
| 653 | - up_w_list = torch.chunk(up_w, getattr(self.args_cmd, 'target_tensor_parallel_size', 1), dim=0) | 653 | + gate_w_list = torch.chunk(gate_w, 1, dim=0) |
| 654 | + up_w_list = torch.chunk(up_w, 1, dim=0) | ||
| 655 | + else: | ||
| 656 | + gate_w_list = torch.chunk(gate_w, getattr(self.args_cmd, 'target_tensor_parallel_size', 1), dim=0) | ||
| 657 | + up_w_list = torch.chunk(up_w, getattr(self.args_cmd, 'target_tensor_parallel_size', 1), dim=0) | ||
| 654 | fc1_weight = torch.cat([torch.cat(weights, dim=0) for weights in zip(gate_w_list, up_w_list)], dim=0) | 658 | fc1_weight = torch.cat([torch.cat(weights, dim=0) for weights in zip(gate_w_list, up_w_list)], dim=0) |
| 655 | experts_linear_fc1_list.append(fc1_weight.t().view(-1)) | 659 | experts_linear_fc1_list.append(fc1_weight.t().view(-1)) |
| 656 | return torch.cat(experts_linear_fc1_list).view(args.hidden_size, -1) | 660 | return torch.cat(experts_linear_fc1_list).view(args.hidden_size, -1) |