File size: 3,321 Bytes
c335050 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 |
import os
import pytest
import torch
import torch.nn.functional as F
from fla.ops.common.chunk_scaled_dot_kkt import chunk_scaled_dot_kkt_fwd
from fla.ops.utils.solve_tril import solve_tril
from fla.utils import assert_close, device, device_platform
@pytest.mark.parametrize(
('B', 'T', 'H', 'chunk_size'),
[
pytest.param(*test, id="B{}-T{}-H{}-chunk_size{}".format(*test))
for test in [
(1, 63, 1, 16),
(2, 500, 4, 32),
(2, 1000, 5, 64),
(3, 1024, 6, 64),
(4, 2048, 8, 64),
]
],
)
@pytest.mark.skipif(
device_platform == 'intel',
reason='Intel Pytorch Failure',
)
def test_solve_tril(B, T, H, chunk_size):
# do not randomly intiialize A otherwise the inverse is not stable
k = F.normalize(torch.randn((B, H, T, 64), dtype=torch.float32, device=device), dim=-1)
# Pad the second-to-last dimension (T) to be a multiple of chunk_size
padding_size = (chunk_size - T % chunk_size) % chunk_size
k_padded = F.pad(k, (0, 0, 0, padding_size, 0, 0, 0, 0))
k_padded = k_padded.reshape(B, H, -1, chunk_size, 64)
A = (k_padded @ k_padded.transpose(-1, -2)).tril(-1)
ref = torch.inverse(A + torch.eye(A.shape[-1], device=A.device)[None, None, None, ...])
ref = ref.reshape(B, H, -1, chunk_size)[:, :, :T, :]
tri = solve_tril(A.reshape(B, H, -1, chunk_size)[:, :, :T, :].transpose(1, 2)).transpose(1, 2)
assert_close('solve_tril', ref, tri, 0.0001)
@pytest.mark.parametrize(
('H', 'D', 'chunk_size', 'cu_seqlens'),
[
pytest.param(*test, id="H{}-D{}-chunk_size{}-cu_seqlens{}".format(*test))
for test in [
(4, 64, 16, [0, 15]),
(4, 64, 32, [0, 256, 500, 1000]),
(4, 100, 64, [0, 15, 100, 300, 1200, 2000]),
(4, 64, 16, [0, 1, 100, 300, 1200, 2048]),
(4, 128, 32, [0, 200, 512, 1200, 2048]),
]
],
)
@pytest.mark.skipif(
os.getenv('SKIP_TEST_CHUNK_VARLEN') == '1',
reason='Skipping test_chunk_varlen because SKIP_TEST_CHUNK_VARLEN is set',
)
@pytest.mark.skipif(
device_platform == 'intel',
reason='Intel Pytorch Failure',
)
def test_solve_tril_varlen(
H: int,
D: int,
chunk_size: int,
cu_seqlens: list[int],
):
T = cu_seqlens[-1]
cu_seqlens = torch.tensor(cu_seqlens, dtype=torch.int32, device=device)
# Construct the input. otherwise inverse's condition number might be too large to measure the error
k = F.normalize(torch.randn((1, T, H, D), dtype=torch.bfloat16, device=device), dim=-1)
beta = torch.randn((1, T, H), dtype=torch.bfloat16, device=device).sigmoid()
A = chunk_scaled_dot_kkt_fwd(k=k, beta=beta, cu_seqlens=cu_seqlens, chunk_size=chunk_size)
ref = torch.zeros_like(A)
for i in range(len(cu_seqlens) - 1):
for j in range(cu_seqlens[i], cu_seqlens[i+1], chunk_size):
actual_size = min(chunk_size, cu_seqlens[i+1] - j)
ref[:, j:j+actual_size, :, :actual_size] = torch.inverse(
A[:, j:j+actual_size, :, :actual_size].transpose(1, 2) +
torch.eye(actual_size, device=A.device, dtype=A.dtype)[None, None, ...],
).transpose(1, 2)
tri = solve_tril(A, cu_seqlens=cu_seqlens)
assert_close('solve_tril_varlen', ref, tri, 0.0001)
|