code: def test_retrieval_decoder_inference(self, model_parallel_config):
      init_method_std = 0.02
  
      # rotary pos emb dim
      batch = 2
      neighbors = 2
      dim = 128
      pad_id = 19999
      num_attention_heads = 8
      chunks = 32
      text_chunk_size = 64
      input_length = chunks * text_chunk_size
      vocab_size = 20000
      # rot_dim = dim // num_attention_heads
      # rotary_pos_emb = RotaryEmbedding(rot_dim).cuda().half()
      hidden = torch.randint(
          0, vocab_size, (batch, input_length)
      ).cuda()  # (seq, batch, dim)
      hidden_mask = (hidden != pad_id).cuda()
  
      hidden_emb = (
          torch.rand(batch, input_length, dim).cuda().half()
      )  # (batch, seq, dim)
  
      # context_chunk_size = 128
      retrieved = torch.randint(
          0, vocab_size, (batch, chunks, neighbors, 2 * text_chunk_size)
      ).cuda()
      # retrieved tokens - (batch, num chunks, num retrieved neighbors, retrieved chunk with continuation)
  
      # context attention mask [b, np, sq, sk]
      pad_id = vocab_size - 1
      context_mask = (retrieved != pad_id).cuda()
      retrieved_emb = (
          torch.rand(batch, chunks, neighbors, 2 * text_chunk_size, dim)
          .cuda()
          .half()
      )
      # retrieved tokens - (batch, num chunks, num retrieved neighbors, retrieved chunk with continuation, hidden)
  
      layer_type = [
          LayerType.encoder,
          LayerType.retrieval_decoder,
          LayerType.encoder,
          LayerType.retrieval_decoder,
      ]
      num_layers = len(layer_type)
  
      init_method = init_method_normal(init_method_std)
      scaled_init_method = scaled_init_method_normal(
          init_method_std, num_layers
      )
      decoder = (
          MegatronRetrievalTransformerDecoderModule(
              config=model_parallel_config,
              init_method=init_method,
              output_layer_init_method=scaled_init_method,
              hidden_size=dim,
              ffn_hidden_size=dim * 4,
              num_layers=num_layers,
              num_attention_heads=num_attention_heads,
              precision=16,
              chunk_size=text_chunk_size,
              layer_type=layer_type,
              hidden_dropout=0.0,
              attention_dropout=0.0,
          )
          .cuda()
          .half()
      )
      out = decoder(
          hidden_emb,
          hidden_mask,
          retrieved_attn_mask=context_mask,
          retrieved_emb=retrieved_emb,
      )
      assert out.shape == torch.Size([input_length, batch, dim])
  
      out_1 = decoder(
          hidden_emb[:, :62],
          hidden_mask[:, :62],
          retrieved_attn_mask=None,
          retrieved_emb=None,
          set_inference_key_value_memory=True,
          inference_max_sequence_len=input_length,
      )
      assert (out[:62] - out_1[:62]).abs().max().item() < 1e-2
      out_1 = decoder(
          hidden_emb[:, 62:63],
          hidden_mask[:, :63],
          retrieved_attn_mask=None,
          retrieved_emb=None,
          set_inference_key_value_memory=False,
          inference_max_sequence_len=input_length,
      )
      assert (out[62] - out_1[0]).abs().max().item() < 1e-2
      out_2 = decoder(
          hidden_emb[:, 63:64],
          hidden_mask[:, :64],
          retrieved_attn_mask=context_mask[:, :1],
          retrieved_emb=retrieved_emb[:, :1],
          set_inference_key_value_memory=False,
          inference_max_sequence_len=input_length,
      )
      assert (out[63] - out_2[0]).abs().max().item() < 1e-2
      for i in range(64, 127):
          out_2 = decoder(
              hidden_emb[:, i : i + 1],
              hidden_mask[:, : i + 1],
              retrieved_attn_mask=context_mask[:, :1],
              retrieved_emb=retrieved_emb[:, :1],
              set_inference_key_value_memory=False,
              inference_max_sequence_len=input_length,
          )
          assert (out[i] - out_2[0]).abs().max().item() < 1e-2
      for i in range(127, 191):
          out_3 = decoder(
              hidden_emb[:, i : i + 1],
              hidden_mask[:, : i + 1],
              retrieved_attn_mask=context_mask[:, :2],
              retrieved_emb=retrieved_emb[:, :2],
              set_inference_key_value_memory=False,
              inference_max_sequence_len=input_length,
          )
          assert (out[i] - out_3[0]).abs().max().item() < 1e-2
  
      out_1 = decoder(
          hidden_emb[:, :130],
          hidden_mask[:, :130],
          retrieved_attn_mask=context_mask[:, :2],
          retrieved_emb=retrieved_emb[:, :2],
          set_inference_key_value_memory=True,
          inference_max_sequence_len=input_length,
      )
      assert (out[:130] - out_1[:130]).abs().max().item() < 1e-2
      for i in range(130, 191):
          out_3 = decoder(
              hidden_emb[:, i : i + 1],
              hidden_mask[:, : i + 1],
              retrieved_attn_mask=context_mask[:, :2],
              retrieved_emb=retrieved_emb[:, :2],
              set_inference_key_value_memory=False,
              inference_max_sequence_len=input_length,
          )
          assert (out[i] - out_3[0]).abs().max().item() < 1e-2
  
file_path: NeMo/tests/collections/nlp/test_retrieval_module_inference.py
imports: import pytest
  import torch
  import torch.nn.functional
  from einops import rearrange
  from pytorch_lightning.trainer.trainer import Trainer
  from nemo.collections.nlp.modules.common.megatron.attention import ParallelChunkedCrossAttention
  from nemo.collections.nlp.modules.common.megatron.layer_type import LayerType
  from nemo.collections.nlp.modules.common.megatron.megatron_init import initialize_model_parallel_for_nemo
  from nemo.collections.nlp.modules.common.megatron.position_embedding import RotaryEmbedding
  from nemo.collections.nlp.modules.common.megatron.retrieval_token_level_encoder_decoder import MegatronRetrievalTokenLevelEncoderDecoderModule
  from nemo.collections.nlp.modules.common.megatron.retrieval_transformer import MegatronRetrievalTransformerDecoderModule, MegatronRetrievalTransformerEncoderModule
  from nemo.collections.nlp.modules.common.megatron.utils import build_attention_mask_3d, init_method_normal, scaled_init_method_normal
  from nemo.collections.nlp.parts.nlp_overrides import NLPDDPStrategy
label_desc: Use image node_id 4 for calling the TestRetrievalModuleInference obj's underlying member method code with example usage: obj.test_retrieval_decoder_inference(model_parallel_config) without return types