Skip to content

llmcompressor.transformers.finetune.callbacks

Training callbacks for compression-aware fine-tuning workflows.

This module provides custom trainer callbacks that integrate compression session management with HuggingFace training loops. Handles precision control, training loop monitoring, and compression lifecycle events during model fine-tuning.

Classes:

DisableHalfPrecisionCallback 

DisableHalfPrecisionCallback(trainer, *args, **kwargs)

Bases: TrainerCallback

TrainerCallback for disabling FP16 training before QAT training begins

Parameters:

  • trainer

    –

    LLM Compressor trainer that will call back into this object

  • args

    –

    args to be passed to base TrainerCallback

  • kwargs

    –

    key word arguments to be passed to base TrainerCallback

Methods:

  • on_epoch_begin –

    Event called at the beginning of an epoch.

  • qat_active –

    :return: True if a quantization modifier is active in the current session

Source code in llmcompressor/transformers/finetune/callbacks.py 
 97
 98
 99
100
101
def __init__(self, trainer, *args, **kwargs):
    super().__init__(*args, **kwargs)
    self.trainer = trainer
    self.on_begin_called = False
    self.quant_start_epoch = math.inf

on_epoch_begin 

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

Event called at the beginning of an epoch.

Source code in llmcompressor/transformers/finetune/callbacks.py 
110
111
112
113
114
115
116
117
118
119
120
121
def on_epoch_begin(
    self,
    args: TrainingArguments,
    state: TrainerState,
    control: TrainerControl,
    **kwargs,
):
    """
    Event called at the beginning of an epoch.
    """
    super().on_epoch_begin(args, state, control, **kwargs)
    self.on_begin_called = True

qat_active 

qat_active() -> bool

Returns:

  • bool –

    True if a quantization modifier is active in the current session

Source code in llmcompressor/transformers/finetune/callbacks.py 
103
104
105
106
107
108
def qat_active(self) -> bool:
    """
    :return: True if a quantization modifier is active in the current session
    """
    session = active_session()
    return session.state.model.qat_active()

TrainingLoopCallbacks 

TrainingLoopCallbacks(trainer, *args, **kwargs)

Bases: TrainerCallback

TrainerCallback for triggering CompressionSession callbacks in the training loop. Used to update the model reference(for running with FSDP) and trigger the post- optim callbacks in each modifier.

Parameters:

  • trainer

    –

    LLM Compressor trainer that will call back into this object

  • args

    –

    args to be passed to base TrainerCallback

  • kwargs

    –

    key word arguments to be passed to base TrainerCallback

Methods:

  • on_step_end –

    Event called at the end of a training step. If using gradient accumulation,

  • on_substep_end –

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

  • on_train_begin –

    Event called at the beginning of training. Update the session reference to the

Source code in llmcompressor/transformers/finetune/callbacks.py 
35
36
37
def __init__(self, trainer, *args, **kwargs):
    super().__init__(*args, **kwargs)
    self.trainer = trainer

on_step_end 

on_step_end(
    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.

Triggers optimizer post_step and batch_end in the active CompressionSession

Source code in llmcompressor/transformers/finetune/callbacks.py 
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
def on_step_end(
    self,
    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.

    Triggers optimizer post_step and batch_end in the active CompressionSession
    """
    super().on_step_end(args, state, control, **kwargs)
    session_callbacks.optim_post_step()
    session_callbacks.batch_end()

on_substep_end 

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

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

Triggers optimizer post_step and batch_end in the active CompressionSession

Source code in llmcompressor/transformers/finetune/callbacks.py 
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
def on_substep_end(
    self,
    args: TrainingArguments,
    state: TrainerState,
    control: TrainerControl,
    **kwargs,
):
    """
    Event called at the end of an substep during gradient accumulation.

    Triggers optimizer post_step and batch_end in the active CompressionSession
    """
    super().on_substep_end(args, state, control, **kwargs)
    session_callbacks.optim_post_step()
    session_callbacks.batch_end()

on_train_begin 

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

Event called at the beginning of training. Update the session reference to the model, as it will have changed to a wrapper if FSDP is enabled

Source code in llmcompressor/transformers/finetune/callbacks.py 
39
40
41
42
43
44
45
46
47
48
49
50
51
52
def on_train_begin(
    self,
    args: TrainingArguments,
    state: TrainerState,
    control: TrainerControl,
    **kwargs,
):
    """
    Event called at the beginning of training. Update the session reference to the
    model, as it will have changed to a wrapper if FSDP is enabled
    """
    super().on_train_begin(args, state, control, **kwargs)
    session = active_session()
    session.state.model = self.trainer.model