geohot/tinygrad
1
0
Fork 0
forked from tinygrad/tinygrad
Code Pull Requests Actions Activity

test scaled dot product attention (#3063)

Browse Source 
* add test

* add initial test for scaled dot product attention

* test pass for scaled dot product attention
This commit is contained in:
Yixiang Gao
2024-01-09 14:30:57 -08:00
committed by GitHub
parent 55ac2a2cf7
commit 73b72b8de2
1 changed files with 23 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/test_multitensor.py
+23
View File
@@ -174,6 +174,29 @@ class TestMultiTensor(unittest.TestCase):
y_shard = layer_norm_sharded(x_sharded).realize()
np.testing.assert_allclose(y.numpy(), y_shard.numpy(), atol=1e-6, rtol=1e-6)
def test_scaled_product_attention(self):
q = Tensor.rand(32, 8, 16, 64).contiguous().realize()
k = Tensor.rand(32, 8, 16, 64).contiguous().realize()
v = Tensor.rand(32, 8, 16, 64).contiguous().realize()
y = Tensor.scaled_dot_product_attention(q, k, v)
q_sharded = q.shard((d0, d1), axis=None).realize()
k_sharded = k.shard((d0, d1), axis=1).realize()
v_sharded = v.shard((d0, d1), axis=1).realize()
y_sharded = Tensor.scaled_dot_product_attention(q_sharded, k_sharded, v_sharded)
np.testing.assert_allclose(y.numpy(), y_sharded.numpy(), atol=1e-6, rtol=1e-6)
m = Tensor.rand(32, 8, 16, 16).contiguous().realize()
y = Tensor.scaled_dot_product_attention(q, k, v, attn_mask=m)
m_sharded = m.shard((d0, d1), axis=None).realize()
y_sharded = Tensor.scaled_dot_product_attention(q_sharded, k_sharded, v_sharded, attn_mask=m_sharded)
np.testing.assert_allclose(y.numpy(), y_sharded.numpy(), atol=1e-6, rtol=1e-6)
y = Tensor.scaled_dot_product_attention(q, k, v, is_causal=True)
y_sharded = Tensor.scaled_dot_product_attention(q_sharded, k_sharded, v_sharded, is_causal=True)
np.testing.assert_allclose(y.numpy(), y_sharded.numpy(), atol=1e-6, rtol=1e-6)
def test_data_parallel_resnet(self):
import sys, pathlib