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

[ROCm] add hip device path (#3301)

Browse source
This commit is contained in:
billishyahao 2025-09-12 17:57:19 +08:00 • committed by GitHub
commit 71ae760aa0
5 changed files with 80 additions and 13 deletions
Download patch file Download diff file Expand all files Collapse all files

17
unsloth/__init__.py
View file

@ -69,8 +69,13 @@ except Exception as exception:
raise exception
pass
def is_hip():
return bool(getattr(getattr(torch, "version", None), "hip", None))
def get_device_type():
if hasattr(torch, "cuda") and torch.cuda.is_available():
if is_hip():
return "hip"
return "cuda"
elif hasattr(torch, "xpu") and torch.xpu.is_available():
return "xpu"
@ -79,7 +84,7 @@ pass
DEVICE_TYPE : str = get_device_type()
def get_device_count():
if DEVICE_TYPE == "cuda":
if DEVICE_TYPE in ("cuda", "hip"):
return torch.cuda.device_count()
elif DEVICE_TYPE == "xpu":
return torch.xpu.device_count()
@ -91,11 +96,12 @@ DEVICE_COUNT : int = get_device_count()
# Reduce VRAM usage by reducing fragmentation
# And optimize pinning of memory
if (DEVICE_TYPE == "cuda") and (os.environ.get("UNSLOTH_VLLM_STANDBY", "0")=="0"):
# 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]"
elif (DEVICE_TYPE == "cuda") and (os.environ.get("UNSLOTH_VLLM_STANDBY", "0")=="1") and \
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"\
@ -153,6 +159,8 @@ if DEVICE_TYPE == "cuda":
def is_bf16_supported(): return SUPPORTS_BFLOAT16
torch.cuda.is_bf16_supported = is_bf16_supported
pass
elif DEVICE_TYPE == "hip":
SUPPORTS_BFLOAT16 = torch.cuda.is_bf16_supported()
elif DEVICE_TYPE == "xpu":
# torch.xpu.is_bf16_supported() does not have including_emulation
# set SUPPORTS_BFLOAT16 as torch.xpu.is_bf16_supported()
@ -218,6 +226,9 @@ if DEVICE_TYPE == "cuda":
"Unsloth will still run for now, but maybe it might crash - let's hope it works!"
)
pass
elif DEVICE_TYPE == "hip":
# NO-OP for rocm device
pass
elif DEVICE_TYPE == "xpu":
# currently intel xpu will not support bnb, will add support in the future
# TODO: check triton for intel installed properly.

6
unsloth/kernels/utils.py
View file

@ -90,7 +90,7 @@ else:
if DEVICE_COUNT > 1:
if DEVICE_TYPE == "cuda":
if DEVICE_TYPE in ("cuda", "hip"):
torch_gpu_device = torch.cuda.device
elif DEVICE_TYPE == "xpu":
torch_gpu_device = torch.xpu.device
@ -312,7 +312,7 @@ if DEVICE_TYPE == "xpu" and HAS_XPU_STREAM:
return out.t() if is_transposed else out
pass
# NVIDIA GPU Default Logic
elif DEVICE_TYPE == "cuda" and HAS_CUDA_STREAM:
elif DEVICE_TYPE in ("cuda", "hip") and HAS_CUDA_STREAM:
@torch.inference_mode
def fast_dequantize(W, quant_state = None, out = None, use_global_buffer = False):
if quant_state is None: return W
@ -513,7 +513,7 @@ if DEVICE_TYPE == "xpu" and HAS_XPU_STREAM:
return out
pass
elif DEVICE_TYPE == "cuda" and HAS_CUDA_STREAM:
elif DEVICE_TYPE in ("cuda", "hip") and HAS_CUDA_STREAM:
def fast_gemv(X, W, quant_state, out = None):
if quant_state is None: return torch_matmul(X, W, out = out)
# For fast X @ W where seq_len == 1

42
unsloth/models/_utils.py
View file

@ -452,7 +452,7 @@ pass
# =============================================
# torch.cuda.amp.custom_fwd is deprecated >= 2.4
torch_version = torch.__version__
if DEVICE_TYPE == "cuda":
if DEVICE_TYPE in ("cuda", "hip"):
if Version(torch_version) < Version("2.4.0"):
torch_amp_custom_fwd = torch.cuda.amp.custom_fwd
torch_amp_custom_bwd = torch.cuda.amp.custom_bwd
@ -506,7 +506,7 @@ pass
# =============================================
# Get Flash Attention v2 if Ampere (RTX 30xx, A100)
if DEVICE_TYPE == "cuda":
if DEVICE_TYPE in ("cuda", "hip"):
import bitsandbytes as bnb
from transformers import AutoTokenizer
@ -565,6 +565,44 @@ if DEVICE_TYPE == "cuda":
# Tri Dao's benchmark shows xformers is faster for now.
HAS_FLASH_ATTENTION = False
pass
elif DEVICE_TYPE == "hip":
SUPPORTS_BFLOAT16 = True
if _is_package_available("flash_attn"):
# Check for CUDA linking errors "undefined symbol: _ZNK3c106SymIntltEl"
try:
try:
# See https://github.com/unslothai/unsloth/issues/1437
from flash_attn.flash_attn_interface import flash_attn_gpu
except:
from flash_attn.flash_attn_interface import flash_attn_cuda
HAS_FLASH_ATTENTION = True
# Also check for softcapping
from flash_attn import __version__ as flash_attn_version
HAS_FLASH_ATTENTION_SOFTCAPPING = Version(flash_attn_version) >= Version("2.6.3")
if not HAS_FLASH_ATTENTION_SOFTCAPPING:
print(
"Unsloth: If you want to finetune Gemma 2, upgrade flash-attn to version 2.6.3 or higher!\n"\
"Newer versions support faster and less memory usage kernels for Gemma 2's attention softcapping!\n"\
"To update flash-attn, do the below:\n"\
'\npip install --no-deps --no-build-isolation --upgrade "flash-attn>=2.6.3"'
)
except:
print(
"Unsloth: Your Flash Attention 2 installation seems to be broken?\n"\
"A possible explanation is you have a new CUDA version which isn't\n"\
"yet compatible with FA2? Please file a ticket to Unsloth or FA2.\n"\
"We shall now use Xformers instead, which does not have any performance hits!\n"\
"We found this negligible impact by benchmarking on 1x A100."
)
# Stop Flash Attention from importing!
import transformers.utils.import_utils
transformers.utils.import_utils.is_flash_attn_2_available = lambda *args, **kwargs: False
import transformers.utils
transformers.utils.is_flash_attn_2_available = lambda *args, **kwargs: False
HAS_FLASH_ATTENTION = False
elif DEVICE_TYPE == "xpu":
SUPPORTS_BFLOAT16 = True

