ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

move PYTORCH_CUDA_ALLOC_CONF into zoo (#3499)

Browse source
This commit is contained in:
DoubleMathew 2025-10-26 21:29:18 -07:00 • committed by GitHub
commit a40dbd2cd5
1 changed files with 0 additions and 47 deletions
Download patch file Download diff file Expand all files Collapse all files

47
unsloth/__init__.py
View file

@ -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