code: def save_ckpt(
      model, optimizer, lr_scheduler, path, epoch, shard, global_step
  ):
      model_path = path + "_pytorch_model.bin"
      optimizer_lr_path = path + ".op_lrs"
      checkpoint = {}
      checkpoint["optimizer"] = optimizer.state_dict()
      checkpoint["lr_scheduler"] = lr_scheduler.state_dict()
      checkpoint["epoch"] = epoch
      checkpoint["shard"] = shard
      checkpoint["global_step"] = global_step
      model_state = (
          model.state_dict()
      )  # each process must run model.state_dict()
      if gpc.get_global_rank() == 0:
          torch.save(checkpoint, optimizer_lr_path)
          torch.save(model_state, model_path)
  
file_path: colossalai/examples/community/roberta/pretraining/pretrain_utils.py
imports: import os
  import sys
  import torch
  import transformers
  from transformers import get_linear_schedule_with_warmup
  from colossalai.legacy.core import global_context
  from colossalai.nn.optimizer import HybridAdam
  from collections import OrderedDict
  import torch.nn
  from model.bert import BertForMaskedLM
  from model.deberta_v2 import DebertaV2ForMaskedLM
label_desc: Use image node_id 5 for calling a global function with example usage: save_ckpt(model, optimizer, lr_scheduler, path, epoch, shard, global_step) without return types