added updated code:

__init__.py

@@ -33,13 +33,14 @@ from .modeling_flexbert import (
     FlexBertForMaskedLM,
     FlexBertForSequenceClassification,
     FlexBertForMultipleChoice,
-    FlexBertForCasualLM,
+    FlexBertForCausalLM,
 )
 from .bert_padding import(
     IndexFirstAxis,
     IndexPutFirstAxis
 )
 
+
 __all__ = [
     "BertAlibiEmbeddings",
     "BertAlibiEncoder",
@@ -69,6 +70,5 @@ __all__ = [
     "FlexBertForMaskedLM",
     "FlexBertForSequenceClassification",
     "FlexBertForMultipleChoice",
-    "IndexFirstAxis",
-    "IndexPutFirstAxis"
+    "FlexBertForCausalLM"
 ]

attention.py

@@ -74,7 +74,7 @@ class BertAlibiUnpadSelfAttention(nn.Module):
                 f"heads ({config.num_attention_heads})"
             )
 
-        self.is_casual = config.casual_mask
+        self.is_causal = config.causal_mask
         self.num_attention_heads = config.num_attention_heads
         self.attention_head_size = int(config.hidden_size / config.num_attention_heads)
         self.all_head_size = self.num_attention_heads * self.attention_head_size
@@ -145,7 +145,7 @@ class BertAlibiUnpadSelfAttention(nn.Module):
                     dropout_p=self.p_dropout,
                     deterministic=self.deterministic_fa2,
                     alibi_slopes=slopes,
-                    casual=self.is_casual
+                    causal=self.is_causal
                 )
                 attention = attention.to(orig_dtype)  # type: ignore
             else:
@@ -156,10 +156,11 @@ class BertAlibiUnpadSelfAttention(nn.Module):
                     dropout_p=self.p_dropout,
                     deterministic=self.deterministic_fa2,
                     alibi_slopes=slopes,
-                    casual = self.is_casual
+                    causal = self.is_causal
                 )
         else:
-            assert not self.is_casual, f"Casual mask not implemented here yet"
+            assert not self.is_causal, f"causal mask not implemented here yet"
+            assert False
             qkv = bert_padding.pad_input(qkv, indices, cu_seqlens.shape[0] - 1, max_seqlen)  # batch, max_seqlen, thd
             unpad_bs, *_ = qkv.shape
             qkv = qkv.view(unpad_bs, -1, 3, self.num_attention_heads, self.attention_head_size)
