Флэш-внимание дает разные результаты для токенов с одинаковыми вложениями? ⇐ Python

[Anonymous]

Я учусь интегрировать Flash Attention в свою модель, чтобы ускорить обучение. Я тестирую функцию, чтобы определить лучший способ ее реализации. Однако я столкнулся с проблемой, из-за которой Flash Attention выдает разные результаты для токенов с идентичными встраиваниями. Я не уверен, совершаю ли я принципиальную ошибку или здесь есть что-то еще.
Вот фрагмент кода, который я использую:

Код: Выделить всё

import torch
from flash_attn.modules.mha import FlashSelfAttention

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
fa_attn = FlashSelfAttention(deterministic=True)
fa_attn.eval()

# Assuming batch_size, seq_len, heads, dim = 1, 4, 1, 4
x = torch.tensor([[0.1, 0.1, 0.1, 0.1],
[0.1, 0.1, 0.1, 0.1],
[0.1, 0.1, 0.1, 0.1],
[0.1, 0.1, 0.1, 0.1]])
q = x.unsqueeze(0).unsqueeze(2)
k = q.clone()
v = q.clone()
qkv = torch.stack([q, k, v], dim=2).half().to(device)
output = fa_attn(qkv)
print(output)

результат:

Код: Выделить всё

tensor([[[[0.1000, 0.1000, 0.1000, 0.1000]],
[[0.0757, 0.0757, 0.0757, 0.0757]],
[[0.1000, 0.1000, 0.1000, 0.1000]],
[[0.0757, 0.0757, 0.0757, 0.0757]]]], device='cuda:0', dtype=torch.float16)

Еще один

Код: Выделить всё

x = torch.tensor([[0.1, 0.1, 0.1, 0.1], [0.1, 0.1, 0.1, 0.1], [0.1, 0.1, 0.1, 0.1]])
q = x.unsqueeze(0).unsqueeze(2)
k = q.clone()
v = q.clone()
qkv = torch.stack([q, k, v], dim=2).half().to(device)
output = fa_attn(qkv)
output

результат:

Код: Выделить всё

tensor([[[[ 0.1000,  0.1000,  0.1000,  0.1000]],

[[-0.5483,  0.5166, -0.5483,  0.5166]],

[[ 0.1000,  0.1000,  0.1000,  0.1000]]]], device='cuda:0',
dtype=torch.float16)

Большое спасибо.

Подробнее здесь: https://stackoverflow.com/questions/791 ... embeddings