slicktune.callbacks

Trainer callbacks used by slick-tune strategies.

Classes

AdaLoRACallback

Call update_and_allocate while AdaLoRA grads are still available.

Module Contents

class slicktune.callbacks.AdaLoRACallback

Bases: transformers.TrainerCallback

Call update_and_allocate while AdaLoRA grads are still available.

PEFT requires update_and_allocate after optimizer.step() and before zero_grad(). Hugging Face Trainer clears grads before on_step_end, so this hooks on_optimizer_step instead.

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

Event called at the beginning of an epoch.

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

Event called at the end of an epoch.

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

Event called after an evaluation phase.

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

Event called at the end of the initialization of the [Trainer].

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

Event called after logging the last logs.

on_optimizer_step(args: transformers.TrainingArguments, state: transformers.TrainerState, control: transformers.TrainerControl, **kwargs: Any) → None

Update AdaLoRA ranks using current parameter gradients.

Parameters:

• args (TrainingArguments) – Trainer args (unused).

• state (TrainerState) – Current trainer state (provides global_step).

• control (TrainerControl) – Trainer control flow (unused).

• **kwargs (Any) – Must include model when available.

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

Event called before the optimizer step but after gradient clipping. Useful for monitoring gradients.

on_predict(args: transformers.training_args.TrainingArguments, state: TrainerState, control: TrainerControl, metrics, **kwargs)

Event called after a successful prediction.

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

Event called after a prediction step.

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

Event called before pushing the model to the hub, at the beginning of Trainer.push_to_hub and Trainer._push_from_checkpoint.

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

Event called after a checkpoint save.

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

Event called at the beginning of a training step. If using gradient accumulation, one training step might take several inputs.

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

Event called at the end of a training step. If using gradient accumulation, one training step might take several inputs.

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

Event called at the end of an substep during gradient accumulation.

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

Event called at the beginning of training.

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

Event called at the end of training.