Skip to content

Commit 98bd7d3

Browse files
cccddd77cccddd77
committed
flash_attention_v2_backward
1 parent 44ad994 commit 98bd7d3

File tree

8 files changed

+745
-24
lines changed

8 files changed

+745
-24
lines changed

‎oneflow/core/autograd/gradient_funcs/scaled_dot_product_attention.cpp

Lines changed: 109 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,109 @@
1+
/*
2+
Copyright 2020 The OneFlow Authors. All rights reserved.
3+
4+
Licensed under the Apache License, Version 2.0 (the "License");
5+
you may not use this file except in compliance with the License.
6+
You may obtain a copy of the License at
7+
8+
http://www.apache.org/licenses/LICENSE-2.0
9+
10+
Unless required by applicable law or agreed to in writing, software
11+
distributed under the License is distributed on an "AS IS" BASIS,
12+
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
See the License for the specific language governing permissions and
14+
limitations under the License.
15+
*/
16+
#include "oneflow/core/common/maybe.h"
17+
#include "oneflow/core/framework/op_expr_grad_function.h"
18+
#include "oneflow/core/framework/op_builder.h"
19+
#include "oneflow/core/framework/op_expr.h"
20+
#include "oneflow/core/framework/op_interpreter/op_interpreter_util.h"
21+
#include "oneflow/core/functional/functional.h"
22+
#include "oneflow/core/common/container_util.h"
23+
24+
namespace oneflow {
25+
26+
namespace one {
27+
28+
struct ScaledDotProductFlashAttentionCaptureState : public AutoGradCaptureState {
29+
bool query_requires_grad = true;
30+
bool key_requires_grad = true;
31+
bool value_requires_grad = true;
32+
size_t query_idx = 0;
33+
size_t key_idx = 0;
34+
size_t value_idx = 0;
35+
size_t out_idx = 0;
36+
size_t softmax_lse_idx = 0;
37+
size_t rng_state_idx = 0;
38+
float p_dropout = .0f;
39+
float softmax_scale = .0f;
40+
bool is_causal = false;
41+
};
42+
43+
class ScaledDotProductFlashAttention
44+
: public OpExprGradFunction<ScaledDotProductFlashAttentionCaptureState> {
45+
public:
46+
Maybe<void> Init(const OpExpr& op) override {
47+
const auto* fw_op_expr = dynamic_cast<const UserOpExpr*>(&op);
48+
CHECK_NOTNULL_OR_RETURN(fw_op_expr) << "fw_op_expr should not be None. ";
49+
base_attrs_ = MakeAttrMapFromUserOpConf(fw_op_expr->proto());
50+
return Maybe<void>::Ok();
51+
}
52+
53+
Maybe<void> Capture(ScaledDotProductFlashAttentionCaptureState* ctx, const TensorTuple& inputs,
54+
const TensorTuple& outputs, const AttrMap& attrs) const override {
55+
CHECK_EQ_OR_RETURN(inputs.size(), 3) << "Input size should be equal to 3. ";
56+
ComposedAttrMap composed_attrs(attrs, base_attrs_);
57+
ctx->p_dropout = JUST(composed_attrs.GetAttr<float>("p_dropout"));
58+
ctx->softmax_scale = JUST(composed_attrs.GetAttr<float>("softmax_scale"));
59+
ctx->is_causal = JUST(composed_attrs.GetAttr<bool>("is_causal"));
60+
ctx->query_requires_grad = JUST(oneflow::VectorAt(inputs, 0))->requires_grad();
61+
ctx->key_requires_grad = JUST(oneflow::VectorAt(inputs, 1))->requires_grad();
62+
ctx->value_requires_grad = JUST(oneflow::VectorAt(inputs, 2))->requires_grad();
63+
ctx->query_idx = ctx->SaveTensorForBackward(JUST(oneflow::VectorAt(inputs, 0)));
64+
ctx->key_idx = ctx->SaveTensorForBackward(JUST(oneflow::VectorAt(inputs, 1)));
65+
ctx->value_idx = ctx->SaveTensorForBackward(JUST(oneflow::VectorAt(inputs, 2)));
66+
ctx->out_idx = ctx->SaveTensorForBackward(JUST(oneflow::VectorAt(outputs, 0)));
67+
ctx->softmax_lse_idx = ctx->SaveTensorForBackward(JUST(oneflow::VectorAt(outputs, 1)));
68+
ctx->rng_state_idx = ctx->SaveTensorForBackward(JUST(oneflow::VectorAt(outputs, 2)));
69+
return Maybe<void>::Ok();
70+
}
71+
72+
Maybe<void> Apply(const ScaledDotProductFlashAttentionCaptureState* ctx,
73+
const TensorTuple& out_grads, TensorTuple* in_grads) const override {
74+
CHECK_EQ_OR_RETURN(out_grads.size(), 3) << "Out grads size should be equal to 3. ";
75+
std::shared_ptr<oneflow::one::TensorTuple> grads;
76+
in_grads->resize(3);
77+
grads = JUST(functional::ScaledDotProductFlashAttentionGrad(
78+
JUST(oneflow::VectorAt(out_grads, 0)),
79+
JUST(oneflow::VectorAt(ctx->SavedTensors(), ctx->query_idx)),
80+
JUST(oneflow::VectorAt(ctx->SavedTensors(), ctx->key_idx)),
81+
JUST(oneflow::VectorAt(ctx->SavedTensors(), ctx->value_idx)),
82+
JUST(oneflow::VectorAt(ctx->SavedTensors(), ctx->out_idx)),
83+
JUST(oneflow::VectorAt(ctx->SavedTensors(), ctx->softmax_lse_idx)),
84+
JUST(oneflow::VectorAt(ctx->SavedTensors(), ctx->rng_state_idx)), ctx->p_dropout,
85+
ctx->is_causal, ctx->softmax_scale));
86+
87+
if (ctx->query_requires_grad) {
88+
JUST(oneflow::VectorAt(*in_grads, 0)) = JUST(oneflow::VectorAt(*grads, 0));
89+
}
90+
if (ctx->key_requires_grad) {
91+
JUST(oneflow::VectorAt(*in_grads, 1)) = JUST(oneflow::VectorAt(*grads, 1));
92+
}
93+
if (ctx->value_requires_grad) {
94+
JUST(oneflow::VectorAt(*in_grads, 2)) = JUST(oneflow::VectorAt(*grads, 2));
95+
}
96+
97+
return Maybe<void>::Ok();
98+
}
99+
100+
private:
101+
AttrMap base_attrs_;
102+
};
103+
104+
REGISTER_OP_EXPR_GRAD_FUNCTION("scaled_dot_product_flash_attention",
105+
ScaledDotProductFlashAttention);
106+
107+
} // namespace one
108+
109+
} // namespace oneflow

