Try out bidirectional_masked prediction

llama_diffusion_model.py

@@ -126,7 +126,7 @@ class BidirectionalLlamaAttention(LlamaAttention):
             attn_mask = base_mask.unsqueeze(0).unsqueeze(1).expand(batch_size, 1, seq_len, seq_len).clone()  # ✅ Copy for each batch
         elif self.masking == 'bidirectional_masked':
             base_mask = torch.ones((seq_len, seq_len), device=hidden_states.device, dtype=torch.bool)
-            base_mask[:, 1:].fill_diagonal_(False)  # ✅ Apply diagonal masking only in 2D
+            base_mask[:, :].fill_diagonal_(False)  # ✅ Apply diagonal masking only in 2D
             attn_mask = base_mask.unsqueeze(0).unsqueeze(1).expand(batch_size, 1, seq_len, seq_len).clone()  # ✅ Copy for each batch
         else: # unidirectional
             # 🚀 Standard autoregressive (causal) mask
@@ -192,7 +192,7 @@ class CustomTransformerModel(PreTrainedModel):
         self.llama.resize_token_embeddings(config.vocab_size)
 
         for i, layer in enumerate(self.llama.model.layers):
-            layer.self_attn = BidirectionalLlamaAttention(layer.self_attn, masking='bidirectional')
+            layer.self_attn = BidirectionalLlamaAttention(layer.self_attn, masking='bidirectional_masked')
 
         # Freeze Llama to retain pre-trained knowledge
         for param in self.llama.parameters():