attn verificiation checks

Author: nitsanluke

fast_llm/layers/transformer/attention.py

@@ -182,11 +182,7 @@ def _attn_fused(
         ).view(b, self._local_head_groups, sq, self._local_heads_per_group, sk)
 
         attn_weights = attn_weights.to(torch.float32) * self._layer_index
-
-        attn_weights = attn_weights.transpose(2, 3)
         attn_weights = torch.where(mask, attn_weights, mask_value)
-        attn_weights = attn_weights.transpose(2, 3)
-
         attn_weights = torch.nn.functional.softmax(attn_weights, dim=-1).to(query.dtype)
 
         with set_generator(self._tensor_space.distributed.tp_generator):
@@ -417,13 +413,14 @@ def forward(self, input_: torch.Tensor, kwargs: dict[str, typing.Any]) -> tuple[
             diff = input_ - flash_input_
             # print(f"Element-wise difference: {diff.shape} {diff}")
             max_diff = diff.abs().max()
-            min_diff = diff.abs().min()
-            print(f"Min element-wise difference: {min_diff.item()}")
+            # min_diff = diff.abs().min()
+            # print(f"Min element-wise difference: {min_diff.item()}")
             print(f"Max element-wise difference: {max_diff.item()}")
-            # if max_diff > 1e-3:
-            #     print("Warning: Max difference exceeds 1e-3")
-            #     import sys
-            #     sys.exit(1)
+            if max_diff > 1e-3:
+                print("Warning: Max difference exceeds 1e-3")
+                import sys
+
+                sys.exit(1)
 
         if self._debug_transformer:
             self._debug_log(query, "query", self._QUERY_DIMS, kwargs)

fast_llm/models/gpt/model.py

@@ -57,7 +57,7 @@ def __init__(
         # TODO: Find a better solution.
         self._preprocessors.append(self._config.transformer.rotary.build(self._tensor_space))
 
-        if not self._config.transformer.diffusion:
+        if self._config.transformer.diffusion is None:
             if self._use_flash_attention:
                 self._preprocessors.append(FlashAttnVarlenPreprocessor(self._config.transformer, self._tensor_space))
             else:
@@ -355,12 +355,21 @@ def preprocess(
 
                         batch_size, seq_len = batch.token_ids.shape
                         seq_len -= 1  # last token is dropped inputs
+                        # attention_mask = torch.ones(
+                        #     (batch_size, 1, seq_len, seq_len),
+                        #     dtype=torch.bool,
+                        #     device=self._tensor_space.distributed.device,
+                        # )
+                        # kwargs[TransformerKwargs.attention_mask] = attention_mask.unsqueeze(1).unsqueeze(1)
                         attention_mask = torch.ones(
-                            (batch_size, 1, seq_len, seq_len),
+                            (seq_len, seq_len),
                             dtype=torch.bool,
                             device=self._tensor_space.distributed.device,
                         )
-                        kwargs[TransformerKwargs.attention_mask] = attention_mask.unsqueeze(1).unsqueeze(1)
+                        kwargs[TransformerKwargs.attention_mask] = attention_mask[
+                            None, None, 0:seq_len, None, :seq_len
+                        ]
+                        print(f"attention_mask: {kwargs[TransformerKwargs.attention_mask]}")
                         # # kwargs[TransformerKwargs.attention_mask_value] = torch.tensor(
                         # #     -10000.0, device=self._tensor_space.distributed.device
                         # # )