def forward(
self,
dec_input,
dec_attn_mask,
enc_output,
enc_attn_mask,
layer_past=None,
get_key_value=False,
):
# convert to Megatron mask
dec_attn_mask_3d = build_attention_mask_3d(
source_mask=dec_attn_mask,
target_mask=dec_attn_mask,
attn_mask_type=self.model_attn_mask_type,
)
enc_dec_attn_mask_3d = build_attention_mask_3d(
source_mask=dec_attn_mask,
target_mask=enc_attn_mask,
attn_mask_type=AttnMaskType.padding,
)
# transformer decoder
dec_output = self.model(
dec_input,
attn_mask_postprocess(dec_attn_mask_3d),
layer_past=layer_past,
get_key_value=get_key_value,
encoder_output=enc_output,
enc_dec_attn_mask=attn_mask_postprocess(enc_dec_attn_mask_3d),
)
return dec_output
def forward(
self,
enc_input,
enc_attn_mask,
context_attn_mask=None,
encoder_output=None,
layer_past=None,
get_key_value=False,
):
# expected enc_input shape [batch, num_chunks, num_neighbors, retrieval_seq_len, dim]
# expected enc_attn_mask shape [batch, num_chunks, num_neighbors, retrieval_seq_len]
# expected encoder_output shape [batch, seq_len, dim]
b, k, r, rn, dim = enc_input.shape
# batch, seq_len, dim
_, n, _ = encoder_output.shape
num_seq_chunks = n // self.chunk_size
assert k == num_seq_chunks, f'sequence requires {num_seq_chunks} retrieved chunks, but only {k} passed in'
seq_index = num_seq_chunks * self.chunk_size
retrieved = rearrange(enc_input, 'b k r n d -> (b k r) n d')
enc_attn_mask = rearrange(enc_attn_mask, 'b k r n -> (b k r) n')
embed_as_context = repeat(encoder_output[:, :seq_index], 'b (k n) d -> (b k r) n d', n=self.chunk_size, r=r)
context_attn_mask = repeat(context_attn_mask[:, :seq_index], 'b (k n) -> (b k r) n', n=self.chunk_size, r=r)
# need to add extra chunk size, since it will be shifted
cross_attn_q_pos_emb = self.rotary_pos_emb(rn, offset=0)
cross_attn_k_pos_emb = self.rotary_pos_emb(self.chunk_size)
attn_pos_emb = (cross_attn_q_pos_emb, cross_attn_q_pos_emb, cross_attn_k_pos_emb)
# # convert to Megatron mask
enc_attn_mask_3d = build_attention_mask_3d(
source_mask=enc_attn_mask, target_mask=enc_attn_mask, attn_mask_type=self.model_attn_mask_type,
)
enc_attn_mask_3d = enc_attn_mask_3d[:, None, :, :]
enc_dec_attn_mask_3d = build_attention_mask_3d(
source_mask=enc_attn_mask, target_mask=context_attn_mask, attn_mask_type=AttnMaskType.padding,
)
enc_dec_attn_mask_3d = enc_dec_attn_mask_3d[:, None, :, :]
# transformer encoder
enc_output = self.model(
retrieved,
enc_attn_mask_3d,
layer_past=layer_past,
get_key_value=get_key_value,
encoder_output=embed_as_context,
enc_dec_attn_mask=enc_dec_attn_mask_3d,
rotary_pos_emb=attn_pos_emb,
)
# revert back to original retrieved shape
enc_output = rearrange(enc_output, '(b k r) n d -> b k r n d', b=b, k=k)
return enc_output
def forward(
self,
dec_input,
dec_attn_mask,
retrieved_attn_mask=None,
retrieved_emb=None,
layer_past=None,
get_key_value=False,
eod_positions=None, # this is a tuple of eod positions returned from tensor.where(tensor == eod_id)
):
# expected dec_input shape [batch, seq_len, dim]
# expected dec_attn_mask shape [batch, seq_len]
# expected retrieved_input shape [batch, num_chunks, num_neighbors, retrival_seq_len, dim]
# expected retrieved_attn_mask shape [batch, num_chunks, num_neighbors, retrival_seq_len]
# batch, seq_len, dim
_, n, _ = dec_input.shape
num_seq_chunks = n // self.chunk_size
if retrieved_emb is not None:
b, k, r, rn, dim = retrieved_emb.shape
assert (
k == num_seq_chunks
), f'sequence requires {num_seq_chunks} retrieved chunks, but only {k} passed in' # need to add extra chunk size, since it will be shifted
self_attn_emb = self.rotary_pos_emb(n)
if retrieved_emb is not None:
cross_attn_q_pos_emb = self.rotary_pos_emb(self.chunk_size * 2 - 1)
cross_attn_k_pos_emb = self.rotary_pos_emb(rn, offset=0)
attn_pos_emb = (self_attn_emb, cross_attn_q_pos_emb, cross_attn_k_pos_emb)
else:
attn_pos_emb = (self_attn_emb, None, None)
dec_attn_mask_3d = self._calculate_dec_att_mask(dec_attn_mask, eod_positions)
if retrieved_emb is not None:
dec_attn_mask = rearrange(dec_attn_mask, 'b (k n) -> (b k) n', k=k)
retrieved_attn_mask = rearrange(retrieved_attn_mask, 'b k r n -> (b k) (r n)')
enc_dec_attn_mask_3d = build_attention_mask_3d(
source_mask=dec_attn_mask, target_mask=retrieved_attn_mask, attn_mask_type=AttnMaskType.padding,
)
enc_dec_attn_mask_3d = enc_dec_attn_mask_3d[:, None, :, :]
else:
enc_dec_attn_mask_3d = None
# transformer encoder
enc_output = self.model(
dec_input,
dec_attn_mask_3d,
layer_past=layer_past,
get_key_value=get_key_value,
encoder_output=None,
retrieved_emb=retrieved_emb,
enc_dec_attn_mask=enc_dec_attn_mask_3d,
rotary_pos_emb=attn_pos_emb,
)
return enc_output
def forward(
self,
enc_input,
enc_attn_mask,
layer_past=None,
get_key_value=False,
):
# convert to Megatron mask
enc_attn_mask_3d = build_attention_mask_3d(
source_mask=enc_attn_mask,
target_mask=enc_attn_mask,
attn_mask_type=self.model_attn_mask_type,
)
# transformer encoder
enc_output = self.model(
enc_input,
attn_mask_postprocess(enc_attn_mask_3d),
layer_past=layer_past,
get_key_value=get_key_value,
)
# we copy input mask for transformer
enc_output_mask = enc_attn_mask
return enc_output, enc_output_mask
def get_attn_mask_3d(hidden_mask, context_mask, chunks):
causal_padding = text_chunk_size - 1
reminder = (text_chunk_size -
(hidden_mask.shape[0] + 1)) % text_chunk_size
hidden_mask = F.pad(hidden_mask, (0, 0, -causal_padding, reminder),
value=False)
dec_attn_mask = rearrange(hidden_mask,
'(k n) b -> (b k) n',
k=chunks)
context_attn_mask = rearrange(context_mask,
'k r n b -> (b k) (r n)')
enc_dec_attn_mask_3d = build_attention_mask_3d(
source_mask=dec_attn_mask,
target_mask=context_attn_mask,
attn_mask_type=AttnMaskType.padding,
)
enc_dec_attn_mask_3d = enc_dec_attn_mask_3d[:, None, :, :]
return enc_dec_attn_mask_3d
def _calculate_dec_att_mask(self, dec_attn_mask, eod_positions):
# # convert to Megatron mask
dec_attn_mask_3d = build_attention_mask_3d(
source_mask=dec_attn_mask, target_mask=dec_attn_mask, attn_mask_type=self.model_attn_mask_type,
)
if eod_positions is not None:
# to mask out the token ids [id, id, eod, id, pad, eod, id, id]
# so attention is not across eod, mask should be:
# [false, true, true, true, true, true, true, true]
# [false, false, true, true, true, true, true, true]
# [false, false, false,true, true, true, true, true]
# [true, true, true, false, true, true, true, true]
# [true, true, true, true, true, true, true, true]
# [true, true, true, false, true, false, true, true]
# [true, true, true, true, true, true, false, true]
# [true, true, true, true, true, true, false, false]
for batch, eod_pos in zip(*eod_positions):
eod_plus_one = eod_pos.item() + 1
dec_attn_mask_3d[batch][eod_plus_one:, :eod_plus_one] = True
dec_attn_mask_3d = dec_attn_mask_3d[:, None, :, :]
return dec_attn_mask_3d
def test_cross_attn(self):
num_layers = 1
init_method_std = 0.02
batch = 2
neighbors = 2
# rotary pos emb dim
dim = 128
pad_id = 19999
num_attention_heads = 8
chunks = 32
text_chunk_size = 64
context_chunk_size = 2 * text_chunk_size
input_length = chunks * text_chunk_size
vocab_size = 20000
rot_dim = dim // num_attention_heads
rotary_pos_emb = RotaryEmbedding(rot_dim).cuda().half()
hidden = torch.randint(0, vocab_size, (input_length, batch)).cuda() # (seq, batch, dim)
hidden_mask = (hidden != pad_id).cuda()
hidden_emb = torch.rand(input_length, batch, dim).cuda().half() # (seq, batch, dim)
retrieved = torch.randint(0, vocab_size, (chunks, neighbors, context_chunk_size, batch)).cuda()
# retrieved tokens - (num chunks, num retrieved neighbors, retrieved chunk with continuation, batch)
# context attention mask [b, np, sq, sk]
context_mask = (retrieved != pad_id).cuda()
retrieved_emb = torch.rand(chunks, neighbors, context_chunk_size, batch, dim).cuda().half()
# retrieved tokens - (num chunks, num retrieved neighbors, retrieved chunk with continuation, batch, hidden)
# need to add extra chunk size, since it will be shifted
cross_attn_q_pos_emb = rotary_pos_emb(text_chunk_size + text_chunk_size - 1, offset=0)
cross_attn_k_pos_emb = rotary_pos_emb(context_chunk_size)
cross_attn_pos_emb = (cross_attn_q_pos_emb, cross_attn_k_pos_emb)
dec_attn_mask = rearrange(hidden_mask, '(k n) b -> (b k) n', k=chunks)
context_attn_mask = rearrange(context_mask, 'k r n b -> (b k) (r n)')
enc_dec_attn_mask_3d = build_attention_mask_3d(
source_mask=dec_attn_mask, target_mask=context_attn_mask, attn_mask_type=AttnMaskType.padding,
)
enc_dec_attn_mask_3d = enc_dec_attn_mask_3d[:, None, :, :]
init_method = init_method_normal(init_method_std)
scaled_init_method = scaled_init_method_normal(init_method_std, num_layers)
cross_attn = (
ParallelChunkedCrossAttention(
init_method=init_method,
output_layer_init_method=scaled_init_method,
layer_number=0,
num_attention_heads=num_attention_heads,
hidden_size=dim,
precision=16,
chunk_size=text_chunk_size,
)
.cuda()
.half()
)
out, bias = cross_attn(
hidden_emb, enc_dec_attn_mask_3d, encoder_output=retrieved_emb, rotary_pos_emb=cross_attn_pos_emb
)
assert out.shape == torch.Size([input_length, batch, dim])
assert bias.shape == torch.Size([dim])
def forward(
self,
dec_input,
dec_attn_mask,
retrieved_attn_mask=None,
retrieved_emb=None,
layer_past=None,
get_key_value=False,
eod_positions=None, # this is a tuple of eod positions returned from tensor.where(tensor == eod_id)
set_inference_key_value_memory=False,
inference_max_sequence_len=None,
):
# expected dec_input shape [batch, seq_len, dim]
# expected dec_attn_mask shape [batch, seq_len]
# expected retrieved_input shape [batch, num_chunks, num_neighbors, retrival_seq_len, dim]
# expected retrieved_attn_mask shape [batch, num_chunks, num_neighbors, retrival_seq_len]
# batch, seq_len, dim
if isinstance(dec_input, tuple):
n, _, _ = dec_input[1].shape
else:
_, n, _ = dec_input.shape
if set_inference_key_value_memory == True:
# seq_index = (n // chunk_size) * chunk_size
self.current_len = n
num_seq_chunks = self.current_len // self.chunk_size
self_attn_emb = self.rotary_pos_emb(self.current_len)
elif inference_max_sequence_len is not None:
# only handles single token increment
assert n == 1
self.current_len += n
self_attn_emb = self.rotary_pos_emb(self.current_len)
num_seq_chunks = self.current_len // self.chunk_size
else:
# this is normal forward without inference
num_seq_chunks = n // self.chunk_size
self_attn_emb = self.rotary_pos_emb(n)
if retrieved_emb is not None:
b, k, r, rn, dim = retrieved_emb.shape
assert (
k == num_seq_chunks
), f'sequence requires {num_seq_chunks} retrieved chunks, but only {k} passed in' # need to add extra chunk size, since it will be shifted
if retrieved_emb is not None:
cross_attn_q_pos_emb = self.rotary_pos_emb(
self.chunk_size * 2 - 1, offset=-self.chunk_size + 1)
cross_attn_k_pos_emb = self.rotary_pos_emb(rn, offset=0)
attn_pos_emb = (self_attn_emb, cross_attn_q_pos_emb,
cross_attn_k_pos_emb)
else:
attn_pos_emb = (self_attn_emb, None, None)
dec_attn_mask_3d = self._calculate_dec_att_mask(
dec_attn_mask, eod_positions)
if retrieved_emb is not None:
# need to shift the dec_attn_mask as first causal_padding elements are ignored
# also pad it to be the multiple of self.chunk_size
causal_padding = self.chunk_size - 1
reminder = (self.chunk_size -
(dec_attn_mask.shape[1] + 1)) % self.chunk_size
dec_attn_mask = F.pad(dec_attn_mask, (-causal_padding, reminder),
value=False)
dec_attn_mask = rearrange(dec_attn_mask, 'b (k n) -> (b k) n', k=k)
retrieved_attn_mask = rearrange(retrieved_attn_mask,
'b k r n -> (b k) (r n)')
enc_dec_attn_mask_3d = build_attention_mask_3d(
source_mask=dec_attn_mask,
target_mask=retrieved_attn_mask,
attn_mask_type=AttnMaskType.padding,
)
enc_dec_attn_mask_3d = enc_dec_attn_mask_3d[:, None, :, :]
else:
enc_dec_attn_mask_3d = None
# transformer encoder
enc_output = self.model(
dec_input,
dec_attn_mask_3d,
layer_past=layer_past,
get_key_value=get_key_value,
encoder_output=None,
retrieved_emb=retrieved_emb,
enc_dec_attn_mask=enc_dec_attn_mask_3d,
rotary_pos_emb=attn_pos_emb,
set_inference_key_value_memory=set_inference_key_value_memory,
inference_max_sequence_len=inference_max_sequence_len,
)
return enc_output
def forward(
self,
enc_input,
enc_attn_mask,
context_attn_mask=None,
encoder_output=None,
layer_past=None,
get_key_value=False,
set_inference_key_value_memory=False, # when doing inference, set this to true to allocate all the cached matrix. later set false to do incremental inference
inference_max_sequence_len=None,
neighbors=2,
):
# expected enc_input shape [batch, num_chunks, num_neighbors, retrieval_seq_len, dim]
# expected enc_attn_mask shape [batch, num_chunks, num_neighbors, retrieval_seq_len]
# expected encoder_output shape [batch, seq_len, dim]
# batch, seq_len, dim
b, n, dim = encoder_output.shape
if set_inference_key_value_memory:
# run once to setup the cache
chunk_start = 0
num_seq_chunks = n // self.chunk_size
num_chunks = inference_max_sequence_len // self.chunk_size
self.cache_output = self._allocate_memory(
b,
num_chunks,
neighbors,
self.chunk_size * 2,
dim,
dtype=encoder_output.dtype)
self.seq_pos_in_chunk = n
self.current_chunk = n // self.chunk_size
self.encoder_output = self._allocate_memory(
b, self.chunk_size, dim, dtype=encoder_output.dtype)
self.context_attn_mask = self._allocate_memory(
b, self.chunk_size, dtype=context_attn_mask.dtype)
self.context_attn_mask
chunk_beg = self.chunk_size * num_seq_chunks
chunk_end = self.chunk_size * num_seq_chunks + self.seq_pos_in_chunk % self.chunk_size
# store the remainders
self.encoder_output[:, :self.seq_pos_in_chunk % self.
chunk_size, :] = encoder_output[:, chunk_beg:
chunk_end, :]
self.context_attn_mask[:, :self.seq_pos_in_chunk % self.
chunk_size] = context_attn_mask[:,
chunk_beg:
chunk_end]
elif inference_max_sequence_len is not None:
# second time of running
# only support one token at a time
assert n == 1
self.seq_pos_in_chunk += n
self.current_chunk = self.seq_pos_in_chunk // self.chunk_size
# if exceed the chunk size
pos_beg = (self.seq_pos_in_chunk - 1) % self.chunk_size
# if self.seq_pos_in_chunk - 1 >= self.chunk_size:
# self.current_chunk += 1
# self.seq_pos_in_chunk -= self.chunk_size
chunk_start = self.current_chunk - 1
self.encoder_output[:, pos_beg:pos_beg + 1, :] = encoder_output
self.context_attn_mask[:, pos_beg:pos_beg +
1] = context_attn_mask[:, self.
seq_pos_in_chunk -
1:self.
seq_pos_in_chunk]
encoder_output = self.encoder_output[:, :pos_beg + 1, :]
context_attn_mask = self.context_attn_mask[:, :pos_beg + 1]
num_seq_chunks = 1
if not self.seq_pos_in_chunk % self.chunk_size == 0:
# still accumulate the encoder_output
# return the cached results
if self.current_chunk == 0:
return None
return self.cache_output[:, :self.current_chunk]
if enc_input is not None:
# only need one chunk for the later calculation
enc_input = enc_input[:, self.current_chunk -
1:self.current_chunk]
enc_attn_mask = enc_attn_mask[:, self.current_chunk -
1:self.current_chunk]
if enc_input is None:
return None
_, k, r, rn, _ = enc_input.shape
assert r == neighbors
if inference_max_sequence_len is None:
num_seq_chunks = n // self.chunk_size
assert k == num_seq_chunks, f'sequence requires {num_seq_chunks} retrieved chunks, but only {k} passed in'
else:
pass
seq_index = num_seq_chunks * self.chunk_size
retrieved = rearrange(enc_input, 'b k r n d -> (b k r) n d')
enc_attn_mask = rearrange(enc_attn_mask, 'b k r n -> (b k r) n')
# embed_as_context = repeat(encoder_output[:, :seq_index], 'b (k n) d -> (b k r) n d', n=self.chunk_size, r=r)
# context_attn_mask = repeat(context_attn_mask[:, :seq_index], 'b (k n) -> (b k r) n', n=self.chunk_size, r=r)
cross_attn_q_pos_emb = self.rotary_pos_emb(rn, offset=0)
if inference_max_sequence_len is not None and not set_inference_key_value_memory:
cross_attn_k_pos_emb = self.rotary_pos_emb(n % self.chunk_size,
offset=pos_beg)
embed_as_context = repeat(encoder_output[:, :seq_index],
'b (k n) d -> (b k r) n d',
n=pos_beg + 1,
r=r)
context_attn_mask = repeat(context_attn_mask[:, :seq_index],
'b (k n) -> (b k r) n',
n=pos_beg + 1,
r=r)
else:
embed_as_context = repeat(encoder_output[:, :seq_index],
'b (k n) d -> (b k r) n d',
n=self.chunk_size,
r=r)
context_attn_mask = repeat(context_attn_mask[:, :seq_index],
'b (k n) -> (b k r) n',
n=self.chunk_size,
r=r)
cross_attn_k_pos_emb = self.rotary_pos_emb(self.chunk_size,
offset=0)
attn_pos_emb = (cross_attn_q_pos_emb, cross_attn_q_pos_emb,
cross_attn_k_pos_emb)
# # convert to Megatron mask
enc_attn_mask_3d = build_attention_mask_3d(
source_mask=enc_attn_mask,
target_mask=enc_attn_mask,
attn_mask_type=self.model_attn_mask_type,
)
enc_attn_mask_3d = enc_attn_mask_3d[:, None, :, :]
enc_dec_attn_mask_3d = build_attention_mask_3d(
source_mask=enc_attn_mask,
target_mask=context_attn_mask,
attn_mask_type=AttnMaskType.padding,
)
enc_dec_attn_mask_3d = enc_dec_attn_mask_3d[:, None, :, :]
# transformer encoder
enc_output = self.model(
retrieved,
enc_attn_mask_3d,
layer_past=layer_past,
get_key_value=get_key_value,
encoder_output=embed_as_context,
enc_dec_attn_mask=enc_dec_attn_mask_3d,
rotary_pos_emb=attn_pos_emb,
)
# revert back to original retrieved shape
enc_output = rearrange(enc_output,
'(b k r) n d -> b k r n d',
b=b,
k=k)
if inference_max_sequence_len is not None:
# update encoded for current chunk
self.cache_output[:, chunk_start:self.
current_chunk, :, :, :] = enc_output
# read all encodings
enc_output = self.cache_output[:, :self.current_chunk]
return enc_output