add some fused kernels #6635

dangkai4u · 2021-10-28T03:33:04Z

Add some fused kernel used in Transformer:

fused_scale_mask_softmax_dropout: scale + mask + softmax + dropout
fused_scale_mask_softmax: scale + mask + softmax

MARD1NO · 2021-11-04T08:44:52Z

oneflow/user/kernels/fused_scale_mask_softmax_dropout.cu

+};
+
+template<typename SRC, typename DST>
+struct DropoutScore {


Score -> Store

已改，顺便问一下，functional_api.yaml中，对于xxx_grad函数也进行了导出，这个是不是没有必要，因为bind_python通常设置为False，而且用户大多数情况下也不会使用

在autograd里面会使用，具体看oneflow/core/autograd/gradient_funcs/

MARD1NO · 2021-11-05T06:08:32Z

oneflow/core/autograd/gradient_funcs/fused_scale_mask_softmax.cpp

+                                         const TensorTuple& out_grads, TensorTuple* in_grads) const {
+  if (!ctx->input_requires_grad) { return Maybe<void>::Ok(); }
+
+  CHECK_EQ_OR_RETURN(out_grads.size(), 2);  // softmax_y, dy


这里的描述可能不太准确

fusedscalemaskSoftmax有两个输出，因此回来的梯度是两个

所以一个是softmax_y的梯度，一个是y的梯度。y的梯度叫dy可以，softmax_y的梯度这里注释需要重新写下，比如softmax_dy

softmax_y 会输出给前向吗得到diff吗，印象中一般是用来传给后向op？

传过来应该都是两个梯度，但我们只要dy的梯度。softmax_y是传给后向op的，这里我只是提醒下这里的注释什么的

MARD1NO · 2021-11-05T06:14:24Z

oneflow/core/functional/impl/nn_functor.cpp

+                                                          .Output("softmax_y")
+                                                          .Build());
+  }
+  Maybe<Tensor> operator()(const std::shared_ptr<one::Tensor>& x, const std::shared_ptr<one::Tensor>& mask,


你最后调用的op是 fused_scale_mask_softmax_dropout_op_，他有两个输出，所以这里返回值应该是一个Tensor Tuple。不然最后输出只有一个，你在求导那里保存不到softmax_y，进而求不了导数

MARD1NO · 2021-11-05T06:23:55Z

oneflow/core/autograd/gradient_funcs/fused_scale_mask_softmax.cpp

+                                         const TensorTuple& out_grads, TensorTuple* in_grads) const {
+  if (!ctx->input_requires_grad) { return Maybe<void>::Ok(); }
+
+  CHECK_EQ_OR_RETURN(out_grads.size(), 2);  // softmax_y, dy


传过来应该都是两个梯度，但我们只要dy的梯度。softmax_y是传给后向op的，这里我只是提醒下这里的注释什么的

MARD1NO · 2021-11-08T06:24:02Z

oneflow/core/functional/functional_api.yaml

 - name: "fused_scale_tril_softmax_mask_scale"
  signature: "TensorTuple (Tensor a, *, Float p=0.5, Int64 diagonal, Float tril_scale_value, Generator generator=None) => FusedScaleTrilSoftmaxMaskScale"
  bind_python: True

 - name: "fused_scale_tril_softmax_mask_scale_grad"
-  signature: "Tensor (Tensor softmax_y, Tensor dy, Tensor mask, Int64 diagonal, Float tril_scale_value, Float mask_scale_value) => FusedScaleTrilSoftmaxMaskScaleGrad"
+  signature: "Tensor (Tensor y, Tensor dy, Tensor mask, Int64 diagonal, Float tril_scale_value, Float mask_scale_value) => FusedScaleTrilSoftmaxMaskScaleGrad"


这个不用改吧，用的就是softmax_y？

好的，这里改了

MARD1NO · 2021-11-08T06:24:20Z

oneflow/core/autograd/gradient_funcs/fused_scale_mask_softmax_dropout.cpp

+  ctx->scale = JUST(composed_attrs.GetAttr<float>("scale_value"));
+  ctx->dropout_scale = JUST(composed_attrs.GetAttr<float>("dropout_scale_value"));
+
+  ctx->SaveTensorForBackward(inputs.at(1));


可以加点注释，表示这里save了什么

这里也要加

dangkai4u · 2021-11-08T13:42:10Z

这个pr我本地测试通过了，你们可以测试一下，之后就合并吧

MARD1NO · 2021-11-10T00:39:40Z