‎oneflow/core/functional/functional_api.yaml

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2688,6 +2688,10 @@
26882688
signature: "Tensor (Tensor query, Tensor key, Tensor value, Tensor attn_mask=None, Float dropout_p=0.0, Bool is_causal=False, Float scale=None, Int64 seed=0) => ScaledDotProductFlashAttention"
26892689
bind_python: True
26902690

2691+
- name: "scaled_dot_product_attention_grad"
2692+
signature: "TensorTuple (Tensor grad_out, Tensor query, Tensor key, Tensor value, Tensor out, Tensor softmax_lse, Tensor rng_state, Float dropout_p=0.0, Bool is_causal=False, Float scale=0.0) => ScaledDotProductFlashAttentionGrad"
2693+
bind_python: False
2694+
26912695
- name: "fused_multi_head_attention_inference"
26922696
signature: "Tensor (Tensor query, Tensor key, Tensor value, Int64 num_heads, Bool causal=False, Int64 query_hidden_slice_start=0, Int64 query_hidden_slice_end=-1, Int64 key_hidden_slice_start=0, Int64 key_hidden_slice_end=-1, Int64 value_hidden_slice_start=0, Int64 value_hidden_slice_end=-1, Tensor attn_bias=None, Int64 causal_diagonal_offset=0) => FusedMultiHeadAttentionInference"
26932697
bind_python: True

