[Callbacks] Remove pre_initialize_structure

src/llmcompressor/entrypoints/oneshot.py

@@ -254,6 +254,7 @@ def _pre_process(self):
         if isinstance(self.model_args.model, (str, PosixPath)):
             self.model_args.model, _ = initialize_model_from_path(self.model_args)
 
+        breakpoint()
         patch_tied_tensors_bug(self.model_args.model)
         modify_save_pretrained(self.model_args.model)
 

src/llmcompressor/transformers/finetune/runner.py

@@ -6,6 +6,7 @@
 import torch
 from loguru import logger
 from torch.utils.data import Dataset
+from transformers import PreTrainedModel
 
 from llmcompressor.args import (
     DatasetArguments,
@@ -154,7 +155,9 @@ def train(self, checkpoint: str, stage: Optional[str] = None):
         # this includes saving the state, optimizer and scheduler
         self.trainer.save_model(output_dir=self._output_dir)
 
-    def run_sequential_stages(self, checkpoint: Optional[str] = None):
+    def run_sequential_stages(
+        self, model: PreTrainedModel, checkpoint: Optional[str] = None
+    ):
         """
         Run the recipe stage by stage, allowing for alternating between one-shot and
         finetuning flows. Optionally save the model output at the end of each stage
@@ -195,7 +198,6 @@ def run_sequential_stages(self, checkpoint: Optional[str] = None):
             if run_type is StageRunType.ONESHOT:
                 from llmcompressor import Oneshot
 
-                model = get_session_model()
                 self._model_args.model = model
 
                 oneshot = Oneshot.from_args(

src/llmcompressor/transformers/finetune/text_generation.py

@@ -422,7 +422,7 @@ def main(
         checkpoint = None
         if last_checkpoint is not None:
             checkpoint = last_checkpoint
-        stage_runner.run_sequential_stages(checkpoint)
+        stage_runner.run_sequential_stages(model, checkpoint)
 
         # exit immediately
         return