huseinzol05
commented
1 month ago This is to replicate the original jax implementation, https://github.com/mesolitica/context-parallelism-xformers/blob/master/playground/blockwise-vanilla-attention-causal.ipynb
I just simply generate global multiplier mask (lazy to do addition),
temp_mask = torch.ones(L, S, dtype=torch.bool).tril(diagonal=0).cuda()
attn_bias_blocks = torch.chunk(temp_mask, chunk_size)
attn_bias_block = attn_bias_blocks[0]
seq_chunk = Q.shape[0] // chunk_size
attn_bias_b = attn_bias_block[:, no * seq_chunk: (no + 1) * seq_chunk]
scores = torch.matmul(Q_block, K_block.T) * attn_bias_bWhile original jax generate addition mask during blockwise,
def _chunk_attention_bias(query_chunk_size, key_chunk_size,
bias, segment_ids, deterministic, attn_dropout, attn_pdrop, causal,
dtype, query_chunk_idx, key_chunk_idx):
query_offset = query_chunk_idx * query_chunk_size
key_offset = key_chunk_idx * key_chunk_size
chunk_bias = jnp.zeros((1, 1, 1, 1), dtype=dtype)
if bias is not None:
chunk_bias = lax.dynamic_slice(
bias,
start_indices=(0, 0, 0, key_offset),
slice_sizes=(*bias.shape[:2], min(bias.shape[-2], query_chunk_size), min(bias.shape[-1], key_chunk_size)),
)
if segment_ids is not None:
q_segment_ids = lax.dynamic_slice(
segment_ids,
start_indices=(0, query_offset),
slice_sizes=(segment_ids.shape[0], query_chunk_size)
)
k_segment_ids = lax.dynamic_slice(
segment_ids,
start_indices=(0, key_offset),
slice_sizes=(segment_ids.shape[0], key_chunk_size)
)
segment_ids_mask = q_segment_ids[:, :, None] != k_segment_ids[:, None, :]
segment_ids_mask = segment_ids_mask[:, None] # B1QK
segment_ids_bias = segment_ids_mask * jnp.finfo(dtype).min
chunk_bias += segment_ids_bias
if causal:
query_idx = lax.broadcasted_iota(dtype=jnp.int32, shape=(query_chunk_size, 1), dimension=0)
key_idx = lax.broadcasted_iota(dtype=jnp.int32, shape=(1, key_chunk_size), dimension=1)
offset = query_offset - key_offset
query_idx += offset
causal_mask_value = (query_idx < key_idx) * jnp.finfo(dtype).min
chunk_bias += causal_mask_value.reshape(1, 1, *causal_mask_value.shape)
if not deterministic and attn_pdrop > 0.0:
attn_dropout_slice = lax.dynamic_slice(
attn_dropout,
start_indices=(0, 0, query_offset, key_offset),
slice_sizes=(
*attn_dropout.shape[:2],
min(attn_dropout.shape[-2], query_chunk_size),
min(attn_dropout.shape[-1], key_chunk_size),
),
)
chunk_bias += attn_dropout_slice * jnp.finfo(dtype).min
return chunk_bias.astype(dtype)
_chunk_bias_fn = partial(
_chunk_attention_bias,
query_chunk_size, key_chunk_size, bias, segment_ids, deterministic,
attn_dropout, attn_pdrop, causal, dtype)
bias_chunk = _chunk_bias_fn(q_chunk_idx_start + q_chunk_idx, k_chunk_idx_start + k_chunk_idx)
zhuzilin
Hi @zhuzilin, follow up from https://github.com/zhuzilin/ring-flash-attention/issues/15
I just wanted to verify the causal, and I simply use loop because I dont have multigpus, but it should be working, when I do causal using your ring logic, the argmax accuracy is super low, but when I do non causal, accuracy is almost perfect 100%, you can check the notebook at https://github.com/mesolitica/context-parallelism-xformers/blob/master/playground/flash-ring-attention-causal.ipynb
From what I understand, let say, I got 2 devices and seqlen of 100k, partitioned to 2, 100k // 2 = 50k 50k, so,
each 50k seq len, device 0: 50k q0k0v0 device 1: 50k q1k1v1
So the blockwise attention calculation, device 0: 50k q0k0v0 + 50k q0k1v1 device 1: 50k q1k0v0 + 50k q1k1v1
(+) denoted as blockwise attention.
For causal base, attention mask is necessary, so the attention mask originally is [100k, 100k] and attention mask we must chunk properly, to become mask0 = [50k, 100k] and mask1 = [50k, 100k], so the blockwise attention calculation,
device 0: 50k (q0k0 mask0[:, 0:50k])v0 + 50k q0k1v1 mask0[:, 50k:100k] device 1: 50k (q1k0 mask1[:, 0:50k])v0 + 50k q1k1v1 mask1[:, 50k:100k]
You can see this slicing from original https://github.com/forhaoliu/ringattention/blob/main/ringattention/ringattention_pallas_tpu.py#L61
Correct me if im wrong here, thanks!