Change LoRA size from 256 to 512, also back to bidirectional_masked

llama_diffusion_model.py

@@ -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():
@@ -202,8 +202,8 @@ class CustomTransformerModel(PreTrainedModel):
             param.requires_grad = True
 
         lora_config = LoraConfig(
-            r=256,
-            lora_alpha=256,
+            r=512,
+            lora_alpha=512,
             lora_dropout=0.0,
             target_modules=["q_proj", "v_proj", "k_proj", "o_proj"],  # Llama-3 uses these attention modules
             bias="none",