[Callbacks] Remove pre_initialize_structure

[kylesayrs]

Purpose

• Remove pre_initialize_structure to simplify codebase
• Consolidate model saving logic
• Remove annoying and misleading logging messages

2025-02-17T17:48:38.477750-0500 | _check_create_state | INFO - State created for compression lifecycle
2025-02-17T17:48:38.478670-0500 | pre_initialize_structure | INFO - Compression lifecycle structure pre-initialized for 0 modifiers
2025-02-17T17:48:38.478836-0500 | pre_initialize_structure | INFO - Compression lifecycle structure pre-initialized for 0 modifiers

Prerequisites

• [Callbacks] Consolidate Saving Methods #1168

Follow-ups

• Remove double initialization

Changes

The preinitialization step used to fulfill a few purposes

• Construct the lifecycle state
  • This is now done by the dataclass directly

- state: Optional[State] = None
+ state: Optional[State] = field(default_factory=State)

• Populate state with model and recipe
  • This is now done (and has always been done) by initialize
  • Some functions such as Trainer.init_model attempt to access the model through the session before initialize is called. In these cases, we can pass the model directly

trainer = Trainer(
-     model_init=get_session_model,
+     model_init=lambda: model,

• Prepend recipes to the recipe.yaml if the model has already been compressed once
  • Move this logic from preinitialization to the save_pretrained function
  • Consolidate all save pathways to use the the same wrapped method

def save_pretrained_wrapper(...):
    update_and_save_recipe(model.name_or_path, save_directory)

• Provide a way for modifiers to influence the model after they have already been applied

  • This can still be a enacted via recipe validation, but likely no longer has a use case and shouldn't be done automatically, at most the LLM Compressor should warn if the recipe configuration is invalid / requires modification

• Create quantization modifier on GPTQ

  • This is now done within the on_initialize function
  • In the future, this should be done by a high-level recipe validation step

def on_initialize(...)
-     self.on_initialize_structure(state, **kwargs)
+     self._maybe_build_quant_modifier(state.model)

• Remove EventType.order() method which is unused

• Extend the Recipe.simplify_recipe class method to support strings

Lifecycle

1. create_session() (doesn't do much and can be hidden behind initialize)
2. initialize(model=..., recipe=...)
   1. Maybe start modifiers
3. LifecycleCallback.event(...)
   1. Maybe start/end modifiers
4. finalize()

Regression Evaluation

Main

vllm (pretrained=/home/kyle/llm-compressor/Meta-Llama-3-8B-Instruct-W4A16-G128,dtype=bfloat16,add_bos_token=True), gen_kwargs: (None), limit: None, num_fewshot: 5, batch_size: 1
|  Tasks   |Version|Filter|n-shot|Metric|   |Value |   |Stderr|
|----------|------:|------|-----:|------|---|-----:|---|-----:|
|winogrande|      1|none  |     5|acc   |↑  |0.7482|±  |0.0122|

This branch

vllm (pretrained=/home/kyle/llm-compressor/Meta-Llama-3-8B-Instruct-W4A16-G128,dtype=bfloat16,add_bos_token=True), gen_kwargs: (None), limit: None, num_fewshot: 5, batch_size: 1
|  Tasks   |Version|Filter|n-shot|Metric|   |Value |   |Stderr|
|----------|------:|------|-----:|------|---|-----:|---|-----:|
|winogrande|      1|none  |     5|acc   |↑  |0.7482|±  |0.0122|

[github-actions]

👋 Hi! Thank you for contributing to llm-compressor. Please add the ready label when the PR is ready for review.

Note: This is required to complete the testing suite, please only add the label once the PR is code complete and local testing has been performed.

[horheynm]

Good job XD