Implement Missing Modifier Methods

[kylesayrs]

Purpose

• Fix model loading when recipe contains a modifier such as OutputDistillationModifier which previously did not implement critical functions like on_initialize_structure (see reload_stories_distill.py test script)
• Make implementing modifiers in the future easier than it is now
• Make existing implementations of modifiers easier to read by removing dummy methods

Traceback

Traceback (most recent call last):
  File "/home/ksayers/llm-compressor/reload_stories_distill.py", line 42, in <module>
    model = SparseAutoModelForCausalLM.from_pretrained(
  File "/home/ksayers/llm-compressor/src/llmcompressor/transformers/sparsification/sparse_model.py", line 152, in from_pretrained
    initialize_recipe(model=model, recipe_path=recipe)
  File "/home/ksayers/llm-compressor/src/llmcompressor/pytorch/model_load/helpers.py", line 77, in initialize_recipe
    pre_initialize_structure(model=model, recipe=recipe_path)
  File "/home/ksayers/llm-compressor/src/llmcompressor/core/session_functions.py", line 70, in pre_initialize_structure
    active_session().pre_initialize_structure(**kwargs)
  File "/home/ksayers/llm-compressor/src/llmcompressor/core/session.py", line 90, in pre_initialize_structure
    mod_data = self._lifecycle.pre_initialize_structure(
  File "/home/ksayers/llm-compressor/src/llmcompressor/core/lifecycle.py", line 94, in pre_initialize_structure
    data = mod.pre_initialize_structure(state=self.state, **extras)
  File "/home/ksayers/llm-compressor/src/llmcompressor/modifiers/stage.py", line 105, in pre_initialize_structure
    modifier.pre_initialize_structure(state, **kwargs)
  File "/home/ksayers/llm-compressor/src/llmcompressor/modifiers/modifier.py", line 87, in pre_initialize_structure
    self.on_initialize_structure(state, **kwargs)
  File "/home/ksayers/llm-compressor/src/llmcompressor/modifiers/modifier.py", line 234, in on_initialize_structure
    raise NotImplementedError()
NotImplementedError

Changes

• Remove requirement for subclass implementations of methods most relevant to training modifiers such as on_finalize, on_start, on_update, on_end, and on_event methods
• Enforce requirement for subclass implementations of methods relevant to training and one shot modifiers such as on_initialize
• Remove requirement for subclass implementations of methods which should be depreciated in the future such as on_initialize_structure

Testing

• Modifiers were found using command grep -P 'class\s+\w+\s*\(.*\bModifier\b.*\)' -r src/ (thank you chatgpt)
• Used below test script to confirm changes.

reload_stories_distill.py

import os, shutil
from llmcompressor.core import create_session
from llmcompressor.transformers import (
    SparseAutoModelForCausalLM,
    oneshot, train
)

output_dir = "./distill_out"
if os.path.exists(output_dir):
    shutil.rmtree(output_dir)
dataset = "open_platypus"
concatenate_data = False
splits = "train[:50%]"
max_steps = 50
num_calibration_samples = 64
recipe_str = "tests/llmcompressor/transformers/finetune/test_finetune_recipe.yaml"

# base
model = SparseAutoModelForCausalLM.from_pretrained(
    "Xenova/llama2.c-stories15M", device_map="auto"
)
distill_teacher = SparseAutoModelForCausalLM.from_pretrained(
    "Xenova/llama2.c-stories15M", device_map="auto"
)

# distill
with create_session():
    train(
        model=model,
        distill_teacher=distill_teacher,
        dataset=dataset,
        output_dir=output_dir,
        num_calibration_samples=num_calibration_samples,
        recipe=recipe_str,
        concatenate_data=concatenate_data,
        splits=splits,
        max_steps=max_steps,
    )

# load
model = SparseAutoModelForCausalLM.from_pretrained(
    output_dir, device_map="auto"
)

[github-actions]

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

[kylesayrs]

Confirmed no new test failures