update

modeling_mitre.py

@@ -433,7 +433,6 @@ class MitreDecoder(MitrePreTrainedModel):
             # ensure contiguous
             expanded_src_tokens = self.check_contiguous(expanded_src_tokens)
             source_tokens = self.check_contiguous(source_tokens)
-            src_length = source_tokens.shape[1]
 
             # get embeds with positions for source tokens (input_ids + registers)
             inputs_embeds = self.source_tokens_embedding_with_positions(expanded_src_tokens, total_token_nums, batch_indices, token_indices)
@@ -444,11 +443,11 @@ class MitreDecoder(MitrePreTrainedModel):
                 decoder_input_ids[:, 0] = source_tokens[:, -1]
             
             tokens = torch.cat([source_tokens, decoder_input_ids], dim=1)
+            src_length = source_tokens.shape[1]
         
         decoder_inputs_embeds = self.embed_tokens(decoder_input_ids)
         decoder_inputs_embeds = decoder_inputs_embeds + self.tgt_embed_positions(decoder_input_ids, past_key_values_length, src_length=src_length)
-        # if past_key_values_length > 0:
-        #     raise ValueError()
+
         if past_key_values_length == 0:
             hidden_states = torch.cat([inputs_embeds, decoder_inputs_embeds], dim=1)
         else:
@@ -759,8 +758,8 @@ class MitreForConditionalGeneration(MitrePreTrainedModel, GenerationMixin):
         return reordered_past
     
     @staticmethod
-    def _reorder_register_nums(register_nums, beam_idx):
-        return register_nums.index_select(0, beam_idx.to(register_nums.device))
+    def _reorder_register_cache(t, beam_idx):
+        return t.index_select(dim=0, index=beam_idx.to(t.device))
     
     @staticmethod
     def _expand_inputs_for_generation(
@@ -844,6 +843,9 @@ class MitreForConditionalGeneration(MitrePreTrainedModel, GenerationMixin):
             if past_key_values is not None:
                 decoder_input_ids_for_generation = decoder_input_ids[:, -1:]
                 attention_mask = registering_cache["attention_mask"]
+                # Get the mask when the first time using kv cache.
+                # After it, we can simply repeat 0. (the last column of mask) to get the next mask.
+                # As a result, we avoid generate the mask from scratch in kv cache and save memory.
                 if attention_mask is not None:
                     attention_mask = torch.cat((attention_mask, attention_mask[..., -1:]), dim=-1)
             else:
@@ -907,7 +909,9 @@ class MitreForConditionalGeneration(MitrePreTrainedModel, GenerationMixin):
             del outputs
 
             past_key_values = self._reorder_cache(past_key_values, beam_idx)
-            registering_cache["register_nums"] = self._reorder_register_nums(registering_cache["register_nums"], beam_idx)
+            registering_cache["register_nums"] = self._reorder_register_cache(registering_cache["register_nums"], beam_idx)
+            if registering_cache["attention_mask"] is not None:
+                registering_cache["attention_mask"] = self._reorder_register_cache(registering_cache["attention_mask"], beam_idx)
 
             cur_len = cur_len + 1