@@ -236,6 +237,7 @@ class BertAlibiUnpadAttention(nn.Module):
             slopes: None or (batch, heads) or (heads,)
         """
         assert (bias is None) == (slopes is None), f"{bias=}, {slopes=}"
+        assert False
         self_output = self.self(input_tensor, cu_seqlens, max_s, indices, attn_mask, bias, slopes)
         if subset_idx is not None:
             return self.output(
@@ -293,7 +295,7 @@ class FlexBertUnpadAttention(FlexBertAttentionBase):
                 f"heads ({config.num_attention_heads})"
             )
 
-        self.is_casual = config.casual_mask
+        self.is_causal = config.causal_mask
         self.num_attention_heads = config.num_attention_heads
         self.attn_head_size = int(config.hidden_size / config.num_attention_heads)
         self.all_head_size = self.num_attention_heads * self.attn_head_size
@@ -402,7 +404,7 @@ class FlexBertUnpadAttention(FlexBertAttentionBase):
                     dropout_p=self.p_dropout,
                     deterministic=self.deterministic_fa2,
                     window_size=self.sliding_window,
-                    casual=self.is_casual
+                    causal=self.is_causal
                 )
                 attn = attn.to(orig_dtype)  # type: ignore
             else:
@@ -413,11 +415,12 @@ class FlexBertUnpadAttention(FlexBertAttentionBase):
                     dropout_p=self.p_dropout,
                     deterministic=self.deterministic_fa2,
                     window_size=self.sliding_window,
-                    casual=self.is_casual
+                    causal=self.is_causal
                 )
             attn = attn.view(bs, dim)
         else:
-            assert not self.is_casual, f"Casual mask not implemented here yet"
+            assert not self.is_causal, f"causal mask not implemented here yet"
+            assert False
             qkv = bert_padding.pad_input(qkv, indices, cu_seqlens.shape[0] - 1, max_seqlen)  # batch, max_seqlen, thd
             unpad_bs, seqlen, _ = qkv.shape
 
@@ -456,7 +459,7 @@ class FlexBertUnpadParallelAttention(FlexBertAttentionBase):
                 f"heads ({config.num_attention_heads})"
             )
 
-        self.is_casual = config.casual_mask
+        self.is_causal = config.causal_mask
         self.num_attention_heads = config.num_attention_heads
         self.attn_head_size = int(config.hidden_size / config.num_attention_heads)
         self.hidden_size = config.hidden_size
@@ -556,7 +559,7 @@ class FlexBertUnpadParallelAttention(FlexBertAttentionBase):
                     dropout_p=self.p_dropout,
                     deterministic=self.deterministic_fa2,
                     window_size=self.sliding_window,
-                    casual=self.is_casual
+                    causal=self.is_causal
                 )
                 attn = attn.to(orig_dtype)  # type: ignore
             else:
@@ -567,11 +570,12 @@ class FlexBertUnpadParallelAttention(FlexBertAttentionBase):
                     dropout_p=self.p_dropout,
                     deterministic=self.deterministic_fa2,
                     window_size=self.sliding_window,
-                    casual=self.is_casual
+                    causal=self.is_causal
                 )
             attn = attn.view(bs, dim)
         else:
-            assert not self.is_casual, f"Casual mask not implemented here yet"
+            assert not self.is_causal, f"causal mask not implemented here yet"
+            assert False
             qkv = bert_padding.pad_input(qkv, indices, cu_seqlens.shape[0] - 1, max_seqlen)  # batch, max_seqlen, thd
             unpad_bs, seqlen, _ = qkv.shape
 
@@ -610,7 +614,7 @@ class FlexBertPaddedAttention(FlexBertAttentionBase):
                 f"heads ({config.num_attention_heads})"
             )
 
-        self.is_casual = config.casual_mask
+        self.is_causal = config.causal_mask
         self.num_attention_heads = config.num_attention_heads
         self.attn_head_size = int(config.hidden_size / config.num_attention_heads)
         self.all_head_size = self.num_attention_heads * self.attn_head_size
@@ -695,7 +699,7 @@ class FlexBertPaddedAttention(FlexBertAttentionBase):
                     dropout_p=self.p_dropout,
                     deterministic=self.deterministic_fa2,
                     window_size=self.sliding_window,
-                    casual=self.is_casual
+                    causal=self.is_causal
                 )
                 attn = attn.to(orig_dtype)  # type: ignore
             else:
@@ -704,10 +708,11 @@ class FlexBertPaddedAttention(FlexBertAttentionBase):
                     dropout_p=self.p_dropout,
                     deterministic=self.deterministic_fa2,
                     window_size=self.sliding_window,
-                    casual=self.is_casual
+                    causal=self.is_causal
                 )
         else:
-            assert not self.is_casual, f"Casual mask not implemented here yet"
+            assert not self.is_causal, f"causal mask not implemented here yet"
+            assert False
             qkv = qkv.view(bs, seqlen, 3, self.num_attention_heads, self.attn_head_size)
 
             q, k, v = qkv.transpose(3, 1).unbind(dim=2)
@@ -743,7 +748,7 @@ class FlexBertUnpadRopeAttention(FlexBertAttentionBase):
                 f"heads ({config.num_attention_heads})"
             )
 
-        self.is_casual = config.casual_mask
+        self.is_causal = config.causal_mask
         self.num_attention_heads = config.num_attention_heads
         self.attn_head_size = int(config.hidden_size / config.num_attention_heads)
         self.all_head_size = self.num_attention_heads * self.attn_head_size
@@ -882,7 +887,7 @@ class FlexBertUnpadRopeAttention(FlexBertAttentionBase):
                     max_seqlen_q=max_seqlen,
                     max_seqlen_k=max_seqlen,
                     deterministic=self.deterministic_fa2,
-                    causal=self.is_casual,
+                    causal=self.is_causal,
                 )
                 attn = attn.to(orig_dtype)  # type: ignore
             else:
@@ -896,7 +901,7 @@ class FlexBertUnpadRopeAttention(FlexBertAttentionBase):
                     max_seqlen_q=max_seqlen,
                     max_seqlen_k=max_seqlen,
                     deterministic=self.deterministic_fa2,
-                    causal=self.is_casual,
+                    causal=self.is_causal,
                 )
             attn = attn.view(bs, dim)
         elif self.use_fa2:
@@ -914,7 +919,7 @@ class FlexBertUnpadRopeAttention(FlexBertAttentionBase):
                     dropout_p=self.p_dropout,
                     deterministic=self.deterministic_fa2,
                     window_size=self.sliding_window,
-                    causal=self.is_casual,
+                    causal=self.is_causal,
                 )
                 attn = attn.to(orig_dtype)  # type: ignore
             else:
@@ -925,11 +930,12 @@ class FlexBertUnpadRopeAttention(FlexBertAttentionBase):
                     dropout_p=self.p_dropout,
                     deterministic=self.deterministic_fa2,
                     window_size=self.sliding_window,
-                    causal=self.is_casual,
+                    causal=self.is_causal,
                 )
             attn = attn.view(bs, dim)
         else:
-            assert not self.is_casual, f"Casual mask not implemented here yet"
+            assert not self.is_causal, f"causal mask not implemented here yet"
+            assert False
             qkv = bert_padding.pad_input(
                 qkv, indices, cu_seqlens.shape[0] - 1, attn_mask.shape[-1]
             )  # batch, max_seqlen, thd
@@ -969,7 +975,7 @@ class FlexBertPaddedRopeAttention(FlexBertAttentionBase):
                 f"heads ({config.num_attention_heads})"
             )
 
-        self.is_casual = config.casual_mask
+        self.is_causal = config.causal_mask
         self.num_attention_heads = config.num_attention_heads
         self.attn_head_size = int(config.hidden_size / config.num_attention_heads)
         self.all_head_size = self.num_attention_heads * self.attn_head_size
@@ -1080,7 +1086,7 @@ class FlexBertPaddedRopeAttention(FlexBertAttentionBase):
                     dropout_p=self.p_dropout,
                     deterministic=self.deterministic_fa2,
                     window_size=self.sliding_window,
-                    casual=self.is_casual,
+                    causal=self.is_causal,
                 )
                 attn = attn.to(orig_dtype)  # type: ignore
             else:
@@ -1089,10 +1095,11 @@ class FlexBertPaddedRopeAttention(FlexBertAttentionBase):
                     dropout_p=self.p_dropout,
                     deterministic=self.deterministic_fa2,
                     window_size=self.sliding_window,
-                    casual=self.is_casual
+                    causal=self.is_causal
                 )
         else:
-            assert not self.is_casual, f"Casual mask not implemented here yet"
+            assert not self.is_causal, f"causal mask not implemented here yet"
+            assert False
             qkv = self.rotary_emb(qkv, seqlen_offset=seqlen_offset, max_seqlen=None)
             q, k, v = qkv.transpose(3, 1).unbind(dim=2)
             attn = F.scaled_dot_product_attention(
@@ -1127,7 +1134,7 @@ class FlexBertUnpadRopeParallelAttention(FlexBertAttentionBase):
                 f"heads ({config.num_attention_heads})"
             )
 
-        self.is_casual = config.casual_mask
+        self.is_causal = config.causal_mask
         self.num_attention_heads = config.num_attention_heads
         self.attn_head_size = int(config.hidden_size / config.num_attention_heads)
         self.hidden_size = config.hidden_size
@@ -1253,7 +1260,7 @@ class FlexBertUnpadRopeParallelAttention(FlexBertAttentionBase):
                     dropout_p=self.p_dropout,
                     deterministic=self.deterministic_fa2,
                     window_size=self.sliding_window,
-                    casual=self.is_casual,
+                    causal=self.is_causal,
                 )
                 attn = attn.to(orig_dtype)  # type: ignore
             else:
@@ -1264,11 +1271,12 @@ class FlexBertUnpadRopeParallelAttention(FlexBertAttentionBase):
                     dropout_p=self.p_dropout,
                     deterministic=self.deterministic_fa2,
                     window_size=self.sliding_window,
-                    casual=self.is_casual,
+                    causal=self.is_causal,
                 )
             attn = attn.view(bs, dim)
         else:
-            assert not self.is_casual, f"Casual mask not implemented here yet"
+            assert not self.is_causal, f"causal mask not implemented here yet"
+            assert False
             qkv = bert_padding.pad_input(
                 qkv, indices, cu_seqlens.shape[0] - 1, attn_mask.shape[-1]
             )  # batch, max_seqlen, thd
@@ -1308,7 +1316,7 @@ class FlexBertPaddedRopeParallelAttention(FlexBertAttentionBase):
                 f"heads ({config.num_attention_heads})"
             )
 
-        self.is_casual = config.casual_mask
+        self.is_causal = config.causal_mask
         self.num_attention_heads = config.num_attention_heads
         self.attn_head_size = int(config.hidden_size / config.num_attention_heads)
         self.hidden_size = config.hidden_size
@@ -1413,7 +1421,7 @@ class FlexBertPaddedRopeParallelAttention(FlexBertAttentionBase):
                     dropout_p=self.p_dropout,
                     deterministic=self.deterministic_fa2,
                     window_size=self.sliding_window,
-                    casual=self.is_casual
+                    causal=self.is_causal
                 )
                 attn = attn.to(orig_dtype)  # type: ignore
             else:
@@ -1422,10 +1430,11 @@ class FlexBertPaddedRopeParallelAttention(FlexBertAttentionBase):
                     dropout_p=self.p_dropout,
                     deterministic=self.deterministic_fa2,
                     window_size=self.sliding_window,
-                    casual=self.is_casual
+                    causal=self.is_causal
                 )
         else:
-            assert not self.is_casual, f"Casual mask not implemented here yet"
+            assert not self.is_causal, f"causal mask not implemented here yet"
+            assert False
             qkv = self.rotary_emb(qkv, seqlen_offset=seqlen_offset, max_seqlen=None)
             q, k, v = qkv.transpose(3, 1).unbind(dim=2)
             attn = F.scaled_dot_product_attention(
@@ -1460,7 +1469,7 @@ class FlexBertPaddedParallelAttention(FlexBertAttentionBase):
                 f"heads ({config.num_attention_heads})"
             )
 
-        self.is_casual = config.casual_mask
+        self.is_causal = config.causal_mask
         self.num_attention_heads = config.num_attention_heads
         self.attn_head_size = int(config.hidden_size / config.num_attention_heads)
         self.hidden_size = config.hidden_size
@@ -1537,7 +1546,7 @@ class FlexBertPaddedParallelAttention(FlexBertAttentionBase):
                     dropout_p=self.p_dropout,
                     deterministic=self.deterministic_fa2,
                     window_size=self.sliding_window,
-                    casual=self.is_casual
+                    causal=self.is_causal
                 )
                 attn = attn.to(orig_dtype)  # type: ignore
             else:
@@ -1546,10 +1555,11 @@ class FlexBertPaddedParallelAttention(FlexBertAttentionBase):
                     dropout_p=self.p_dropout,
                     deterministic=self.deterministic_fa2,
                     window_size=self.sliding_window,
-                    casual=self.is_casual
+                    causal=self.is_causal
                 )
         else:
-            assert not self.is_casual, f"Casual attention mask not yet implemented here"
+            assert not self.is_causal, f"causal attention mask not yet implemented here"
+            assert False
             qkv = qkv.view(bs, seqlen, 3, self.num_attention_heads, self.attn_head_size)
             q, k, v = qkv.transpose(3, 1).unbind(dim=2)  # b h s d
             attn = F.scaled_dot_product_attention(

config.json

@@ -88,4 +88,4 @@
   "use_sdpa_attn_mask": false,
   "vocab_size": 50368,
   "is_casual": true
-}
+}

configuration_bert.py

@@ -97,7 +97,7 @@ class FlexBertConfig(TransformersBertConfig):
         pad_logits: bool = False,
         compile_model: bool = False,
         masked_prediction: bool = False,
-        casual_mask: bool = False,
+        causal_mask: bool = False,
         **kwargs,
     ):
         """
