Fix causal mode

modeling_gptbert.py

@@ -333,7 +333,7 @@ class MaskedSoftmax(torch.autograd.Function):
 
 
 class SelfAttention(nn.Module):
-    def __init__(self, config: GptBertConfig, layer_idx: int):
+    def __init__(self, config: GptBertConfig, layer_idx: int, is_decoder: bool):
         super().__init__()
 
         self.config = config
@@ -349,6 +349,8 @@ class SelfAttention(nn.Module):
         self.k_out_dim = self.d_qk * self.num_kv_heads
         self.v_out_dim = self.d_v * self.num_kv_heads
 
+        self.is_causal = is_decoder
+
         self.qk_proj = MultiCastedLinearOrthoIn(self.hidden_size, [self.q_out_dim, self.k_out_dim], bias=False)
         self.v_proj = CastedLinearIn(self.hidden_size, self.v_out_dim, bias=False)
         self.out_proj = CastedLinearIn(self.d_v*self.num_attention_heads, self.hidden_size, bias=False)
@@ -376,7 +378,6 @@ class SelfAttention(nn.Module):
         self.lambdas = nn.Parameter(torch.tensor([0.5]))
 
         self.sequence_length = config.max_sequence_length
-        self.is_causal = config.is_decoder
         self.window_length = None
 
     def set_window_length(self, window_length: int):
@@ -526,10 +527,10 @@ class FeedForward(nn.Module):
 
 
 class Layer(nn.Module):
-    def __init__(self, config: GptBertConfig, layer_idx: int):
+    def __init__(self, config: GptBertConfig, layer_idx: int, is_decoder: bool):
         super().__init__()
 
-        self.attention = SelfAttention(config, layer_idx)
+        self.attention = SelfAttention(config, layer_idx, is_decoder)
         self.mlp = FeedForward(config)
         self.lambdas = nn.Parameter(torch.tensor([0., 0., 1., 0., 1., 0.]))
 
@@ -550,9 +551,9 @@ class Layer(nn.Module):
 
 
 class Encoder(nn.Module):
-    def __init__(self, config: GptBertConfig):
+    def __init__(self, config: GptBertConfig, is_decoder: bool):
         super().__init__()
-        self.layers = nn.ModuleList([Layer(config, i) for i in range(config.num_layers)])
+        self.layers = nn.ModuleList([Layer(config, i, is_decoder) for i in range(config.num_layers)])
         self.local_global_ratio = config.local_global_ratio
 
     def set_window_length(self, config: GptBertConfig):
@@ -613,9 +614,10 @@ class GptBertModel(GptBertPreTrainedModel):
         super().__init__(config, **kwargs)
         self.config = config
         self.hidden_size = config.hidden_size
+        self.is_decoder = self.is_decoder if hasattr(self, "is_decoder") else False
 
         self.embedding = Embedding(config)
-        self.encoder = Encoder(config)
+        self.encoder = Encoder(config, self.is_decoder)
         self.classifier = LMClassifier(config, config.vocab_size) if add_mlm_layer else None
         self.set_window_length(config)
         self.gradient_checkpointing = False
@@ -718,6 +720,7 @@ class GptBertForMaskedLM(GptBertModel):
     _tied_weights_keys = ["classifier.emb2vocab.weight"]
 
     def __init__(self, config: GptBertConfig, **kwargs):
+        self.is_decoder = False
         super().__init__(config, add_mlm_layer=True, **kwargs)
 
     def get_output_embeddings(self):
@@ -769,7 +772,7 @@ class GptBertForCausalLM(GptBertModel):
     _tied_weights_keys = ["classifier.emb2vocab.weight"]
 
     def __init__(self, config: GptBertConfig, **kwargs):
-        config.is_decoder = True
+        self.is_decoder = True
         super().__init__(config, add_mlm_layer=True, **kwargs)
 
     def get_output_embeddings(self):
@@ -886,6 +889,7 @@ class GptBertForSequenceClassification(GptBertModel):
     _keys_to_ignore_on_load_unexpected = ["classifier.emb2vocab.weight", "classifier.emb2vocab.bias"]
 
     def __init__(self, config: GptBertConfig, **kwargs):
+        self.is_decoder = False
         super().__init__(config, add_mlm_layer=False, **kwargs)
 
         self.num_labels = config.num_labels
@@ -941,6 +945,7 @@ class GptBertForTokenClassification(GptBertModel):
     _keys_to_ignore_on_load_unexpected = ["classifier.emb2vocab.weight", "classifier.emb2vocab.bias"]
 
     def __init__(self, config: GptBertConfig, **kwargs):
+        self.is_decoder = False
         super().__init__(config, add_mlm_layer=False, **kwargs)
 
         self.num_labels = config.num_labels
@@ -978,6 +983,7 @@ class GptBertForQuestionAnswering(GptBertModel):
     _keys_to_ignore_on_load_unexpected = ["classifier.emb2vocab.weight", "classifier.emb2vocab.bias"]
 
     def __init__(self, config: GptBertConfig, **kwargs):
+        self.is_decoder = False
         super().__init__(config, add_mlm_layer=False, **kwargs)
 
         self.num_labels = config.num_labels
@@ -1034,6 +1040,7 @@ class GptBertForMultipleChoice(GptBertModel):
     _keys_to_ignore_on_load_unexpected = ["classifier.emb2vocab.weight", "classifier.emb2vocab.bias"]
 
     def __init__(self, config: GptBertConfig, **kwargs):
+        self.is_decoder = False
         super().__init__(config, add_mlm_layer=False, **kwargs)
 
         self.num_labels = getattr(config, "num_labels", 2)