code: def test_encoder_decoder_module(self, model_parallel_config):
      # 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
      enc_num_layers = 4
      dec_num_layers = 6
      enc_cross_attention = [3]  # layer numbers for cross attention
      dec_cross_attention = [3, 5]  # layer numbers for cross attention
  
      all_tokens = torch.randint(
          0, vocab_size, (batch, input_length + 1)
      ).cuda()  # (seq, batch, dim)
      hidden = all_tokens[:, :-1]
      labels = all_tokens[:, 1:]
  
      hidden_mask = (hidden != pad_id).cuda()
      retrieved = torch.randint(
          0, vocab_size, (batch, chunks, neighbors, 2 * text_chunk_size)
      ).cuda()
  
      pad_id = vocab_size - 1
      context_mask = (retrieved != pad_id).cuda()
  
      class FakeTokenizer:
          eos_id = vocab_size - 2
  
      tokenizer = FakeTokenizer()
  
      encoder_decoder = (
          MegatronRetrievalTokenLevelEncoderDecoderModule(
              config=model_parallel_config,
              vocab_size=vocab_size,
              hidden_size=dim,
              max_position_embeddings=input_length,
              num_attention_heads=num_attention_heads,
              ffn_hidden_size=dim * 4,
              precision=16,
              chunk_size=text_chunk_size,
              enc_num_layers=enc_num_layers,
              dec_num_layers=dec_num_layers,
              enc_cross_attention=enc_cross_attention,
              dec_cross_attention=dec_cross_attention,
              add_position_embedding=False,
              tokenizer=tokenizer,
          )
          .cuda()
          .half()
      )
  
      out = encoder_decoder(
          hidden,
          hidden_mask,
          retrieved_ids=retrieved,
          retrieved_attn_mask=context_mask,
          labels=labels,
      )
      assert out.shape == torch.Size([batch, input_length])
  
      # verify the attention mask matrix is correct
  
      all_tokens = torch.tensor(
          [
              [
                  1,
                  2,
                  vocab_size - 2,
                  3,
                  vocab_size - 1,
                  vocab_size - 2,
                  3,
                  4,
                  5,
              ]
          ]
      ).cuda()
  
      encoder_decoder = (
          MegatronRetrievalTokenLevelEncoderDecoderModule(
              config=model_parallel_config,
              vocab_size=vocab_size,
              hidden_size=dim,
              max_position_embeddings=8,
              num_attention_heads=num_attention_heads,
              ffn_hidden_size=dim * 4,
              precision=16,
              chunk_size=4,
              enc_num_layers=enc_num_layers,
              dec_num_layers=dec_num_layers,
              enc_cross_attention=enc_cross_attention,
              dec_cross_attention=dec_cross_attention,
              add_position_embedding=False,
              tokenizer=tokenizer,
          )
          .cuda()
          .half()
      )
  
      hidden = all_tokens[:, :-1]
      labels = all_tokens[:, 1:]
  
      hidden_mask = (hidden != pad_id).cuda()
      retrieved = torch.randint(
          0, vocab_size, (1, 2, neighbors, 8)
      ).cuda()
  
      pad_id = vocab_size - 1
      context_mask = (retrieved != pad_id).cuda()
  
      out = encoder_decoder(
          hidden,
          hidden_mask,
          retrieved_ids=retrieved,
          retrieved_attn_mask=context_mask,
          labels=labels,
      )
  
      mask3d = encoder_decoder.pre_decoder._calculate_dec_att_mask(
          hidden_mask, torch.where(hidden == vocab_size - 2)
      )
      expected = torch.tensor(
          [
              [
                  [
                      [False, True, True, True, True, True, True, True],
                      [
                          False,
                          False,
                          True,
                          True,
                          True,
                          True,
                          True,
                          True,
                      ],
                      [
                          False,
                          False,
                          False,
                          True,
                          True,
                          True,
                          True,
                          True,
                      ],
                      [True, True, True, False, True, True, True, True],
                      [True, True, True, True, True, True, True, True],
                      [
                          True,
                          True,
                          True,
                          False,
                          True,
                          False,
                          True,
                          True,
                      ],
                      [True, True, True, True, True, True, False, True],
                      [
                          True,
                          True,
                          True,
                          True,
                          True,
                          True,
                          False,
                          False,
                      ],
                  ]
              ]
          ]
      ).cuda()
      assert (mask3d == expected).all()
  
file_path: NeMo/tests/collections/nlp/test_retrieval_module.py
imports: import pytest
  import torch
  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 5 for calling the TestRetrievalModule obj's underlying member method code with example usage: obj.test_encoder_decoder_module(model_parallel_config) without return types