fix

attention.py

@@ -863,7 +863,7 @@ class FlexBertUnpadRopeAttention(FlexBertAttentionBase):
         qkv = self.Wqkv(hidden_states)
 
         # only needed for inference when we have KV cache
-        seqlen_offset = max_seqlen * (cu_seqlens[0].item() // max_seqlen)
+        seqlen_offset = 0
 
         # (total_seqlen, 3, nheads, headdim)
         qkv = qkv.view(-1, 3, self.num_attention_heads, self.attn_head_size)

config.json

@@ -74,7 +74,7 @@
   "padding": "unpadded",
   "pooling_type": "cls",
   "position_embedding_type": "absolute",
-  "rotary_emb_base": 10000.0,
+  "rotary_emb_base": 1000.0,
   "rotary_emb_dim": 64,
   "rotary_emb_interleaved": false,
   "rotary_emb_scale_base": null,

modeling_flexbert.py

@@ -1733,7 +1733,6 @@ class FlexBertForCausalLM(FlexBertPreTrainedModel):
         input_ids, indices, cu_seqlens, max_seqlen, position_ids, _ = self.unpad_inputs(
             input_ids, attention_mask, position_ids, None
         )
-        breakpoint()
         return {
             "input_ids": input_ids,
             "attention_mask": attention_mask,