@@ -157,7 +157,7 @@ class FlexBertConfig(TransformersBertConfig):
             pad_logits (bool): Pad logits after the calculating the loss.
             compile_model (bool): Compile the subset of the model which can be compiled.
             masked_prediction (bool): Use only pass the masked tokens throught the final MLM layers
-            casual_mask (bool): Use a casual mask, defaulting to false.
+            causal (bool): Use a causal mask, defaulting to false.
             **kwargs: Additional keyword arguments.
         """
         super().__init__(attention_probs_dropout_prob=attention_probs_dropout_prob, **kwargs)
@@ -215,7 +215,7 @@ class FlexBertConfig(TransformersBertConfig):
         self.pad_logits = pad_logits
         self.compile_model = compile_model
         self.masked_prediction = masked_prediction
-        self.casual_mask = casual_mask
+        self.causal_mask = causal_mask
 
         if loss_kwargs.get("return_z_loss", False):
             if loss_function != "fa_cross_entropy":

modeling_flexbert.py

@@ -125,7 +125,6 @@ from .rotary import UnpaddedRotaryEmbedding
 
 logger = logging.getLogger(__name__)
 
-
 def _count_parameters(model: nn.Module, trainable: bool = True) -> int:
     if trainable:
         return sum(p.numel() for p in model.parameters() if p.requires_grad)
@@ -873,7 +872,7 @@ class FlexBertPreTrainedModel(BertPreTrainedModel):
 
     def _init_module_weights(self, module: nn.Module):
         """
