same as training code

modeling_flexbert.py

@@ -1529,16 +1529,13 @@ class FlexBertForCausalLM(FlexBertPreTrainedModel):
         self.unpad_embeddings = config.unpad_embeddings
         self.pad_logits = config.pad_logits
         self.compile_model = config.compile_model
-        self.vocab_size = config.vocab_size
         # self.masked_prediction = config.masked_prediction
 
         # Initialize weights and apply final processing
         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)):
@@ -1552,7 +1549,7 @@ class FlexBertForCausalLM(FlexBertPreTrainedModel):
             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)
 
@@ -1640,27 +1637,22 @@ class FlexBertForCausalLM(FlexBertPreTrainedModel):
         #
         # Prediction scores are only computed for masked tokens and the (bs,
         # 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]        
-            if attention_mask is None:  
-                # unpad expects a encoder-like mask where all non-padding are ones
-                attention_mask = torch.ones_like(input_ids)
-                attention_mask[input_ids == 50283] = 0  # zero out pad tokens
+        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(
                 input_ids, attention_mask, position_ids, labels
             )
 
-        
         hidden_states = self.bert(
             input_ids,
-            attention_mask=None, # let FA handle it
+            attention_mask=None, # let FA do this
             position_ids=position_ids,
             indices=indices,
             cu_seqlens=cu_seqlens,
             max_seqlen=max_seqlen,
         )
-        # print(hidden_states.shape)
 
         if self.compile_model:
             logits = self.compiled_lm_head(hidden_states)
@@ -1673,26 +1665,24 @@ class FlexBertForCausalLM(FlexBertPreTrainedModel):
                 shift_labels = torch.full_like(input_ids, -100)
                 shift_labels[:-1] = input_ids[1:]
 
-                # Mask boundaries
+                # Mask boundaries, so eos doesn't predict bos
                 for i in range(len(cu_seqlens) - 1):
                     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
+                # NOTE: no padding or mask in there for now
+                assert 50283 not in shift_labels, f"PAD token found in shift_labels: {shift_labels}"
+                assert 50284 not in shift_labels, f"MASK token found in shift_labels: {shift_labels}"
+                assert shift_labels.shape == logits.shape[:-1] # Verify shapes align
                     
             else:
                 # Padded case: simple shift
                 shift_labels = input_ids[..., 1:].contiguous()
                 logits = logits[..., :-1, :].contiguous()
+                # mask out PAD tokens in the shift_labels
+                mask = (shift_labels == 50283)
+                shift_labels = torch.where(mask, torch.tensor(-100, device=shift_labels.device), shift_labels)
+                assert shift_labels.shape == logits.shape[:-1] # Verify shapes align
 
             # For both cases, we'll use the shifted input_ids as our labels
             labels = shift_labels
@@ -1703,26 +1693,14 @@ class FlexBertForCausalLM(FlexBertPreTrainedModel):
                 shift_labels.view(-1)
             )
 
-        if self.unpad_embeddings: # revert back to normal logits
-            logits = logits.view(batch_size, -1, self.vocab_size)
-
         if self.pad_logits:
-            # print(f"Padding logits: {logits.shape}")
-            new_logits = self.pad_inputs(logits, indices, batch_size, seq_len-1)[0]
-            # print(f"New logits: {new_logits.shape}")
-            # print(new_logits.shape)
-            # if new_logits.dim() == 2:
-            #     new_logits = new_logits.unsqueeze(0)
             return CausalLMOutput(
                 loss=loss,
-                logits=new_logits,
+                logits=self.pad_inputs(logits, indices, batch_size, seq_len)[0],
                 hidden_states=None,
                 attentions=None,
             )
         else:
-            # print(f"Non-padding logits: {logits.shape}")
-            # if logits.dim() == 2:
-            #     logits = logits.unsqueeze(0)
             return CausalLMOutput(
                 loss=loss,
                 logits=logits,
@@ -1947,4 +1925,4 @@ def init_mlm_model_from_pretrained(
             pretrained_model.decoder, new_model.decoder, linear_type=TileLinear.default, mode=mode, bias_only=True
         )
     else:
-        tile_linear(pretrained_model.decoder, new_model.decoder, linear_type=TileLinear.default, mode=mode)
+        tile_linear(pretrained_model.decoder, new_model.decoder, linear_type=TileLinear.default, mode=mode)

padding.py

@@ -84,4 +84,4 @@ def pad_input(
         padded_labels[indices] = labels
         padded_labels = padded_labels.view(batch, seqlen)
 
-    return padded_inputs, padded_labels
+    return padded_inputs, padded_labels