class_docstr: Transformer class.
class_docstr_tok: ["Transformer","class","."]
docstr: Forward method with activation checkpointing.
docstr_tok: ["Forward","method","with","activation","checkpointing","."]
code: def _checkpointed_forward(
      self,
      hidden_states: Tensor,
      attention_mask: Tensor,
      context: Tensor,
      context_mask: Tensor,
      rotary_pos_emb: Tensor,
  ):
      """Forward method with activation checkpointing."""
  
      def custom(start: int, end: int):
          def custom_forward(
              hidden_states,
              attention_mask,
              context,
              context_mask,
              rotary_pos_emb,
          ):
              for index in range(start, end):
                  layer = self._get_layer(index)
                  hidden_states, context = layer(
                      hidden_states=hidden_states,
                      attention_mask=attention_mask,
                      context=context,
                      context_mask=context_mask,
                      rotary_pos_emb=rotary_pos_emb,
                      inference_params=None,
                  )
              return hidden_states, context
  
          return custom_forward
  
      if self.config.recompute_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_per_pipeline_rank:
              hidden_states, context = tensor_parallel.checkpoint(
                  custom(l, l + self.config.recompute_num_layers),
                  self.config.distribute_saved_activations,
                  hidden_states,
                  attention_mask,
                  context,
                  context_mask,
                  rotary_pos_emb,
              )
  
              l += self.config.recompute_num_layers
  
      elif self.config.recompute_method == "block":
          # 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_per_pipeline_rank):
              if l < self.config.recompute_num_layers:
                  hidden_states, context = tensor_parallel.checkpoint(
                      custom(l, l + 1),
                      self.config.distribute_saved_activations,
                      hidden_states,
                      attention_mask,
                      context,
                      context_mask,
                      rotary_pos_emb,
                  )
              else:
                  hidden_states, context = custom(l, l + 1)(
                      hidden_states,
                      attention_mask,
                      context,
                      context_mask,
                      rotary_pos_emb,
                  )
      else:
          raise ValueError("Invalid activation recompute method.")
  
      return hidden_states
  
imports: import re
  from contextlib import nullcontext
  from dataclasses import dataclass
  from typing import List, Union
  import torch
  from torch import Tensor
  from megatron.core import InferenceParams, parallel_state, tensor_parallel
  from megatron.core.fusions.fused_layer_norm import FusedLayerNorm
  from megatron.core.transformer.custom_layers.transformer_engine import TENorm
  from megatron.core.transformer.enums import AttnMaskType
  from megatron.core.transformer.module import MegatronModule
  from megatron.core.transformer.spec_utils import ModuleSpec, build_module
  from megatron.core.transformer.transformer_config import TransformerConfig
  from megatron.core.transformer.transformer_layer import TransformerLayer
  from megatron.core.utils import make_sharded_tensor_for_checkpoint, make_viewless_tensor
label_desc: Use image node_id 4 for calling the TransformerBlock obj's underlying member method code with example usage: obj._checkpointed_forward(hidden_states, attention_mask, context, context_mask, rotary_pos_emb) and returns: hidden_states, custom_forward, hidden_states,  context