-        Custom weight init of modules using .bert_layers.initialization.init_weights
+        Custom weight init of modules using src.bert_layers.initialization.init_weights
         Currently only supports init of embedding modules
         """
         assert isinstance(module, nn.Module)
@@ -1126,7 +1125,6 @@ class FlexBertForMaskedLM(FlexBertPreTrainedModel):
         # seqlen) dimensions are flattened
 
         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]
             input_ids, indices, cu_seqlens, max_seqlen, position_ids, labels = self.unpad_inputs(
@@ -1506,9 +1504,7 @@ class FlexBertForMultipleChoice(FlexBertPreTrainedModel):
         return params
 
 
-class FlexBertForCasualLM(FlexBertPreTrainedModel):
-    config_class = FlexBertConfig 
-    
+class FlexBertForCausalLM(FlexBertPreTrainedModel):
     """Bert Model transformer with a LM head.
 
     This head is just a standard LM head module. Used for causal language modeling tasks.
@@ -1538,23 +1534,14 @@ class FlexBertForCasualLM(FlexBertPreTrainedModel):
         self._init_weights(reset_params=False)
 
     def _init_weights(self, module: Optional[nn.Module] = None, reset_params: Optional[bool] = None):
-        # Handle the XOR condition
         assert (module is None) != (reset_params is None), "arg module xor reset_params must be specified"
