add on lifecycle methods

Signed-off-by: Brian Dellabetta <bdellabe@redhat.com>

Author: brian-dellabetta

examples/transform/llama3_example.py

@@ -3,14 +3,13 @@
 
 from llmcompressor.modifiers.quantization import GPTQModifier
 from llmcompressor.modifiers.transform import TransformModifier
-from llmcompressor.transformers import oneshot
+from llmcompressor import oneshot
 
 # Select model and load it.
 MODEL_ID = "meta-llama/Meta-Llama-3-8B-Instruct"
 
 model = AutoModelForCausalLM.from_pretrained(
     MODEL_ID,
-    device_map="auto",
     torch_dtype="auto",
 )
 tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
@@ -66,6 +65,7 @@ def tokenize(sample):
     model=model,
     dataset=ds,
     recipe=recipe,
+    pipeline="sequential",
     max_seq_length=MAX_SEQUENCE_LENGTH,
     num_calibration_samples=NUM_CALIBRATION_SAMPLES,
 )

src/llmcompressor/modifiers/transform/transform.py

@@ -2,7 +2,7 @@
 
 from compressed_tensors.transform import TransformScheme, apply_transform_config
 
-from llmcompressor.core import State
+from llmcompressor.core import Event, EventType, State
 from llmcompressor.modifiers import Modifier
 
 from .template.quip import QUIP
@@ -12,9 +12,9 @@ class TransformModifier(Modifier):
     preset_config: Optional[str] = None
     config_groups: Optional[Dict[str, TransformScheme]] = None
 
-    # model validator to validate both preset and config gropus are not provided
+    # model validator to validate both preset and config groups are not provided
 
-    def on_initialize(self, state: State, **kwargs):
+    def on_initialize(self, state: State, **kwargs) -> bool:
         if self.preset_config is not None:
             # import config template and customize to model
             pass
@@ -23,4 +23,29 @@ def on_initialize(self, state: State, **kwargs):
         config = QUIP
 
         apply_transform_config(state.model, config)
-        breakpoint()
+
+        return True
+
+    def on_start(self, state: State, event: Event, **kwargs):
+        self.started_ = True
+
+    def on_event(self, state: State, event: Event, **kwargs):
+        if event.type_ == EventType.CALIBRATION_EPOCH_START:
+            if not self.started_:
+                self.on_start(state, None)
+
+        elif event.type_ == EventType.SEQUENTIAL_EPOCH_END:
+            pass
+
+        elif event.type_ == EventType.CALIBRATION_EPOCH_END:
+            if not self.ended_:
+                self.on_end(state, None)
+
+    def on_end(self, state: State, event: Event, **kwargs):
+        self.ended_ = True
+
+    def on_finalize(self, state: State, **kwargs) -> bool:
+        if not self.ended_:
+            self.on_end(state, None)
+
+        return True