allow double initialization

Signed-off-by: Kyle Sayers <kylesayrs@gmail.com>

Author: kylesayrs

src/llmcompressor/core/lifecycle.py

@@ -84,14 +84,8 @@ def initialize(
         :return: List of data returned from initialization of modifiers
         :rtype: List[Any]
         """
-        if self.initialized_:
-            raise ValueError(
-                "Initialize was called twice. To update state values prior to "
-                "initialization, please use `active_session().state.update()`"
-            )
-
-        self.state.update(**kwargs)
         logger.debug("Initializing compression lifecycle")
+        self.state.update(**kwargs)
         self.recipe_container.append(recipe, recipe_stage, recipe_args)
         self.modifiers = self.recipe_container.get_modifiers()
         self._set_model_layer_prefix()

src/llmcompressor/transformers/finetune/session_mixin.py

@@ -143,25 +143,6 @@ 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):
             active_session().initialize(
@@ -175,7 +156,6 @@ def initialize_session(
                 copy_data=False,
                 attach_optim_callbacks=True,
                 fsdp_active=self.is_fsdp_enabled,
-                steps_per_epoch=self.total_steps_per_epoch,
                 metadata=self.metadata,
             )
 
@@ -219,6 +199,29 @@ 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().initialize(
+            optimizer=self.optimizer, steps_per_epoch=self.total_steps_per_epoch
+        )
+
         return self.optimizer
 
     def create_scheduler(
@@ -255,7 +258,7 @@ def training_step(
         """
         self._check_super_defined("training_step")
 
-        callbacks.batch_start(batch_data=inputs)
+        callbacks.batch_start(batch_data=inputs, global_step=self.state.epoch)
         model_outputs = super().training_step(
             model=model, inputs=inputs, num_items_in_batch=num_items_in_batch
         )