code: def finetune(
      args,
      tokenizer: CPMBeeTokenizer,
      model: CPMBee,
      optimizer: bmt.optim.AdamOffloadOptimizer,
      lr_scheduler: bmt.lr_scheduler.WarmupLRScheduler,
      optim_manager: bmt.optim.OptimManager,
  ):
      average_time = bmt.utils.AverageRecorder()
      if model.config.dtype == torch.half:
          loss_func = bmt.loss.FusedCrossEntropy(ignore_index=-100)
      else:
          loss_func = torch.nn.CrossEntropyLoss(ignore_index=-100)
  
      if args.tensorboard is not None and bmt.rank() == 0:
          from torch.utils.tensorboard import SummaryWriter
          import distutils.version  # noqa: F401
  
          if not os.path.exists(args.tensorboard):
              os.makedirs(args.tensorboard)
          writer = SummaryWriter(log_dir=args.tensorboard)
  
      best_eval_loss, eval_loss_increase = 1e9, 0
      global_token_pass = 0.0
      global_steps = 0
      global_world_size = bmt.world_size()
      dataloader = FinetuneDataset(
          args.dataset,
          args.batch_size,
          args.max_length,
          tokenizer,
          max_depth=8,
          task_name=args.task_name,
          drop_last=args.drop_last,
      )
  
      for epoch in range(args.epoch):
          epoch = epoch + 1
          last_data = None
          for iteration, data in enumerate(dataloader):
              iteration = iteration + 1
              global_steps = global_steps + 1
              skip_this_batch = False
              if data is None:
                  if last_data is None:
                      raise RuntimeError(
                          "Dataset is too small, please use a smaller batch size or sequence length!"
                      )
                  data = last_data  # use last data
                  skip_this_batch = True
              else:
                  last_data = data
  
              input_ids = (
                  torch.from_numpy(data["inputs"])
                  .cuda()
                  .to(torch.int32)
              )
              input_ids_sub = (
                  torch.from_numpy(data["inputs_sub"])
                  .cuda()
                  .to(torch.int32)
              )
              input_length = (
                  torch.from_numpy(data["length"])
                  .cuda()
                  .to(torch.int32)
              )
              input_context = (
                  torch.from_numpy(data["context"]).cuda().bool()
              )
              input_sample_ids = (
                  torch.from_numpy(data["sample_ids"])
                  .cuda()
                  .to(torch.int32)
              )
              input_num_segments = (
                  torch.from_numpy(data["num_segments"])
                  .cuda()
                  .to(torch.int32)
              )
              input_segment_ids = (
                  torch.from_numpy(data["segment_ids"])
                  .cuda()
                  .to(torch.int32)
              )
              input_segment_rel_offset = (
                  torch.from_numpy(data["segment_rel_offset"])
                  .cuda()
                  .to(torch.int32)
              )
              input_segment_rel = (
                  torch.from_numpy(data["segment_rel"])
                  .cuda()
                  .to(torch.int32)
              )
              input_span = (
                  torch.from_numpy(data["spans"]).cuda().to(torch.int32)
              )
              targets = (
                  torch.from_numpy(data["target"])
                  .cuda()
                  .to(torch.int32)
              )
              ext_table_ids = (
                  torch.from_numpy(data["ext_ids"])
                  .cuda()
                  .to(torch.int32)
              )
              ext_table_sub = (
                  torch.from_numpy(data["ext_sub"])
                  .cuda()
                  .to(torch.int32)
              )
              task_ids = (
                  torch.from_numpy(data["task_ids"])
                  .cuda()
                  .to(torch.int32)
              )
              task_names = data["task_names"]
              # ===========
              optim_manager.zero_grad()
              mem_usage = {}
              tim_usage = {}
              mem_usage, tim_usage = add_mem_time(
                  "init", mem_usage, tim_usage
              )
  
              # ===========
              logits, _ = model(
                  input_ids,
                  input_ids_sub,
                  input_length,
                  input_context,
                  input_sample_ids,
                  input_num_segments,
                  input_segment_ids,
                  input_segment_rel_offset,
                  input_segment_rel,
                  input_span,
                  ext_table_ids,
                  ext_table_sub,
              )
              loss = loss_func(
                  logits.view(-1, logits.size(-1)),
                  targets.long().view(-1),
              )
              if skip_this_batch:
                  loss = loss * 0
  
              mem_usage, tim_usage = add_mem_time(
                  "forward", mem_usage, tim_usage
              )
  
              # ===========
              optim_manager.backward(loss)
              mem_usage, tim_usage = add_mem_time(
                  "backward", mem_usage, tim_usage
              )
  
              # ===========
              grad_norm = optim_manager.clip_grad_norm(
                  optimizer.param_groups, max_norm=1.0
              )
              optim_manager.step()
              mem_usage, tim_usage = add_mem_time(
                  "optim", mem_usage, tim_usage
              )
  
              # ==========
              iteration_time = tim_usage["optim"] - tim_usage["init"]
              average_time.record(iteration_time)
  
              with torch.no_grad():
                  task_num = len(task_names)
                  targets_tmp = targets.expand(task_num, -1, -1)
                  task = torch.arange(
                      task_num, dtype=torch.int32, device="cuda"
                  )[:, None, None]
                  targets_tmp = torch.where(
                      task_ids == task,
                      targets_tmp,
                      torch.scalar_tensor(
                          -100, dtype=torch.int32, device="cuda"
                      ),
                  )
  
                  task_loss_map: Dict[str, float] = {}
                  if not skip_this_batch:
                      for i in range(task_num):
                          task_loss = loss_func(
                              logits.view(-1, logits.size(-1)),
                              targets_tmp[i, :].long().view(-1),
                          )
                          task_loss_map[
                              task_names[i]
                          ] = task_loss.item()
                  gatherd_task_loss_map: List[
                      Dict[str, float]
                  ] = allgather_objects(task_loss_map)
  
                  global_task_loss_map: Dict[
                      str, Union[List[float], float]
                  ] = {}
                  for local_task_loss_map in gatherd_task_loss_map:
                      for (
                          task_name,
                          task_loss,
                      ) in local_task_loss_map.items():
                          if task_name not in global_task_loss_map:
                              global_task_loss_map[task_name] = []
                          global_task_loss_map[task_name].append(
                              task_loss
                          )
  
                  task_loss_map = {}
                  for task_name in sorted(
                      list(global_task_loss_map.keys())
                  ):
                      avg_loss = sum(
                          global_task_loss_map[task_name]
                      ) / len(global_task_loss_map[task_name])
                      task_loss_map[task_name] = avg_loss
  
              local_total_rate = torch.Tensor(
                  [input_length.float().mean() / args.max_length]
              ).cuda()
              local_total_rate = bmt.sum_loss(local_total_rate).item()
              global_token_pass += (
                  global_world_size
                  * local_total_rate
                  * args.max_length
                  * args.batch_size
              )
              avg_time = average_time.value
  
              train_info = {
                  "time": tim_usage["init"],
                  "epoch": epoch,
                  "iteration": iteration,
                  "loss": task_loss_map[args.task_name],
                  "lr": lr_scheduler.current_lr,
                  "lr_scale": int(optim_manager.loss_scale),
                  "time_usage": tim_usage,
                  "mem_usage": mem_usage,
                  "avg_time": avg_time,
                  "token_max": local_total_rate,
                  "token_pass": global_token_pass,
                  "throughout": args.max_length
                  * args.batch_size
                  * local_total_rate
                  / avg_time,
                  "grad_norm": grad_norm.item(),
                  "mask_max": (
                      (targets >= 0).sum(-1).float().mean()
                      / args.max_length
                  ).item(),
                  "num_gpus": global_world_size,
                  "task_loss": task_loss_map,
              }
  
              bmt.print_rank(
                  (
                      "| Epoch: {:3d} | Iter: {:6d} | loss: {:.4f} "
                      + "| lr: {:.4e}, scale: {:10.4f} | time: {:.4f} |"
                      + " token/max: {:.4f} | mask/max: {:.4f} | grad_norm: {:.4f}"
                  ).format(
                      epoch,
                      iteration,
                      task_loss_map[args.task_name],
                      lr_scheduler.current_lr,
                      int(optim_manager.loss_scale),
                      avg_time,
                      input_length.float().mean() / args.max_length,
                      (targets >= 0).sum(-1).float().mean()
                      / args.max_length,
                      grad_norm,
                  )
              )
              bmt.print_rank(
                  "| "
                  + " | ".join(
                      [
                          "{} loss: {:.4f}".format(task_name, loss)
                          for task_name, loss in task_loss_map.items()
                      ]
                  )
              )
              if iteration % args.inspect_iters == 0:
                  model_inspect = bmt.inspect.inspect_model(model, "*")
                  bmt.print_rank(
                      bmt.inspect.format_summary(model_inspect)
                  )
                  train_info["model_inspect"] = model_inspect
  
              # write log here
              if args.tensorboard is not None and bmt.rank() == 0:
                  writer.add_scalar(
                      "Loss/train",
                      task_loss_map[args.task_name],
                      global_steps,
                  )
                  for task_name, loss in task_loss_map.items():
                      writer.add_scalar(
                          "Loss/train/{}".format(task_name),
                          loss,
                          global_steps,
                      )
  
              # evaluation
              if global_steps % args.eval_interval == 0:
                  eval_loss = evaluation(
                      model, args, tokenizer, loss_func
                  )
                  if args.tensorboard is not None and bmt.rank() == 0:
                      writer.add_scalar(
                          "Loss/eval", eval_loss, global_steps
                      )
                  if eval_loss < best_eval_loss:
                      best_eval_loss = eval_loss
                      eval_loss_increase = 0
                      if args.save is not None:
                          if not args.use_delta:
                              bmt.save(
                                  model,
                                  os.path.join(
                                      args.save,
                                      args.save_name + "-best.pt",
                                  ),
                              )
                          else:
                              state_dict = model.state_dict()
                              if bmt.rank() == 0:
                                  torch.save(
                                      state_dict,
                                      os.path.join(
                                          args.save,
                                          args.save_name
                                          + "-delta-best.pt",
                                      ),
                                  )
                  else:
                      eval_loss_increase += 1
                  bmt.print_rank(
                      "| Eval loss: {:.4f} | Increase: {:2d}".format(
                          eval_loss, eval_loss_increase
                      )
                  )
                  if eval_loss_increase == args.early_stop_patience:
                      bmt.print_rank(
                          "Eval loss has increased {:d} times, the finetune loop early stopped.".format(
                              eval_loss_increase
                          )
                      )
                      return
  
imports: import time
  from typing import Dict, List, Union
  import torch
  import bmtrain
  import os
  from opendelta import LoraModel
  from cpm_live.arguments import get_args
  from cpm_live.models import CPMBee, CPMBeeConfig
  from cpm_live.tokenizers import CPMBeeTokenizer
  from cpm_live.utils import allgather_objects
  from cpm_live.training_tasks.bee import FinetuneDataset
label_desc: Use image node_id 10 for calling a global function with example usage: finetune(args, tokenizer, model, optimizer, lr_scheduler, optim_manager) without return types