|
from functools import partial |
|
from time import time |
|
import os |
|
import numpy as np |
|
import jax |
|
import jax.flatten_util |
|
import jax.numpy as jnp |
|
import mlxu |
|
from EasyLM.bpt import blockwise_attn |
|
from EasyLM.jax_utils import ( |
|
get_float_dtype_by_name, set_random_seed, next_rng, JaxRNG |
|
) |
|
|
|
|
|
FLAGS, _ = mlxu.define_flags_with_default( |
|
seed=42, |
|
dtype='fp32', |
|
embed_dim=2048, |
|
n_heads=16, |
|
ref_attn_seq_len=2048, |
|
eff_attn_seq_len=16384, |
|
batch_size=1, |
|
query_chunk_size=2048, |
|
key_chunk_size=2048, |
|
warmup_steps=40, |
|
steps=200, |
|
) |
|
|
|
|
|
def main(argv): |
|
|
|
def random_kqv(rng_key, seq_len): |
|
rng_generator = JaxRNG(rng_key) |
|
kqv = [] |
|
for i in range(3): |
|
kqv.append( |
|
jax.random.normal( |
|
rng_generator(), |
|
(FLAGS.batch_size, seq_len, FLAGS.n_heads, FLAGS.embed_dim // FLAGS.n_heads), |
|
dtype=get_float_dtype_by_name(FLAGS.dtype) |
|
) |
|
) |
|
return tuple(kqv) |
|
|
|
def reference_attn(query, key, value): |
|
dtype = get_float_dtype_by_name(FLAGS.dtype) |
|
query = query / jnp.sqrt(query.shape[-1]).astype(dtype) |
|
logits = jnp.einsum("bqhc,bkhc->bhqk", query, key) |
|
mask_value = jnp.finfo(logits.dtype).min |
|
_, q_seq_len, _, _ = query.shape |
|
_, kv_seq_len, _, _ = key.shape |
|
mask_shape = (q_seq_len, kv_seq_len) |
|
row_ids = jax.lax.broadcasted_iota(jnp.int32, mask_shape, 0) |
|
col_ids = jax.lax.broadcasted_iota(jnp.int32, mask_shape, 1) |
|
causal_mask = (row_ids < col_ids)[None, None, :, :] |
|
logits = logits + jnp.where(causal_mask, mask_value, 0.0) |
|
weights = jax.nn.softmax(logits, axis=-1) |
|
out = jnp.einsum("bhqk,bkhc->bqhc", weights, value) |
|
return out |
|
|
|
def efficient_attention(query, key, value): |
|
dtype = get_float_dtype_by_name(FLAGS.dtype) |
|
return blockwise_attn( |
|
query, key, value, |
|
bias=None, |
|
deterministic=True, |
|
dropout_rng=None, |
|
attn_pdrop=0.0, |
|
causal=True, |
|
query_chunk_size=FLAGS.query_chunk_size, |
|
key_chunk_size=FLAGS.key_chunk_size, |
|
dtype=get_float_dtype_by_name(FLAGS.dtype), |
|
policy=jax.checkpoint_policies.nothing_saveable(), |
|
precision=None, |
|
float32_logits=True, |
|
prevent_cse=True, |
|
) |
|
|
|
|
|
@partial(jax.jit, static_argnums=(1,)) |
|
def reference_attn_forward_backward(rng_key, seq_len): |
|
@partial(jax.grad, argnums=(0, 1, 2)) |
|
@partial(jax.checkpoint, policy=jax.checkpoint_policies.nothing_saveable()) |
|
def grad_fn(query, key, value): |
|
out = reference_attn(query, key, value) |
|
return jnp.mean(out) |
|
|
|
query, key, value = random_kqv(rng_key, seq_len) |
|
return jax.flatten_util.ravel_pytree( |
|
grad_fn(query, key, value)[1] |
|
)[0].mean() |
|
|
|
@partial(jax.jit, static_argnums=(1,)) |
|
def efficient_attn_forward_backward(rng_key, seq_len): |
|
@partial(jax.grad, argnums=(0, 1, 2)) |
|
def grad_fn(query, key, value): |
|
out = efficient_attention(query, key, value) |
|
return jnp.mean(out) |
|
|
|
query, key, value = random_kqv(rng_key, seq_len) |
|
return jax.flatten_util.ravel_pytree( |
|
grad_fn(query, key, value)[1] |
|
)[0].mean() |
|
|
|
|
|
set_random_seed(FLAGS.seed) |
|
|
|
jax.block_until_ready(reference_attn_forward_backward(next_rng(), FLAGS.ref_attn_seq_len)) |
|
jax.block_until_ready(efficient_attn_forward_backward(next_rng(), FLAGS.eff_attn_seq_len)) |
|
|
|
all_results = [] |
|
for i in range(FLAGS.warmup_steps): |
|
all_results.append(reference_attn_forward_backward(next_rng(), FLAGS.ref_attn_seq_len)) |
|
jax.block_until_ready(all_results) |
|
|
|
start_time = time() |
|
all_results = [] |
|
for i in range(FLAGS.steps): |
|
all_results.append(reference_attn_forward_backward(next_rng(), FLAGS.ref_attn_seq_len)) |
|
|
|
jax.block_until_ready(all_results) |
|
elapsed_time_ref_attn = time() - start_time |
|
print(f'Reference attention: {elapsed_time_ref_attn:.3f} seconds') |
|
|
|
|
|
all_results = [] |
|
for i in range(FLAGS.warmup_steps): |
|
all_results.append(efficient_attn_forward_backward(next_rng(), FLAGS.eff_attn_seq_len)) |
|
jax.block_until_ready(all_results) |
|
|
|
|
|
start_time = time() |
|
all_results = [] |
|
for i in range(FLAGS.steps): |
|
all_results.append(efficient_attn_forward_backward(next_rng(), FLAGS.eff_attn_seq_len)) |
|
|
|
jax.block_until_ready(all_results) |
|
elapsed_time_efficient_attn = time() - start_time |
|
print(f'Efficient attention: {elapsed_time_efficient_attn:.3f} seconds') |
|
|
|
flops_ratio = (FLAGS.eff_attn_seq_len / FLAGS.ref_attn_seq_len) ** 2 |
|
efficiency = elapsed_time_ref_attn / elapsed_time_efficient_attn * flops_ratio |
|
print(f'Efficiency: {efficiency:.3f}') |
|
|
|
|
|
if __name__ == '__main__': |
|
mlxu.run(main) |
|
|
|
|
|
|
|
|