10
unsloth/models/llama.py
View file

@ -1853,6 +1853,8 @@ class FastLlamaModel:
if major_version < 7:
print("Unsloth: vLLM does not work on older GPUs - will switch to Unsloth inference!")
fast_inference = False
elif DEVICE_TYPE == "hip":
fast_inference = True
if unsloth_vllm_standby and os.environ.get("UNSLOTH_VLLM_STANDBY", "0") == "0":
raise RuntimeError("Unsloth: `unsloth_vllm_standby` is True, but environment variable `UNSLOTH_VLLM_STANDBY` is not set to 1!")
pass
@ -1866,6 +1868,14 @@ class FastLlamaModel:
gpu_version = torch.version.cuda
gpu_stats_snippet = f"CUDA: {gpu_stats.major}.{gpu_stats.minor}. CUDA Toolkit: {gpu_version}."
from importlib.metadata import version as importlib_version
try: vllm_version = f" vLLM: {importlib_version('vllm')}."
except: vllm_version = ""
elif DEVICE_TYPE == "hip":
gpu_stats = torch.cuda.get_device_properties(0)
gpu_version = torch.version.hip
gpu_stats_snippet = f"ROCm Toolkit: {gpu_version}."
from importlib.metadata import version as importlib_version
try: vllm_version = f" vLLM: {importlib_version('vllm')}."
except: vllm_version = ""

18
unsloth/models/vision.py
View file

@ -278,6 +278,14 @@ class FastBaseModel:
gpu_version = torch.version.cuda
gpu_stats_snippet = f"CUDA: {gpu_stats.major}.{gpu_stats.minor}. CUDA Toolkit: {gpu_version}."
from importlib.metadata import version as importlib_version
try: vllm_version = f" vLLM: {importlib_version('vllm')}."
except: vllm_version = ""
elif DEVICE_TYPE == "hip":
gpu_stats = torch.cuda.get_device_properties(0)
gpu_version = torch.version.hip
gpu_stats_snippet = f"ROCm Toolkit: {gpu_version}."
from importlib.metadata import version as importlib_version
try: vllm_version = f" vLLM: {importlib_version('vllm')}."
except: vllm_version = ""
@ -463,7 +471,7 @@ class FastBaseModel:
# Clear deleted GPU items
for _ in range(3):
gc.collect()
if DEVICE_TYPE == "cuda": torch.cuda.empty_cache()
if DEVICE_TYPE in ("cuda", "hip"): torch.cuda.empty_cache()
elif DEVICE_TYPE == "xpu": torch.xpu.empty_cache()
pass
@ -558,7 +566,7 @@ class FastBaseModel:
# Clear deleted GPU items
for _ in range(3):
gc.collect()
if DEVICE_TYPE == "cuda":
if DEVICE_TYPE in ("cuda", "hip"):
torch.cuda.empty_cache()
elif DEVICE_TYPE == "xpu":
torch.xpu.empty_cache()
@ -627,7 +635,7 @@ class FastBaseModel:
# Clear deleted GPU items
for _ in range(3):
gc.collect()
if DEVICE_TYPE == "cuda":
if DEVICE_TYPE in ("cuda", "hip"):
torch.cuda.empty_cache()
elif DEVICE_TYPE == "xpu":
torch.xpu.empty_cache()
@ -663,7 +671,7 @@ class FastBaseModel:
# Clear deleted GPU items
for _ in range(3):
gc.collect()
if DEVICE_TYPE == "cuda":
if DEVICE_TYPE in ("cuda", "hip"):
torch.cuda.empty_cache()
elif DEVICE_TYPE == "xpu":
torch.xpu.empty_cache()
@ -728,7 +736,7 @@ class FastBaseModel:
# Clear deleted GPU items
for _ in range(3):
gc.collect()
if DEVICE_TYPE == "cuda":
if DEVICE_TYPE in ("cuda", "hip"):
torch.cuda.empty_cache()
elif DEVICE_TYPE == "xpu":
torch.xpu.empty_cache()