code: def compute_activation_memory(args, num_microbatches, verbose=False):
      # Using formula in Table 2 of https://arxiv.org/pdf/2205.05198.pdf.
      # We are trying to compute the maximum activation footprint, so all calculations in this function
      # are for the first pipeline stage.
  
      # Memory footprint from transformer layer (self-attention and MLP).
      activation_memory = (
          args.seq_length * args.micro_batch_size * args.hidden_size
      ) * 34
      if verbose:
          print(
              f"Activation memory footprint per transformer layer: "
              f"{activation_memory / NUM_BYTES_IN_MEGABYTE / args.tensor_model_parallel_size:.1f} MB"
          )
      activation_memory *= args.num_layers
  
      # Now add activation memory required for input embeddings, last LayerNorm and output layer.
  
      # Input to embedding (pp_size microbatches in flight).
      activation_memory += (
          8
          * args.seq_length
          * args.micro_batch_size
          * args.pipeline_model_parallel_size
      )
      # Dropout in embedding layer (pp_size microbatches in flight).
      activation_memory += (
          args.seq_length
          * args.micro_batch_size
          * args.hidden_size
          * args.pipeline_model_parallel_size
      )
  
      # Multiply by interleaved PP memory factor.
      if args.virtual_pipeline_model_parallel_size is not None:
          interleaved_schedule_memory_penalty = 1 + (
              (args.pipeline_model_parallel_size - 1)
              / (
                  args.pipeline_model_parallel_size
                  * args.virtual_pipeline_model_parallel_size
              )
          )
          in_flight_microbatches = math.ceil(
              interleaved_schedule_memory_penalty
              * args.pipeline_model_parallel_size
          )
          if verbose:
              print(
                  f"Memory penalty from interleaved schedule: {interleaved_schedule_memory_penalty:.2f}"
              )
              print(
                  f"Number of in-flight microbatches: {in_flight_microbatches}"
              )
          activation_memory *= interleaved_schedule_memory_penalty
  
      # If using non-interleaved schedule, number of microbatches in pipeline can be less than pp_size,
      # so discount accordingly.
      if (
          args.virtual_pipeline_model_parallel_size is None
          and args.pipeline_model_parallel_size > 1
      ):
          if num_microbatches is not None:
              activation_memory *= min(
                  1,
                  num_microbatches / args.pipeline_model_parallel_size,
              )
              in_flight_microbatches = min(
                  num_microbatches, args.pipeline_model_parallel_size
              )
          else:
              in_flight_microbatches = args.pipeline_model_parallel_size
          if verbose:
              print(
                  f"Number of in-flight microbatches: {in_flight_microbatches}"
              )
  
      if args.pipeline_model_parallel_size == 1:
          # Inputs to output layer and CE loss.
          activation_memory += (
              args.seq_length
              * args.micro_batch_size
              * args.hidden_size
              * 4
              * (1 + (args.padded_vocab_size / args.hidden_size))
          )
  
      # Activation memory is partitioned by TP size due to tensor and sequence model parallelism.
      return activation_memory / args.tensor_model_parallel_size
  
imports: import math
label_desc: Use image node_id 2 for calling a global function with example usage: compute_activation_memory(args, num_microbatches, verbose) and returns: unknown