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

Fix fused LoRA dtype mismatch on fp32 activations into fp16/bf16 weights

Browse source 
The fused LoRA ops are decorated with @custom_fwd, which disables
autocast. When an fp32 activation (for example the fp32 output of
fast_rms_layernorm on an fp32 hidden/residual stream) reaches a fused
LoRA path whose base weight is fp16/bf16, nothing downcasts the
activation, so torch.matmul raises:

    expected mat1 and mat2 to have the same dtype: float != c10::Half

This is not torch-version specific; torch.matmul never auto-promotes a
mixed fp16/fp32 matmul.

Reconcile the dtypes at the matmul_lora choke point: after dequantizing
W (both the plain and the fast_dequantize branches, not the fp8 branch),
downcast X to the base weight compute dtype the way autocast would for a
plain Linear, and keep the LoRA A/B cast dtype consistent with it.

Keep the backward pass dtype-consistent: LoRA_MLP, LoRA_QKV and LoRA_W
now record the original input dtype in ctx.input_dtype, pin the saved
activation to the compute dtype, cast the incoming gradients to that
dtype, and cast the returned dX back to the original input dtype to honor
the autograd grad-dtype contract. Every added branch is a no-op when the
dtypes already match, so normal fp16/bf16/fp32 training is unchanged.
This commit is contained in:
Daniel Han 2026-07-09 06:07:43 +00:00
commit 08a985d28a
2 changed files with 56 additions and 3 deletions
Download patch file Download diff file Expand all files Collapse all files

48
unsloth/kernels/fast_lora.py
View file

@ -95,6 +95,15 @@ class LoRA_MLP(torch.autograd.Function):
h = _forward_function(e, g)
i = matmul_lora(h, downW, downW_quant, downA, downB, downS)
# custom_fwd disables autocast, so X may arrive in a different dtype than
# the fused-op compute dtype (e.g. fp32 hidden states from
# fast_rms_layernorm meeting fp16/bf16 base weights). matmul_lora computes
# in the base weight dtype (== e.dtype); keep the saved activation in that
# same dtype so the backward pass stays dtype-consistent. The incoming
# dtype is remembered in ctx.input_dtype and restored on the returned dX.
ctx.input_dtype = dtype
X = X.to(e.dtype)
ctx.custom_saved_tensors = (
gateW,
gateW_quant,
@ -134,6 +143,9 @@ class LoRA_MLP(torch.autograd.Function):
e = e.view(-1, e.shape[-1])
g = g.view(-1, g.shape[-1])
dtype = X.dtype
# X (and thus dtype) is the compute dtype pinned in forward; align dY too.
if dY.dtype != dtype:
dY = dY.to(dtype)
gateA, gateB, upA, upB, downA, downB = (
gateA.to(dtype),
@ -206,8 +218,11 @@ class LoRA_MLP(torch.autograd.Function):
# gateW, gateW_quant, gateA, gateB, gateS,
# upW, upW_quant, upA, upB, upS,
# downW, downW_quant, downA, downB, downS,
dX = dX.view(batch, seq_len, hd)
if dX.dtype != ctx.input_dtype:
dX = dX.to(ctx.input_dtype)
return (
dX.view(batch, seq_len, hd),
dX,
None,
None,
d_gateA.t(),
@ -404,6 +419,12 @@ class LoRA_QKV(torch.autograd.Function):
K = K.view(orig_shape[0], orig_shape[1], -1)
V = V.view(orig_shape[0], orig_shape[1], -1)
# custom_fwd disables autocast; matmul_lora computed in the base weight
# (compute) dtype == Q.dtype. Pin the saved activation to it so backward
# is dtype-consistent, and remember the incoming dtype for the dX grad.
ctx.input_dtype = dtype
X = X.to(Q.dtype)
ctx.custom_saved_tensors = (
QW,
QW_quant,
@ -447,6 +468,13 @@ class LoRA_QKV(torch.autograd.Function):
dV = dV.view(-1, dV.shape[-1])
X = X.view(-1, X.shape[-1])
dtype = X.dtype
# X (and thus dtype) is the compute dtype pinned in forward; align grads.
if dQ.dtype != dtype:
dQ = dQ.to(dtype)
if dK.dtype != dtype:
dK = dK.to(dtype)
if dV.dtype != dtype:
dV = dV.to(dtype)
QA, QB, KA, KB, VA, VB = (
QA.to(dtype),
@ -519,8 +547,11 @@ class LoRA_QKV(torch.autograd.Function):
# QW, QW_quant, QA, QB, QS,
# KW, KW_quant, KA, KB, KS,
# VW, VW_quant, VA, VB, VS,
dX = dX.view(batch, seq_len, hd)
if dX.dtype != ctx.input_dtype:
dX = dX.to(ctx.input_dtype)
return (
dX.view(batch, seq_len, hd),
dX,
None,
None,
d_QA.t(),
@ -604,6 +635,11 @@ class LoRA_W(torch.autograd.Function):
def forward(ctx, X: torch.Tensor, W, W_quant, A, B, S):
dtype = X.dtype
XW = matmul_lora(X, W, W_quant, A, B, S)
# custom_fwd disables autocast; pin the saved activation to the compute
# dtype (== XW.dtype) so backward is dtype-consistent, and remember the
# incoming dtype for the returned dX grad.
ctx.input_dtype = dtype
X = X.to(XW.dtype)
ctx.custom_saved_tensors = (
W,
W_quant,
@ -622,6 +658,9 @@ class LoRA_W(torch.autograd.Function):
dY = dY.reshape(-1, dY.shape[-1]) # Must be reshape
X = X.reshape(-1, X.shape[-1]) # Must be reshape
dtype = X.dtype
# X (and thus dtype) is the compute dtype pinned in forward; align dY too.
if dY.dtype != dtype:
dY = dY.to(dtype)
A, B = A.to(dtype), B.to(dtype)
@ -647,7 +686,10 @@ class LoRA_W(torch.autograd.Function):
dX.addmm_(dY @ B.t(), A.t(), alpha = S)
# W, W_quant, A, B, S
return dX.view(batch, seq_len, hd), None, None, d_A.t(), d_B.t(), None
dX = dX.view(batch, seq_len, hd)
if dX.dtype != ctx.input_dtype:
dX = dX.to(ctx.input_dtype)
return dX, None, None, d_A.t(), d_B.t(), None
def apply_lora_o(self, X):

11
unsloth/kernels/utils.py
View file

@ -1091,11 +1091,22 @@ def matmul_lora(
W = W.dequantize()
else:
W = W.contiguous()
# custom_fwd disables autocast, so reconcile the activation dtype to the
# weight (compute) dtype the way autocast would for a plain Linear. This
# covers fp32 hidden states (e.g. fp32 fast_rms_layernorm output) meeting
# fp16/bf16 base weights.
if X.dtype != W.dtype:
X = X.to(W.dtype)
dtype = W.dtype
out = torch_matmul(X, W.t(), out = out)
elif W.dtype == torch.float8_e4m3fn:
out = fp8_linear(X, W, W_quant)
else:
W = fast_dequantize(W, W_quant, use_global_buffer = True)
# See note above: align the activation dtype to the base weight dtype.
if X.dtype != W.dtype:
X = X.to(W.dtype)
dtype = W.dtype
out = torch_matmul(X, W.t(), out = out)
if W_quant is not None:
del W