class_docstr: Transformer class.
class_docstr_tok: ["Transformer","class","."]
docstr: Forward method with activation checkpointing.
code: def _checkpointed_forward(
      self,
      hidden_states,
      attention_mask,
      encoder_output,
      enc_dec_attn_mask,
      rotary_pos_emb,
      self_attention_relative_position_bias,
      cross_attention_relative_position_bias,
      checkpoint_activations_all_layers,
  ):
      """Forward method with activation checkpointing."""
  
      def custom(start, end):
          if self.transformer_engine:
  
              def custom_forward(*inputs):
                  hidden_states = inputs[0]
                  attention_mask = inputs[1]
                  encoder_output = inputs[2]
                  enc_dec_attn_mask = inputs[3]
                  # Cache FP8 weight and transpose at (1) the first micro-batch in each global-batch
                  # in training, (2) the first micro-batch in each validation and test routine.
                  # The caching happens in TransformerEngine when passing `is_first_microbatch=True`.
                  is_first_microbatch = (
                      self.is_first_train_microbatch and self.training
                  ) or (
                      self.is_prev_microbatch_training
                      and not self.training
                  )
                  for index in range(start, end):
                      layer = self._get_layer(index)
                      hidden_states = layer(
                          hidden_states,
                          attention_mask,
                          encoder_output=encoder_output,
                          enc_dec_attn_mask=enc_dec_attn_mask,
                          inference_params=None,
                          is_first_microbatch=is_first_microbatch,
                          checkpoint_core_attention=False,
                      )
  
                  return hidden_states
  
          else:
  
              def custom_forward(*inputs):
                  if len(inputs) == 9:
                      hidden_states = inputs[0]
                      attention_mask = inputs[1]
                      encoder_output = inputs[2]
                      enc_dec_attn_mask = inputs[3]
                      rotary_pos_emb = (inputs[4], inputs[5], inputs[6])
                      self_attention_relative_position_bias = inputs[7]
                      cross_attention_relative_position_bias = inputs[8]
                  elif len(inputs) == 10:
                      hidden_states = (inputs[0], inputs[1])
                      attention_mask = inputs[2]
                      encoder_output = inputs[3]
                      enc_dec_attn_mask = inputs[4]
                      rotary_pos_emb = (inputs[5], inputs[6], inputs[7])
                      self_attention_relative_position_bias = inputs[8]
                      cross_attention_relative_position_bias = inputs[9]
                  else:
                      hidden_states = inputs[0]
                      attention_mask = inputs[1]
                      encoder_output = inputs[2]
                      enc_dec_attn_mask = inputs[3]
                      rotary_pos_emb = inputs[4]
                      self_attention_relative_position_bias = inputs[5]
                      cross_attention_relative_position_bias = inputs[6]
                  for index in range(start, end):
                      layer = self._get_layer(index)
                      hidden_states = layer(
                          hidden_states=hidden_states,
                          attention_mask=attention_mask,
                          encoder_output=encoder_output,
                          enc_dec_attn_mask=enc_dec_attn_mask,
                          rotary_pos_emb=rotary_pos_emb,
                          self_attention_relative_position_bias=self_attention_relative_position_bias,
                          cross_attention_relative_position_bias=cross_attention_relative_position_bias,
                      )
                      if isinstance(hidden_states, tuple):
                          pass
                      else:
                          hidden_states = hidden_states.contiguous()
                  return hidden_states
  
          return custom_forward
  
      if self.activations_checkpoint_method == "uniform":
          # Uniformly divide the total number of Transformer layers and checkpoint
          # the input activation of each divided chunk.
          # A method to further reduce memory usage reducing checkpoints.
          l = 0
          while l < self.num_layers:
              if isinstance(hidden_states, tuple):
                  hidden_tuple = (hidden_states[0], hidden_states[1])
              else:
                  hidden_tuple = (hidden_states,)
              middle_tuple = (
                  attention_mask,
                  encoder_output,
                  enc_dec_attn_mask,
              )
  
              if rotary_pos_emb is None:
                  rot_tuple = (rotary_pos_emb,)
              else:
                  rot_tuple = (
                      rotary_pos_emb[0],
                      rotary_pos_emb[1],
                      rotary_pos_emb[2],
                  )
  
              final_tuple = (
                  self_attention_relative_position_bias,
                  cross_attention_relative_position_bias,
              )
              arg_tuple = (
                  hidden_tuple + middle_tuple + rot_tuple + final_tuple
              )
  
              if self.transformer_engine:
                  hidden_states = te_checkpoint(
                      custom(
                          l, l + self.activations_checkpoint_num_layers
                      ),
                      False,
                      tensor_parallel.random.get_cuda_rng_tracker,
                      parallel_state.get_tensor_model_parallel_group(),
                      *arg_tuple,
                  )
              else:
                  hidden_states = tensor_parallel.checkpoint(
                      custom(
                          l, l + self.activations_checkpoint_num_layers
                      ),
                      False,
                      *arg_tuple,
                  )
              l += self.activations_checkpoint_num_layers
      elif self.activations_checkpoint_method == "block":
          # When pipeline-parallel size > 1 and 'num_micro_batches_with_partial_activation_checkpoints' = int,
          # pipeline scheduling can force to checkpoint all layers or partial layers in a micro-batch.
          if checkpoint_activations_all_layers:
              activations_checkpoint_num_layers = self.num_layers
          else:
              activations_checkpoint_num_layers = (
                  self.activations_checkpoint_num_layers
              )
              if (
                  parallel_state.get_pipeline_model_parallel_world_size()
                  > 0
                  and self.activations_checkpoint_layers_per_pipeline
                  is not None
              ):
                  # Decrease the number of layers to checkpoint at later pipeline stages
                  activations_checkpoint_num_layers -= int(
                      parallel_state.get_pipeline_model_parallel_rank()
                      * self.activations_checkpoint_layers_per_pipeline
                  )
          # Checkpoint the input activation of only a set number of individual
          # Transformer layers and skip the rest.
          # A method fully use the device memory removing redundant re-computation.
          for l in range(self.num_layers):
              if isinstance(hidden_states, tuple):
                  hidden_tuple = (hidden_states[0], hidden_states[1])
              else:
                  hidden_tuple = (hidden_states,)
              middle_tuple = (
                  attention_mask,
                  encoder_output,
                  enc_dec_attn_mask,
              )
  
              if rotary_pos_emb is None:
                  rot_tuple = (rotary_pos_emb,)
              else:
                  rot_tuple = (
                      rotary_pos_emb[0],
                      rotary_pos_emb[1],
                      rotary_pos_emb[2],
                  )
  
              final_tuple = (
                  self_attention_relative_position_bias,
                  cross_attention_relative_position_bias,
              )
              arg_tuple = (
                  hidden_tuple + middle_tuple + rot_tuple + final_tuple
              )
  
              if l < activations_checkpoint_num_layers:
                  if self.transformer_engine:
                      hidden_states = te_checkpoint(
                          custom(l, l + 1),
                          False,
                          tensor_parallel.random.get_cuda_rng_tracker,
                          parallel_state.get_tensor_model_parallel_group(),
                          *arg_tuple,
                      )
                  else:
                      hidden_states = tensor_parallel.checkpoint(
                          custom(l, l + 1), False, *arg_tuple
                      )
              else:
                  hidden_states = custom(l, l + 1)(*arg_tuple)
      else:
          raise ValueError("Invalid activation checkpoint method.")
  
      return hidden_states
  
