fix

modeling_flexbert.py

@@ -1722,10 +1722,12 @@ class FlexBertForCausalLM(FlexBertPreTrainedModel):
         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)
+        # Calculate sequence-local positions
+        seqlens = attention_mask.sum(dim=-1)  # Get length of each sequence
+        position_ids = torch.zeros_like(input_ids)
+        for i in range(len(seqlens)):
+            position_ids[i, :seqlens[i]] = torch.arange(seqlens[i], device=input_ids.device)
+            
             
         batch_size, seq_len = input_ids.shape[:2]
         input_ids, indices, cu_seqlens, max_seqlen, position_ids, _ = self.unpad_inputs(