Class CheckpointingCallback

utils/callback_utils.py:72–84

class CheckpointingCallback(TrainerCallback):
    def __init__(self, steps_to_save):
        self.steps_to_save = steps_to_save

        self.rank = torch.distributed.get_rank()

    def on_step_end(self, args: TrainingArguments, state: TrainerState, control: TrainerControl, **kwargs):
        if state.global_step in self.steps_to_save:
            control.should_save = True
            control.should_evaluate = True
            if self.rank == 0:
                print(f"Saving checkpoint at step {state.global_step}")
        return control