start on initialize

Signed-off-by: Kyle Sayers <[email protected]>

Author: kylesayrs

src/llmcompressor/modifiers/modifier.py

@@ -89,6 +89,12 @@ def initialize(self, state: State, **kwargs):
 
         self.initialized_ = self.on_initialize(state=state, **kwargs)
 
+        # trigger start
+        fake_start_event = Event(type_=EventType.BATCH_START, global_step=0)
+        if self.should_start(fake_start_event):
+            self.on_start(state, fake_start_event, **kwargs)
+            self.started_ = True
+
     def finalize(self, state: State, **kwargs):
         """
         Finalize the modifier for the given model and state.

tests/llmcompressor/pytorch/modifiers/pruning/constant/test_pytorch.py

@@ -81,7 +81,7 @@ def test_constant_pruning_modifier_e2e(model, optimizer):
         end=1,
         update=0.5,
     )
-    modifier.initialize(state, start=0)
+    modifier.initialize(state)
 
     # check mask is added and has correct sparsity