Fix modeling_gpt1.py

Fix an issue with the attention mask, where its size would not be
correct during training and inference.

modeling_gpt1.py

@@ -154,6 +154,7 @@ class GPT1Model(GPT1PreTrainedModel):
         self.register_buffer('causal_mask',
                              torch.triu(causal_mask, diagonal=1),
                              persistent=False)
+        self.mask_cache_len = config.max_position_embeddings
 
         self.post_init()
 
@@ -172,12 +173,18 @@ class GPT1Model(GPT1PreTrainedModel):
         position_embeds = self.pos_emb(position_ids)
         hidden_state = self.embs_dropout(input_embeds) + position_embeds
 
-        if attention_mask is not None:
-            causal_mask = attention_mask.to(dtype=input_embeds.dtype,
-                                            device=input_embeds.device)
-        else:
-            causal_mask = self.causal_mask.to(dtype=input_embeds.dtype,
-                                              device=input_embeds.device)
+        if attention_mask is not None and attention_mask.size(1) > self.mask_cache_len:
+            seq_len = attention_mask.size(1)
+            self.mask_cache_len = seq_len
+
+            causal_mask = torch.full((seq_len, seq_len),
+                                     fill_value=float('-inf'))
+            self.register_buffer('causal_mask',
+                                 torch.triu(causal_mask, diagonal=1),
+                                 persistent=False)
+
+        causal_mask = self.causal_mask.to(dtype=input_embeds.dtype,
+                                          device=input_embeds.device)
 
         for layer in self.layers:
             hidden_state = layer(hidden_state, attn_mask=causal_mask)
@@ -240,10 +247,13 @@ class GPT1ForCausalLM(GPT1PreTrainedModel):
             logits=logits
         )
 
-    def prepare_inputs_for_generation(self, input_ids, *args, **kwargs):
-        seq_len = input_ids.size(1)
+    def prepare_inputs_for_generation(self, input_ids, attention_mask,
+                                      *args, **kwargs):
+        assert attention_mask.size(1) == input_ids.size(1)
+
+        seq_len = attention_mask.size(1)
 
-        attn_mask = torch.full((1, seq_len, seq_len), fill_value=float('-inf'))
+        attn_mask = torch.full((seq_len, seq_len), fill_value=float('-inf'))
         attn_mask = torch.triu(attn_mask, diagonal=1)
 
         return {