code: def on_step_end(
      self,
      args: TrainingArguments,
      state: TrainerState,
      control: TrainerControl,
      **kwargs,
  ):
      if (state.global_step + 1) % self.save_steps == 0:
          self.trainer.accelerator.wait_for_everyone()
          state_dict = self.trainer.accelerator.get_state_dict(
              self.trainer.deepspeed
          )
          unwrapped_model = self.trainer.accelerator.unwrap_model(
              self.trainer.deepspeed
          )
          if self.trainer.accelerator.is_main_process:
              unwrapped_model.save_pretrained(
                  args.output_dir, state_dict=state_dict
              )
          self.trainer.accelerator.wait_for_everyone()
      return control
  
file_path: autotrain-advanced/src/autotrain/trainers/clm/callbacks.py
imports: import os
  import torch
  from peft import set_peft_model_state_dict
  from transformers import TrainerCallback, TrainerControl, TrainerState, TrainingArguments
  from transformers.trainer_utils import PREFIX_CHECKPOINT_DIR
label_desc: Use image node_id 2 for calling the SaveDeepSpeedPeftModelCallback obj's underlying member method code with example usage: obj.on_step_end(args, state, control) and returns: control