From 51257d37a82acf42b6fee2af63695728d7c3bdf3 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Wed, 1 Jan 2025 22:32:58 -0800 Subject: [PATCH] Update _utils.py --- unsloth/models/_utils.py | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/unsloth/models/_utils.py b/unsloth/models/_utils.py index 386d71354d..9bd3598b19 100644 --- a/unsloth/models/_utils.py +++ b/unsloth/models/_utils.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -__version__ = "2024.12.12" +__version__ = "2025.1.1" __all__ = [ "prepare_model_for_kbit_training", @@ -110,6 +110,9 @@ from unsloth_zoo.compiler import ( get_transformers_model_type, unsloth_compile_transformers as _unsloth_compile_transformers, ) +from unsloth_zoo.peft_utils import ( + requires_grad_for_gradient_checkpointing, +) # ============================================= # Disable some warnings which can get annoying @@ -557,6 +560,10 @@ def prepare_model_for_kbit_training( output.requires_grad_(True) model.get_input_embeddings().register_forward_hook(make_inputs_require_grad) + # Enable grads on non language models as well + requires_grad_for_gradient_checkpointing() + pass + return model pass