class_docstr: This class mirrors deepspeed functionality. Not all changes are implemented in this version.
  
  See TorchEngineFull for more modifications.
class_docstr_tok: ["This","class","mirrors","deepspeed","functionality",".","Not","all","changes","are","implemented","in","this","version",".","See","TorchEngineFull","for","more","modifications","."]
docstr: This checkpoint can be used for downstream tasks.
  The default behavior is to save this checkpoint to a checkpoints folder under base_directory/name/checkpoints
docstr_tok: ["This","checkpoint","can","be","used","for","downstream","tasks",".","The","default","behavior","is","to","save","this","checkpoint","to","a","checkpoints","folder","under","base_directory\/name\/checkpoints"]
code: def save_final_model(
      self,
      base_directory,
      identifier,
      tokenizer,
      cfg_arch,
      dryrun=False,
  ):
      """This checkpoint can be used for downstream tasks.
      The default behavior is to save this checkpoint to a checkpoints folder under base_directory/name/checkpoints
      """
      try:
          identifier_str = f"{identifier:2.4f}"
      except ValueError:
          identifier_str = str(identifier)
      full_path = os.path.join(
          base_directory, "checkpoints", identifier_str
      )
      os.makedirs(full_path, exist_ok=True)
      # This saves tokenizer_config.json, tokenizer.json and special_tokens_map.json to this folder
      if not dryrun:
          tokenizer.save_pretrained(full_path)
          # Save model.safetensors, model_config.json
          save_file(
              self.retrieve_model_state_dict(),
              os.path.join(full_path, "model.safetensors"),
          )
          # legacy save: torch.save(self.retrieve_model_state_dict(), os.path.join(full_path, "model.pth"))
          with open(
              os.path.join(full_path, "model_config.json"), "w"
          ) as file:
              json.dump(
                  OmegaConf.to_container(cfg_arch, resolve=True), file
              )
  
imports: import torch
  import torch._inductor.utils
  import os
  import json
  from omegaconf import OmegaConf
  from functools import partial
  from contextlib import nullcontext
  import time
  import logging
  import transformers
  from safetensors.torch import load_file, save_file
  from transformers.utils.generic import working_or_temp_dir
  from torch.distributed.optim import ZeroRedundancyOptimizer
  from .utils import group_parameters, prepare_pretraining_dataloader, update_ema, updated_latest_weight_average
  from .optimizers.schedulers import get_schedule_fn
  from .optimizers import Adahessian, AdamWScale, Shampoo, LARS, SAM, ProgressiveBatching, AGD, Sophia
  import warnings
label_desc: Use image node_id 17 for calling the TorchEngineMinimal obj's underlying member method code with example usage: obj.save_final_model(base_directory, identifier, tokenizer, cfg_arch, dryrun) without return types