revert train

vllm-project · Mar 6, 2025 · 8bf0dec · 8bf0dec
1 parent 110b91d
commit 8bf0dec
Showing 1 changed file with 9 additions and 0 deletions.
diff --git a/src/llmcompressor/transformers/finetune/text_generation.py b/src/llmcompressor/transformers/finetune/text_generation.py
@@ -263,6 +263,15 @@ def main(
         # exit immediately
         return
 
+    # Training
+    if training_args.do_train:
+        checkpoint = None
+        if training_args.resume_from_checkpoint is not None:
+            checkpoint = training_args.resume_from_checkpoint
+        elif last_checkpoint is not None:
+            checkpoint = last_checkpoint
+        stage_runner.train(checkpoint)
+
     # save if model was provided as a string or custom output_dir was set
     if isinstance(model_args.model, str) or (
         training_args.output_dir