‎oneflow/core/functional/impl/nn_functor.cpp

Lines changed: 135 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5538,6 +5538,139 @@ class ScaledDotProductFlashAttentionFunctor {
55385538
#endif // CUDA_VERSION >= 11070
55395539
};
55405540

5541+
class ScaledDotProductFlashAttentionGradFunctor {
5542+
public:
5543+
ScaledDotProductFlashAttentionGradFunctor() {
5544+
#if CUDA_VERSION >= 11070
5545+
op_ = CHECK_JUST(one::OpBuilder("scaled_dot_product_flash_attention_grad")
5546+
.Input("grad_out")
5547+
.Input("query")
5548+
.Input("key")
5549+
.Input("value")
5550+
.Input("out")
5551+
.Input("softmax_lse")
5552+
.Input("rng_state")
5553+
.Output("grad_q")
5554+
.Output("grad_k")
5555+
.Output("grad_v")
5556+
.Build());
5557+
#endif
5558+
}
5559+
5560+
Maybe<TensorTuple> operator()(
5561+
const std::shared_ptr<one::Tensor>& grad_out, const std::shared_ptr<one::Tensor>& query,
5562+
const std::shared_ptr<one::Tensor>& key, const std::shared_ptr<one::Tensor>& value,
5563+
const std::shared_ptr<one::Tensor>& out, const std::shared_ptr<one::Tensor>& softmax_lse,
5564+
const std::shared_ptr<one::Tensor>& rng_state, const float& dropout_p, const bool& is_causal,
5565+
const float& scale) const {
5566+
#if CUDA_VERSION >= 11070
5567+
// grad_out(batch x q_sqe_len x num_heads x head_size)
5568+
// query (batch x q_seq_len x num_heads x head_size_padded)
5569+
// key (batch x kv_seq_len x num_heads_k x head_size_padded)
5570+
// value (batch x kv_seq_len x num_heads_k x head_size_padded)
5571+
// out (batch x kv_seq_len x num_heads x head_size_padded)
5572+
// softmax_lse (batch x num_heads x q_seq_len)
5573+
const auto head_size = grad_out->shape()->At(3);
5574+
const auto head_size_padded = query->shape()->At(3);
5575+
const auto batch_size = query->shape()->At(0);
5576+
const auto seqlen_q = query->shape()->At(1);
5577+
const auto seqlen_k = key->shape()->At(1);
5578+
const auto num_heads = query->shape()->At(2);
5579+
const auto num_heads_k = key->shape()->At(2);
5580+
CHECK_EQ_OR_RETURN(batch_size, key->shape()->At(0))
5581+
<< " key has different batch size from query.";
5582+
CHECK_EQ_OR_RETURN(batch_size, value->shape()->At(0))
5583+
<< " value has different batch size from query.";
5584+
CHECK_EQ_OR_RETURN(batch_size, grad_out->shape()->At(0))
5585+
<< " grad_out has different batch size from query.";
5586+
CHECK_EQ_OR_RETURN(batch_size, out->shape()->At(0))
5587+
<< " out has different batch size from query.";
5588+
CHECK_EQ_OR_RETURN(batch_size, softmax_lse->shape()->At(0))
5589+
<< " softmax_lse has different batch size from query.";
5590+
CHECK_EQ_OR_RETURN(num_heads, grad_out->shape()->At(2))
5591+
<< " grad_out has different num_heads from query.";
5592+
CHECK_EQ_OR_RETURN(num_heads, softmax_lse->shape()->At(1))
5593+
<< " softmax_lse has different num_heads from query.";
5594+
CHECK_EQ_OR_RETURN(num_heads_k, value->shape()->At(2))
5595+
<< " value has different num_heads from key.";
5596+
CHECK_EQ_OR_RETURN(seqlen_q, grad_out->shape()->At(1))
5597+
<< " grad_out has different seq_len from query.";
5598+
CHECK_EQ_OR_RETURN(seqlen_q, softmax_lse->shape()->At(2))
5599+
<< " softmax_lse has different seq_len from query.";
5600+
CHECK_EQ_OR_RETURN(head_size_padded, key->shape()->At(3))
5601+
<< " key has different head dims from query.";
5602+
CHECK_EQ_OR_RETURN(head_size_padded, value->shape()->At(3))
5603+
<< " key has different head dims from query.";
5604+
CHECK_EQ_OR_RETURN(head_size_padded, out->shape()->At(3))
5605+
<< " out has different head dims from query.";
5606+
5607+
bool padded = head_size % 8;
5608+
5609+
auto grad_out_ = padded ? JUST(pad_last_dim<8>(grad_out)) : grad_out;
5610+
5611+
auto& attrs = THREAD_CACHED_MUTABLE_ATTR_MAP("p_dropout", "softmax_scale", "is_causal",
5612+
"window_size_left", "window_size_right");
5613+
attrs.SetAllAttrs(dropout_p, scale, is_causal, -1, -1);
5614+
5615+
auto output = std::make_shared<TensorTuple>(3);
5616+
auto output_ = JUST(OpInterpUtil::Dispatch<TensorTuple>(
5617+
*op_, {grad_out_, query, key, value, out, softmax_lse, rng_state}, attrs));
5618+
CHECK_EQ(output_->size(), 3);
5619+
auto grad_q_ = (*output_)[0];
5620+
auto grad_k_ = (*output_)[1];
5621+
auto grad_v_ = (*output_)[2];
5622+
5623+
// auto grad_q_padded = JUST(functional::Transpose(grad_q_, {0, 2, 1, 3}));
5624+
// auto grad_k_padded_expanded = JUST(functional::Transpose(grad_k_, {0, 2, 1, 3}));
5625+
// auto grad_v_padded_expanded = JUST(functional::Transpose(grad_v_, {0, 2, 1, 3}));
5626+
5627+
std::shared_ptr<Tensor> grad_q_padded, grad_k_padded, grad_v_padded;
5628+
5629+
bool expanded = num_heads != num_heads_k;
5630+
5631+
grad_q_padded = grad_q_;
5632+
if (expanded) {
5633+
grad_k_padded = JUST(functional::ReduceSum(
5634+
JUST(functional::Reshape(grad_k_, {batch_size, seqlen_k, num_heads_k,
5635+
num_heads / num_heads_k, head_size_padded})),
5636+
{3}, false, grad_k_->dtype()));
5637+
grad_v_padded = JUST(functional::ReduceSum(
5638+
JUST(functional::Reshape(grad_v_, {batch_size, seqlen_k, num_heads_k,
5639+
num_heads / num_heads_k, head_size_padded})),
5640+
{3}, false, grad_v_->dtype()));
5641+
} else {
5642+
grad_k_padded = grad_k_;
5643+
grad_v_padded = grad_v_;
5644+
}
5645+
5646+
auto grad_q = padded ? JUST(functional::Slice(grad_q_padded, {0, 0, 0, 0},
5647+
{batch_size, seqlen_q, num_heads, head_size},
5648+
{1, 1, 1, 1}, false))
5649+
: grad_q_padded;
5650+
auto grad_k = padded ? JUST(functional::Slice(grad_k_padded, {0, 0, 0, 0},
5651+
{batch_size, seqlen_k, num_heads_k, head_size},
5652+
{1, 1, 1, 1}, false))
5653+
: grad_k_padded;
5654+
auto grad_v = padded ? JUST(functional::Slice(grad_v_padded, {0, 0, 0, 0},
5655+
{batch_size, seqlen_k, num_heads_k, head_size},
5656+
{1, 1, 1, 1}, false))
5657+
: grad_v_padded;
5658+
5659+
(*output)[0] = grad_q;
5660+
(*output)[1] = grad_k;
5661+
(*output)[2] = grad_v;
5662+
return output;
5663+
5664+
#endif // CUDA_VERSION >= 11070
5665+
5666+
UNIMPLEMENTED_THEN_RETURN() << "only support CUDA_VERSION >= 11070.";
5667+
}
5668+
5669+
private:
5670+
#if CUDA_VERSION >= 11070
5671+
std::shared_ptr<OpExpr> op_;
5672+
#endif // CUDA_VERSION >= 11070
5673+
};
55415674
} // namespace impl
55425675

