[PIR AMP]Fix some error for bert amp (PaddlePaddle#64497)

0x45f · co63oc · commit 36afd973622a · 2024-05-23T15:30:24.000+08:00
diff --git a/paddle/fluid/operators/fused/fused_feedforward_op.cu b/paddle/fluid/operators/fused/fused_feedforward_op.cu
@@ -518,14 +518,14 @@ void FusedFeedForwardGradKernel(
     bool add_residual,
     int ring_id,
     DenseTensor* x_grad,
-    DenseTensor* ln1_scale_grad,
-    DenseTensor* ln1_bias_grad,
-    DenseTensor* ln2_scale_grad,
-    DenseTensor* ln2_bias_grad,
     DenseTensor* linear1_weight_grad,
     DenseTensor* linear1_bias_grad,
     DenseTensor* linear2_weight_grad,
-    DenseTensor* linear2_bias_grad) {
+    DenseTensor* linear2_bias_grad,
+    DenseTensor* ln1_scale_grad,
+    DenseTensor* ln1_bias_grad,
+    DenseTensor* ln2_scale_grad,
+    DenseTensor* ln2_bias_grad) {
   using U = phi::funcs::LayerNormParamType<T>;
 
   auto* ln1_out_ptr = pre_layer_norm ? ln1_out.get_ptr() : nullptr;
@@ -672,9 +672,9 @@ PD_REGISTER_KERNEL(fused_feedforward_grad,
                    double,
                    phi::dtype::float16) {
   if (kernel_key.dtype() == phi::DataType::FLOAT16) {
-    kernel->OutputAt(1).SetDataType(phi::DataType::FLOAT32);
-    kernel->OutputAt(2).SetDataType(phi::DataType::FLOAT32);
-    kernel->OutputAt(3).SetDataType(phi::DataType::FLOAT32);
-    kernel->OutputAt(4).SetDataType(phi::DataType::FLOAT32);
+    kernel->OutputAt(5).SetDataType(phi::DataType::FLOAT32);
+    kernel->OutputAt(6).SetDataType(phi::DataType::FLOAT32);
+    kernel->OutputAt(7).SetDataType(phi::DataType::FLOAT32);
+    kernel->OutputAt(8).SetDataType(phi::DataType::FLOAT32);
   }
 }
diff --git a/paddle/fluid/operators/generator/parse_utils.py b/paddle/fluid/operators/generator/parse_utils.py
@@ -643,6 +643,8 @@ def validate_backward_inputs(
 
 
 def validate_backward_outputs(op, forward_inputs, backward_outputs):
+    if op in ['fused_attention_grad']:
+        return
     assert len(backward_outputs) <= len(
         forward_inputs
     ), f"{op } has too many outputs"
diff --git a/paddle/fluid/operators/ops_signature/fused_feedforward_sig.cc b/paddle/fluid/operators/ops_signature/fused_feedforward_sig.cc
@@ -84,14 +84,14 @@ KernelSignature FeedForwardGradFuseOpArgumentMapping(
                           "add_residual",
                           "ring_id"},
                          {"X@GRAD",
-                          "Ln1Scale@GRAD",
-                          "Ln1Bias@GRAD",
-                          "Ln2Scale@GRAD",
-                          "Ln2Bias@GRAD",
                           "Linear1Weight@GRAD",
                           "Linear1Bias@GRAD",
                           "Linear2Weight@GRAD",
-                          "Linear2Bias@GRAD"});
+                          "Linear2Bias@GRAD",
+                          "Ln1Scale@GRAD",
+                          "Ln1Bias@GRAD",
+                          "Ln2Scale@GRAD",
+                          "Ln2Bias@GRAD"});
 }
 }  // namespace phi
 
diff --git a/paddle/fluid/pir/dialect/op_generator/op_gen.py b/paddle/fluid/pir/dialect/op_generator/op_gen.py
@@ -1119,10 +1119,11 @@ def get_input_grad_semantic(op_info, op_info_items):
 
     # get backward op
     bwd_op_name = op_info.backward_name
