def test_small_forward(self):
N = 12
H = 8
L = 2000
S = 2000
E = 32
k = 32
C = 100
I = 10
B = 32
Q = torch.randn(N, H, L, E).to(self.device)
K = torch.randn(N, H, S, E).to(self.device)
lengths = torch.full((N, ), L, dtype=torch.int32).to(self.device)
groups, counts = cluster_queries(Q, lengths, C, I, B)
sorted_g, sorted_gi = torch.sort(groups.view(N * H, -1), dim=-1)
sorted_rev_gi = torch.argsort(sorted_gi, dim=-1)
q_offset = torch.arange(N * H, device=Q.device).unsqueeze(-1) * L
q_flat = (sorted_gi + q_offset).reshape(-1)
s_queries = Q.reshape(-1, E).index_select(0, q_flat).view(N, H, L, E)
Q_grouped = aggregate(s_queries, sorted_g.view(N, H, L),
1 / counts.float())
QK = torch.einsum("nhle,nhse->nhls", Q_grouped, K)
_, topk = torch.topk(QK, k, dim=-1)
topk = topk.contiguous()
topk_broadcast = broadcast(
topk.float(), groups, torch.ones_like(counts, dtype=torch.float32),
torch.zeros((N, H, L, k), device=Q.device))
weights_sorted = clustered_sparse_dot_product(s_queries, K, topk,
groups, counts, lengths)
q_rev_flat = (sorted_rev_gi + q_offset).reshape(-1)
weights = weights_sorted.reshape(-1, k).index_select(0,
q_rev_flat).view(
N, H, L, k)
values = torch.randn(N, H, S, E).to(self.device)
for i in range(2000):
output_hat = clustered_sparse_weighted_average(
weights_sorted, values, topk, groups, counts)
s = torch.cuda.Event(enable_timing=True)
e = torch.cuda.Event(enable_timing=True)
s.record()
output_hat = clustered_sparse_weighted_average(weights, values, topk,
groups, counts)
e.record()
torch.cuda.synchronize()
t_sparse = s.elapsed_time(e)
print('T_sparse Forward:{}'.format(t_sparse))
def test_forward(self):
N = 5
H = 2
L = 100
S = 100
E = 32
C = 10
I = 10
B = 32
k = 5
Q = torch.randn(N, H, L, E).to(self.device)
K = torch.randn(N, H, S, E).to(self.device)
lengths = torch.full((N, ), L, dtype=torch.int32).to(self.device)
groups, counts = cluster_queries(Q, lengths, C, I, B)
Q_grouped = aggregate(Q, groups, 1 / counts.float())
QK = torch.einsum("nhle,nhse->nhls", Q_grouped, K)
_, topk = torch.topk(QK, k, dim=-1)
topk = topk.contiguous()
topk_broadcast = broadcast(
topk.float(), groups, torch.ones_like(counts, dtype=torch.float32),
torch.zeros((N, H, L, k), device=Q.device))
weights = torch.rand(N, H, L, k).to(self.device).requires_grad_(True)
values = torch.randn(N, H, S, E).to(self.device).requires_grad_(True)
values_selected = values[
torch.arange(N).view(N, 1, 1, 1).to(self.device),
torch.arange(H).view(1, H, 1, 1).to(self.device),
topk_broadcast.long()]
output = (weights.unsqueeze(-1) * values_selected).sum(-2)
output_hat = clustered_sparse_weighted_average(weights, values, topk,
groups)
self.assertLess(torch.abs(output - output_hat).max(), 1e-4)
def test_small_forward_backward(self):
N = 12
H = 8
L = 2000
S = 2000
E = 32
k = 32
C = 100
I = 10
B = 32
Q = torch.randn(N, H, L, E).to(self.device)
K = torch.randn(N, H, S, E).to(self.device)
lengths = torch.full((N, ), L, dtype=torch.int32).to(self.device)
groups, counts = cluster_queries(Q, lengths, C, I, B)
Q_grouped = aggregate(Q, groups, 1 / counts.float())
QK = torch.einsum("nhle,nhse->nhls", Q_grouped, K)
_, topk = torch.topk(QK, k, dim=-1)
topk = topk.contiguous()
topk_broadcast = broadcast(
topk.float(), groups, torch.ones_like(counts, dtype=torch.float32),
torch.zeros((N, H, L, k), device=Q.device))
weights = torch.rand(N, H, L, k).to(self.device).requires_grad_(True)
values = torch.randn(N, H, S, E).to(self.device).requires_grad_(True)
self._zero_grad(weights, values)
n_runs = 20
s = time.time()
for i in range(n_runs):
output_hat = clustered_sparse_weighted_average(
weights, values, topk, groups, counts)
output_hat.sum().backward()
e = time.time()
t_sparse = (e - s) / n_runs
print('T_sparse Forward Backward:{}'.format(t_sparse))
def test_small_forward(self):
N = 12
H = 8
L = 2000
S = 2000
E = 32
k = 32
C = 100
I = 10
B = 32
Q = torch.randn(N, H, L, E).to(self.device)
K = torch.randn(N, H, S, E).to(self.device)
lengths = torch.full((N, ), L, dtype=torch.int32).to(self.device)
groups, counts = cluster_queries(Q, lengths, C, I, B)
Q_grouped = aggregate(Q, groups, 1 / counts.float())
QK = torch.einsum("nhle,nhse->nhls", Q_grouped, K)
_, topk = torch.topk(QK, k, dim=-1)
topk = topk.contiguous()
topk_broadcast = broadcast(
topk.float(), groups, torch.ones_like(counts, dtype=torch.float32),
torch.zeros((N, H, L, k), device=Q.device))
weights = torch.rand(N, H, L, k).to(self.device).requires_grad_(True)
values = torch.randn(N, H, S, E).to(self.device).requires_grad_(True)
for i in range(2000):
output_hat = clustered_sparse_weighted_average(
weights, values, topk, groups)
s = torch.cuda.Event(enable_timing=True)
e = torch.cuda.Event(enable_timing=True)
s.record()
output_hat = clustered_sparse_weighted_average(weights, values, topk,
groups)
e.record()
torch.cuda.synchronize()
t_sparse = s.elapsed_time(e)
print('T_sparse Forward:{}'.format(t_sparse))
def test_correctness_masked(self):
N = 12
H = 6
L = 1000
S = 1000
E = 32
k = 32
C = 100
I = 10
B = 32
for exp in range(30):
N = np.random.randint(1, 6)
H = np.random.randint(1, 8)
C = np.random.randint(10, 500)
L = np.random.randint(C, 2000)
E = np.random.randint(10, 128)
S = np.random.randint(100, 1000)
k = np.random.randint(10, 64)
if os.getenv("VERBOSE_TESTS", ""):
print(("Testing Masked: N H L S E C k: "
"{} {} {} {} {} {} {}").format(N, H, L, S, E, C, k))
Q = torch.randn(N, H, L, E).to(self.device)
K = torch.randn(N, H, S, E).to(self.device)
lengths = np.random.randint(C, L + 1, N)
lengths = torch.tensor(lengths, dtype=torch.int32).to(self.device)
lengths[0] = L
query_lengths = LengthMask(lengths, max_len=L)
groups, counts = cluster_queries(Q, lengths, C, I, B)
sorted_g, sorted_gi = torch.sort(groups.view(N * H, -1), dim=-1)
sorted_rev_gi = torch.argsort(sorted_gi, dim=-1)
q_offset = torch.arange(N * H, device=Q.device).unsqueeze(-1) * L
q_flat = (sorted_gi + q_offset).reshape(-1)
s_queries = Q.reshape(-1, E).index_select(0,
q_flat).view(N, H, L, E)
Q_grouped = aggregate(s_queries, sorted_g.view(N, H, L),
1 / counts.float())
QK = torch.einsum("nhle,nhse->nhls", Q_grouped, K)
_, topk = torch.topk(QK, k, dim=-1)
topk = topk.contiguous()
topk_broadcast = broadcast(
topk.float(), groups,
torch.ones_like(counts, dtype=torch.float32),
torch.zeros((N, H, L, k), device=Q.device))
weights_sorted = torch.rand(N, H, L,
k).to(self.device).requires_grad_(True)
weights_sorted.retain_grad()
q_rev_flat = (sorted_rev_gi + q_offset).reshape(-1)
weights = torch.clone(
weights_sorted.reshape(-1, k).index_select(0, q_rev_flat).view(
N, H, L, k))
weights.retain_grad()
values = torch.randn(N, H, S,
E).to(self.device).requires_grad_(True)
self._zero_grad(weights, values)
values_selected = values[
torch.arange(N).view(N, 1, 1, 1).to(self.device),
torch.arange(H).view(1, H, 1, 1).to(self.device),
topk_broadcast.long()]
output = (weights.unsqueeze(-1) * values_selected).sum(-2)
output = output * query_lengths.float_matrix[:, None, :, None]
output.sum().backward()
grad = [torch.clone(weights.grad), torch.clone(values.grad)]
self._zero_grad(weights_sorted, values)
self._zero_grad(weights, values)
output_hat_sorted = clustered_sparse_weighted_average(
weights_sorted, values, topk, groups, counts)
output_hat = output_hat_sorted.reshape(-1, E).index_select(
0, q_rev_flat).view(N, H, L, E)
self.assertLess(torch.abs(output - output_hat).max(), 1e-4)
output_hat.sum().backward()
weights_grad_sorted = torch.clone(weights_sorted.grad)
weights_grad = torch.clone(
weights_grad_sorted.reshape(-1, k).index_select(
0, q_rev_flat).view(N, H, L, k))
grad_hat = [weights_grad, torch.clone(values.grad)]
for g1, g2 in zip(grad, grad_hat):
self.assertLess(torch.abs(g1 - g2).max(), 1e-3)
def test_forward(self):
N = 6
H = 5
L = 100
S = 100
E = 32
C = 10
I = 10
B = 32
k = 5
for exp in range(30):
C = np.random.randint(10, 500)
L = np.random.randint(C, 2000)
E = np.random.randint(10, 128)
S = np.random.randint(100, 1000)
k = np.random.randint(10, 64)
if os.getenv("VERBOSE_TESTS", ""):
print(("Testing: N H L S E C k: "
"{} {} {} {} {} {} {}").format(N, H, L, S, E, C, k))
Q = torch.randn(N, H, L, E).to(self.device)
K = torch.randn(N, H, S, E).to(self.device)
lengths = torch.full((N, ), L, dtype=torch.int32).to(self.device)
groups, counts = cluster_queries(Q, lengths, C, I, B)
sorted_g, sorted_gi = torch.sort(groups.view(N * H, -1), dim=-1)
sorted_rev_gi = torch.argsort(sorted_gi, dim=-1)
q_offset = torch.arange(N * H, device=Q.device).unsqueeze(-1) * L
q_flat = (sorted_gi + q_offset).reshape(-1)
s_queries = Q.reshape(-1, E).index_select(0,
q_flat).view(N, H, L, E)
Q_grouped = aggregate(s_queries, sorted_g.view(N, H, L),
1 / counts.float())
QK = torch.einsum("nhle,nhse->nhls", Q_grouped, K)
_, topk = torch.topk(QK, k, dim=-1)
topk = topk.contiguous()
topk_broadcast = broadcast(
topk.float(), groups,
torch.ones_like(counts, dtype=torch.float32),
torch.zeros((N, H, L, k), device=Q.device))
weights_sorted = clustered_sparse_dot_product(
s_queries, K, topk, groups, counts, lengths)
weights = torch.softmax(weights_sorted, dim=-1)
q_rev_flat = (sorted_rev_gi + q_offset).reshape(-1)
weights = weights_sorted.reshape(-1, k).index_select(
0, q_rev_flat).view(N, H, L, k)
values = torch.randn(N, H, S,
E).to(self.device).requires_grad_(True)
values_selected = values[
torch.arange(N).view(N, 1, 1, 1).to(self.device),
torch.arange(H).view(1, H, 1, 1).to(self.device),
topk_broadcast.long()]
output = (weights.unsqueeze(-1) * values_selected).sum(-2)
output_hat_sorted = clustered_sparse_weighted_average(
weights_sorted, values, topk, groups, counts)
output_hat = output_hat_sorted.reshape(-1, E).index_select(
0, q_rev_flat).view(N, H, L, E)
self.assertLess(torch.abs(output_hat - output).max(), 1e-3)