55435676
ONEFLOW_FUNCTION_LIBRARY(m) {
@@ -5676,6 +5809,8 @@ ONEFLOW_FUNCTION_LIBRARY(m) {
56765809
m.add_functor<impl::MultiTensorYoloV5WeightUpdateFunctor>("MultiTensorYoloV5WeightUpdate");
56775810
m.add_functor<impl::FusedClipGradFunctor>("FusedClipGrad");
56785811
m.add_functor<impl::ScaledDotProductFlashAttentionFunctor>("ScaledDotProductFlashAttention");
5812+
m.add_functor<impl::ScaledDotProductFlashAttentionGradFunctor>(
5813+
"ScaledDotProductFlashAttentionGrad");
56795814
}
56805815

56815816
} // namespace functional

‎oneflow/ir/include/OneFlow/OneFlowUserOps.td

Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2903,6 +2903,35 @@ def OneFlow_ScaledDotProductFlashAttentionOp : OneFlow_BaseOp<"scaled_dot_produc
29032903
let has_data_type_infer_fn = 1;
29042904
}
29052905

2906+
def OneFlow_ScaledDotProductFlashAttentionGradOp : OneFlow_BaseOp<"scaled_dot_product_flash_attention_grad", [NoMemoryEffect, NoGrad, DeclareOpInterfaceMethods<UserOpCompatibleInterface>]> {
2907+
let input = (ins
2908+
OneFlow_Tensor:$grad_out,
2909+
OneFlow_Tensor:$query,
2910+
OneFlow_Tensor:$key,
2911+
OneFlow_Tensor:$value,
2912+
OneFlow_Tensor:$out,
2913+
OneFlow_Tensor:$softmax_lse,
2914+
OneFlow_Tensor:$rng_state,
2915+
Optional<OneFlow_Tensor>:$alibi_slopes_
2916+
);
2917+
let output = (outs
2918+
OneFlow_Tensor:$grad_q,
2919+
OneFlow_Tensor:$grad_k,
2920+
OneFlow_Tensor:$grad_v
2921+
);
2922+
let attrs = (ins
2923+
DefaultValuedAttr<F32Attr, "0.">:$p_dropout,
2924+
DefaultValuedAttr<F32Attr, "0.">:$softmax_scale,
2925+
DefaultValuedAttr<BoolAttr, "false">:$is_causal,
2926+
SI32Attr:$window_size_left,
2927+
SI32Attr:$window_size_right
2928+
);
2929+
let has_logical_tensor_desc_infer_fn = 1;
2930+
let has_physical_tensor_desc_infer_fn = 1;
2931+
let has_get_sbp_fn = 1;
2932+
let has_data_type_infer_fn = 1;
2933+
}
2934+
29062935
def OneFlow_FusedMultiHeadAttentionInferenceOp : OneFlow_BaseOp<"fused_multi_head_attention_inference", [NoMemoryEffect, AttrSizedOperandSegments, DeclareOpInterfaceMethods<UserOpCompatibleInterface>]> {
29072936
let input = (ins
29082937
OneFlow_Tensor:$query,

0 commit comments

Comments
 (0)