fix non-fusion opt

flexflow · Oct 8, 2024 · cf85d60 · cf85d60
1 parent 7ea8bd4
commit cf85d60
Show file tree

Hide file tree

Showing 4 changed files with 30 additions and 9 deletions.
diff --git a/src/ops/add_bias_residual_layer_norm.cc b/src/ops/add_bias_residual_layer_norm.cc
@@ -670,8 +670,18 @@ void AddBiasResidualLayerNorm::inference_task(
   AddBiasResidualLayerNormMeta *m =
       *((AddBiasResidualLayerNormMeta **)task->local_args);
 
-  assert(regions.size() ==
-         4 + (m->elementwise_affine ? (m->use_bias ? 2 : 1) : 0));
+  int expected_regions =
+      5; // input, attn_bias, residual (input), added_output, output
+  if (m->inplace_residual) {
+    expected_regions--; // input == added_output
+  }
+  if (m->elementwise_affine) {
+    expected_regions += 1; // gamma
+    if (m->use_bias) {
+      expected_regions += 1; // beta
+    }
+  }
+  assert(regions.size() == expected_regions);
 
   int rid = 0, tid = 0, did = 0;
   GenericTensorAccessorR input =

diff --git a/src/ops/fused.cu b/src/ops/fused.cu
@@ -645,7 +645,7 @@ __host__ void
         assert(false && "Fusion currently does not support type");
       }
     }
-    if (metas->meta[op]->inference_debugging ) {
+    if (metas->meta[op]->inference_debugging) {
       std::vector<GenericTensorAccessorR> input_accessors_to_save;
       std::vector<GenericTensorAccessorR> weight_accessors_to_save;
       std::vector<GenericTensorAccessorR> output_accessors_to_save;

diff --git a/src/ops/linear.cc b/src/ops/linear.cc
@@ -668,11 +668,11 @@ void Linear::inference_task(Task const *task,
     }
     Linear::save_inference_tensors_to_file(
         m, shard_id, bc, {input}, weights_accessors, {output});
-    printf("\tin=[%i,%i].T @ w=[%i,%i] -> out=[%i,%i]\n",
-           in_dim,
-           bc->num_tokens,
+    printf("\tw=[%i,%i].T @ in=[%i,%i] -> out=[%i,%i]\n",
            in_dim,
            out_dim,
+           in_dim,
+           bc->num_tokens,
            out_dim,
            bc->num_tokens);
   }

diff --git a/src/ops/residual_layer_norm.cc b/src/ops/residual_layer_norm.cc
@@ -988,9 +988,20 @@ void ResidualLayerNorm::inference_task(
     return;
   }
 
-  assert(regions.size() ==
-         3 + m->use_two_residuals +
-             (m->elementwise_affine ? (m->use_bias ? 2 : 1) : 0));
+  int expected_num_regions = 4; // input, residual1, added_output, output
+  if (m->use_two_residuals) {
+    expected_num_regions++; // residual2
+  }
+  if (m->inplace_residual) {
+    expected_num_regions--; // added_output = input
+  }
+  if (m->elementwise_affine) {
+    expected_num_regions += 1; // gamma
+    if (m->use_bias) {
+      expected_num_regions += 1; // beta
+    }
+  }
+  assert(regions.size() == expected_num_regions);
 
   int region_idx = 0, task_region_idx = 0;
   GenericTensorAccessorR input =