+    sparse_op_name_suffix = '_sp' if op_info.is_sparse_op else ''
     if (bwd_op_name is None) or (bwd_op_name not in op_info_items.keys()):
         input_grad_semantics = ["false" for i in range(num_inputs)]
     else:
-        bwd_op_info = op_info_items[bwd_op_name]
+        bwd_op_info = op_info_items[bwd_op_name + sparse_op_name_suffix]
 
         # cut "_grad" of each output of bwd_op, and then compare each modified output with corresponding input
         # thus determine whether each input has grad semantic
@@ -1153,12 +1154,13 @@ def get_mutable_attribute_grad_semantic(op_info, op_info_items):
 
     # get backward op
     bwd_op_name = op_info.backward_name
+    sparse_op_name_suffix = '_sp' if op_info.is_sparse_op else ''
     if (bwd_op_name is None) or (bwd_op_name not in op_info_items.keys()):
         mutable_attribute_grad_semantics = [
             "false" for i in range(len(fwd_mutable_attribute_list))
         ]
     else:
-        bwd_op_info = op_info_items[bwd_op_name]
+        bwd_op_info = op_info_items[bwd_op_name + sparse_op_name_suffix]
 
         # cut "_grad" of each output of bwd_op, and then compare each modified output with corresponding attribute
         # thus determine whether each attribute has grad semantic