imports: from contextlib import nullcontext
  from typing import Any, Callable, Optional
  import torch
  import torch.nn
  from einops import rearrange
  from nemo.collections.common.parts.adapter_modules import LinearAdapterConfig
  from nemo.collections.nlp.modules.common.megatron.adapters.parallel_adapters import AdapterName, ParallelLinearAdapterConfig, ParallelLinearAdapterWeightTyingConfig
  from nemo.collections.nlp.modules.common.megatron.attention import ParallelAttention, ParallelChunkedCrossAttention
  from nemo.collections.nlp.modules.common.megatron.fused_bias_dropout_add import bias_dropout_add, bias_dropout_add_fused_inference, bias_dropout_add_fused_train, dropout_add
  from nemo.collections.nlp.modules.common.megatron.fused_layer_norm import get_layer_norm
  from nemo.collections.nlp.modules.common.megatron.layer_norm_1p import LayerNorm1P, LPLayerNorm
  from nemo.collections.nlp.modules.common.megatron.layer_type import LayerType
  from nemo.collections.nlp.modules.common.megatron.mlp import ParallelMLP, SwitchMLP
  from nemo.collections.nlp.modules.common.megatron.module import MegatronModule
  from nemo.collections.nlp.modules.common.megatron.utils import ApexGuardDefaults
  from nemo.collections.nlp.parts import utils_funcs
  from nemo.core import adapter_mixins
  from nemo.utils import logging
label_desc: Use image node_id 4 for calling the ParallelTransformer obj's underlying member method code with example usage: obj._checkpointed_forward(hidden_states, attention_mask, encoder_output, enc_dec_attn_mask, rotary_pos_emb, _attention_relative_position_bias, cross_attention_relative_position_bias, checkpoint_activations_all_layers) and returns: hidden_states, custom_forward, hidden_states, hidden_states