On 1/26/22, k <[email protected]> wrote:
> the AminRezaei0x443 implementation also produces the same data, attached
> again.
>
> the aminrezaei implementation does the square root, provides for
> optional mask and bias tensors, is on pypi, and has both a jax and
> torch implementation, so it seems the way to go.
>
> next i'll be timing it compared to the paper's implementation that i
> noted as speedy. just on my raspberry pi, though. i'm guessing it's
> roughly the same on good hardware with large models, where the core
> batches dominate everything. sometimes i mostly engage stuff i bump
> into.
>
> maybe it would be good just to quickly run through the source and
> verify that aminrezai does checkpointing and lax mapping like in the
> paper.
>
import jax, torch, math
import memory_efficient_attention
def moskomule_attention(queries, keys, values, query_chunk_size, key_chunk_size):
queries /= math.sqrt(keys.shape[-1])
return memory_efficient_attention.efficient_attention(queries, keys, values, chunk_size=key_chunk_size, checkpointing = True, out_of_place = False)
def aminrezaei_attention(queries, keys, values, query_chunk_size, key_chunk_size):
return memory_efficient_attention.efficient_dot_product_attention_pt(queries, keys, values, None, None, query_chunk_Size = query_chunk_Size, key_chunk_size = key_chunk_size)
def attention(queries, keys, values, query_chunk_size, key_chunk_size):
queries = queries.permute(0, 2, 1, 3)
keys = keys.permute(0, 2, 1, 3)
values = values.permute(0, 2, 1, 3)
if hasattr(memory_efficient_attention, 'efficient_attention'):
return moskomule_attention(queries, keys, values, query_chunk_size, key_chunk_size)
elif hasattr(memory_efficient_attention, 'efficient_dot_product_attention_pt'):
return aminrezaei_attention(queries, keys, values, query_chunk_size, key_chunk_size)
if __name__ == '__main__':
queries, keys, values = torch.from_numpy(jax.random.normal(jax.random.PRNGKey(0), (3, 1, 64, 8, 16)).to_py())
out = attention(queries, keys, values, query_chunk_size=4, key_chunk_size=4)
print(out)
