Expose flag for should_save_peft_only (#1357)

mosaicml · Jul 15, 2024 · 54746bf · 54746bf
1 parent d271210
commit 54746bf
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 1 deletion.
diff --git a/llmfoundry/models/hf/hf_causal_lm.py b/llmfoundry/models/hf/hf_causal_lm.py
@@ -90,6 +90,7 @@ def __init__(
         use_train_metrics: bool = True,
         additional_train_metrics: Optional[List] = None,
         additional_eval_metrics: Optional[List] = None,
+        should_save_peft_only: bool = True,
     ):
 
         config_overrides = config_overrides or {}
@@ -131,6 +132,7 @@ def __init__(
             eval_metrics=eval_metrics,
             init_device=init_device,
             peft_config=peft_config_object,
+            should_save_peft_only=should_save_peft_only,
         )
 
     @staticmethod

diff --git a/llmfoundry/models/hf/model_wrapper.py b/llmfoundry/models/hf/model_wrapper.py
@@ -40,6 +40,7 @@ def __init__(
         shift_labels: bool = False,
         init_device: Optional[str] = None,
         peft_config: Optional['PeftConfig'] = None,
+        should_save_peft_only: bool = True,
     ):
         super().__init__(
             model,
@@ -49,7 +50,7 @@ def __init__(
             eval_metrics=eval_metrics,
             shift_labels=shift_labels,
             peft_config=peft_config,
-            should_save_peft_only=True,
+            should_save_peft_only=should_save_peft_only,
         )
 
         self.prepare_inner_model(self.model, init_device)