Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

add num_splist for flash_attn_bwd and FlashAttnUnpaddedGradKernel #56363

Merged
merged 10 commits into from
Sep 4, 2023
Merged

add num_splist for flash_attn_bwd and FlashAttnUnpaddedGradKernel #56363

Show file tree
Hide file tree
Changes from 4 commits
Commits
Show all changes
10 commits
Select commit Hold shift + click to select a range
b742184
add num_splist for flash_attn_bwd and FlashAttnUnpaddedGradKernel
AnnaTrainingG Aug 16, 2023
5a5c307
update
AnnaTrainingG Aug 16, 2023
7d9ce2f
update
AnnaTrainingG Aug 17, 2023
814b9c5
Add assertTrue
AnnaTrainingG Aug 21, 2023
ded6af8
Update submodule to a specific commit
AnnaTrainingG Aug 30, 2023
063f8f3
update test
AnnaTrainingG Aug 31, 2023
51817ec
update
AnnaTrainingG Aug 31, 2023
7dd1f71
update
AnnaTrainingG Sep 1, 2023
75ff6f8
update
AnnaTrainingG Sep 1, 2023
8b585ab
update
AnnaTrainingG Sep 1, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
16 changes: 11 additions & 5 deletions paddle/phi/kernels/gpu/flash_attn_grad_kernel.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -236,11 +236,10 @@ void FlashAttnUnpaddedGradKernel(const Context& ctx,
const int64_t total_k = k.dims()[0];
const int64_t num_heads_k = k.dims()[1];

// TODO(umiswing): add deterministic in fa2.
// int num_splits = 0; // 0 for an internal heuristic, which is optimal
// if (FLAGS_cudnn_deterministic) {
// num_splits = 1;
// }
int num_splits = 0; // 0 for an internal heuristic, which is optimal
if (FLAGS_cudnn_deterministic) {
num_splits = 1;
}
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

封装成一个函数

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

done


// TODO(umiswing): add shape check
PADDLE_ENFORCE_EQ(
Expand Down Expand Up @@ -294,6 +293,7 @@ void FlashAttnUnpaddedGradKernel(const Context& ctx,
params.scale,
params.causal,
params.is_bf16,
num_splits,
stream,
params.seed,
params.offset);
Expand Down Expand Up @@ -401,6 +401,11 @@ void FlashAttnGradKernel(const Context& ctx,
VLOG(10) << "FlashAttn bwd seed: " << params.seed
<< ", offset: " << params.offset;

int num_splits = 0; // 0 for an internal heuristic, which is optimal
if (FLAGS_cudnn_deterministic) {
num_splits = 1;
}

bool succ = phi::dynload::flash_attn_bwd(dout.data(),
q.data(),
k.data(),
Expand All @@ -426,6 +431,7 @@ void FlashAttnGradKernel(const Context& ctx,
params.scale,
params.causal,
params.is_bf16,
num_splits,
stream,
params.seed,
params.offset);
Expand Down
29 changes: 24 additions & 5 deletions test/legacy_test/test_flash_attention.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -191,17 +191,13 @@ def test_unpadded(self):
fetches_result[0], out_, rtol=5e-03, atol=1e-03
)

def test_all(self):
def flash_attn_compute(self, query, key, value):
print(
f"Test case shape {self.shape} dtype {self.dtype} causal {self.causal}"
)
# test dynamic
paddle.disable_static()

query = np.random.random(self.shape)
key = np.random.random(self.shape)
value = np.random.random(self.shape)

q = paddle.to_tensor(
query, place=self.place, dtype=self.dtype, stop_gradient=False
)
Expand Down Expand Up @@ -306,6 +302,29 @@ def test_all(self):
np.testing.assert_allclose(
fetches_result[0], out_, rtol=5e-03, atol=1e-03
)
return out, out_, fetches_result[0]
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

fetches_result[0]是静态图执行的前向输出吧，确定性实现可以只比较测试动态图，但是需要检查前向out和反向的dq、dk、dv，保证两次执行的结果完全一样。

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

done


def test_all(self):
query = np.random.random(self.shape)
key = np.random.random(self.shape)
value = np.random.random(self.shape)
out, out_, _ = self.flash_attn_compute(query, key, value)

def test_all_flag(self):
paddle.set_flags({'FLAGS_cudnn_deterministic': 1})
query = np.random.random(self.shape)
key = np.random.random(self.shape)
value = np.random.random(self.shape)

out1, out1_, fetches_result1 = self.flash_attn_compute(
query, key, value
)
out2, out2_, fetches_result2 = self.flash_attn_compute(
query, key, value
)
self.assertTrue(np.equal(out1.numpy(), out2.numpy()).all())
self.assertTrue(np.equal(fetches_result1, fetches_result2).all())
paddle.set_flags({'FLAGS_cudnn_deterministic': 0})


@unittest.skipIf(
Expand Down