support load from lora checkpoint

README.md

@@ -31,7 +31,7 @@ An extensible, convenient, and efficient toolbox for finetuning large machine le
 
 ## Latest News
 
-* [2025-03-18] With full support for Accelerate and lots of streamlining, LMFlow-nightly is now available! Feel free to try out the latest features and improvements by `git checkout lmflow-nightly`.
+* [2025-03-18] With full support for Accelerate and lots of streamlining, LMFlow-nightly is now available! Feel free to try out the latest features and improvements by `git checkout lmflow-nightly`, or see [here](https://github.com/OptimalScale/LMFlow/tree/lmflow-nightly).
 * [2024-12-02] Support [Hymba](https://github.com/NVlabs/hymba), a new family of small language models featuring a hybrid-head parallel architecture. Check out [Post-training Hymba](https://github.com/OptimalScale/LMFlow/tree/main/experimental/Hymba) for more details.
 * [2024-07-01] 🏆 LMFlow receives the [**Best Demo Paper Award**](https://docs.google.com/presentation/d/1TVDooAZqkNObz5ysVhDFtqnnVHR-u8wqYvgix-gzPMs/edit#slide=id.g2e55907bbcc_0_70) at **NAACL 2024**! 🎉
 * [2024-06-30] Expanding Optimization Options! We now support custom optimizer training with a variety of optimizers. Dive into the details and try out the new features with our updated script at [custom_optimizers](https://github.com/OptimalScale/LMFlow/blob/main/scripts/run_finetune_with_custom_optim.sh).

src/lmflow/models/hf_model_mixin.py

@@ -365,16 +365,7 @@ def __prepare_model_for_training(
 
         if model_args.use_lora or model_args.use_dora:
             model.enable_input_require_grads()
-            if model_args.lora_model_path is not None:
-                # Load model from LoRA weights
-                model = PeftModel.from_pretrained(
-                    model,
-                    model_args.lora_model_path,
-                    is_trainable=True,
-                )
-            else:
-                # New LoRA Finetuning
-                model = get_peft_model(model, self.peft_config)
+            model = get_peft_model(model, self.peft_config)
             model.print_trainable_parameters()
 
         # We resize the embeddings only when necessary to avoid index errors.

src/lmflow/pipeline/finetuner.py

@@ -39,7 +39,6 @@
 from lmflow.models.hf_encoder_decoder_model import HFEncoderDecoderModel
 from lmflow.models.hf_text_regression_model import HFTextRegressionModel
 from lmflow.pipeline.base_tuner import BaseTuner
-from lmflow.pipeline.utils.peft_trainer import PeftTrainer, PeftSavingCallback
 
 
 logger = logging.getLogger(__name__)
@@ -509,13 +508,9 @@ def compute_metrics(eval_preds):
 
         # Initialize our Trainer
         training_args = finetuner_args
-
-        if model_args.use_lora:
-            FinetuningTrainer = PeftTrainer
-            trainer_callbacks = [PeftSavingCallback]
-        else:
-            FinetuningTrainer = Trainer
-            trainer_callbacks = []
+        FinetuningTrainer = Trainer
+        trainer_callbacks = []
+
         if data_collator is None:
             data_collator = default_data_collator
 
@@ -603,6 +598,7 @@ def switch_active_layers(self):
             checkpoint = None
             last_checkpoint = self.last_checkpoint
             if training_args.resume_from_checkpoint is not None:
+                # this also supports load from lora checkpoint
                 checkpoint = training_args.resume_from_checkpoint
             elif last_checkpoint is not None:
                 checkpoint = last_checkpoint

src/lmflow/pipeline/rm_tuner.py

@@ -15,8 +15,7 @@
 from lmflow.datasets import Dataset
 from lmflow.models.hf_text_regression_model import HFTextRegressionModel
 from lmflow.pipeline.finetuner import Finetuner
-from lmflow.pipeline.utils.rm_trainer import compute_metrics, RewardTrainer, PeftRewardTrainer
-from lmflow.pipeline.utils.peft_trainer import PeftSavingCallback
+from lmflow.pipeline.utils.rm_trainer import compute_metrics, RewardTrainer
 from lmflow.pipeline.utils.rm_dataprocessor import RewardDataCollatorWithPadding
 
 
@@ -106,12 +105,8 @@ def tune(
             )
 
         # 2. prepare trainer
-        if self.model_args.use_lora:
-            RewardModelingTrainer = PeftRewardTrainer
-            trainer_callbacks = [PeftSavingCallback]
-        else:
-            RewardModelingTrainer = RewardTrainer
-            trainer_callbacks = []
+        RewardModelingTrainer = RewardTrainer
+        trainer_callbacks = []
 
         if self.finetuner_args.use_lisa:
             class DynamicLayerActivationCallback(TrainerCallback):