code: def main():
      disable_existing_loggers()
      launch_from_torch(config={})
      logger = get_dist_logger()
      config = transformers.GPT2Config(
          n_position=SEQ_LENGTH,
          n_layer=NUM_LAYERS,
          n_head=NUM_HEADS,
          n_embd=HIDDEN_DIM,
      )
      if FP16:
          model = GPT2LMHeadModel(config=config).half().to("cuda")
      else:
          model = GPT2LMHeadModel(config=config).to("cuda")
      global_numel = sum([p.numel() for p in model.parameters()])
  
      meta_input_sample = {
          "input_ids": torch.zeros(
              (BATCH_SIZE, SEQ_LENGTH), dtype=torch.int64
          ).to("meta"),
          "attention_mask": torch.zeros(
              (BATCH_SIZE, SEQ_LENGTH), dtype=torch.int64
          ).to("meta"),
      }
  
      gm, solution = autoparallelize(
          model, meta_input_sample, return_solution=True
      )
  
      # print solution on rank 0
      if gpc.get_global_rank() == 0:
          for node_strategy in solution:
              print(node_strategy)
  
      # build criterion
      criterion = GPTLMLoss()
  
      optimizer = torch.optim.Adam(gm.parameters(), lr=0.01)
      logger.info(get_mem_info(prefix="After init model, "), ranks=[0])
      get_tflops_func = partial(
          get_tflops, global_numel, BATCH_SIZE, SEQ_LENGTH
      )
      torch.cuda.synchronize()
      model.train()
  
      for n in range(10):
          # we just use randomly generated data here
          input_ids, attn_mask = get_data(
              BATCH_SIZE, SEQ_LENGTH, VOCAB_SIZE
          )
          optimizer.zero_grad()
          start = time()
          outputs = gm(input_ids, attn_mask)
          loss = criterion(outputs, input_ids)
          loss.backward()
          optimizer.step()
          torch.cuda.synchronize()
          step_time = time() - start
          logger.info(
              f"[{n+1}/{NUM_STEPS}] Loss:{loss.item():.3f}, Step time: {step_time:.3f}s, TFLOPS: {get_tflops_func(step_time):.3f}",
              ranks=[0],
          )
      torch.cuda.synchronize()
  
file_path: ColossalAI/examples/language/gpt/experiments/auto_parallel/auto_parallel_with_gpt.py
imports: from functools import partial
  from time import time
  import psutil
  import torch
  import transformers
  from gpt_modules import GPT2LMHeadModel, GPTLMLoss
  from colossalai.auto_parallel.tensor_shard.initialize import autoparallelize
  from colossalai.initialize import launch_from_torch
  from colossalai.legacy.core import global_context
  from colossalai.logging import disable_existing_loggers, get_dist_logger
label_desc: Use image node_id 6 for calling a global function with example usage: main() without return types