From fb9ba349b5842bc45bed4dbc85ac5ee64b116183 Mon Sep 17 00:00:00 2001 From: Yevhen Bochkov Date: Wed, 12 Aug 2026 17:50:57 +0300 Subject: [PATCH] Fix AIM row alignment across batch and head boundaries get_aim_states and get_aim_star_states flattened (batch, head) rows and detected word boundaries with roll() and a global cumsum(). Word ids touching a row boundary (or the wrap-around between the tensor's ends) could collide, merging word pairs from unrelated rows or dropping final words. Teacher and student AIM states are built independently, so any merged or dropped row shifted every following row and broke their one-to-one correspondence in the loss. Assign pair ids per (batch row, head) slice instead, and return an empty state tensor when a batch contains no complete word pair rather than crashing on max() of an empty tensor. Add regression tests that keep identical word ids separate across batch rows and across attention heads for both AIM and AIM*. Co-Authored-By: Claude Fable 5 --- matt/training/aim_impl.py | 122 ++++++++++++++++++++++++-------------- tests/test_aim_impl.py | 47 +++++++++++++++ 2 files changed, 124 insertions(+), 45 deletions(-) create mode 100644 tests/test_aim_impl.py 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()