calculate total_steps_per_epoch earlier, remove compression_ready

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

Author: kylesayrs

src/llmcompressor/core/lifecycle.py

@@ -218,12 +218,6 @@ def _check_setup_event_lifecycle(self, event_type: EventType):
                 "Cannot invoke event before recipe, model, and start are set"
             )
 
-        if not self.state.compression_ready:
-            logger.error("Cannot invoke event before recipe, model, and start are set")
-            raise ValueError(
-                "Cannot invoke event before recipe, model, and start are set"
-            )
-
         logger.debug("Setting up event lifecycle for event type: {}", event_type)
 
         for mod in self.modifiers:

src/llmcompressor/core/state.py

@@ -119,18 +119,6 @@ class State:
     model_log_cadence: Optional[float] = None
     _last_log_step: Union[float, int, None] = None
 
-    @property
-    def compression_ready(self) -> bool:
-        """
-        Check if the model and optimizer are set for compression.
-
-        :return: True if model and optimizer are set, False otherwise
-        :rtype: bool
-        """
-        ready = self.model is not None and self.optimizer is not None
-        logger.debug("Compression ready: {}", ready)
-        return ready
-
     def update(
         self,
         model: Any = None,

src/llmcompressor/transformers/finetune/session_mixin.py

@@ -149,6 +149,25 @@ def initialize_session(
 
         train_data = self.get_train_dataloader()
 
+        # calculate total_steps_per_epoch
+        # n_gpu handled internally by dataloader
+        total_batch_size = (
+            self.args.per_device_train_batch_size
+            * self.args.gradient_accumulation_steps
+        )
+        if isinstance(self.train_dataset, IterableDataset):
+            logger.warning(
+                "Training is being run with a streamed dataset, "
+                "steps_per_epoch cannot be determined and will default to "
+                "1. LLM Compressor modifiers utilizing this statistic may not "
+                "behave as expected. "
+            )
+            self.total_steps_per_epoch = 1
+        else:
+            self.total_steps_per_epoch = math.ceil(
+                len(self.train_dataset) / total_batch_size
+            )
+
         self.accelerator.wait_for_everyone()
         with summon_full_params_context(self.model, offload_to_cpu=True):
             initialize(
@@ -161,6 +180,7 @@ def initialize_session(
                 start=epoch,
                 copy_data=False,
                 fsdp_active=self.is_fsdp_enabled,
+                steps_per_epoch=self.total_steps_per_epoch,
                 metadata=self.metadata,
             )
         self.accelerator.wait_for_everyone()
@@ -203,29 +223,6 @@ def create_optimizer(self):
         self._check_super_defined("create_optimizer")
         super().create_optimizer()
 
-        # n_gpu handled internally by dataloader
-        total_batch_size = (
-            self.args.per_device_train_batch_size
-            * self.args.gradient_accumulation_steps
-        )
-
-        if isinstance(self.train_dataset, IterableDataset):
-            logger.warning(
-                "Training is being run with a streamed dataset, "
-                "steps_per_epoch cannot be determined and will default to "
-                "1. LLM Compressor modifiers utilizing this statistic may not "
-                "behave as expected. "
-            )
-            self.total_steps_per_epoch = 1
-        else:
-            self.total_steps_per_epoch = math.ceil(
-                len(self.train_dataset) / total_batch_size
-            )
-
-        active_session().state.update(
-            optimizer=self.optimizer, steps_per_epoch=self.total_steps_per_epoch
-        )
-
         return self.optimizer
 
     def create_scheduler(

tests/unit/core/test_state.py

@@ -67,16 +67,6 @@ def test_state_update():
     assert state.model_log_cadence == 2
 
 
-@pytest.mark.regression
-def test_state_sparsification_ready():
-    state = State()
-    assert not state.compression_ready
-
-    state.model = "model"
-    state.optimizer = "optimizer"
-    assert state.compression_ready
-
-
 @pytest.mark.regression
 def test_state_update_loggers():
     state = State()