From 4f596958105c0cc1fa2dcfca66bab3c7c190b0c6 Mon Sep 17 00:00:00 2001 From: Fizza-Mukhtar Date: Tue, 30 Dec 2025 07:08:10 -0800 Subject: [PATCH] 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,