From a40dbd2cd5ee974a3f92888bd492556f7710f053 Mon Sep 17 00:00:00 2001 From: DoubleMathew Date: Sun, 26 Oct 2025 21:29:18 -0700 Subject: [PATCH] move PYTORCH_CUDA_ALLOC_CONF into zoo (#3499) --- unsloth/__init__.py | 47 --------------------------------------------- 1 file changed, 47 deletions(-) diff --git a/unsloth/__init__.py b/unsloth/__init__.py index 0469a09999..26f7da23f8 100644 --- a/unsloth/__init__.py +++ b/unsloth/__init__.py @@ -101,53 +101,6 @@ from unsloth_zoo.device_type import ( DEVICE_COUNT, ALLOW_PREQUANTIZED_MODELS, ) -# Reduce VRAM usage by reducing fragmentation -# And optimize pinning of memory -# TODO(billishyahao): need to add hip related optimization... -if (DEVICE_TYPE in ("cuda", "hip")) and (os.environ.get("UNSLOTH_VLLM_STANDBY", "0")=="0"): - os.environ["PYTORCH_CUDA_ALLOC_CONF"] = \ - "expandable_segments:True,"\ - "roundup_power2_divisions:[32:256,64:128,256:64,>:32]" - os.environ["PYTORCH_HIP_ALLOC_CONF"] = "expandable_segments:True" -elif (DEVICE_TYPE in ("cuda", "hip")) and (os.environ.get("UNSLOTH_VLLM_STANDBY", "0")=="1") and \ - ("expandable_segments:True" in os.environ.get("PYTORCH_CUDA_ALLOC_CONF", "")): - warnings.warn( - "Unsloth: `UNSLOTH_VLLM_STANDBY` is on, but requires `expandable_segments` to be off.\n"\ - "We will remove `expandable_segments`.", - stacklevel = 2, - ) - os.environ["PYTORCH_CUDA_ALLOC_CONF"] = re.sub( - r"expandable\_segments\:True\,?", - "", - os.environ["PYTORCH_CUDA_ALLOC_CONF"], - ) - os.environ["PYTORCH_HIP_ALLOC_CONF"] = re.sub( - r"expandable\_segments\:True\,?", - "", - os.environ["PYTORCH_HIP_ALLOC_CONF"], - ) -pass -# We support Pytorch 2 -# Fixes https://github.com/unslothai/unsloth/issues/38 -torch_version = str(re.match(r"[0-9\.]{3,}", str(torch.__version__)).group(0)).split(".") -major_torch, minor_torch = torch_version[0], torch_version[1] -major_torch, minor_torch = int(major_torch), int(minor_torch) -if (major_torch < 2): - raise ImportError("Unsloth only supports Pytorch 2 for now. Please update your Pytorch to 2.1.\n"\ - "We have some installation instructions on our Github page.") -elif (major_torch == 2) and (minor_torch < 2): - # Disable expandable_segments - del os.environ["PYTORCH_CUDA_ALLOC_CONF"] - del os.environ["PYTORCH_HIP_ALLOC_CONF"] -pass - -# CCE fails on Torch 2.8 and above -# OutOfResources: out of resource: shared memory, Required: 98304, Hardware limit: 65536. Reducing block sizes or `num_stages` -if (major_torch >= 2 and minor_torch >= 8) or (major_torch > 2): - os.environ["UNSLOTH_ENABLE_CCE"] = "0" -elif DEVICE_TYPE == "hip": - # CCE also fails in HIP / AMD - os.environ["UNSLOTH_ENABLE_CCE"] = "0" # Fix other issues from .import_fixes import fix_xformers_performance_issue