code: def main():
      args = set_args()
      config = AutoConfig.from_pretrained(
          args.model_dir, trust_remote_code=True
      )
      config.pre_seq_len = args.pre_seq_len
      config.prefix_projection = args.prefix_projection
  
      model = AutoModel.from_pretrained(
          args.model_dir, config=config, trust_remote_code=True
      )
      tokenizer = AutoTokenizer.from_pretrained(
          args.model_dir, trust_remote_code=True
      )
      model = model.half().cuda()
      if args.prefix_projection:
          model.gradient_checkpointing_enable()
  
      conf = {
          "train_micro_batch_size_per_gpu": args.train_batch_size,
          "gradient_accumulation_steps": args.gradient_accumulation_steps,
          "optimizer": {
              "type": "Adam",
              "params": {
                  "lr": 1e-5,
                  "betas": [0.9, 0.95],
                  "eps": 1e-8,
                  "weight_decay": 5e-4,
              },
          },
          "fp16": {"enabled": True},
          "zero_optimization": {
              "stage": 1,
              "offload_optimizer": {
                  "device": "cpu",
                  "pin_memory": True,
              },
              "allgather_partitions": True,
              "allgather_bucket_size": 2e8,
              "overlap_comm": True,
              "reduce_scatter": True,
              "reduce_bucket_size": 2e8,
              "contiguous_gradients": True,
          },
          "steps_per_print": args.log_steps,
      }
  
      for name, param in model.named_parameters():
          if not any(nd in name for nd in ["prefix_encoder"]):
              param.requires_grad = False
  
      print_trainable_parameters(model)
  
      for name, param in model.named_parameters():
          if param.requires_grad == True:
              print(name)
  
      train_dataset = Seq2SeqDataSet(
          args.train_path, tokenizer, args.max_len, args.max_src_len
      )
      train_dataloader = DataLoader(
          train_dataset,
          batch_size=conf["train_micro_batch_size_per_gpu"],
          sampler=RandomSampler(train_dataset),
          collate_fn=coll_fn,
          drop_last=True,
          num_workers=0,
      )
  
      model_engine, optimizer, _, _ = deepspeed.initialize(
          config=conf, model=model, model_parameters=model.parameters()
      )
      model_engine.train()
      global_step = 0
      for i_epoch in range(args.num_train_epochs):
          train_iter = iter(train_dataloader)
          for step, batch in enumerate(train_iter):
              input_ids = batch["input_ids"].cuda()
              labels = batch["labels"].cuda()
              outputs = model_engine.forward(
                  input_ids=input_ids, labels=labels, use_cache=False
              )
              loss = outputs[0]
              if conf["gradient_accumulation_steps"] > 1:
                  loss = loss / conf["gradient_accumulation_steps"]
              model_engine.backward(loss)
              torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
              if (step + 1) % conf["gradient_accumulation_steps"] == 0:
                  model_engine.step()
                  global_step += 1
              if global_step % args.log_steps == 0:
                  print(
                      "loss:{}, global_step:{}".format(
                          float(loss.item()), global_step
                      )
                  )
          save_dir = os.path.join(
              args.output_dir, f"global_step-{global_step}"
          )
          model.save_pretrained(save_dir)
          copy(
              os.path.join(args.model_dir, "tokenizer_config.json"),
              os.path.join(save_dir, "tokenizer_config.json"),
          )
          copy(
              os.path.join(args.model_dir, "ice_text.model"),
              os.path.join(save_dir, "ice_text.model"),
          )
  
file_path: DeepKE/example/llm/InstructKGC/src/finetuning_pt.py
imports: import argparse
  import os
  from shutil import copy
  import json
  import torch
  from torch.nn.utils.rnn import pad_sequence
  from torch.utils.data import Dataset, RandomSampler, DataLoader
  import deepspeed
  from transformers import AutoConfig, AutoTokenizer, AutoModel
label_desc: Use image node_id 4 for calling a global function with example usage: main() without return types