This is an automated email from the ASF dual-hosted git repository.
masahi pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/main by this push:
new 97ddd667c8 [Relay][Pytorch] Add support for
`aten::scaled_dot_product_attention` (#16143)
97ddd667c8 is described below
commit 97ddd667c821cad1d034f9411d1c4aa1cda225bb
Author: Masahiro Hiramori <[email protected]>
AuthorDate: Wed Nov 29 12:50:45 2023 -0800
[Relay][Pytorch] Add support for `aten::scaled_dot_product_attention`
(#16143)
* add support for `aten::scaled_dot_product_attention`
* convert `_expr.Constant` into `int` in the unflattened_size
* enable `test_transformer`
* explicit scale not support with `torch==2.0.0`
* fix error on gpu
---
python/tvm/relay/frontend/pytorch.py | 116 +++++++++++++++++++++++++-
tests/python/frontend/pytorch/test_forward.py | 84 ++++++++++++++++++-
2 files changed, 196 insertions(+), 4 deletions(-)
diff --git a/python/tvm/relay/frontend/pytorch.py
b/python/tvm/relay/frontend/pytorch.py
index 30711de0a7..18e628d565 100644
--- a/python/tvm/relay/frontend/pytorch.py
+++ b/python/tvm/relay/frontend/pytorch.py
@@ -1555,15 +1555,21 @@ class PyTorchOpConverter:
dim = dim if dim >= 0 else len(dshape) + dim
assert len(dshape) > dim >= 0
- assert unflattened_size.count(-1) <= 1
+ new_unflattened_size = []
+ for s in unflattened_size:
+ if isinstance(s, _expr.Constant):
+ s = s.data.numpy().item()
+ new_unflattened_size.append(s)
+
+ assert new_unflattened_size.count(-1) <= 1
- mult = np.multiply.reduce(unflattened_size)
+ mult = np.multiply.reduce(new_unflattened_size)
if mult < 0:
assert dshape[dim] % mult == 0
else:
assert dshape[dim] == mult
- new_shape = dshape[:dim] + unflattened_size + dshape[dim + 1 :]
+ new_shape = dshape[:dim] + tuple(new_unflattened_size) + dshape[dim +
1 :]
out = _op.reshape(data, new_shape)
return out
@@ -3891,6 +3897,109 @@ class PyTorchOpConverter:
reci_ord,
)
+ def scaled_dot_product_attention(self, inputs, input_types):
+ query = inputs[0]
+ key = inputs[1]
+ value = inputs[2]
+ attn_mask = inputs[3]
+ dropout_p = inputs[4]
+ is_causal = inputs[5]
+
+ # Explicit scale can be used from torch>=2.1.0
+ if len(inputs) == 7:
+ scale = inputs[6]
+ else:
+ scale = None
+
+ assert (
+ input_types[0] == input_types[1] == input_types[2]
+ ), "Expected query, key, and value to have the same dtype"
+
+ dtype = input_types[0]
+ assert dtype == "float32" or dtype == "float64", "Data type can be
float32 or float64"
+
+ query_shape = self.infer_shape_with_prelude(query)
+ key_shape = self.infer_shape_with_prelude(key)
+ value_shape = self.infer_shape_with_prelude(value)
+ assert 3 <= len(query_shape) <= 4, "Only 3D or 4D query supported"
+ assert 3 <= len(key_shape) <= 4, "Only 3D or 4D key supported"
+ assert 3 <= len(value_shape) <= 4, "Only 3D or 4D value supported"
+
+ assert dropout_p == 0.0, "Only dropout_p==0.0 supported"
+
+ L, S = query_shape[-2], key_shape[-2]
+
+ if scale is None:
+ scale_factor = _expr.const(1 / math.sqrt(query_shape[-1]),
dtype=dtype)
+ else:
+ scale_factor = _expr.const(scale, dtype=dtype)
+
+ attn_bias = _op.full(_expr.const(0.0, dtype=dtype), (L, S))
+
+ if is_causal:
+ assert attn_mask is None, "Explicit attn_mask shouldn't be set
when is_causal=True"
+ temp_mask = _op.full(_expr.const(True), [L, S], dtype="bool")
+ temp_mask = _op.trilu(temp_mask, 0, upper=False)
+ temp_mask = _op.cast(temp_mask, dtype="bool")
+ temp_mask = _op.logical_not(temp_mask)
+ fill_value = _op.cast(_expr.const(float("-inf")), dtype=dtype)
+ attn_bias = _op.where(temp_mask, fill_value, attn_bias)
+ attn_bias = _op.cast(attn_bias, dtype)
+
+ if attn_mask is not None:
+ if input_types[3] == "bool":
+ attn_mask = _op.logical_not(attn_mask)
+ fill_value = _op.cast(_expr.const(float("-inf")), dtype=dtype)
+ attn_bias = _op.where(attn_mask, fill_value, attn_bias)
+ else:
+ attn_bias = _op.add(attn_bias, attn_mask)
+
+ if len(query_shape) < len(key_shape):
+ batch_size = key_shape[0]
+ else:
+ batch_size = query_shape[0]
+ if len(query_shape) == 4 and len(key_shape) == 4:
+ query = _op.reshape(query, newshape=[-3, -2])
+ key = _op.reshape(key, newshape=[-3, -2])
+ if len(query_shape) == 3 and len(key_shape) == 4:
+ query = _op.broadcast_to(query, shape=(batch_size,) + query_shape)
+ query = _op.reshape(query, newshape=[-3, -2])
+ key = _op.reshape(key, newshape=[-3, -2])
+ if len(query_shape) == 4 and len(key_shape) == 3:
+ query = _op.reshape(query, newshape=[-3, -2])
+ key = _op.broadcast_to(key, shape=(batch_size,) + key_shape)
+ key = _op.reshape(key, newshape=[-3, -2])
+ attn_weight = _op.nn.batch_matmul(query, key)
+ if len(query_shape) == 4 or len(key_shape) == 4:
+ attn_weight = _op.reshape(attn_weight, newshape=[-4, batch_size,
-1, -2])
+ attn_weight = _op.squeeze(attn_weight, axis=[])
+
+ attn_weight = _op.multiply(attn_weight, scale_factor)
+ attn_weight = _op.add(attn_weight, attn_bias)
+ attn_weight = _op.nn.softmax(attn_weight)
+ attn_weight = _op.nn.dropout(attn_weight, rate=dropout_p)
+
+ aw_shape = self.infer_shape_with_prelude(attn_weight)
+ if len(aw_shape) < len(value_shape):
+ batch_size = value_shape[0]
+ else:
+ batch_size = aw_shape[0]
+ if len(aw_shape) == 4 and len(value_shape) == 4:
+ attn_weight = _op.reshape(attn_weight, newshape=[-3, -2])
+ value = _op.reshape(value, newshape=[-3, -2])
+ if len(aw_shape) == 3 and len(value_shape) == 4:
+ attn_weight = _op.broadcast_to(attn_weight, shape=(batch_size,) +
aw_shape)
+ attn_weight = _op.reshape(attn_weight, newshape=[-3, -2])
+ value = _op.reshape(value, newshape=[-3, -2])
+ if len(aw_shape) == 4 and len(value_shape) == 3:
+ attn_weight = _op.reshape(attn_weight, newshape=[-3, -2])
+ value = _op.broadcast_to(value, shape=(batch_size,) + value_shape)
+ value = _op.reshape(value, newshape=[-3, -2])
+ attn_weight = _op.nn.batch_matmul(attn_weight, value,
transpose_b=False)
+ if len(aw_shape) == 4 or len(value_shape) == 4:
+ attn_weight = _op.reshape(attn_weight, newshape=[-4, batch_size,
-1, -2])
+ return attn_weight
+
# Operator mappings
def create_convert_map(self):
self.convert_map = {
@@ -4167,6 +4276,7 @@ class PyTorchOpConverter:
"aten::copy_": self.inplace_copy,
"aten::swapaxes": self.transpose,
"aten::linalg_vector_norm": self.linalg_vector_norm,
+ "aten::scaled_dot_product_attention":
self.scaled_dot_product_attention,
}
def update_convert_map(self, custom_map):
diff --git a/tests/python/frontend/pytorch/test_forward.py
b/tests/python/frontend/pytorch/test_forward.py
index d9ecbce265..6d4ed4539b 100644
--- a/tests/python/frontend/pytorch/test_forward.py
+++ b/tests/python/frontend/pytorch/test_forward.py
@@ -4776,7 +4776,6 @@ def test_masked_fill():
verify_model(test_fn, [inp.to(torch.float64), inp > 0.5])
[email protected](reason="unsupported op:
'aten::scaled_dot_product_attention'")
def test_transformer():
"""test_transformer"""
model = torch.nn.Transformer(d_model=256, nhead=8, num_encoder_layers=6,
num_decoder_layers=6)
@@ -5489,6 +5488,89 @@ def test_linalg_vector_norm():
verify_model(test_fn(order=0), input_data=input_data)
[email protected]_gpu
+def test_scaled_dot_product_attention():
+ """test_scaled_dot_product_attention"""
+ torch.set_grad_enabled(False)
+
+ def test_fn(attn_mask=None, is_causal=False):
+ return lambda query, key, value:
torch.nn.functional.scaled_dot_product_attention(
+ query, key, value, attn_mask=attn_mask, is_causal=is_causal
+ )
+
+ L, S, E, Ev = 5, 7, 11, 13
+ query_4d = torch.randn(2, 3, L, E)
+ query_3d = torch.randn(3, L, E)
+ key_4d = torch.randn(2, 3, S, E)
+ key_3d = torch.randn(3, S, E)
+ value_4d = torch.randn(2, 3, S, Ev)
+ value_3d = torch.randn(3, S, Ev)
+
+ verify_model(test_fn(), [query_4d, key_4d, value_4d])
+ verify_model(test_fn(), [query_4d, key_4d, value_3d])
+ verify_model(test_fn(), [query_4d, key_3d, value_4d])
+ verify_model(test_fn(), [query_4d, key_3d, value_3d])
+ verify_model(test_fn(), [query_3d, key_4d, value_4d])
+ verify_model(test_fn(), [query_3d, key_4d, value_3d])
+ verify_model(test_fn(), [query_3d, key_3d, value_4d])
+ verify_model(test_fn(), [query_3d, key_3d, value_3d])
+
+ verify_model(test_fn(is_causal=True), [query_4d, key_4d, value_4d])
+ verify_model(test_fn(is_causal=True), [query_4d, key_4d, value_3d])
+ verify_model(test_fn(is_causal=True), [query_4d, key_3d, value_4d])
+ verify_model(test_fn(is_causal=True), [query_4d, key_3d, value_3d])
+ verify_model(test_fn(is_causal=True), [query_3d, key_4d, value_4d])
+ verify_model(test_fn(is_causal=True), [query_3d, key_4d, value_3d])
+ verify_model(test_fn(is_causal=True), [query_3d, key_3d, value_4d])
+ verify_model(test_fn(is_causal=True), [query_3d, key_3d, value_3d])
+
+ # Test with explicit attn_mask
+ attn_mask = torch.ones((L, S), dtype=torch.bool).tril(diagonal=0)
+ if torch.cuda.is_available():
+ attn_mask = attn_mask.cuda()
+ verify_model(test_fn(attn_mask=attn_mask), [query_4d, key_4d, value_4d])
+ verify_model(test_fn(attn_mask=attn_mask), [query_4d, key_4d, value_3d])
+ verify_model(test_fn(attn_mask=attn_mask), [query_4d, key_3d, value_4d])
+ verify_model(test_fn(attn_mask=attn_mask), [query_4d, key_3d, value_3d])
+ verify_model(test_fn(attn_mask=attn_mask), [query_3d, key_4d, value_4d])
+ verify_model(test_fn(attn_mask=attn_mask), [query_3d, key_4d, value_3d])
+ verify_model(test_fn(attn_mask=attn_mask), [query_3d, key_3d, value_4d])
+ verify_model(test_fn(attn_mask=attn_mask), [query_3d, key_3d, value_3d])
+
+ # Test with float64
+ query_4d = torch.randn(2, 3, L, E, dtype=torch.float64)
+ query_3d = torch.randn(3, L, E, dtype=torch.float64)
+ key_4d = torch.randn(2, 3, S, E, dtype=torch.float64)
+ key_3d = torch.randn(3, S, E, dtype=torch.float64)
+ value_4d = torch.randn(2, 3, S, Ev, dtype=torch.float64)
+ value_3d = torch.randn(3, S, Ev, dtype=torch.float64)
+ verify_model(test_fn(), [query_4d, key_4d, value_4d])
+ verify_model(test_fn(), [query_4d, key_4d, value_3d])
+ verify_model(test_fn(), [query_4d, key_3d, value_4d])
+ verify_model(test_fn(), [query_4d, key_3d, value_3d])
+ verify_model(test_fn(), [query_3d, key_4d, value_4d])
+ verify_model(test_fn(), [query_3d, key_4d, value_3d])
+ verify_model(test_fn(), [query_3d, key_3d, value_4d])
+ verify_model(test_fn(), [query_3d, key_3d, value_3d])
+
+ # Test with larger tensors
+ L, S, E, Ev = 128, 128, 64, 64
+ query_4d = torch.randn(32, 8, L, E)
+ query_3d = torch.randn(8, L, E)
+ key_4d = torch.randn(32, 8, S, E)
+ key_3d = torch.randn(8, S, E)
+ value_4d = torch.randn(32, 8, S, Ev)
+ value_3d = torch.randn(8, S, Ev)
+ verify_model(test_fn(), [query_4d, key_4d, value_4d])
+ verify_model(test_fn(), [query_4d, key_4d, value_3d])
+ verify_model(test_fn(), [query_4d, key_3d, value_4d])
+ verify_model(test_fn(), [query_4d, key_3d, value_3d])
+ verify_model(test_fn(), [query_3d, key_4d, value_4d])
+ verify_model(test_fn(), [query_3d, key_4d, value_3d])
+ verify_model(test_fn(), [query_3d, key_3d, value_4d])
+ verify_model(test_fn(), [query_3d, key_3d, value_3d])
+
+
class TestSetSpan:
"""test structural equal between translated / hand-crafted relay IR with
span tagged."""
