add unpad back in with created attn_mask

config.json

@@ -82,7 +82,7 @@
   "sliding_window": 128,
   "transformers_version": "4.44.1",
   "type_vocab_size": 2,
-  "unpad_embeddings": false,
+  "unpad_embeddings": true,
   "use_cache": true,
   "use_fa2": true,
   "use_sdpa_attn_mask": false,

modeling_flexbert.py

@@ -1643,6 +1643,8 @@ class FlexBertForCausalLM(FlexBertPreTrainedModel):
         return_dict = return_dict if return_dict is not None else self.config.use_return_dict
         if self.unpad_embeddings and (indices is None and cu_seqlens is None and max_seqlen is None):
             batch_size, seq_len = input_ids.shape[:2]
+            if attention_mask is None:  # Create causal mask (lower triangular)
+                attention_mask = torch.tril(torch.ones(batch, seqlen), diagonal=0)
             input_ids, indices, cu_seqlens, max_seqlen, position_ids, labels = self.unpad_inputs(
                 input_ids, attention_mask, position_ids, labels
             )