Update modeling_llada.py

modeling_llada.py

@@ -30,7 +30,6 @@ from transformers import PreTrainedModel
 from transformers.modeling_outputs import CausalLMOutputWithPast
 from transformers.models.auto import AutoModel
 from transformers.cache_utils import Cache
-from transformers import AutoConfig
 
 from .configuration_llada import (
     LLaDAConfig,
@@ -637,7 +636,7 @@ class LLaDABlock(nn.Module):
         """
         if self.flash_attn_func is not None and attn_mask is None:
             r = self.flash_attn_func(
-                q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2), dropout_p=dropout_p, causal=is_causal
+                q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2), dropout_p=dropout_p, causal=False
             )
             return r.transpose(1, 2)
         else:
@@ -657,7 +656,7 @@ class LLaDABlock(nn.Module):
                 v,
                 attn_mask=None,
                 dropout_p=dropout_p,
-                # is_causal=False,
+                is_causal=False,
             )
 
     def attention(
@@ -713,9 +712,9 @@ class LLaDABlock(nn.Module):
             q,
             k,
             v,
-            attn_mask=attention_bias,
+            attn_mask=None,
             dropout_p=0.0 if not self.training else self.config.attention_dropout,
-            is_causal=attention_bias is None,
+            is_causal=False,
         )
 
         # Re-assemble all head outputs side-by-side.
@@ -1491,4 +1490,4 @@ class LLaDAModelLM(PreTrainedModel):
             self.model.transformer.ff_out = self.model.transformer.wte
 
 # Register the model so that it is available for transformer pipelines, auto-loading, etc.
-AutoModel.register(LLaDAConfig, LLaDAModelLM)
+AutoModel.register(LLaDAConfig, LLaDAModelLM)