From f3a4d082fd76e3d6b48ba12849062029133bfabc Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 2 Feb 2025 17:12:21 -0800 Subject: [PATCH] Update utils.py --- unsloth/kernels/utils.py | 59 +++++++++++++++------------------------- 1 file changed, 22 insertions(+), 37 deletions(-) diff --git a/unsloth/kernels/utils.py b/unsloth/kernels/utils.py index 3b0c1d3919..ac468e43ad 100644 --- a/unsloth/kernels/utils.py +++ b/unsloth/kernels/utils.py @@ -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,)