From 97ed0b46d73a1668c9d35d5e84eb81531aad5e85 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Wed, 5 Mar 2025 03:49:38 -0800 Subject: [PATCH] bug fix --- unsloth/kernels/utils.py | 33 +++++++++++++++++---------------- unsloth/models/llama.py | 2 +- 2 files changed, 18 insertions(+), 17 deletions(-) diff --git a/unsloth/kernels/utils.py b/unsloth/kernels/utils.py index 8b66b1769e..db1d73c340 100644 --- a/unsloth/kernels/utils.py +++ b/unsloth/kernels/utils.py @@ -104,6 +104,11 @@ cdequantize_blockwise_fp16_nf4 = bnb.functional.lib.cdequantize_blockwise_fp16_ cdequantize_blockwise_bf16_nf4 = bnb.functional.lib.cdequantize_blockwise_bf16_nf4 cgemm_4bit_inference_naive_fp16 = bnb.functional.lib.cgemm_4bit_inference_naive_fp16 cgemm_4bit_inference_naive_bf16 = bnb.functional.lib.cgemm_4bit_inference_naive_bf16 +torch_mm = torch.mm +torch_mv = torch.mv +torch_matmul = torch.matmul +torch_addmm = torch.addmm +torch_empty = torch.empty def QUANT_STATE(W): return getattr(W, "quant_state", None) @@ -194,8 +199,8 @@ if HAS_CUDA_STREAM: WEIGHT_BUFFER = WEIGHT_BUFFERS[device_index] ABSMAX_BUFFER = ABSMAX_BUFFERS[device_index] if WEIGHT_BUFFER is None: - WEIGHT_BUFFERS[device_index] = WEIGHT_BUFFER = torch.empty(size, dtype = dtype, device = device, requires_grad = False) - ABSMAX_BUFFERS[device_index] = ABSMAX_BUFFER = torch.empty(n_elements_absmax, dtype = torch.float32, device = device, requires_grad = False) + WEIGHT_BUFFERS[device_index] = WEIGHT_BUFFER = torch_empty(size, dtype = dtype, device = device, requires_grad = False) + ABSMAX_BUFFERS[device_index] = ABSMAX_BUFFER = torch_empty(n_elements_absmax, dtype = torch.float32, device = device, requires_grad = False) if size > WEIGHT_BUFFER.numel(): WEIGHT_BUFFER.resize_(size) if n_elements_absmax > ABSMAX_BUFFER.numel(): ABSMAX_BUFFER.resize_(n_elements_absmax) @@ -204,11 +209,11 @@ if HAS_CUDA_STREAM: out_absmax = ABSMAX_BUFFER[:n_elements_absmax] else: if out is None: - out = torch.empty(shape, dtype = dtype, device = device, requires_grad = False) + out = torch_empty(shape, dtype = dtype, device = device, requires_grad = False) else: assert(out.shape == shape) assert(out.dtype == dtype) - out_absmax = torch.empty(n_elements_absmax, dtype = torch.float32, device = device, requires_grad = False) + out_absmax = torch_empty(n_elements_absmax, dtype = torch.float32, device = device, requires_grad = False) pass # NF4 dequantization of statistics @@ -258,11 +263,11 @@ else: # Create weight matrix if out is None: - out = torch.empty(shape, dtype = dtype, device = device, requires_grad = False) + out = torch_empty(shape, dtype = dtype, device = device, requires_grad = False) else: assert(out.shape == shape) assert(out.dtype == dtype) - out_absmax = torch.empty(n_elements_absmax, dtype = torch.float32, device = device, requires_grad = False) + out_absmax = torch_empty(n_elements_absmax, dtype = torch.float32, device = device, requires_grad = False) # Do dequantization ptr_out_absmax = get_ptr(out_absmax) @@ -286,7 +291,7 @@ pass if HAS_CUDA_STREAM: def fast_gemv(X, W, quant_state, out = None): - if quant_state is None: return torch.matmul(X, W, out = out) + if quant_state is None: return torch_matmul(X, W, out = out) # For fast X @ W where seq_len == 1 # From https://github.com/TimDettmers/bitsandbytes/blob/main/bitsandbytes/functional.py#L1469 _, q_len, hd = X.shape @@ -318,7 +323,7 @@ if HAS_CUDA_STREAM: bout = shape[0] if out is None: - out = torch.empty((1, 1, bout,), dtype = dtype, device = device) + out = torch_empty((1, 1, bout,), dtype = dtype, device = device) # else: # assert(out.shape == (1, 1, bout,)) # pass @@ -336,7 +341,7 @@ if HAS_CUDA_STREAM: ldb = ctypes_c_int32(ldb) ldc = ctypes_c_int32(ldc) - df = torch.empty(absmax.shape, dtype = torch.float32, device = device) + df = torch_empty(absmax.shape, dtype = torch.float32, device = device) with torch_cuda_device(device): cdequantize_blockwise_fp32( get_ptr(code2), get_ptr(absmax), get_ptr(absmax2), get_ptr(df), @@ -385,7 +390,7 @@ else: device = W.device if out is None: - out = torch.empty((1, 1, bout,), dtype = dtype, device = device) + out = torch_empty((1, 1, bout,), dtype = dtype, device = device) # else: # assert(out.shape == (1, 1, bout,)) # pass @@ -403,7 +408,7 @@ else: ldb = ctypes_c_int32(ldb) ldc = ctypes_c_int32(ldc) - df = torch.empty(absmax.shape, dtype = torch.float32, device = device) + df = torch_empty(absmax.shape, dtype = torch.float32, device = device) cdequantize_blockwise_fp32( get_ptr(code2), get_ptr(absmax), get_ptr(absmax2), get_ptr(df), ctypes_c_int(blocksize2), ctypes_c_int(df.numel()), @@ -423,10 +428,6 @@ else: pass -torch_mm = torch.mm -torch_mv = torch.mv -torch_matmul = torch.matmul -torch_addmm = torch.addmm def fast_linear_forward(proj, X, temp_lora = None, out = None): W, W_quant, lora_A, lora_B, lora_S, bias = get_lora_parameters_bias(proj) @@ -438,7 +439,7 @@ def fast_linear_forward(proj, X, temp_lora = None, out = None): elif bsz == 1 and q_len == 1: out = fast_gemv(X, W, W_quant, out = out) else: - W = fast_dequantize(W.t(), W_quant, use_global_buffer = False) + W = fast_dequantize(W.t(), W_quant, use_global_buffer = True) out = torch_matmul(X, W, out = out) pass diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py index 8ba7c45368..356e81a018 100644 --- a/unsloth/models/llama.py +++ b/unsloth/models/llama.py @@ -261,7 +261,7 @@ def LlamaAttention_fast_forward_inference( # pass # Attention - if bsz == 1: + if True:#bsz == 1: Qn *= self.scalar # See https://github.com/ggerganov/llama.cpp/issues/7805#issuecomment-2153349963 # It seems like doing (Q * scalar) @ K is better than (Q @ K) * scalar to stop overflows A = torch_matmul(Qn, Knn.transpose(2, 3), out = self.attention[:,:,:,:cached_len])