Add KV-Cache for MegatronLMEncoderDecoderModel
State: closed | Pull request: no
Labels: name: stale

**Is your feature request related to a problem? Please describe.**

I have been using the MegatronLMEncoderDecoderModel which seems to be the closest implementation for my need (encoder-decoder model), but when generating results (with the `decode` method), the generation becomes slower as it generates more tokens. My guess is that KV-Cache is not used for this model.

After a lot of digging, I found that the MegatronGPTModel uses for inference `set_inference_key_value_memory` and `inference_max_sequence_len`. I have tried to set `set_inference_key_value_memory=True` during the first decoding iteration, then to False, as well as providing only 1 token at a time, similar to `GPTModelTextGenerationStrategy` https://github.com/NVIDIA/NeMo/blob/900ca0b8bac3d038b9854bd498967b0525aaeb35/nemo/collections/nlp/modules/common/text_generation_strategy.py#L312-L327 , but I am faced with many issues, and the usage isn't straightforward.

To notice the generation speed decrease, you can wrap the for loop with tqdm:
```python
from tqdm.auto import tqdm

def decode(
        self,
        tokens_enc,
        enc_mask,
        num_tokens_to_generate,
        encoder_input=None,
        tokenizer=None,
        enc_output=None,
        enc_output_attn_mask=None,
        ignore_ids=[],
        bos_id=None,  # If bos=None, will use tokenizer.bos_id unless explicitly set to something else.
        predicted_tokens_dec=None,
        batch_data=None,
        sampling_method: str = "greedy-search",
        sampling_kwargs: dict = {},
    ):
   ...

        for i in tqdm(range(num_tokens_to_generate)):
```

**Describe the solution you'd like**

It would be nice if KV-Cache could be used for the MegatronLMEncoderDecoderModel, or provide a snippet of code to use the `decode` method without any slowdown during the generation.


**Describe alternatives you've considered**

`MegatronGPTModel` seems to uses the KV-Cache, but is a decoder only model.  Since `MegatronGPTModel` uses internally (to the lower level) a `ParallelTransformer`, like the `MegatronTransformerEncoderModule` and `MegatronTransformerDecoderModule`, a parameter could be specified to `MegatronGPTModel` to enable cross attention inside it to be used like a decoder, with an encoder output specified
```python
def forward(
        self,
        input_ids,
        position_ids,
        attention_mask,
       enc_output=None, # Add this input
        loss_mask=None,
        labels=None,
        token_type_ids=None,
        layer_past=None,
        get_key_value=False,
        forward_method_parallel_output=None,
        encoder_input=None,
        set_inference_key_value_memory=False,
        inference_max_sequence_len=None,
        checkpoint_activations_all_layers=None,
    ):
       lm_output = self.language_model(
            input_ids,
            position_ids,
            attention_mask,
            encoder_output=encoder_output, # Add this
            layer_past=layer_past,
            get_key_value=get_key_value,
            encoder_input=encoder_input,
            set_inference_key_value_memory=set_inference_key_value_memory,
            inference_max_sequence_len=inference_max_sequence_len,
            checkpoint_activations_all_layers=checkpoint_activations_all_layers,
        )
```
So that GPT model can be used with CrossAttention like [HuggingFace GPT2 model](https://github.com/huggingface/transformers/blob/6c1d0b069de22d7ed8aa83f733c25045eea0585d/src/transformers/models/gpt2/modeling_gpt2.py#L1091), because KV-Cache is implemented in `MegatronGPTModel.generate`



This issue is stale because it has been open for 30 days with no activity. Remove stale label or comment or this will be closed in 7 days.

This issue was closed because it has been inactive for 7 days since being marked as stale.

Hi @Kurokabe,
Thank you for pointing out. Indeed, T5 inference is not supporting KV_cache yet. 
We have plans to support it in the future. However, as of now, we are having other priorities.