diff --git a/paddle/fluid/primitive/codegen/gen.py b/paddle/fluid/primitive/codegen/gen.py
@@ -40,6 +40,7 @@
 VJPS_BLACK_LIST = [
     'reshape_grad',
     'add_n_grad',
+    'fused_attention_grad',
 ]
 
 BACKENDS_BLACK_LIST = [
diff --git a/paddle/fluid/primitive/rule/vjp/manual/manual_vjp.cc b/paddle/fluid/primitive/rule/vjp/manual/manual_vjp.cc
@@ -67,5 +67,124 @@ std::vector<std::vector<paddle::Tensor>> reshape_vjp(
   return vjp_res;
 }
 
+std::vector<std::vector<paddle::Tensor>> fused_attention_vjp(
+    const Tensor& out_grad,
+    const Tensor& x,
+    const Tensor& qkv_weight,
+    const paddle::optional<Tensor>& qkv_bias,
+    const paddle::optional<Tensor>& qkv_bias_out,
+    const paddle::optional<Tensor>& src_mask,
+    const paddle::optional<Tensor>& src_mask_out,
+    const Tensor& out_linear_weight,
+    const paddle::optional<Tensor>& out_linear_bias,
+    const paddle::optional<Tensor>& ln_scale,
+    const paddle::optional<Tensor>& ln_bias,
+    const paddle::optional<Tensor>& ln_scale_2,
+    const paddle::optional<Tensor>& ln_bias_2,
+    const paddle::optional<Tensor>& ln_out,
+    const paddle::optional<Tensor>& ln_mean,
+    const paddle::optional<Tensor>& ln_var,
+    const paddle::optional<Tensor>& ln_mean_2,
+    const paddle::optional<Tensor>& ln_var_2,
+    const paddle::optional<Tensor>& bias_dropout_residual_out,
+    const Tensor& qkv_out,
+    const Tensor& transpose_out_2,
+    const Tensor& qk_out,
+    const Tensor& qktv_out,
+    const Tensor& softmax_out,
+    const Tensor& attn_dropout_mask_out,
+    const Tensor& attn_dropout_out,
+    const Tensor& fmha_out,
+    const Tensor& out_linear_out,
+    const Tensor& dropout_mask_out,
+    int num_heads,
+    bool transpose_qkv_wb,
+    bool pre_layer_norm,
+    float epsilon,
+    float attn_dropout_rate,
+    bool is_test,
+    bool attn_dropout_fix_seed,
+    int attn_dropout_seed,
+    const std::string& attn_dropout_implementation,
+    float dropout_rate,
+    bool dropout_fix_seed,
+    int dropout_seed,
+    const std::string& dropout_implementation,
+    float ln_epsilon,
+    bool add_residual,
+    int ring_id,
+    const std::vector<std::vector<bool>>& stop_gradients) {
+  std::vector<std::vector<paddle::Tensor>> vjp_res;
+  for (auto arg : stop_gradients) {
+    vjp_res.push_back(std::vector<paddle::Tensor>(arg.size()));
+  }
+  auto op_res =
+      backend::fused_attention_grad<LazyTensor>(out_grad,
+                                                x,
+                                                qkv_weight,
+                                                qkv_bias,
+                                                qkv_bias_out,
+                                                src_mask,
+                                                src_mask_out,
+                                                out_linear_weight,
+                                                out_linear_bias,
+                                                ln_scale,
+                                                ln_bias,
+                                                ln_scale_2,
+                                                ln_bias_2,
+                                                ln_out,
+                                                ln_mean,
+                                                ln_var,
+                                                ln_mean_2,
+                                                ln_var_2,
+                                                bias_dropout_residual_out,
+                                                qkv_out,
+                                                transpose_out_2,
+                                                qk_out,
+                                                qktv_out,
+                                                softmax_out,
+                                                attn_dropout_mask_out,
+                                                attn_dropout_out,
+                                                fmha_out,
+                                                out_linear_out,
+                                                dropout_mask_out,
+                                                num_heads,
+                                                transpose_qkv_wb,
+                                                pre_layer_norm,
+                                                epsilon,
+                                                attn_dropout_rate,
+                                                is_test,
+                                                attn_dropout_fix_seed,
+                                                attn_dropout_seed,
+                                                attn_dropout_implementation,
+                                                dropout_rate,
+                                                dropout_fix_seed,
+                                                dropout_seed,
+                                                dropout_implementation,
+                                                ln_epsilon,
+                                                add_residual,
+                                                ring_id);
+  // x_grad
+  vjp_res[0][0] = std::get<8>(op_res);
+  // ln_scale_grad
+  vjp_res[1][0] = std::get<4>(op_res);
+  // ln_bias_grad
+  vjp_res[2][0] = std::get<5>(op_res);
+  // qkv_weight_grad
+  vjp_res[3][0] = std::get<9>(op_res);
+  // qkv_bias_grad
+  vjp_res[4][0] = std::get<0>(op_res);
+  // out_linear_weight_grad
+  vjp_res[5][0] = std::get<10>(op_res);
+  // out_linear_bias_grad
+  vjp_res[6][0] = std::get<3>(op_res);
+  // ln_scale_2_grad
+  vjp_res[7][0] = std::get<6>(op_res);
+  // ln_bias_2_grad
+  vjp_res[8][0] = std::get<7>(op_res);
+  vjp_res = ConstructVjpResultByStopGradients(vjp_res, stop_gradients);
+  return vjp_res;
+}
+
 }  // namespace primitive
 }  // namespace paddle
diff --git a/paddle/fluid/primitive/rule/vjp/manual/manual_vjp.h b/paddle/fluid/primitive/rule/vjp/manual/manual_vjp.h
@@ -33,5 +33,53 @@ std::vector<std::vector<paddle::Tensor>> reshape_vjp(
     const Tensor& out_grad,
     const std::vector<std::vector<bool>>& stop_gradients);
 
