From 44b3006678cd7fd52b1fc8a8460163702c767a74 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Tue, 5 Nov 2024 01:36:11 -0800 Subject: [PATCH] Forward hook --- unsloth/models/_utils.py | 14 +++++++------- unsloth/models/llama.py | 1 + 2 files changed, 8 insertions(+), 7 deletions(-) diff --git a/unsloth/models/_utils.py b/unsloth/models/_utils.py index 22aa4d3d80..72994fa623 100644 --- a/unsloth/models/_utils.py +++ b/unsloth/models/_utils.py @@ -467,13 +467,13 @@ def prepare_model_for_kbit_training( pass # If use_reentrant = True which is the Pytorch default, we just make the input requires_grad. - if use_reentrant: - if hasattr(model, "enable_input_require_grads"): - model.enable_input_require_grads() - else: - def make_inputs_require_grad(module, input, output): - output.requires_grad_(True) - model.get_input_embeddings().register_forward_hook(make_inputs_require_grad) + # if use_reentrant: + # if hasattr(model, "enable_input_require_grads"): + # model.enable_input_require_grads() + # else: + # def make_inputs_require_grad(module, input, output): + # output.requires_grad_(True) + # model.get_input_embeddings().register_forward_hook(make_inputs_require_grad) return model pass diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py index c8bf5f8894..61857431b9 100644 --- a/unsloth/models/llama.py +++ b/unsloth/models/llama.py @@ -606,6 +606,7 @@ def LlamaModel_fast_forward( # Embed positions if inputs_embeds is None: inputs_embeds = self.embed_tokens(input_ids) + inputs_embeds.requires_grad_(True) # inputs_embeds = inputs_embeds.to(self.config.torch_dtype) torch_dtype = __DTYPE_MAP.get(self.config.torch_dtype, None)