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-02-02 17:12:21 -08:00
commit f3a4d082fd
1 changed files with 22 additions and 37 deletions
Download patch file Download diff file Expand all files Collapse all files

59
unsloth/kernels/utils.py
View file

@ -67,6 +67,7 @@ global CUDA_STREAM
CUDA_STREAM = None
get_ptr = bnb.functional.get_ptr
import ctypes
ctypes_c_int = ctypes.c_int
cdequantize_blockwise_fp32 = bnb.functional.lib.cdequantize_blockwise_fp32
cdequantize_blockwise_fp16_nf4 = bnb.functional.lib.cdequantize_blockwise_fp16_nf4
cdequantize_blockwise_bf16_nf4 = bnb.functional.lib.cdequantize_blockwise_bf16_nf4
@ -121,8 +122,6 @@ WEIGHT_BUFFER = None
global ABSMAX_BUFFER
ABSMAX_BUFFER = None
ctypes_c_int = ctypes.c_int
if HAS_CUDA_STREAM:
def fast_dequantize(W, quant_state = None, out = None, use_global_buffer = False):
if quant_state is None: return W
@ -159,14 +158,12 @@ if HAS_CUDA_STREAM:
if WEIGHT_BUFFER is None:
WEIGHT_BUFFER = torch.empty(size, dtype = dtype, device = "cuda:0")
ABSMAX_BUFFER = torch.empty(n_elements_absmax, dtype = torch.float32, device = "cuda:0")
ABSMAX_BUFFER.ptr_out_absmax = get_ptr(ABSMAX_BUFFER)
if size > WEIGHT_BUFFER.numel(): WEIGHT_BUFFER.resize_(size)
if n_elements_absmax > ABSMAX_BUFFER.numel(): ABSMAX_BUFFER.resize_(n_elements_absmax)
out = WEIGHT_BUFFER[:size].view(shape)
out_absmax = ABSMAX_BUFFER[:n_elements_absmax]
ptr_out_absmax = ABSMAX_BUFFER.ptr_out_absmax
else:
if out is None:
out = torch.empty(shape, dtype = dtype, device = "cuda:0")
@ -174,33 +171,21 @@ if HAS_CUDA_STREAM:
assert(out.shape == shape)
assert(out.dtype == dtype)
out_absmax = torch.empty(n_elements_absmax, dtype = torch.float32, device = "cuda:0")
ptr_out_absmax = get_ptr(out_absmax)
pass
# NF4 dequantization of statistics
ptr_out_absmax = get_ptr(out_absmax)
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),
CUDA_STREAM,
get_ptr(code2), get_ptr(absmax), get_ptr(absmax2), ptr_out_absmax,
ctypes_c_int(blocksize2), ctypes_c_int(n_elements_absmax), CUDA_STREAM,
)
print(offset, out_absmax)
out_absmax += offset
# Dequantize W
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()),
CUDA_STREAM,)
fx(get_ptr(None), get_ptr(W), ptr_out_absmax, get_ptr(out),
ctypes_c_int(blocksize), ctypes_c_int(out.numel()), CUDA_STREAM,)
# Careful returning transposed data
is_transposed = (True if W.shape[0] == 1 else False)
@ -318,17 +303,17 @@ if HAS_CUDA_STREAM:
lda = shape[0]
ldc = shape[0]
ldb = (hd+1)//2
m = ctypes.c_int32(m)
n = ctypes.c_int32(n)
k = ctypes.c_int32(k)
lda = ctypes.c_int32(lda)
ldb = ctypes.c_int32(ldb)
ldc = ctypes.c_int32(ldc)
m = ctypes_c_int32(m)
n = ctypes_c_int32(n)
k = ctypes_c_int32(k)
lda = ctypes_c_int32(lda)
ldb = ctypes_c_int32(ldb)
ldc = ctypes_c_int32(ldc)
df = torch.empty(absmax.shape, dtype = torch.float32, device = "cuda:0")
cdequantize_blockwise_fp32(
get_ptr(code2), get_ptr(absmax), get_ptr(absmax2), get_ptr(df),
ctypes.c_int(blocksize2), ctypes.c_int(df.numel()), CUDA_STREAM,
ctypes_c_int(blocksize2), ctypes_c_int(df.numel()), CUDA_STREAM,
)
df += offset
absmax = df
@ -336,7 +321,7 @@ if HAS_CUDA_STREAM:
fx = cgemm_4bit_inference_naive_fp16 if dtype == torch.float16 else \
cgemm_4bit_inference_naive_bf16
blocksize = ctypes.c_int32(blocksize)
blocksize = ctypes_c_int32(blocksize)
fx(m, n, k, get_ptr(X), get_ptr(W), get_ptr(absmax), get_ptr(stats), get_ptr(out),
lda, ldb, ldc, blocksize, CUDA_STREAM,)
@ -382,17 +367,17 @@ else:
lda = shape[0]
ldc = shape[0]
ldb = (hd+1)//2
m = ctypes.c_int32(m)
n = ctypes.c_int32(n)
k = ctypes.c_int32(k)
lda = ctypes.c_int32(lda)
ldb = ctypes.c_int32(ldb)
ldc = ctypes.c_int32(ldc)
m = ctypes_c_int32(m)
n = ctypes_c_int32(n)
k = ctypes_c_int32(k)
lda = ctypes_c_int32(lda)
ldb = ctypes_c_int32(ldb)
ldc = ctypes_c_int32(ldc)
df = torch.empty(absmax.shape, dtype = torch.float32, device = "cuda:0")
cdequantize_blockwise_fp32(
get_ptr(code2), get_ptr(absmax), get_ptr(absmax2), get_ptr(df),
ctypes.c_int(blocksize2), ctypes.c_int(df.numel()),
ctypes_c_int(blocksize2), ctypes_c_int(df.numel()),
)
df += offset
absmax = df
@ -400,7 +385,7 @@ else:
fx = cgemm_4bit_inference_naive_fp16 if dtype == torch.float16 else \
cgemm_4bit_inference_naive_bf16
blocksize = ctypes.c_int32(blocksize)
blocksize = ctypes_c_int32(blocksize)
fx(m, n, k, get_ptr(X), get_ptr(W), get_ptr(absmax), get_ptr(stats), get_ptr(out),
lda, ldb, ldc, blocksize,)