fix

flexflow · Feb 19, 2024 · 2cbc0b7 · 2cbc0b7
1 parent fcbeea0
commit 2cbc0b7
Showing 1 changed file with 10 additions and 3 deletions.
diff --git a/src/runtime/model.cc b/src/runtime/model.cc
@@ -3798,9 +3798,16 @@ bool FFModel::check_operators_integrity(
         }
         for (int i = 0; i < fused->op_num_outputs[op]; i++) {
           int my_off = fused->op_output_idx[i + ooff];
-          assert(fused->op_output_source[i + ooff] == FusedOp::SOURCE_OUTPUT);
-          assert(FusedOp::use_same_regions(
-              fused->outputs[my_off], old_op->outputs[i], pt_mapping));
+          assert(fused->op_output_source[i + ooff] == FusedOp::SOURCE_OUTPUT ||
+                (fused->op_output_source[i + ooff] == FusedOp::SOURCE_INPUT &&
+                 (old_op->op_type == OP_RESIDUAL_LAYERNORM ||
+                  old_op->op_type == OP_RESIDUAL_RMS_NORM ||
+                  old_op->op_type == OP_ADD_BIAS_RESIDUAL_LAYERNORM)));
+          if (fused->op_output_source[i + ooff] == FusedOp::SOURCE_OUTPUT) {
+            assert(FusedOp::use_same_regions(fused->outputs[my_off], old_op->outputs[i], pt_mapping));
+          } else {
+            assert(FusedOp::use_same_regions(fused->inputs[my_off], old_op->outputs[i], pt_mapping));
+          }
         }
         ioff += fused->op_num_inputs[op];
         woff += fused->op_num_weights[op];