From 73b72b8de2fa9b63023c1e67c19c93a7f24b480c Mon Sep 17 00:00:00 2001 From: Yixiang Gao Date: Tue, 9 Jan 2024 14:30:57 -0800 Subject: [PATCH] test scaled dot product attention (#3063) * add test * add initial test for scaled dot product attention * test pass for scaled dot product attention --- test/test_multitensor.py | 23 +++++++++++++++++++++++ 1 file changed, 23 insertions(+) diff --git a/test/test_multitensor.py b/test/test_multitensor.py index 78818ec144..2ec42dfd7e 100644 --- a/test/test_multitensor.py +++ b/test/test_multitensor.py @@ -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