integrations.llm_compressor.plugin

integrations.llm_compressor.plugin

Sparse Finetuning plugin for Axolotl — enables handling of sparse neural networks by maintaining masks for zero weights during training.

Classes

Name	Description
LLMCompressorCallbackHandler	Trainer callback for Sparse Finetuning.
LLMCompressorPlugin	Sparse Finetuning plugin for Axolotl integration.

LLMCompressorCallbackHandler

integrations.llm_compressor.plugin.LLMCompressorCallbackHandler(trainer, recipe)

Trainer callback for Sparse Finetuning. Maintains sparsity patterns during training by applying masks after optimization steps, ensuring zero-weight updates are canceled out.

Methods

Name	Description
on_step_begin	Called at the beginning of a training step. Triggers batch_start callback.
on_step_end	Called at the end of a training step. Triggers optimizer and batch_end callbacks.
on_train_begin	Called at the beginning of training. Initializes the compression session.
on_train_end	Called at the end of training. Finalizes the compression session.

on_step_begin

integrations.llm_compressor.plugin.LLMCompressorCallbackHandler.on_step_begin(
    args,
    state,
    control,
    **kwargs,
)

Called at the beginning of a training step. Triggers batch_start callback.

on_step_end

integrations.llm_compressor.plugin.LLMCompressorCallbackHandler.on_step_end(
    args,
    state,
    control,
    **kwargs,
)

Called at the end of a training step. Triggers optimizer and batch_end callbacks.

on_train_begin

integrations.llm_compressor.plugin.LLMCompressorCallbackHandler.on_train_begin(
    args,
    state,
    control,
    **kwargs,
)

Called at the beginning of training. Initializes the compression session.

Parameters

Name	Type	Description	Default
args	TrainingArguments	Training arguments.	required
state	TrainerState	Trainer state.	required
control	TrainerControl	Trainer control.	required

on_train_end

integrations.llm_compressor.plugin.LLMCompressorCallbackHandler.on_train_end(
    args,
    state,
    control,
    **kwargs,
)

Called at the end of training. Finalizes the compression session.

LLMCompressorPlugin

integrations.llm_compressor.plugin.LLMCompressorPlugin()

Sparse Finetuning plugin for Axolotl integration.

Methods

Name	Description
add_callbacks_post_trainer	Adds Sparse Finetuning callback to the Trainer instance.
get_input_args	Returns the path to the plugin’s argument definition.

add_callbacks_post_trainer

integrations.llm_compressor.plugin.LLMCompressorPlugin.add_callbacks_post_trainer(
    cfg,
    trainer,
)

Adds Sparse Finetuning callback to the Trainer instance.

Parameters

Name	Type	Description	Default
cfg	Any	Configuration object containing the sparse recipe.	required
trainer	Trainer	Huggingface Trainer instance.	required

Returns

Name	Type	Description
list	list	List containing the configured callback instances.

get_input_args

integrations.llm_compressor.plugin.LLMCompressorPlugin.get_input_args()

Returns the path to the plugin’s argument definition.

Returns

Name	Type	Description
str	str	Dotted path to the LLMCompressorArgs class.

Functions

Name	Description
compute_loss_wrapper	Wraps the loss computation function to trigger the loss_calculated callback.

compute_loss_wrapper

integrations.llm_compressor.plugin.compute_loss_wrapper(compute_loss_func)

Wraps the loss computation function to trigger the loss_calculated callback.

Parameters

Name	Type	Description	Default
compute_loss_func	Callable	Original loss computation function.	required

Returns

Name	Type	Description
Callable	Callable[Concatenate[Module, P], R]	Wrapped function that also invokes the loss_calculated callback.