-        
-        if module is not None:
-            # Add basic initialization for common module types
-            if isinstance(module, (nn.Linear, nn.Embedding)):
-                module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
-                if isinstance(module, nn.Linear) and module.bias is not None:
-                    module.bias.data.zero_()
-            elif isinstance(module, nn.LayerNorm):
-                module.bias.data.zero_()
-                module.weight.data.fill_(1.0)
+        if module:
+            self._init_module_weights(module)
         else:
             assert isinstance(reset_params, bool)
             self.bert._init_weights(reset_params=reset_params)
             self.lm_head._init_weights(reset_params=reset_params)
-            
+
             if not self.config.tie_word_embeddings:
                 init_weights(self.config, self.decoder, self.config.hidden_size, type_of_module=ModuleType.final_out)
 
@@ -1644,7 +1631,6 @@ class FlexBertForCasualLM(FlexBertPreTrainedModel):
         # seqlen) dimensions are flattened
 
         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]
             input_ids, indices, cu_seqlens, max_seqlen, position_ids, labels = self.unpad_inputs(
@@ -1664,29 +1650,28 @@ class FlexBertForCasualLM(FlexBertPreTrainedModel):
             logits = self.compiled_lm_head(hidden_states)
         else:
             logits = self.lm_head(hidden_states)
-        
+
         loss = None
         if labels is not None:
-            if indices is not None:                
-                # Unpadded case: shift within each sequence using input_ids
-                # Initialize shifted labels from input_ids
+            if cu_seqlens is not None:                
                 shift_labels = torch.full_like(input_ids, -100)
-                
-                # For each sequence, shift the input_ids to create labels
+                shift_labels[:-1] = input_ids[1:]
+
+                # Mask boundaries
                 for i in range(len(cu_seqlens) - 1):
-                    start = cu_seqlens[i]
-                    end = cu_seqlens[i + 1]
-                    # Input: [A, B, C, D] -> Labels: [B, C, D, -100]
-                    shift_labels[start:end-1] = input_ids[start+1:end]
-                        
-                # Debug prints
-                # print(f"input_ids slice: {input_ids[:20]}")  # Show first 20 tokens
-                # print(f"shift_labels slice: {shift_labels[:20]}")  # Show first 20 token
-                        
-                # # Debug prints
-                # print(f"input_ids slice: {input_ids[:20]}")  # Show first 20 tokens
-                # print(f"shift_labels slice: {shift_labels[:20]}")  # Show first 20 tokens
-                # print(f"First sequence length: {cu_seqlens[1] - cu_seqlens[0]}")
+                    boundary_pos = cu_seqlens[i+1] - 1
+                    shift_labels[boundary_pos] = -100
+                
+                # Mask out PAD tokens
+                mask = (shift_labels == 50283)
+                shift_labels = torch.where(mask, torch.tensor(-100, device=shift_labels.device), shift_labels)
+            
+
+            # print input_ids[(cu_seqlens[2]+1)-5:(cu_seqlens[2]+1)+5]
+            # print shift_labels[(cu_seqlens[2]+1)-5:(cu_seqlens[2]+1)+5]
+            # print input_ids[(cu_seqlens[-2]+1)-5:(cu_seqlens[-2]+1)+5]
+            # print shift_labels[(cu_seqlens[-2]+1)-5:(cu_seqlens[-2]+1)+5]
+            # breakpoint() # pkill -u oweller2 -f wandb
                     
             else:
                 # Padded case: simple shift
@@ -1703,7 +1688,7 @@ class FlexBertForCasualLM(FlexBertPreTrainedModel):
             )
 
         if self.pad_logits:
-            print(f"Padding logits: {logits.shape}")
+            # print(f"Padding logits: {logits.shape}")
             new_logits = self.pad_inputs(logits, indices, batch_size, seq_len)[0]
             if len(new_logits.shape) == 2:
                 new_logits = new_logits.unsqueeze(0)
@@ -1714,7 +1699,7 @@ class FlexBertForCasualLM(FlexBertPreTrainedModel):
                 attentions=None,
             )
         else:
-            print(f"Non-padding logits: {logits.shape}")
+            # print(f"Non-padding logits: {logits.shape}")
             if len(logits.shape) == 2:
                 logits = logits.unsqueeze(0)
             return CausalLMOutput(
@@ -1757,7 +1742,6 @@ class FlexBertForCasualLM(FlexBertPreTrainedModel):
         params += _count_parameters(self.lm_head, trainable)
         return params
 
-FlexBertForCasualLM.register_for_auto_class("AutoModelForCausalLM")
 
 def init_model_from_pretrained(
     pretrained_model: FlexBertModel,