Skip to content

Commit bed652d

Browse files
limin2021limin2021
authored
[Cherry pick 2.2] fix bugs to support bias add none for fused_attention op. (#37411) (#37483)
Add support for bias is none for fused_attention op.
1 parent f873d3a commit bed652d

File tree

4 files changed

+173
-71
lines changed

4 files changed

+173
-71
lines changed

‎paddle/fluid/operators/fused/fused_attention_op.cc‎

Lines changed: 47 additions & 30 deletions
Original file line numberDiff line numberDiff line change
@@ -28,12 +28,8 @@ class FusedAttentionOp : public framework::OperatorWithKernel {
2828
void InferShape(framework::InferShapeContext *ctx) const override {
2929
OP_INOUT_CHECK(ctx->HasInput("X"), "Input", "X", "FusedAttentionOp");
3030
OP_INOUT_CHECK(ctx->HasInput("QKVW"), "Input", "QKVW", "FusedAttentionOp");
31-
OP_INOUT_CHECK(ctx->HasInput("QKVBias"), "Input", "QKVBias",
32-
"FusedAttentionOp");
3331
OP_INOUT_CHECK(ctx->HasInput("OutLinearW"), "Input", "OutLinearW",
3432
"FusedAttentionOp");
35-
OP_INOUT_CHECK(ctx->HasInput("OutLinearBias"), "Input", "OutLinearBias",
36-
"FusedAttentionOp");
3733

3834
if (ctx->Attrs().Get<bool>("pre_layer_norm") == true) {
3935
OP_INOUT_CHECK(ctx->HasOutput("LnMean"), "Output", "LnMean",
@@ -54,8 +50,10 @@ class FusedAttentionOp : public framework::OperatorWithKernel {
5450
// qkv_out: [batch_size, seq_len, 3, num_head, dim_head]
5551
OP_INOUT_CHECK(ctx->HasOutput("QKVOut"), "Output", "QKVOut",
5652
"FusedAttentionOp");
57-
OP_INOUT_CHECK(ctx->HasOutput("QKVBiasOut"), "Output", "QKVBiasOut",
58-
"FusedAttentionOp");
53+
if (ctx->HasInput("QKVBias")) {
54+
OP_INOUT_CHECK(ctx->HasOutput("QKVBiasOut"), "Output", "QKVBiasOut",
55+
"FusedAttentionOp");
56+
}
5957
OP_INOUT_CHECK(ctx->HasOutput("TransposeOut2"), "Output", "TransposeOut2",
6058
"FusedAttentionOp");
6159
OP_INOUT_CHECK(ctx->HasOutput("QKOut"), "Output", "QKOut",
@@ -107,6 +105,13 @@ class FusedAttentionOp : public framework::OperatorWithKernel {
107105
"input qkv_weight = [%s]",
108106
x_dim, y_dim));
109107

108+
PADDLE_ENFORCE_EQ(y_dim[1] * y_dim[2], y_dim[3],
109+
platform::errors::InvalidArgument(
110+
"The dimensions of qkv_weight must be 4"
111+
"(3, num_head, dim_head, dim_embed),"
112+
"and must satisfy the limitations: "
113+
"(num_head * dim_head == dim_embed)"));
114+
110115
if (ctx->Attrs().Get<bool>("pre_layer_norm") == true) {
111116
ctx->SetOutputDim("LnMean", {x_dim[0] * x_dim[1]});
112117
ctx->SetOutputDim("LnVariance", {x_dim[0] * x_dim[1]});
@@ -119,8 +124,11 @@ class FusedAttentionOp : public framework::OperatorWithKernel {
119124
// [batch_size, seq_len, 3, num_head, head_size]
120125
ctx->SetOutputDim("QKVOut",
121126
{x_dim[0], x_dim[1], y_dim[0], y_dim[1], y_dim[2]});
122-
ctx->SetOutputDim("QKVBiasOut",
123-
{x_dim[0], x_dim[1], y_dim[0], y_dim[1], y_dim[2]});
127+
128+
if (ctx->HasInput("QKVBias")) {
129+
ctx->SetOutputDim("QKVBiasOut",
130+
{x_dim[0], x_dim[1], y_dim[0], y_dim[1], y_dim[2]});
131+
}
124132
// [3, batch_size, num_head, seq_len, head_size]
125133
ctx->SetOutputDim("TransposeOut2",
126134
{y_dim[0], x_dim[0], y_dim[1], x_dim[1], y_dim[2]});
@@ -173,11 +181,11 @@ class FusedAttentionOpMaker : public framework::OpProtoAndCheckerMaker {
173181
"H. Here, H represents the last dimension of its input tensor.")
174182
.AsDispensable();
175183
AddInput("QKVW", "The qkv weight tensor.");
176-
AddInput("QKVBias", "The qkv bias tensor.");
184+
AddInput("QKVBias", "The qkv bias tensor.").AsDispensable();
177185
AddInput("SrcMask", "(optional) The attention mask tensor in fmha.")
178186
.AsDispensable();
179187
AddInput("OutLinearW", "The out_linear weight tensor.");
180-
AddInput("OutLinearBias", "The out_linear bias tensor.");
188+
AddInput("OutLinearBias", "The out_linear bias tensor.").AsDispensable();
181189
AddInput("Ln2Scale",
182190
"(optional) Scale is a 1-dimensional tensor of size "
183191
"H. Here, H represents the last dimension of its input tensor.")
@@ -379,12 +387,8 @@ class FusedAttentionGradOp : public framework::OperatorWithKernel {
379387
OP_INOUT_CHECK(ctx->HasInput("X"), "Input", "X", "FusedAttentionGrad");
380388
OP_INOUT_CHECK(ctx->HasInput("QKVW"), "Input", "QKVW",
381389
"FusedAttentionGrad");
382-
OP_INOUT_CHECK(ctx->HasInput("QKVBias"), "Input", "QKVBias",
383-
"FusedAttentionGrad");
384390
OP_INOUT_CHECK(ctx->HasInput("OutLinearW"), "Input", "OutLinearW",
385391
"FusedAttentionGrad");
386-
OP_INOUT_CHECK(ctx->HasInput("OutLinearBias"), "Input", "OutLinearBias",
387-
"FusedAttentionGrad");
388392

389393
if (ctx->Attrs().Get<bool>("pre_layer_norm") == true) {
390394
if (ctx->HasOutput(framework::GradVarName("LnScale"))) {
@@ -399,14 +403,17 @@ class FusedAttentionGradOp : public framework::OperatorWithKernel {
399403
if (ctx->HasOutput(framework::GradVarName("X"))) {
400404
ctx->SetOutputDim(framework::GradVarName("X"), ctx->GetInputDim("X"));
401405
}
402-
403-
ctx->SetOutputDim(framework::GradVarName("OutLinearBias"),
404-
ctx->GetInputDim("OutLinearBias"));
406+
if (ctx->HasOutput(framework::GradVarName("OutLinearBias"))) {
407+
ctx->SetOutputDim(framework::GradVarName("OutLinearBias"),
408+
ctx->GetInputDim("OutLinearBias"));
409+
}
405410
ctx->SetOutputDim(framework::GradVarName("OutLinearW"),
406411
ctx->GetInputDim("OutLinearW"));
407412
ctx->SetOutputDim(framework::GradVarName("QKVW"), ctx->GetInputDim("QKVW"));
408-
ctx->SetOutputDim(framework::GradVarName("QKVBias"),
409-
ctx->GetInputDim("QKVBias"));
413+
if (ctx->HasOutput(framework::GradVarName("QKVBias"))) {
414+
ctx->SetOutputDim(framework::GradVarName("QKVBias"),
415+
ctx->GetInputDim("QKVBias"));
416+
}
410417

411418
if (ctx->Attrs().Get<bool>("pre_layer_norm") == true) {
412419
ctx->SetOutputDim(framework::GradVarName("LnOut"),
@@ -434,8 +441,10 @@ class FusedAttentionGradOp : public framework::OperatorWithKernel {
434441
}
435442
ctx->SetOutputDim(framework::GradVarName("QKVOut"),
436443
ctx->GetInputDim("QKVOut"));
437-
ctx->SetOutputDim(framework::GradVarName("QKVBiasOut"),
438-
ctx->GetInputDim("QKVBiasOut"));
444+
if (ctx->HasOutput(framework::GradVarName("QKVBiasOut"))) {
445+
ctx->SetOutputDim(framework::GradVarName("QKVBiasOut"),
446+
ctx->GetInputDim("QKVBiasOut"));
447+
}
439448
ctx->SetOutputDim(framework::GradVarName("OutLinearOut"),
440449
ctx->GetInputDim("OutLinearOut"));
441450
}
@@ -462,7 +471,15 @@ class FusedAttentionGradOpMaker : public framework::SingleGradOpMaker<T> {
462471
// inputs x, parameters and their grad.
463472
op->SetInput("X", this->Input("X"));
464473
op->SetInput("QKVW", this->Input("QKVW"));
465-
op->SetInput("QKVBias", this->Input("QKVBias"));
474+
475+
if (this->HasInput("QKVBias")) {
476+
op->SetInput("QKVBias", this->Input("QKVBias"));
477+
op->SetOutput(framework::GradVarName("QKVBias"),
478+
this->InputGrad("QKVBias"));
479+
op->SetInput("QKVBiasOut", this->Output("QKVBiasOut"));
480+
op->SetOutput(framework::GradVarName("QKVBiasOut"),
481+
this->OutputGrad("QKVBiasOut"));
482+
}
466483

467484
if (this->HasInput("SrcMask")) {
468485
op->SetInput("SrcMask", this->Input("SrcMask"));
@@ -472,7 +489,11 @@ class FusedAttentionGradOpMaker : public framework::SingleGradOpMaker<T> {
472489
}
473490

474491
op->SetInput("OutLinearW", this->Input("OutLinearW"));
475-
op->SetInput("OutLinearBias", this->Input("OutLinearBias"));
492+
if (this->HasInput("OutLinearBias")) {
493+
op->SetInput("OutLinearBias", this->Input("OutLinearBias"));
494+
op->SetOutput(framework::GradVarName("OutLinearBias"),
495+
this->InputGrad("OutLinearBias"));
496+
}
476497

477498
op->SetAttrMap(this->Attrs());
478499
bool is_pre_layer_norm =
@@ -503,10 +524,7 @@ class FusedAttentionGradOpMaker : public framework::SingleGradOpMaker<T> {
503524

504525
op->SetOutput(framework::GradVarName("X"), this->InputGrad("X"));
505526
op->SetOutput(framework::GradVarName("QKVW"), this->InputGrad("QKVW"));
506-
op->SetOutput(framework::GradVarName("QKVBias"),
507-
this->InputGrad("QKVBias"));
508-
op->SetOutput(framework::GradVarName("OutLinearBias"),
509-
this->InputGrad("OutLinearBias"));
527+
510528
op->SetOutput(framework::GradVarName("OutLinearW"),
511529
this->InputGrad("OutLinearW"));
512530

@@ -528,7 +546,7 @@ class FusedAttentionGradOpMaker : public framework::SingleGradOpMaker<T> {
528546
this->Output("BiasDropoutResidualOut"));
529547
}
530548
op->SetInput("QKVOut", this->Output("QKVOut"));
531-
op->SetInput("QKVBiasOut", this->Output("QKVBiasOut"));
549+
532550
op->SetInput("TransposeOut2", this->Output("TransposeOut2"));
533551
op->SetInput("QKOut", this->Output("QKOut"));
534552
op->SetInput("QKTVOut", this->Output("QKTVOut"));
@@ -553,8 +571,7 @@ class FusedAttentionGradOpMaker : public framework::SingleGradOpMaker<T> {
553571
}
554572

555573
op->SetOutput(framework::GradVarName("QKVOut"), this->OutputGrad("QKVOut"));
556-
op->SetOutput(framework::GradVarName("QKVBiasOut"),
557-
this->OutputGrad("QKVBiasOut"));
574+
558575
op->SetOutput(framework::GradVarName("QKTVOut"),
559576
this->OutputGrad("QKTVOut"));
560577
op->SetOutput(framework::GradVarName("TransposeOut2"),

‎paddle/fluid/operators/fused/fused_attention_op.cu‎

Lines changed: 76 additions & 31 deletions
Original file line numberDiff line numberDiff line change
@@ -96,9 +96,11 @@ class FusedAttentionOpKernel : public framework::OpKernel<T> {
9696

9797
auto *x_data = input_x->data<T>();
9898
auto *qkv_weight_data = qkv_weight->data<T>();
99-
auto *qkv_bias_data = qkv_bias->data<T>();
99+
auto *qkv_bias_data = (qkv_bias == nullptr) ? nullptr : qkv_bias->data<T>();
100100
auto *qkv_out_data = qkv_out->mutable_data<T>(ctx.GetPlace());
101-
auto *qkv_bias_out_data = qkv_bias_out->mutable_data<T>(ctx.GetPlace());
101+
auto *qkv_bias_out_data =
102+
(qkv_bias == nullptr) ? nullptr
103+
: qkv_bias_out->mutable_data<T>(ctx.GetPlace());
102104

103105
// get data ptr for FMHA.
104106
auto *transpose_out_2_data =
@@ -117,7 +119,8 @@ class FusedAttentionOpKernel : public framework::OpKernel<T> {
117119

118120
// get data ptr for out_linear.
119121
auto *out_linear_weight_data = out_linear_weight->data<T>();
120-
auto *out_linear_bias_data = out_linear_bias->data<T>();
122+
auto *out_linear_bias_data =
123+
(out_linear_bias == nullptr) ? nullptr : out_linear_bias->data<T>();
121124
auto *out_linear_out_data = out_linear_out->mutable_data<T>(ctx.GetPlace());
122125

123126
// get data ptr for bias+dropout+residual+layernorm
@@ -139,9 +142,15 @@ class FusedAttentionOpKernel : public framework::OpKernel<T> {
139142

140143
auto layer_norm_compute = AttnLayerNorm<T>(ctx.cuda_device_context(),
141144
epsilon, bsz_seq, dim_embed);
145+
146+
bool compute_bias = true;
147+
if (qkv_bias == nullptr) {
148+
compute_bias = false;
149+
}
142150
// (transA, transB, compute_bias) = (false, true, true)
143-
auto qkv_compute = AttnMatMul<T>(ctx.cuda_device_context(), false, true,
144-
bsz_seq, output_size, input_size, true);
151+
auto qkv_compute =
152+
AttnMatMul<T>(ctx.cuda_device_context(), false, true, bsz_seq,
153+
output_size, input_size, compute_bias);
145154

146155
AttnDropoutParam attn_dropout_param(
147156
is_test_1, dropout_implementation_1, attn_dropout_rate,
@@ -176,10 +185,17 @@ class FusedAttentionOpKernel : public framework::OpKernel<T> {
176185
qkv_compute.ComputeForward(qkv_weight, input_x, qkv_bias, qkv_out,
177186
qkv_bias_out);
178187
}
179-
fmha_ref_compute.ComputeForward(*qkv_bias_out, src_mask, transpose_out_2,
180-
qk_out, src_mask_out, softmax_out,
181-
attn_dropout_mask_out, attn_dropout_out,
182-
qktv_out, fmha_out);
188+
if (qkv_bias == nullptr) {
189+
fmha_ref_compute.ComputeForward(*qkv_out, src_mask, transpose_out_2,
190+
qk_out, src_mask_out, softmax_out,
191+
attn_dropout_mask_out, attn_dropout_out,
192+
qktv_out, fmha_out);
193+
} else {
194+
fmha_ref_compute.ComputeForward(*qkv_bias_out, src_mask, transpose_out_2,
195+
qk_out, src_mask_out, softmax_out,
196+
attn_dropout_mask_out, attn_dropout_out,
197+
qktv_out, fmha_out);
198+
}
183199

184200
// fmha_out: [batch_size, seq_len, num_head, head_dim]
185201
// weight: [embed_dim, embed_dim]
@@ -249,9 +265,10 @@ class FusedAttentionGradKernel : public framework::OpKernel<T> {
249265
auto *out_linear_bias = ctx.Input<Tensor>("OutLinearBias");
250266
auto *src_mask_data = (src_mask == nullptr ? nullptr : src_mask->data<T>());
251267
auto *qkv_weight_data = qkv_weight->data<T>();
252-
auto *qkv_bias_data = qkv_bias->data<T>();
268+
auto *qkv_bias_data = (qkv_bias == nullptr) ? nullptr : qkv_bias->data<T>();
253269
auto *out_linear_weight_data = out_linear_weight->data<T>();
254-
auto *out_linear_bias_data = out_linear_bias->data<T>();
270+
auto *out_linear_bias_data =
271+
(out_linear_bias == nullptr) ? nullptr : out_linear_bias->data<T>();
255272

256273
// fw output
257274
auto *fmha_out = ctx.Input<Tensor>("FMHAOut");
@@ -299,8 +316,15 @@ class FusedAttentionGradKernel : public framework::OpKernel<T> {
299316
auto *d_bias_dropout_residual_out =
300317
ctx.Output<Tensor>(framework::GradVarName("BiasDropoutResidualOut"));
301318
auto *d_x_data = d_x->mutable_data<T>(ctx.GetPlace());
302-
auto *d_qkv_out_data = d_qkv_out->mutable_data<T>(ctx.GetPlace());
303-
auto *d_qkv_bias_out_data = d_qkv_bias_out->mutable_data<T>(ctx.GetPlace());
319+
// when qkv_bias is not nullptr, d_qkv_out is equals to d_qkv_bias_out, the
320+
// space can be reused.
321+
auto *d_qkv_out_data = (d_qkv_bias_out != nullptr)
322+
? nullptr
323+
: d_qkv_out->mutable_data<T>(ctx.GetPlace());
324+
auto *d_qkv_bias_out_data =
325+
(d_qkv_bias_out == nullptr)
326+
? nullptr
327+
: d_qkv_bias_out->mutable_data<T>(ctx.GetPlace());
304328
auto *d_qktv_out_data = d_qktv_out->mutable_data<T>(ctx.GetPlace());
305329
auto *d_transpose_out_2_data =
306330
d_transpose_out_2->mutable_data<T>(ctx.GetPlace());
@@ -326,11 +350,15 @@ class FusedAttentionGradKernel : public framework::OpKernel<T> {
326350
auto *d_ln_2_bias = ctx.Output<Tensor>(framework::GradVarName("Ln2Bias"));
327351

328352
auto *d_qkv_weight_data = d_qkv_weight->mutable_data<T>(ctx.GetPlace());
329-
auto *d_qkv_bias_data = d_qkv_bias->mutable_data<T>(ctx.GetPlace());
353+
auto *d_qkv_bias_data = (d_qkv_bias == nullptr)
354+
? nullptr
355+
: d_qkv_bias->mutable_data<T>(ctx.GetPlace());
330356
auto *d_out_linear_weight_data =
331357
d_out_linear_weight->mutable_data<T>(ctx.GetPlace());
332358
auto *d_out_linear_bias_data =
333-
d_out_linear_bias->mutable_data<T>(ctx.GetPlace());
359+
(d_out_linear_bias == nullptr)
360+
? nullptr
361+
: d_out_linear_bias->mutable_data<T>(ctx.GetPlace());
334362

335363
const auto input_x_dims = input_x->dims();
336364
const auto qkv_w_dims = qkv_weight->dims();
@@ -352,12 +380,15 @@ class FusedAttentionGradKernel : public framework::OpKernel<T> {
352380

353381
bool transA = false;
354382
bool transB = true;
355-
bool compute_bias = true;
383+
bool compute_qkv_bias = true;
384+
if (qkv_bias == nullptr) {
385+
compute_qkv_bias = false;
386+
}
356387
auto layer_norm_compute = AttnLayerNorm<T>(ctx.cuda_device_context(),
357388
epsilon, bsz_seq, dim_embed);
358389
auto qkv_compute =
359390
AttnMatMul<T>(ctx.cuda_device_context(), transA, transB, bsz_seq,
360-
output_size, input_size, compute_bias);
391+
output_size, input_size, compute_qkv_bias);
361392
AttnDropoutParam attn_dropout_param(
362393
is_test_1, dropout_implementation_1, attn_dropout_prob,
363394
is_upscale_in_train_1, is_fix_seed_1, seed_val_1, seed_1);
@@ -367,7 +398,7 @@ class FusedAttentionGradKernel : public framework::OpKernel<T> {
367398
output_size = hidden_size;
368399
transA = false;
369400
transB = false;
370-
compute_bias = false;
401+
bool compute_bias = false;
371402
auto out_linear_compute =
372403
AttnMatMul<T>(ctx.cuda_device_context(), transA, transB, bsz_seq,
373404
output_size, input_size, compute_bias);
@@ -405,14 +436,19 @@ class FusedAttentionGradKernel : public framework::OpKernel<T> {
405436
d_out_linear_out, d_fmha_out,
406437
d_out_linear_weight, nullptr);
407438

408-
fmha_ref_compute.ComputeBackward(
409-
*transpose_out_2, src_mask, *softmax_out, *attn_dropout_mask_out,
410-
*attn_dropout_out, *qk_out, *src_mask_out, *d_fmha_out, d_qktv_out,
411-
d_attn_dropout_out, d_softmax_out, d_src_mask_out, d_qk_out,
412-
d_transpose_out_2, nullptr, d_qkv_bias_out);
413-
cudaMemcpyAsync(d_qkv_out_data, d_qkv_bias_out_data,
414-
bsz_seq * 3 * num_head * dim_head * sizeof(T),
415-
cudaMemcpyDeviceToDevice);
439+
if (qkv_bias != nullptr) {
440+
fmha_ref_compute.ComputeBackward(
441+
*transpose_out_2, src_mask, *softmax_out, *attn_dropout_mask_out,
442+
*attn_dropout_out, *qk_out, *src_mask_out, *d_fmha_out, d_qktv_out,
443+
d_attn_dropout_out, d_softmax_out, d_src_mask_out, d_qk_out,
444+
d_transpose_out_2, nullptr, d_qkv_bias_out);
445+
} else {
446+
fmha_ref_compute.ComputeBackward(
447+
*transpose_out_2, src_mask, *softmax_out, *attn_dropout_mask_out,
448+
*attn_dropout_out, *qk_out, *src_mask_out, *d_fmha_out, d_qktv_out,
449+
d_attn_dropout_out, d_softmax_out, d_src_mask_out, d_qk_out,
450+
d_transpose_out_2, nullptr, d_qkv_out);
451+
}
416452

417453
if (pre_layer_norm) {
418454
auto *ln_mean = ctx.Input<Tensor>("LnMean");
@@ -432,15 +468,24 @@ class FusedAttentionGradKernel : public framework::OpKernel<T> {
432468
auto *d_ln_bias_data =
433469
(d_ln_bias == nullptr ? nullptr
434470
: d_ln_bias->mutable_data<U>(ctx.GetPlace()));
435-
436-
qkv_compute.ComputeBackward(ln_out, qkv_weight, d_qkv_bias_out, d_ln_out,
437-
d_qkv_weight, d_qkv_bias);
471+
if (qkv_bias != nullptr) {
472+
qkv_compute.ComputeBackward(ln_out, qkv_weight, d_qkv_bias_out,
473+
d_ln_out, d_qkv_weight, d_qkv_bias);
474+
} else {
475+
qkv_compute.ComputeBackward(ln_out, qkv_weight, d_qkv_out, d_ln_out,
476+
d_qkv_weight, d_qkv_bias);
477+
}
438478
layer_norm_compute.ComputeBackward(x_data, d_ln_out_data, ln_scale_data,
439479
ln_mean_data, ln_var_data, d_x_data,
440480
d_ln_scale_data, d_ln_bias_data);
441481
} else {
442-
qkv_compute.ComputeBackward(input_x, qkv_weight, d_qkv_bias_out, d_x,
443-
d_qkv_weight, d_qkv_bias);
482+
if (qkv_bias != nullptr) {
483+
qkv_compute.ComputeBackward(input_x, qkv_weight, d_qkv_bias_out, d_x,
484+
d_qkv_weight, d_qkv_bias);
485+
} else {
486+
qkv_compute.ComputeBackward(input_x, qkv_weight, d_qkv_out, d_x,
487+
d_qkv_weight, d_qkv_bias);
488+
}
444489
}
445490
// gradient accumulation
446491
std::vector<const Tensor *> ins;

0 commit comments

Comments
 (0)