diff --git a/matt/training/aim_impl.py b/matt/training/aim_impl.py index 2cdae6c..84ed6d1 100644 --- a/matt/training/aim_impl.py +++ b/matt/training/aim_impl.py @@ -111,6 +111,69 @@ ] +def last_token_mask(word_ids: torch.Tensor) -> torch.Tensor: + """Return a mask for the final token of each non-special segment.""" + if word_ids.ndim != 2: + raise ValueError(f"word_ids must be two-dimensional, got {tuple(word_ids.shape)}") + + next_word_ids = torch.full_like(word_ids, -100) + next_word_ids[:, :-1] = word_ids[:, 1:] + return (word_ids != -100) & (word_ids != next_word_ids) + + +def causal_word_pair_ids( + word_ids: torch.Tensor, + num_heads: int, +) -> tuple[torch.Tensor, int]: + """Build flattened AIM pair IDs without crossing batch/head boundaries. + + The returned tensor has the same flattened order as + ``attn_weights.unsqueeze(-1) * value_states.unsqueeze(-3)``. + + Word ids restart in every sequence, so pair ids must be derived per + (batch row, head) slice. An earlier implementation compared neighbours + with ``roll()`` and numbered pairs with a global ``cumsum()`` over the + flattened batch: whenever the word ids touching a row boundary (or the + wrap-around from the last token back to the first) happened to be equal, + word segments from different rows or heads were merged into one pair row + or dropped altogether. Teacher and student AIM states are built + independently, so any merged or dropped row shifts every following row + and breaks their one-to-one correspondence in the loss. Keeping the + row/head structure until the very end makes such collisions impossible. + """ + if word_ids.ndim != 2: + raise ValueError(f"word_ids must be two-dimensional, got {tuple(word_ids.shape)}") + if num_heads <= 0: + raise ValueError(f"num_heads must be positive, got {num_heads}") + + device = word_ids.device + _, seq_len = word_ids.shape + repeated_word_ids = word_ids.repeat_interleave(num_heads, dim=0) + valid_query_mask = last_token_mask(repeated_word_ids) + + causal_mask = torch.tril(torch.ones(seq_len, seq_len, device=device, dtype=torch.bool)) + causal_word_ids = ( + repeated_word_ids + .unsqueeze(-2) + .expand(-1, seq_len, -1) + .masked_fill(~causal_mask, -100) + .masked_fill(~valid_query_mask.unsqueeze(-1), -100) + ) + + previous_word_ids = torch.full_like(causal_word_ids, -100) + previous_word_ids[..., 1:] = causal_word_ids[..., :-1] + pair_start_mask = (causal_word_ids != -100) & (causal_word_ids != previous_word_ids) + + pair_counts = pair_start_mask.sum(dim=-1) + flat_pair_counts = pair_counts.reshape(-1) + pair_offsets = (flat_pair_counts.cumsum(0) - flat_pair_counts).reshape_as(pair_counts) + pair_ids = pair_start_mask.cumsum(dim=-1) - 1 + pair_ids = pair_ids + pair_offsets.unsqueeze(-1) + pair_ids = pair_ids.masked_fill(causal_word_ids == -100, -100) + + return pair_ids.reshape(-1), int(flat_pair_counts.sum().item()) + + def aim_impl( teacher_attn_weights: torch.Tensor, teacher_value_states: torch.Tensor, @@ -163,53 +226,26 @@ def get_aim_states( (num_pairs, hidden_size) """ - device = word_ids.device - - batch_size, seq_len = word_ids.shape num_heads = attn_weights.size(1) hidden_size = value_states.size(-1) - # (batch_size * num_heads, seq_len) -> (batch_size * num_heads * seq_len) - rep_word_ids = word_ids.repeat_interleave(num_heads, dim=0) - rep_word_ids_flat = rep_word_ids.view(-1) - valid_word_mask = torch.logical_and( - rep_word_ids_flat != -100, # no padding or special tokens - rep_word_ids_flat != rep_word_ids_flat.roll(-1), # last token of the word - ) - - causal_mask = torch.tril(torch.ones(seq_len, seq_len, device=device, dtype=torch.bool)) - rep_full_word_ids = ( - # (batch_size * num_heads, seq_len) - rep_word_ids - .unsqueeze(-2) - # (batch_size * num_heads, seq_len, seq_len) - .repeat(1, seq_len, 1) - .masked_fill(~causal_mask, -100) - # (batch_size * num_heads * seq_len, seq_len) - .view(-1, seq_len) - .masked_fill(~valid_word_mask.unsqueeze(-1), -100) - # (batch_size * num_heads * seq_len * seq_len) - .view(-1) - ) - - # (batch_size * num_heads * seq_len * seq_len) - full_word_ids = torch.where( - rep_full_word_ids != -100, - # mask of where the words change, cumsum to get word ids across the batch - torch.logical_and( - rep_full_word_ids != -100, - rep_full_word_ids != rep_full_word_ids.roll(1), - ).cumsum(0) - 1, # -1 because we want to start from 0 - -100, - ) + # Pair ids are assigned per (batch row, head) slice so that identical + # word ids in neighbouring rows or heads never share a pair row. + full_word_ids, num_pairs = causal_word_pair_ids(word_ids, num_heads) valid_word_ids_mask = full_word_ids != -100 valid_word_ids = full_word_ids[valid_word_ids_mask] # (batch_size * num_heads * valid_rows * valid_cols, hidden_size) attv = attn_weights.unsqueeze(-1) * value_states.unsqueeze(-3) - attv = attv.view(-1, hidden_size)[valid_word_ids_mask, :] - - num_pairs = valid_word_ids.max() + 1 + attv = attv.reshape(-1, hidden_size)[valid_word_ids_mask, :] + + if num_pairs == 0: + return torch.empty( + 0, + hidden_size, + device=word_ids.device, + dtype=attv.dtype, + ) # (num_pairs, hidden_size) attv = torch.zeros( @@ -285,16 +321,12 @@ def get_aim_star_states( # (batch_size * num_heads, seq_len) -> (batch_size * num_heads * seq_len) rep_word_ids = word_ids.repeat_interleave(num_heads, dim=0) - rep_word_ids_flat = rep_word_ids.view(-1) - valid_word_mask = torch.logical_and( - rep_word_ids_flat != -100, # no padding or special tokens - rep_word_ids_flat != rep_word_ids_flat.roll(-1), # last token of the word - ) + valid_word_mask = last_token_mask(rep_word_ids) # (batch_size * num_heads * valid_rows, hidden_size) word_states = ( torch.matmul(attn_weights, value_states) - .view(-1, hidden_size)[valid_word_mask] + .reshape(-1, hidden_size)[valid_word_mask.reshape(-1)] ) return word_states diff --git a/tests/test_aim_impl.py b/tests/test_aim_impl.py new file mode 100644 index 0000000..6d95cc0 --- /dev/null +++ b/tests/test_aim_impl.py @@ -0,0 +1,47 @@ +import unittest + +import torch + +from matt.training.aim_impl import get_aim_states, get_aim_star_states + + +class AIMStatesTest(unittest.TestCase): + def test_aim_states_keep_same_word_ids_separate_across_batch_rows(self) -> None: + word_ids = torch.tensor([[0], [0]]) + attn_weights = torch.ones(2, 1, 1, 1) + value_states = torch.tensor([[[[1.0]]], [[[2.0]]]]) + + actual = get_aim_states(word_ids, attn_weights, value_states) + + torch.testing.assert_close(actual, torch.tensor([[1.0], [2.0]])) + + def test_aim_states_keep_same_word_ids_separate_across_heads(self) -> None: + word_ids = torch.tensor([[0]]) + attn_weights = torch.ones(1, 2, 1, 1) + value_states = torch.tensor([[[[1.0]], [[2.0]]]]) + + actual = get_aim_states(word_ids, attn_weights, value_states) + + torch.testing.assert_close(actual, torch.tensor([[1.0], [2.0]])) + + def test_aim_star_states_keep_same_word_ids_separate_across_batch_rows(self) -> None: + word_ids = torch.tensor([[0], [0]]) + attn_weights = torch.ones(2, 1, 1, 1) + value_states = torch.tensor([[[[1.0]]], [[[2.0]]]]) + + actual = get_aim_star_states(word_ids, attn_weights, value_states) + + torch.testing.assert_close(actual, torch.tensor([[1.0], [2.0]])) + + def test_aim_star_states_keep_same_word_ids_separate_across_heads(self) -> None: + word_ids = torch.tensor([[0]]) + attn_weights = torch.ones(1, 2, 1, 1) + value_states = torch.tensor([[[[1.0]], [[2.0]]]]) + + actual = get_aim_star_states(word_ids, attn_weights, value_states) + + torch.testing.assert_close(actual, torch.tensor([[1.0], [2.0]])) + + +if __name__ == "__main__": + unittest.main()