+std::vector<std::vector<paddle::Tensor>> fused_attention_vjp(
+    const Tensor& out_grad,
+    const Tensor& x,
+    const Tensor& qkv_weight,
+    const paddle::optional<Tensor>& qkv_bias,
+    const paddle::optional<Tensor>& qkv_bias_out,
+    const paddle::optional<Tensor>& src_mask,
+    const paddle::optional<Tensor>& src_mask_out,
+    const Tensor& out_linear_weight,
+    const paddle::optional<Tensor>& out_linear_bias,
+    const paddle::optional<Tensor>& ln_scale,
+    const paddle::optional<Tensor>& ln_bias,
+    const paddle::optional<Tensor>& ln_scale_2,
+    const paddle::optional<Tensor>& ln_bias_2,
+    const paddle::optional<Tensor>& ln_out,
+    const paddle::optional<Tensor>& ln_mean,
+    const paddle::optional<Tensor>& ln_var,
+    const paddle::optional<Tensor>& ln_mean_2,
+    const paddle::optional<Tensor>& ln_var_2,
+    const paddle::optional<Tensor>& bias_dropout_residual_out,
+    const Tensor& qkv_out,
+    const Tensor& transpose_out_2,
+    const Tensor& qk_out,
+    const Tensor& qktv_out,
+    const Tensor& softmax_out,
+    const Tensor& attn_dropout_mask_out,
+    const Tensor& attn_dropout_out,
+    const Tensor& fmha_out,
+    const Tensor& out_linear_out,
+    const Tensor& dropout_mask_out,
+    int num_heads,
+    bool transpose_qkv_wb,
+    bool pre_layer_norm,
+    float epsilon,
+    float attn_dropout_rate,
+    bool is_test,
+    bool attn_dropout_fix_seed,
+    int attn_dropout_seed,
+    const std::string& attn_dropout_implementation,
+    float dropout_rate,
+    bool dropout_fix_seed,
+    int dropout_seed,
+    const std::string& dropout_implementation,
+    float ln_epsilon,
+    bool add_residual,
+    int ring_id,
+    const std::vector<std::vector<bool>>& stop_gradients);
+
 }  // namespace primitive
 }  // namespace paddle
