code: def main():
      runtimes = {}
      nlist = (
          list(range(2000, 10000, 2000))
          + list(range(10000, 50000, 5000))
          + list(range(50000, 100000, 10000))
      )
      print("N-List:", nlist)
      for data_parr, tens_parr, pipe_parr in [
          (8, 1, 1),
          (4, 2, 1),
          (2, 1, 4),
          (1, 2, 4),
      ]:
          for offload in [True, False]:
              dist_setting = (
                  "ddp="
                  + str(data_parr)
                  + ", tensor_parr="
                  + str(tens_parr)
                  + ", pipe_parr="
                  + str(pipe_parr)
                  + ", offload="
                  + str(offload)
              )
              runtimes[dist_setting] = {}
              print("Beginning Testing for", dist_setting)
              for n in nlist:
                  cmd = "python3 -m torch.distributed.launch --nproc_per_node=8 run_gpt_minimal_test.py"
                  cmd += (
                      " --micro-batch-size 1 --num-layers "
                      + str(n)
                      + " --hidden-size 128 --num-attention-heads 16"
                  )
                  cmd += (
                      " --max-position-embeddings 128 --seq-length 128 --tensor-model-parallel-size "
                      + str(tens_parr)
                  )
                  cmd += (
                      " --pipeline-model-parallel-size "
                      + str(pipe_parr)
                      + (" --cpu-offload" if offload else "")
                  )
                  print(cmd)
                  runtime, bill_params, success, errs = run_gpt(cmd)
                  if success:
                      runtimes[dist_setting][bill_params] = runtime
                      print(
                          str(runtime) + "s per training iter for",
                          str(bill_params) + "B parameter GPT-2",
                      )
                      if n >= 10000:
                          plot(runtimes)
                  else:
                      print(
                          "GPT-2 w/",
                          n,
                          "layers failed using",
                          dist_setting,
                      )
                      print(
                          "Moving on to the next distributed setting..."
                      )
                      print("#" * (25))
                      print()
                      plot(runtimes)
                      break
      print(runtimes)
      plot(runtimes)
  
imports: import subprocess
  import os
  from apex.transformer.testing.commons import TEST_SUCCESS_MESSAGE
label_desc: Use image node_id 3 for calling a global function with example usage: main() without return types