From e2c04a2d79858fe102a795a48fbdae228bb442f4 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Tue, 4 Mar 2025 19:11:54 -0800 Subject: [PATCH] Update llama.py --- unsloth/models/llama.py | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py index 297c2984c3..9e764197e1 100644 --- a/unsloth/models/llama.py +++ b/unsloth/models/llama.py @@ -1535,7 +1535,6 @@ class LongRopeRotaryEmbedding(torch.nn.Module): pass -@torch.inference_mode def unsloth_fast_generate( self, inputs: Optional[torch.Tensor] = None, @@ -1577,7 +1576,7 @@ def unsloth_fast_generate( kwargs["pad_token_id"] = kwargs.pop("pad_token_id", model_eos_token_id) # Mixed precision autocast - with torch.autocast(device_type = "cuda", dtype = dtype): + with torch.inference_mode(), torch.autocast(device_type = "cuda", dtype = dtype): output = self._old_generate(self, inputs, *args, **kwargs) pass @@ -1972,7 +1971,7 @@ class FastLlamaModel: # Patch generate if model.generate.__name__ != "unsloth_fast_generate": model._old_generate = model.generate - model.generate = unsloth_fast_generate + model.generate = types.MethodType(unsloth_fast_generate, model) model.generate.__doc__ = model._old_generate.__doc__ return model, tokenizer pass @@ -2410,7 +2409,7 @@ class FastLlamaModel: # Patch generate if model.generate.__name__ != "unsloth_fast_generate": model._old_generate = model.generate - model.generate = unsloth_fast_generate + model.generate = types.MethodType(unsloth_fast_generate, model) model.generate.__doc__ = model._old_generate.__doc__ return model pass