!24586 model parallel communication recompute rule change

Merge pull request !24586 from yao_yf/model_parallel_recompute_change
This commit is contained in:
i-robot 2021-10-09 09:04:02 +00:00 committed by Gitee
commit e66b0db835
1 changed files with 4 additions and 4 deletions

View File

@ -351,10 +351,10 @@ void InsertRedistribution(const RedistributionOpListPtr &redistribution_oplist_p
if (prim_out != nullptr && prim_in != nullptr) {
auto prim_out_attr = prim_out->attrs();
auto prim_in_attr = prim_in->attrs();
if (prim_out_attr.find(RECOMPUTE_COMM_OP) != prim_out_attr.end() &&
!GetValue<bool>(prim_out_attr[RECOMPUTE_COMM_OP]) &&
prim_in_attr.find(RECOMPUTE_COMM_OP) != prim_in_attr.end() &&
!GetValue<bool>(prim_in_attr[RECOMPUTE_COMM_OP]) &&
if (((prim_out_attr.find(RECOMPUTE_COMM_OP) != prim_out_attr.end() &&
!GetValue<bool>(prim_out_attr[RECOMPUTE_COMM_OP])) ||
(prim_in_attr.find(RECOMPUTE_COMM_OP) != prim_in_attr.end() &&
!GetValue<bool>(prim_in_attr[RECOMPUTE_COMM_OP]))) &&
COMMUNICATION_OPS.find(op_name) != COMMUNICATION_OPS.end()) {
MS_LOG(INFO) << "The redistribution node would not be recomputed.";
instance_name = instance_name + "_" + NOT_RECOMPUTE;