ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Fix 3D tensor support for bitsandbytes 8-bit matmul in forward pass

Browse source
This commit is contained in:
Fizza-Mukhtar 2025-12-30 07:56:01 -08:00
commit f2e87251c7
1 changed files with 5 additions and 5 deletions
Download patch file Download diff file Expand all files Collapse all files

10
unsloth/kernels/fast_lora.py
View file

@ -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: