class_docstr: Simple model saver to filesystem
class_docstr_tok: ["Simple","model","saver","to","filesystem"]
code: def _st_save(self, step, model):
      try:
          from safetensors.torch import save_file
      except ImportError:
          raise ImportError(
              "run: pip install safetensors, to use safetensors"
          )
      if (
          hasattr(self.model_opt, "lora_layers")
          and len(self.model_opt.lora_layers) > 0
      ) or (
          hasattr(self.model_opt, "lora_embedding")
          and self.model_opt.lora_embedding
      ):
          model_state_dict = lora_state_dict(model, bias="lora_only")
      else:
          model_state_dict = model.state_dict()
  
      if torch.distributed.is_initialized():
          ws = torch.distributed.get_world_size()
      else:
          ws = 1
      if ws > 1:
          full_model = [None for _ in range(ws)]
          for key, value in model_state_dict.items():
              model_state_dict[key] = value.cpu()
          torch.distributed.all_gather_object(
              full_model, model_state_dict
          )
          fm_sd = {}
          for key in full_model[0].keys():
              if key.split(".")[-1] == "lora_A":
                  if key.split(".")[-2] in [
                      "linear_keys",
                      "linear_values",
                      "linear_query",
                      "w_1",
                      "w_3",
                  ]:
                      fm_sd[key] = (
                          sum(
                              [
                                  full_model[i][key].cpu()
                                  for i in range(ws)
                              ]
                          )
                          / ws
                      )
                  elif key.split(".")[-2] in ["final_linear", "w_2"]:
                      fm_sd[key] = torch.cat(
                          [full_model[i][key].cpu() for i in range(ws)],
                          1,
                      )
              elif key.split(".")[-1] == "lora_B":
                  if key.split(".")[-2] in [
                      "linear_keys",
                      "linear_values",
                      "linear_query",
                      "w_1",
                      "w_3",
                  ]:
                      fm_sd[key] = torch.cat(
                          [full_model[i][key].cpu() for i in range(ws)],
                          0,
                      )
                  elif key.split(".")[-2] in ["final_linear", "w_2"]:
                      fm_sd[key] = (
                          sum(
                              [
                                  full_model[i][key].cpu()
                                  for i in range(ws)
                              ]
                          )
                          / ws
                      )
              elif key.split(".")[-1] in [
                  "linear_keys",
                  "linear_values",
                  "linear_query",
                  "w_1",
                  "w_3",
              ]:
                  fm_sd[key] = torch.cat(
                      [full_model[i][key].cpu() for i in range(ws)], 0
                  )
              elif key.split(".")[-1] in ["final_linear", "w_2"]:
                  fm_sd[key] = torch.cat(
                      [full_model[i][key].cpu() for i in range(ws)], 1
                  )
              else:
                  fm_sd[key] = full_model[0][key]
          model_state_dict = fm_sd
  
      checkpoint = {
          "vocab": vocabs_to_dict(self.vocabs),
          "opt": self.model_opt,
          "optim": self.optim.state_dict(),
      }
  
      if (
          not torch.distributed.is_initialized()
          or torch.distributed.get_rank() == 0
      ):
          logger.info(
              "Saving checkpoint %s_step_%d.pt" % (self.base_path, step)
          )
          ckpt_path = "%s_step_%d.pt" % (self.base_path, step)
          torch.save(checkpoint, ckpt_path)
          logger.info(
              "Saving safetensors %s_step_%d.pt"
              % (self.base_path, step)
          )
          model_path = "%s_step_%d.safetensors" % (self.base_path, step)
          save_file(model_state_dict, model_path)
      else:
          ckpt_path = None
          model_path = None
      if torch.distributed.is_initialized():
          torch.distributed.barrier()
  
      return ckpt_path, model_path
  
file_path: OpenNMT-py/onmt/models/model_saver.py
imports: import os
  import torch
  import re
  from collections import deque
  from onmt.utils.logging import logger
  from onmt.inputters.inputter import vocabs_to_dict
  from onmt.modules.lora import lora_state_dict
label_desc: Use image node_id 2 for calling the ModelSaver obj's underlying member method code with example usage: obj._st_save(step, model) and returns: ckpt_path,  model_path