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 2024-09-08 19:47:21 -07:00
commit d27d992f58
1 changed files with 7 additions and 1 deletions
Download patch file Download diff file Expand all files Collapse all files

8
unsloth/kernels/utils.py
View file

@ -56,7 +56,8 @@ pass
import bitsandbytes as bnb
# https://github.com/bitsandbytes-foundation/bitsandbytes/pull/1330/files
HAS_CUDA_STREAM = Version(bnb.__version__) > Version("0.43.3")
CUDA_STREAM = torch.cuda.current_stream("cuda:0")
global CUDA_STREAM
CUDA_STREAM = None
get_ptr = bnb.functional.get_ptr
import ctypes
cdequantize_blockwise_fp32 = bnb.functional.lib.cdequantize_blockwise_fp32
@ -129,6 +130,8 @@ if HAS_CUDA_STREAM:
offset, state2 = compressed_stats
absmax2, code2, blocksize2, _, _, _, _ = state2
pass
global CUDA_STREAM
if CUDA_STREAM is None: CUDA_STREAM = torch.cuda.current_stream("cuda:0")
# Create weight matrix
if out is None:
@ -236,6 +239,9 @@ if HAS_CUDA_STREAM:
offset, state2 = compressed_stats
absmax2, code2, blocksize2, _, _, _, _ = state2
pass
global CUDA_STREAM
if CUDA_STREAM is None: CUDA_STREAM = torch.cuda.current_stream("cuda:0")
# assert(dtype == X.dtype)
bout = shape[0]