code: def main():
      # version check
      # this example is supposed to work for versions greater than 0.2.0
      assert version.parse(CAI_VERSION) >= version.parse("0.2.0")
  
      set_cpu_maximum_parallelism()
      args = parse_args()
  
      # if args.distplan not in ["colossalai", "torch_ddp", "torch_zero", "zero1", "zero2"]:
      if args.distplan not in [
          "CAI_ZeRO1",
          "CAI_ZeRO2",
          "CAI_Gemini",
          "Pytorch_DDP",
          "Pytorch_ZeRO",
      ]:
          raise TypeError(f"{args.distplan} is error")
  
      # batch size per DP degree
      BATCH_SIZE = args.batch_size
      SEQ_LEN = 1024
      VOCAB_SIZE = 50257
  
      NUM_STEPS = args.train_step
  
      WARMUP_STEPS = 1
      assert (
          WARMUP_STEPS < NUM_STEPS
      ), "warmup steps should smaller than the total steps"
      assert (
          NUM_STEPS - WARMUP_STEPS
      ) % 2 == 1, (
          "the number of valid steps should be odd to take the median"
      )
      PROF_FLAG = False  # The flag of profiling, False by default
  
      disable_existing_loggers()
      colossalai.launch_from_torch(config={})
  
      logger = get_dist_logger()
      logger.info(
          f"{args.model_type}, {args.distplan}, batch size {BATCH_SIZE}",
          ranks=[0],
      )
  
      # build criterion
      criterion = GPTLMLoss()
      torch.manual_seed(123)
      if args.distplan.startswith("CAI"):
          ctx = (
              LazyInitContext(default_device=get_current_device())
              if args.distplan == "CAI_Gemini"
              else nullcontext()
          )
          # build GPT model
          with ctx:
              model = model_builder(args.model_type)(checkpoint=True)
  
          # assign running configurations
          if args.distplan == "CAI_ZeRO1":
              zero_stage = 1
          elif args.distplan == "CAI_ZeRO2":
              zero_stage = 2
          elif args.distplan == "CAI_Gemini":
              zero_stage = 3
          else:
              raise RuntimeError
  
          plugin = None
          if args.distplan.startswith("CAI_ZeRO"):
              plugin = LowLevelZeroPlugin(
                  stage=zero_stage,
                  reduce_bucket_size_in_m=12,
                  overlap_communication=True,
                  verbose=True,
              )
          elif args.distplan == "CAI_Gemini":
              plugin = GeminiPlugin(
                  search_range_m=128, hidden_dim=model.config.n_embd
              )
          else:
              raise RuntimeError
  
          # build a highly optimized gpu/cpu optimizer
          optimizer = HybridAdam(model.parameters(), lr=1e-3)
  
          logger.info(
              get_mem_info(prefix="After init optim, "), ranks=[0]
          )
      elif args.distplan.startswith("Pytorch"):
          assert (
              args.tp_degree == 1
          ), "The degree of TP should be 1 for DDP examples."
          model = model_builder(args.model_type)(checkpoint=True).cuda()
          plugin = TorchDDPPlugin()
          if args.distplan.endswith("DDP"):
              optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
          elif args.distplan.endswith("ZeRO"):
              from torch.distributed.optim import (
                  ZeroRedundancyOptimizer,
              )
  
              optimizer = ZeroRedundancyOptimizer(
                  model.parameters(),
                  optimizer_class=torch.optim.Adam,
                  lr=1e-3,
              )
  
      else:
          raise RuntimeError
      # wrap your model and optimizer
      booster = Booster(plugin=plugin)
      model, optimizer, criterion, _, _ = booster.boost(
          model, optimizer, criterion
      )
  
      # model is shared after TP
      numel = get_model_size(model)
      logger.info(
          f"the size of testing model size is {model_size_formatter(numel)}."
      )
      logger.info(get_mem_info(prefix="After init model, "), ranks=[0])
  
      # Tflops_per_GPU = global_batch * global_numel * seq_len * 8 / #gpu
      # = (batch_per_DP_group * dp_degree) * (numel * tp_degree) * seq_len * 8 / (tp_degree * dp_degree)
      # = batch_per_DP_group * numel * seq_len * 8
      get_tflops_func = partial(get_tflops, numel, BATCH_SIZE, SEQ_LEN)
  
      torch.cuda.synchronize()
      model.train()
      tflops_list = []
  
      def train_step():
          # we just use randomly generated data here
          input_ids, attn_mask = get_data(
              BATCH_SIZE, SEQ_LEN, VOCAB_SIZE
          )
          optimizer.zero_grad()
  
          start = time()
          outputs = model(input_ids, attn_mask)
          loss = criterion(outputs, input_ids)
          torch.cuda.synchronize()
          fwd_end = time()
          fwd_time = fwd_end - start
          logger.info(
              get_mem_info(prefix=f"[{n + 1}/{NUM_STEPS}] Forward "),
              ranks=[0],
          )
          booster.backward(loss, optimizer)
  
          torch.cuda.synchronize()
          bwd_end = time()
          bwd_time = bwd_end - fwd_end
          logger.info(
              get_mem_info(prefix=f"[{n + 1}/{NUM_STEPS}] Backward "),
              ranks=[0],
          )
  
          optimizer.step()
          torch.cuda.synchronize()
          optim_time = time() - bwd_end
          step_time = time() - start
          logger.info(
              get_mem_info(
                  prefix=f"[{n + 1}/{NUM_STEPS}] Optimizer step "
              ),
              ranks=[0],
          )
  
          step_tflops = get_tflops_func(step_time)
          logger.info(
              f"[{n + 1}/{NUM_STEPS}] Loss:{loss.item():.3f}, Step time: {step_time:.3f}s, TFLOPS: {get_tflops_func(step_time):.3f}, FWD time: {fwd_time:.3f}s, BWD time: {bwd_time:.3f}s, OPTIM time: {optim_time:.3f}s",
              ranks=[0],
          )
          if n >= WARMUP_STEPS:
              tflops_list.append(step_tflops)
  
      demo_profiler = get_profile_context(
          PROF_FLAG,
          WARMUP_STEPS,
          NUM_STEPS - WARMUP_STEPS,
          save_dir=f"profile/{get_time_stamp()}-demo",
      )
  
      with demo_profiler as prof:
          for n in range(NUM_STEPS):
              train_step()
              prof.step()
  
      tflops_list.sort()
      median_index = ((NUM_STEPS - WARMUP_STEPS) >> 1) + WARMUP_STEPS
      logger.info(f"Median TFLOPS is {tflops_list[median_index]:.3f}")
      torch.cuda.synchronize()
  
imports: import argparse
  import os
  from contextlib import nullcontext
  from functools import partial
  from time import time
  import psutil
  import torch
  import torch.nn
  from commons.model_zoo import model_builder
  from commons.utils import get_data, get_profile_context, get_tflops, get_time_stamp
  from packaging import version
  import colossalai
  from colossalai.booster import Booster
  from colossalai.booster.plugin import GeminiPlugin, LowLevelZeroPlugin, TorchDDPPlugin
  from colossalai.lazy import LazyInitContext
  from colossalai.logging import disable_existing_loggers, get_dist_logger
  from colossalai.nn.optimizer import HybridAdam
  from colossalai.utils import get_current_device
label_desc: Use image node_id 8 for calling a global function with example usage: main() without return types