reset session in between stages

Signed-off-by: Kyle Sayers <[email protected]>
vllm-project · Mar 10, 2025 · c2456f0 · c2456f0
1 parent 2a59554
commit c2456f0
Showing 1 changed file with 2 additions and 3 deletions.
diff --git a/src/llmcompressor/transformers/finetune/runner.py b/src/llmcompressor/transformers/finetune/runner.py
@@ -17,7 +17,6 @@
 from llmcompressor.core import active_session
 from llmcompressor.pytorch.model_load.helpers import (
     get_completed_stages,
-    get_session_model,
     save_checkpoint,
     save_completed_stages,
 )
@@ -222,6 +221,7 @@ def run_sequential_stages(
                     recipe_stage=stage_name,
                 )
             elif run_type is StageRunType.TRAIN:
+                self.trainer.model = model
                 self.train(checkpoint=checkpoint, stage=stage_name)
 
             checkpoint = None
@@ -248,11 +248,10 @@ def run_sequential_stages(
 
             # setup for next stage
             session = active_session()
-            session.reset_stage()
+            session.reset()
 
             # synchronize and clean up memory
             self.trainer.accelerator.wait_for_everyone()
-            self.trainer.model = get_session_model()
             torch.cuda.empty_cache()
             self.trainer.accelerator.free_memory()
             self.trainer.accelerator.wait_for_everyone()