lmflow.pipeline.utils.peft_trainer

Trainer for Peft models

Module Contents

Classes

PeftTrainer
PeftSavingCallback	Correctly save PEFT model and not full model

class lmflow.pipeline.utils.peft_trainer.PeftTrainer

Bases: transformers.Trainer

_save_checkpoint(_, trial, metrics=None)

Don’t save base model, optimizer etc. but create checkpoint folder (needed for saving adapter)

class lmflow.pipeline.utils.peft_trainer.PeftSavingCallback

Bases: transformers.trainer_callback.TrainerCallback

Correctly save PEFT model and not full model

_save(model, folder)

on_train_end(args: transformers.training_args.TrainingArguments, state: transformers.trainer_callback.TrainerState, control: transformers.trainer_callback.TrainerControl, **kwargs)

Save final best model adapter

on_epoch_end(args: transformers.training_args.TrainingArguments, state: transformers.trainer_callback.TrainerState, control: transformers.trainer_callback.TrainerControl, **kwargs)

Save intermediate model adapters in case of interrupted training

on_save(args: transformers.training_args.TrainingArguments, state: transformers.trainer_callback.TrainerState, control: transformers.trainer_callback.TrainerControl, **kwargs)