From c452eb13f54dc572bb0b12c63ae47746543c4654 Mon Sep 17 00:00:00 2001 From: Fizza-Mukhtar Date: Tue, 30 Dec 2025 07:08:10 -0800 Subject: [PATCH 1/3] Fix 3D tensor support for bitsandbytes 8-bit matmul in forward pass --- unsloth/kernels/fast_lora.py | 13 +++++++++++++ 1 file changed, 13 insertions(+) diff --git a/unsloth/kernels/fast_lora.py b/unsloth/kernels/fast_lora.py index 60d0c318c3..16b679d005 100644 --- a/unsloth/kernels/fast_lora.py +++ b/unsloth/kernels/fast_lora.py @@ -379,9 +379,22 @@ class LoRA_QKV(torch.autograd.Function): ): dtype = X.dtype + # bitsandbytes 8-bit matmul expects 2D inputs. + # TorchInductor/AOTAutograd fails on 3D tensors during backward, + # so we explicitly flatten the sequence dimension. + orig_shape = X.shape + if X.dim() == 3: + X = X.view(-1, X.shape[-1]) + Q = matmul_lora(X, QW, QW_quant, QA, QB, QS) K = matmul_lora(X, KW, KW_quant, KA, KB, KS) V = matmul_lora(X, VW, VW_quant, VA, VB, VS) + + # Restore original shape after matmul + if len(orig_shape) == 3: + Q = Q.view(orig_shape[0], orig_shape[1], -1) + K = K.view(orig_shape[0], orig_shape[1], -1) + V = V.view(orig_shape[0], orig_shape[1], -1) ctx.custom_saved_tensors = ( QW, From f2e87251c721482d05bf8dd21452e4eb5c20ba02 Mon Sep 17 00:00:00 2001 From: Fizza-Mukhtar Date: Tue, 30 Dec 2025 07:56:01 -0800 Subject: [PATCH 2/3] Fix 3D tensor support for bitsandbytes 8-bit matmul in forward pass --- unsloth/kernels/fast_lora.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/unsloth/kernels/fast_lora.py b/unsloth/kernels/fast_lora.py index 16b679d005..fbb18c3a15 100644 --- a/unsloth/kernels/fast_lora.py +++ b/unsloth/kernels/fast_lora.py @@ -383,12 +383,12 @@ class LoRA_QKV(torch.autograd.Function): # TorchInductor/AOTAutograd fails on 3D tensors during backward, # so we explicitly flatten the sequence dimension. orig_shape = X.shape + X_for_matmul = X if X.dim() == 3: - X = X.view(-1, X.shape[-1]) - - Q = matmul_lora(X, QW, QW_quant, QA, QB, QS) - K = matmul_lora(X, KW, KW_quant, KA, KB, KS) - V = matmul_lora(X, VW, VW_quant, VA, VB, VS) + X_for_matmul = X.view(-1, X.shape[-1]) + Q = matmul_lora(X_for_matmul, QW, QW_quant, QA, QB, QS) + K = matmul_lora(X_for_matmul, KW, KW_quant, KA, KB, KS) + V = matmul_lora(X_for_matmul, VW, VW_quant, VA, VB, VS) # Restore original shape after matmul if len(orig_shape) == 3: From e43e67cb18e3f9842ca34416db8f42b21f7154f2 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Tue, 30 Dec 2025 15:58:40 +0000 Subject: [PATCH 3/3] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- unsloth/kernels/fast_lora.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/unsloth/kernels/fast_lora.py b/unsloth/kernels/fast_lora.py index fbb18c3a15..f1c0e298d9 100644 --- a/unsloth/kernels/fast_lora.py +++ b/unsloth/kernels/fast_lora.py @@ -389,7 +389,7 @@ class LoRA_QKV(torch.autograd.Function): Q = matmul_lora(X_for_matmul, QW, QW_quant, QA, QB, QS) K = matmul_lora(X_for_matmul, KW, KW_quant, KA, KB, KS) V = matmul_lora(X_for_matmul, VW, VW_quant, VA, VB, VS) - + # Restore original shape after matmul if len(orig_shape) == 3: Q = Q.view(orig_shape[0], orig_shape[1], -1)