code: def compute_weight_and_optimizer_memory(args, verbose=False):
      if not args.group_query_attention:
          args.num_query_groups = args.num_attention_heads
      num_parameters_in_transformer_layers = (
          10
          * args.num_layers
          * args.hidden_size
          * args.hidden_size
          * (
              1
              + (
                  args.num_query_groups
                  / (5.0 * args.num_attention_heads)
              )
              + (2 / (5 * args.hidden_size))
              + (1 / (5 * args.num_layers * args.hidden_size))
          )
      )
      embedding_size = args.hidden_size * args.padded_vocab_size
      if args.untie_embeddings_and_output_weights:
          num_total_parameters_with_embeddings = (
              num_parameters_in_transformer_layers
              + (2 * embedding_size)
          )
      else:
          num_total_parameters_with_embeddings = (
              num_parameters_in_transformer_layers + embedding_size
          )
      if verbose:
          print(
              f"Number of parameters in billions: {num_total_parameters_with_embeddings / 10**9:.2f}"
          )
  
      # Most loaded model shard has (1/pp_size transformer layers + 1 embedding layer) / tp_size.
      num_parameters_on_most_loaded_model_shard = (
          (
              num_parameters_in_transformer_layers
              / args.pipeline_model_parallel_size
          )
          + embedding_size
      ) / args.tensor_model_parallel_size
      if (
          args.untie_embeddings_and_output_weights
          and args.pipeline_model_parallel_size == 1
      ):
          num_parameters_on_most_loaded_model_shard += (
              embedding_size / args.tensor_model_parallel_size
          )
      if verbose:
          print(
              f"Number of parameters in most loaded shard in billions: {num_parameters_on_most_loaded_model_shard / 10**9:.4f}"
          )
  
      if args.pipeline_model_parallel_size > 1:
          # Other shards just have (1/pp_size transformer layers) / tp_size.
          num_parameters_on_other_model_shards = (
              num_parameters_in_transformer_layers
              / (
                  args.pipeline_model_parallel_size
                  * args.tensor_model_parallel_size
              )
          )
          if verbose:
              print(
                  f"Number of parameters in other shards in billions: {num_parameters_on_other_model_shards / 10**9:.4f}"
              )
  
      num_bytes_per_parameter = (
          18
          if not args.use_distributed_optimizer
          else 6 + (12 / args.data_parallel_size)
      )
      weight_and_optimizer_memory = (
          num_parameters_on_most_loaded_model_shard
          * num_bytes_per_parameter
      )
  
      return weight_and_optimizer_memory
  
imports: import math
label_desc: Use image node_id 1 for calling a global function with example usage: compute_weight_and_optimizer_memory(args, verbose) and returns: weight_and_optimizer_memory