pad

modeling_flexbert.py

@@ -1721,6 +1721,11 @@ class FlexBertForCausalLM(FlexBertPreTrainedModel):
     ) -> dict:
         if attention_mask is None:
             attention_mask = torch.ones_like(input_ids)
+
+        # Calculate positions before unpadding
+        if position_ids is None:
+            position_ids = attention_mask.long().cumsum(-1) - 1
+            position_ids.masked_fill_(attention_mask == 0, 1)
             
         batch_size, seq_len = input_ids.shape[:2]
         input_ids, indices, cu_seqlens, max_seqlen, position_ids, _ = self.unpad_inputs(