class_docstr: Callback to merge LoRA weights into the base model and save full-weight checkpoints
code: def on_step_begin(
      self,
      args: TrainingArguments,
      state: TrainerState,
      control: TrainerControl,
      model: peft.LoraModel,
      optimizer: torch.optim.Optimizer,
      **_kwargs,
  ):
      if (
          state.global_step > 0
          and state.global_step % self.relora_steps == 0
      ):
          checkpoint_folder = os.path.join(
              args.output_dir,
              f"{PREFIX_CHECKPOINT_DIR}-{state.global_step}",
              "relora",
          )
  
          with torch.no_grad():
              merge_and_save(
                  model,
                  self.last_full_model,
                  checkpoint_folder,
                  reinit=True,
                  quantized=self.quantized,
                  actually_save=is_main_process(),
                  cpu_offload=self.cpu_offload,
              )
              reset_optimizer(optimizer)
  
          if self.quantized:
              self.last_full_model = checkpoint_folder
          self.num_lora_restarts += 1
  
      return control
  
imports: import glob
  import json
  import logging
  import os.path
  import shutil
  from pathlib import Path
  from typing import Dict, List, Sequence
  import bitsandbytes
  import peft
  import safetensors.torch
  import torch
  from huggingface_hub import snapshot_download
  from torch.optim.lr_scheduler import LRScheduler
  from torch.optim.optimizer import Optimizer
  from transformers import TrainerCallback, TrainerControl, TrainerState, TrainingArguments
  from transformers.trainer_utils import PREFIX_CHECKPOINT_DIR
  from axolotl.utils.dict import DictDefault
  from axolotl.utils.distributed import is_main_process
label_desc: Use image node_id 3 for calling the ReLoRACallback obj's underlying member method code with example usage: obj.on_step_begin(args, state, control, model, optimizer) and returns: control