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

Update utils.py

Browse source
This commit is contained in:
Daniel Han 2025-03-04 02:35:19 -08:00
commit 5ee394758d
1 changed files with 5 additions and 6 deletions
Download patch file Download diff file Expand all files Collapse all files

11
unsloth/kernels/utils.py
View file

@ -28,7 +28,6 @@ else:
torch_amp_custom_fwd = torch.amp.custom_fwd(device_type = "cuda")
torch_amp_custom_bwd = torch.amp.custom_bwd(device_type = "cuda")
pass
torch_cuda_device = torch.cuda.device
# tl.math.tanh now is libdevice.tanh
@ -70,10 +69,10 @@ HAS_CUDA_STREAM = Version(bnb.__version__) > Version("0.43.3")
get_ptr = bnb.functional.get_ptr
if torch.cuda.device_count() > 1:
def _cuda_device_of(a: torch_Tensor): return torch.cuda.device_of(a)
torch_cuda_device = torch.cuda.device
else:
from contextlib import nullcontext
def _cuda_device_of(a: torch_Tensor): return nullcontext()
def torch_cuda_device(device): return nullcontext()
pass
_cuda_getCurrentRawStream = torch._C._cuda_getCurrentRawStream
c_void_p = ctypes.c_void_p
@ -215,10 +214,10 @@ if HAS_CUDA_STREAM:
# NF4 dequantization of statistics
ptr_out_absmax = get_ptr(out_absmax)
with _cuda_device_of(absmax):
with torch_cuda_device(device):
cdequantize_blockwise_fp32(
get_ptr(code2), get_ptr(absmax), get_ptr(absmax2), ptr_out_absmax,
ctypes_c_int(blocksize2), ctypes_c_int(n_elements_absmax), _get_tensor_stream(absmax),
ctypes_c_int(blocksize2), ctypes_c_int(n_elements_absmax), CUDA_STREAM
)
out_absmax += offset
@ -226,7 +225,7 @@ if HAS_CUDA_STREAM:
fx = cdequantize_blockwise_fp16_nf4 if dtype == torch.float16 else \
cdequantize_blockwise_bf16_nf4
fx(get_ptr(None), get_ptr(W), ptr_out_absmax, get_ptr(out),
ctypes_c_int(blocksize), ctypes_c_int(out.numel()), _get_tensor_stream(absmax),)
ctypes_c_int(blocksize), ctypes_c_int(out.numel()), CUDA_STREAM,)
pass
# Careful returning transposed data
is_transposed = (True if W.shape[0] == 1 else False)