diff --git a/paddle/phi/infermeta/backward.cc b/paddle/phi/infermeta/backward.cc
@@ -510,7 +510,7 @@ void GeneralTernaryGradInferMeta(const MetaTensor& x,
   if (dx) {
     dx->share_meta(x);
   }
-  if (dy) {
+  if (dy && y) {
     dy->share_meta(y);
   }
   if (dz) {
diff --git a/paddle/phi/infermeta/fusion.cc b/paddle/phi/infermeta/fusion.cc
@@ -924,6 +924,7 @@ void FusedAttentionInferMeta(const MetaTensor& x,
   }
 
   out->set_dims(x.dims());
+  out->set_dtype(x.dtype());
 }
 
 void FusedAttentionGradInferMeta(const MetaTensor& out_grad,
@@ -998,19 +999,19 @@ void FusedAttentionGradInferMeta(const MetaTensor& out_grad,
                         "GradOp is only callable when is_test is false"));
 
   if (!pre_layer_norm) {
-    if (ln_scale_2_grad) {
+    if (ln_scale_2_grad && ln_scale_2) {
       ln_scale_2_grad->set_dims(ln_scale_2.dims());
     }
-    if (ln_bias_2_grad) {
+    if (ln_bias_2_grad && ln_bias_2) {
       ln_bias_2_grad->set_dims(ln_bias_2.dims());
     }
   }
 
-  if (pre_layer_norm) {
+  if (pre_layer_norm && ln_scale) {
     if (ln_scale_grad) {
       ln_scale_grad->set_dims(ln_scale.dims());
     }
-    if (ln_bias_grad) {
+    if (ln_bias_grad && ln_bias) {
       ln_bias_grad->set_dims(ln_bias.dims());
     }
   }
@@ -1019,7 +1020,7 @@ void FusedAttentionGradInferMeta(const MetaTensor& out_grad,
     x_grad->set_dims(x.dims());
   }
 
-  if (out_linear_bias_grad) {
+  if (out_linear_bias_grad && out_linear_bias) {
     out_linear_bias_grad->set_dims(out_linear_bias.dims());
   }
 
@@ -1031,7 +1032,7 @@ void FusedAttentionGradInferMeta(const MetaTensor& out_grad,
     qkv_weight_grad->set_dims(qkv_weight.dims());
   }
 
-  if (qkv_bias_grad) {
+  if (qkv_bias_grad && qkv_bias) {
     qkv_bias_grad->set_dims(qkv_bias.dims());
   }
 
@@ -1040,7 +1041,7 @@ void FusedAttentionGradInferMeta(const MetaTensor& out_grad,
       ln_out_grad->set_dims(ln_out.dims());
     }
   } else {
-    if (bias_dropout_residual_out_grad) {
+    if (bias_dropout_residual_out_grad && bias_dropout_residual_out) {
       bias_dropout_residual_out_grad->set_dims(
           bias_dropout_residual_out.dims());
     }
@@ -1556,36 +1557,36 @@ void FusedFeedForwardGradInferMeta(const MetaTensor& out_grad,
                                    bool add_residual,
                                    int ring_id,
                                    MetaTensor* x_grad,
-                                   MetaTensor* ln1_scale_grad,
-                                   MetaTensor* ln1_bias_grad,
-                                   MetaTensor* ln2_scale_grad,
-                                   MetaTensor* ln2_bias_grad,
                                    MetaTensor* linear1_weight_grad,
                                    MetaTensor* linear1_bias_grad,
                                    MetaTensor* linear2_weight_grad,
-                                   MetaTensor* linear2_bias_grad) {
+                                   MetaTensor* linear2_bias_grad,
+                                   MetaTensor* ln1_scale_grad,
+                                   MetaTensor* ln1_bias_grad,
+                                   MetaTensor* ln2_scale_grad,
+                                   MetaTensor* ln2_bias_grad) {
   auto d_out_dim = out_grad.dims();
   x_grad->set_dims(d_out_dim);
-  if (ln1_scale_grad) {
+  if (ln1_scale_grad && ln1_scale) {
     ln1_scale_grad->set_dims(ln1_scale.dims());
   }
-  if (ln1_bias_grad) {
+  if (ln1_bias_grad && ln1_bias) {
     ln1_bias_grad->set_dims(ln1_bias.dims());
   }
-  if (ln2_scale_grad) {
+  if (ln2_scale_grad && ln2_scale) {
     ln2_scale_grad->set_dims(ln2_scale.dims());
   }
-  if (ln2_bias_grad) {
+  if (ln2_bias_grad && ln2_bias) {
     ln2_bias_grad->set_dims(ln2_bias.dims());
   }
 
   linear1_weight_grad->set_dims(linear1_weight.dims());
-  if (linear1_bias_grad) {
+  if (linear1_bias_grad && linear1_bias) {
     linear1_bias_grad->set_dims(linear1_bias.dims());
   }
 
   linear2_weight_grad->set_dims(linear2_weight.dims());
-  if (linear2_bias_grad) {
+  if (linear2_bias_grad && linear2_bias) {
     linear2_bias_grad->set_dims(linear2_bias.dims());
   }
 }
diff --git a/paddle/phi/infermeta/fusion.h b/paddle/phi/infermeta/fusion.h
diff --git a/paddle/phi/kernels/fused_feedforward_grad_kernel.h b/paddle/phi/kernels/fused_feedforward_grad_kernel.h
diff --git a/paddle/phi/kernels/fusion/xpu/fused_feedforward_grad_kernel.cc b/paddle/phi/kernels/fusion/xpu/fused_feedforward_grad_kernel.cc
diff --git a/paddle/phi/ops/yaml/inconsistent/static_backward.yaml b/paddle/phi/ops/yaml/inconsistent/static_backward.yaml

Original file line number	Diff line number	Diff line change
`@@ -40,6 +40,7 @@`
`40`	`40`	`VJPS_BLACK_LIST = [`
`41`	`41`	`'reshape_grad',`
`42`	`42`	`'add_n_grad',`
	`43`	`+ 'fused_attention_grad',`
`43`	`44`	`]`
`44`	`45`
`45`	`46`	`BACKENDS_BLACK_LIST = [`
Original file line number	Diff line number	Diff line change
`@@ -510,7 +510,7 @@ void GeneralTernaryGradInferMeta(const MetaTensor& x,`
`510`	`510`	`if (dx) {`
`511`	`511`	`dx->share_meta(x);`
`512`	`512`	`}`
`513`		`- if (dy) {`
	`513`	`+ if (dy && y) {`
`514`	`514`	`dy->share_meta(y);`
`515`	`515`	`}`
`516`	`516`	`if (dz) {`
Original file line number	Diff line number	Diff line change
`@@ -924,6 +924,7 @@ void FusedAttentionInferMeta(const MetaTensor& x,`
`924`	`924`	`}`
`925`	`925`
`926`	`926`	`out->set_dims(x.dims());`
	`927`	`+ out->set_dtype(x.dtype());`
`927`	`928`	`}`
`928`	`929`
`929`	`930`	`void FusedAttentionGradInferMeta(const MetaTensor& out_grad,`
`@@ -998,19 +999,19 @@ void FusedAttentionGradInferMeta(const MetaTensor& out_grad,`
`998`	`999`	`"GradOp is only callable when is_test is false"));`
`999`	`1000`
`1000`	`1001`	`if (!pre_layer_norm) {`
`1001`		`- if (ln_scale_2_grad) {`
	`1002`	`+ if (ln_scale_2_grad && ln_scale_2) {`
`1002`	`1003`	`ln_scale_2_grad->set_dims(ln_scale_2.dims());`
`1003`	`1004`	`}`
`1004`		`- if (ln_bias_2_grad) {`
	`1005`	`+ if (ln_bias_2_grad && ln_bias_2) {`
`1005`	`1006`	`ln_bias_2_grad->set_dims(ln_bias_2.dims());`
`1006`	`1007`	`}`
`1007`	`1008`	`}`
`1008`	`1009`
`1009`		`- if (pre_layer_norm) {`
	`1010`	`+ if (pre_layer_norm && ln_scale) {`
`1010`	`1011`	`if (ln_scale_grad) {`
`1011`	`1012`	`ln_scale_grad->set_dims(ln_scale.dims());`
`1012`	`1013`	`}`
`1013`		`- if (ln_bias_grad) {`
	`1014`	`+ if (ln_bias_grad && ln_bias) {`
`1014`	`1015`	`ln_bias_grad->set_dims(ln_bias.dims());`
`1015`	`1016`	`}`
`1016`	`1017`	`}`
`@@ -1019,7 +1020,7 @@ void FusedAttentionGradInferMeta(const MetaTensor& out_grad,`
`1019`	`1020`	`x_grad->set_dims(x.dims());`
`1020`	`1021`	`}`
`1021`	`1022`
`1022`		`- if (out_linear_bias_grad) {`
	`1023`	`+ if (out_linear_bias_grad && out_linear_bias) {`
`1023`	`1024`	`out_linear_bias_grad->set_dims(out_linear_bias.dims());`
`1024`	`1025`	`}`
`1025`	`1026`
`@@ -1031,7 +1032,7 @@ void FusedAttentionGradInferMeta(const MetaTensor& out_grad,`
`1031`	`1032`	`qkv_weight_grad->set_dims(qkv_weight.dims());`
`1032`	`1033`	`}`
`1033`	`1034`
`1034`		`- if (qkv_bias_grad) {`
	`1035`	`+ if (qkv_bias_grad && qkv_bias) {`
`1035`	`1036`	`qkv_bias_grad->set_dims(qkv_bias.dims());`
`1036`	`1037`	`}`
`1037`	`1038`
`@@ -1040,7 +1041,7 @@ void FusedAttentionGradInferMeta(const MetaTensor& out_grad,`
`1040`	`1041`	`ln_out_grad->set_dims(ln_out.dims());`
`1041`	`1042`	`}`
`1042`	`1043`	`} else {`
`1043`		`- if (bias_dropout_residual_out_grad) {`
	`1044`	`+ if (bias_dropout_residual_out_grad && bias_dropout_residual_out) {`
`1044`	`1045`	`bias_dropout_residual_out_grad->set_dims(`
`1045`	`1046`	`bias_dropout_residual_out.dims());`
`1046`	`1047`	`}`
`@@ -1556,36 +1557,36 @@ void FusedFeedForwardGradInferMeta(const MetaTensor& out_grad,`
`1556`	`1557`	`bool add_residual,`
`1557`	`1558`	`int ring_id,`
`1558`	`1559`	`MetaTensor* x_grad,`
`1559`		`- MetaTensor* ln1_scale_grad,`
`1560`		`- MetaTensor* ln1_bias_grad,`
`1561`		`- MetaTensor* ln2_scale_grad,`
`1562`		`- MetaTensor* ln2_bias_grad,`
`1563`	`1560`	`MetaTensor* linear1_weight_grad,`
`1564`	`1561`	`MetaTensor* linear1_bias_grad,`
`1565`	`1562`	`MetaTensor* linear2_weight_grad,`
`1566`		`- MetaTensor* linear2_bias_grad) {`
	`1563`	`+ MetaTensor* linear2_bias_grad,`
	`1564`	`+ MetaTensor* ln1_scale_grad,`
	`1565`	`+ MetaTensor* ln1_bias_grad,`
	`1566`	`+ MetaTensor* ln2_scale_grad,`
	`1567`	`+ MetaTensor* ln2_bias_grad) {`
`1567`	`1568`	`auto d_out_dim = out_grad.dims();`
`1568`	`1569`	`x_grad->set_dims(d_out_dim);`
`1569`		`- if (ln1_scale_grad) {`
	`1570`	`+ if (ln1_scale_grad && ln1_scale) {`
`1570`	`1571`	`ln1_scale_grad->set_dims(ln1_scale.dims());`
`1571`	`1572`	`}`
`1572`		`- if (ln1_bias_grad) {`
	`1573`	`+ if (ln1_bias_grad && ln1_bias) {`
`1573`	`1574`	`ln1_bias_grad->set_dims(ln1_bias.dims());`
`1574`	`1575`	`}`
`1575`		`- if (ln2_scale_grad) {`
	`1576`	`+ if (ln2_scale_grad && ln2_scale) {`
`1576`	`1577`	`ln2_scale_grad->set_dims(ln2_scale.dims());`
`1577`	`1578`	`}`
`1578`		`- if (ln2_bias_grad) {`
	`1579`	`+ if (ln2_bias_grad && ln2_bias) {`
`1579`	`1580`	`ln2_bias_grad->set_dims(ln2_bias.dims());`
`1580`	`1581`	`}`
`1581`	`1582`
`1582`	`1583`	`linear1_weight_grad->set_dims(linear1_weight.dims());`
`1583`		`- if (linear1_bias_grad) {`
	`1584`	`+ if (linear1_bias_grad && linear1_bias) {`
`1584`	`1585`	`linear1_bias_grad->set_dims(linear1_bias.dims());`
`1585`	`1586`	`}`
`1586`	`1587`
`1587`	`1588`	`linear2_weight_grad->set_dims(linear2_weight.dims());`
`1588`		`- if (linear2_bias_grad) {`
	`1589`	`+ if (linear2_bias_grad && linear2_bias) {`
`1589`	`1590`	`linear2_bias_grad->set_dims(linear2_bias.dims());`
`1590`	`1591`	`}`
`1591`	`1592`	`}`