code: def bench_shardformer(
      BATCH,
      N_CTX,
      provider,
      model_func,
      dtype=torch.float32,
      device="cuda",
  ):
      warmup = 10
      rep = 100
      # prepare data
      data = data_gen_for_sequence_classification(BATCH, N_CTX)
      data = {k: v.cuda() for k, v in data.items()}
      model = model_func().to(device)
      model.train()
      if provider == "org_model":
          fn = lambda: train(model, data)
          ms = triton.testing.do_bench(fn, warmup=warmup, rep=rep)
          return ms
      if provider == "shard_model":
          shard_config = ShardConfig(
              enable_fused_normalization=True,
              enable_tensor_parallelism=True,
          )
          shard_former = ShardFormer(shard_config=shard_config)
          sharded_model, _ = shard_former.optimize(model)
          sharded_model = sharded_model.cuda()
          fn = lambda: train(sharded_model, data)
          ms = triton.testing.do_bench(fn, warmup=warmup, rep=rep)
          return ms
  
file_path: colossalai/colossalai/shardformer/examples/performance_benchmark.py
imports: import torch
  import torch.distributed
  import transformers
  import triton
  import colossalai
  from colossalai.shardformer import ShardConfig, ShardFormer
label_desc: Use image node_id 4 for calling a global function with example usage: bench_shardformer(BATCH, N_CTX, provider, model_func, dtype, device) and returns: ms, ms