Update modeling_dbrx.py

modeling_dbrx.py

@@ -1,3 +1,4 @@
+# code adapted from https://huggingface.co/fahadh4ilyas
 """PyTorch Dbrx model."""
 
 import math
@@ -244,28 +245,13 @@ def resolve_ffn_act_fn(
 # Copied from LLaMaAttention
 #############################################################################
 
-def get_max_seqlen_in_batch(attention_mask):
-    max_num = torch.max(attention_mask)
-    # attention_mask: B x N
-    counts = []
-    for i in range(1, max_num + 1):
-        counts.append(
-            torch.sum(attention_mask == i, axis=-1)
-        )  # shape: B, count length of data point maksed with i
-    result = torch.stack(counts, axis=1)
-    result = result.flatten()
-    return result[result.nonzero()].squeeze(-1).to(dtype=torch.int32)
-
-
-def _get_unpad_data(attention_mask):
-    seqlens_in_batch = get_max_seqlen_in_batch(
-        attention_mask
-    )  # attention_mask.sum(dim=-1, dtype=torch.int32)
+
+def _get_unpad_data(attention_mask: torch.Tensor):
+    seqlens_in_batch = attention_mask.sum(dim=-1, dtype=torch.int32)
     indices = torch.nonzero(attention_mask.flatten(), as_tuple=False).flatten()
     max_seqlen_in_batch = seqlens_in_batch.max().item()
-    cu_seqlens = F.pad(
-        torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.torch.int32), (1, 0)
-    )
+    cu_seqlens = F.pad(torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.int32),
+                       (1, 0))
     return (
         indices,
         cu_seqlens,
@@ -426,7 +412,7 @@ class DbrxFlashAttention2(DbrxAttention):
         **kwargs: Any,
     ) -> Tuple[torch.Tensor, Optional[torch.Tensor],
                Optional[Tuple[torch.Tensor]]]:
-        logger.info(
+        logger.debug(
             'Implicitly setting `output_attentions` to False as it is not supported in Flash Attention.'
         )
         output_attentions = False
@@ -1459,4 +1445,4 @@ class DbrxForCausalLM(DbrxPreTrainedModel):
             reordered_past += (tuple(
                 past_state.index_select(0, beam_idx.to(past_state.device))
                 for past_state in layer_past),)
-        return reordered_past
+        return reordered_past