code: def run_gpu_forward(
      self, recompute_granularity, use_transformer_engine
  ):
      config = self.get_config()
      config.recompute_granularity = recompute_granularity
      modules = self.get_modules(
          config, use_transformer_engine, use_gpu=True
      )
  
      seq_length = 32
      micro_batch_size = 2
      n_chunks_per_sample = (
          seq_length // config.retro_preprocess.retro_gpt_chunk_length
      )
  
      # Init tensors.
      hidden_states = torch.ones(
          (
              seq_length,
              micro_batch_size,
              config.hidden_size,
          )
      ).cuda()
      attention_mask = None
      decoder_context = torch.ones(
          (
              config.retro_preprocess.retro_gpt_retrieved_length,
              config.retro_num_neighbors
              * micro_batch_size
              * n_chunks_per_sample,
              config.hidden_size,
          )
      ).cuda()
      encoder_context = torch.ones(
          (
              config.retro_preprocess.retro_gpt_chunk_length,
              micro_batch_size * n_chunks_per_sample,
              config.hidden_size,
          )
      ).cuda()
  
      # Forward decoder.
      decoder_attn_output = modules.decoder_attn(
          hidden_states,
          attention_mask,
          decoder_context,
      )
      with torch.enable_grad():
          decoder_bda_output = modules.decoder_bda(True, True)(
              decoder_attn_output,
              hidden_states,
              config.hidden_dropout,
          )
  
      # Forward encoder.
      encoder_attn_output_tuples = modules.encoder_attn(
          decoder_context,
          None,
          encoder_context,
      )
      with torch.enable_grad():
          encoder_bda_output = modules.encoder_bda(True, True)(
              encoder_attn_output_tuples,
              decoder_context,
              config.retro_encoder_hidden_dropout,
          )
      encoder_norm_output = modules.encoder_norm(encoder_bda_output)
  
      # Verify decoder.
      assert set(decoder_attn_output.keys()) == set(
          [
              "ns",
              "bs",
              "d",
              "l",
              "pad",
              "attention_output",
              "attention_bias",
              "context",
          ]
      )
      assert decoder_attn_output["ns"] == seq_length
      assert decoder_attn_output["bs"] == micro_batch_size
      assert decoder_attn_output["d"] == config.hidden_size
      assert decoder_attn_output["l"] == n_chunks_per_sample
      assert decoder_attn_output["pad"] == 3
      assert tuple(decoder_attn_output["attention_output"].shape) == (
          config.retro_preprocess.retro_gpt_chunk_length,
          micro_batch_size * n_chunks_per_sample,
          config.hidden_size,
      )
      assert tuple(decoder_attn_output["attention_bias"].shape) == (
          config.hidden_size,
      )
      assert decoder_attn_output["context"].shape == (
          config.retro_preprocess.retro_gpt_retrieved_length
          * config.retro_num_neighbors,
          micro_batch_size * n_chunks_per_sample,
          config.hidden_size,
      )
      assert decoder_bda_output.shape == hidden_states.shape
  
      # Verify encoder.
      assert (
          len(encoder_attn_output_tuples) == config.retro_num_neighbors
      )
      for output, bias, residual in encoder_attn_output_tuples:
          assert tuple(output.shape) == (
              config.retro_preprocess.retro_gpt_retrieved_length,
              micro_batch_size * n_chunks_per_sample,
              config.hidden_size,
          )
          assert tuple(bias.shape) == (config.hidden_size,)
          assert tuple(residual.shape) == (
              config.retro_preprocess.retro_gpt_retrieved_length,
              micro_batch_size * n_chunks_per_sample,
              config.hidden_size,
          )
      assert encoder_bda_output.shape == (
          config.retro_preprocess.retro_gpt_retrieved_length,
          config.retro_num_neighbors
          * micro_batch_size
          * n_chunks_per_sample,
          config.hidden_size,
      )
      assert encoder_norm_output.shape == (
          config.retro_preprocess.retro_gpt_retrieved_length,
          config.retro_num_neighbors
          * micro_batch_size
          * n_chunks_per_sample,
          config.hidden_size,
      )
  
file_path: megatron-lm/tests/unit_tests/transformer/test_retro_attention.py
imports: import torch
  import types
  from megatron.core.models.retro import RetroConfig, get_retro_decoder_block_spec
  from megatron.core.models.retro.decoder_attention import RetroDecoderCrossAttention, RetroDecoderBiasDropoutAdd
  from megatron.core.models.retro.encoder_attention import RetroEncoderCrossAttention, RetroEncoderBiasDropoutAdd, RetroEncoderLayerNorm
  from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed
  from megatron.core.transformer.transformer_block import TransformerBlock
  from tests.unit_tests.test_utilities import Utils
label_desc: Use image node_id 7 for calling the TestRetroAttention obj's underlying member method code with example usage: obj.run_gpu_forward(recompute_granularity, use_transformer_engine) without return types