From 648ca59185637a540dea2c233f2d067bad83eb3c Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Tue, 5 Nov 2024 12:05:17 -0800 Subject: [PATCH] Update _utils.py --- unsloth/models/_utils.py | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/unsloth/models/_utils.py b/unsloth/models/_utils.py index 5a7d0b005f..94cf1b74e0 100644 --- a/unsloth/models/_utils.py +++ b/unsloth/models/_utils.py @@ -321,8 +321,7 @@ except Exception as error: ) pass import xformers.ops.fmha as xformers -# [TODO] Unsure why Xformers is also breaking as well -xformers_attention = torch.compiler.disable(xformers.memory_efficient_attention) +xformers_attention = xformers.memory_efficient_attention # Check TRL version from trl import __version__ as trl_version @@ -375,8 +374,8 @@ pass # ============================================= # Torch compile settings -UNSLOTH_COMPILE_DEBUG = True #"UNSLOTH_COMPILE_DEBUG" in os.environ -UNSLOTH_COMPILE_MAXIMUM = True #"UNSLOTH_COMPILE_MAXIMUM" in os.environ +UNSLOTH_COMPILE_DEBUG = "UNSLOTH_COMPILE_DEBUG" in os.environ +UNSLOTH_COMPILE_MAXIMUM = "UNSLOTH_COMPILE_MAXIMUM" in os.environ # Just remove max_autotune_gemm warning import functools @functools.lru_cache(None)