update model

modeling_flexbert.py

@@ -1654,7 +1654,7 @@ class FlexBertForCausalLM(FlexBertPreTrainedModel):
         
         hidden_states = self.bert(
             input_ids,
-            attention_mask=None,
+            attention_mask=None, # let FA handle it
             position_ids=position_ids,
             indices=indices,
             cu_seqlens=cu_seqlens,
@@ -1703,11 +1703,8 @@ class FlexBertForCausalLM(FlexBertPreTrainedModel):
                 shift_labels.view(-1)
             )
 
-        if self.unpad_embeddings:
-            try:
-                logits = logits.view(batch_size, -1, self.vocab_size)
-            except Exception as e:
-                breakpoint()
+        if self.unpad_embeddings: # revert back to normal logits
+            logits = logits.view(batch_size, -1, self.vocab_size)
 
         if self.pad_logits:
             # print(f"Padding logits: {logits.shape}")

pytorch_model.bin

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:fea155de40c6fd0d7f58f431f493c5f614a64dabe168b72cbc74421a9bf17baf
+oid sha256:7863cc4c58494c661ffd3c77af90796f5caa8217917c2f6e7c99cc28d65b58c2
 size 598685038