From 08a985d28abcb4d7186b1000ff15c287bc4cc7d0 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Thu, 9 Jul 2026 06:07:43 +0000 Subject: [PATCH 1/3] Fix fused LoRA dtype mismatch on fp32 activations into fp16/bf16 weights 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. --- unsloth/kernels/fast_lora.py | 48 +++++++++++++++++++++++++++++++++--- unsloth/kernels/utils.py | 11 +++++++++ 2 files changed, 56 insertions(+), 3 deletions(-) diff --git a/unsloth/kernels/fast_lora.py b/unsloth/kernels/fast_lora.py index f5d85e8088..265e728862 100644 --- a/unsloth/kernels/fast_lora.py +++ b/unsloth/kernels/fast_lora.py @@ -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): diff --git a/unsloth/kernels/utils.py b/unsloth/kernels/utils.py index 43ed198a4a..89b610f5d0 100644 --- a/unsloth/kernels/utils.py +++ b/unsloth/kernels/utils.py @@ -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 From f0705ef6e5148cd600e5f0d5d1197705748f9f02 Mon Sep 17 00:00:00 2001 From: danielhanchen Date: Thu, 9 Jul 2026 07:47:23 +0000 Subject: [PATCH 2/3] Tighten fused-LoRA dtype-fix comments --- unsloth/kernels/fast_lora.py | 10 ++++------ unsloth/kernels/utils.py | 5 ++--- 2 files changed, 6 insertions(+), 9 deletions(-) diff --git a/unsloth/kernels/fast_lora.py b/unsloth/kernels/fast_lora.py index 265e728862..027dd7b8e9 100644 --- a/unsloth/kernels/fast_lora.py +++ b/unsloth/kernels/fast_lora.py @@ -95,12 +95,10 @@ 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. + # custom_fwd disables autocast, so X may mismatch the compute dtype (e.g. + # fp32 fast_rms_layernorm output meeting fp16/bf16 base weights). Pin the + # saved activation to the compute dtype (== e.dtype) so backward stays + # consistent; remember the incoming dtype to restore it on the dX grad. ctx.input_dtype = dtype X = X.to(e.dtype) diff --git a/unsloth/kernels/utils.py b/unsloth/kernels/utils.py index 89b610f5d0..5bc5ab9a74 100644 --- a/unsloth/kernels/utils.py +++ b/unsloth/kernels/utils.py @@ -1092,9 +1092,8 @@ def matmul_lora( 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. + # weight (compute) dtype as autocast would (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 From cb1ef59b52f0a525c7ef290e279438fbbad19653 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 13 Jul 2026 11:29:30 +0000 Subject: [PATCH 3/3] Avoid recomputing LoRA B-projected gradient in fused-LoRA backward Each fused-LoRA backward computed dY @ B.t() (the LoRA-B projected gradient, shape (batch*seq, r)) twice per projection: once for the A weight gradient and once for the dX accumulation. Hoist it into a single temporary and reuse it in both, removing 6 small matmuls per layer (3 in LoRA_QKV, 2 in LoRA_MLP, 1 in LoRA_W). Results are bitwise identical; measured about 2 to 3 percent faster on the LoRA forward plus backward for qkv and mlp in both 16-bit and 4-bit. --- unsloth/kernels/fast_lora.py | 36 ++++++++++++++++++++++++------------ 1 file changed, 24 insertions(+), 12 deletions(-) diff --git a/unsloth/kernels/fast_lora.py b/unsloth/kernels/fast_lora.py index 027dd7b8e9..b1206e4a2c 100644 --- a/unsloth/kernels/fast_lora.py +++ b/unsloth/kernels/fast_lora.py @@ -182,12 +182,16 @@ class LoRA_MLP(torch.autograd.Function): d_downA.addmm_(h.t(), dY @ downB.t(), alpha = downS, beta = 0) d_downB.addmm_(downA.t() @ h.t(), dY, alpha = downS, beta = 0) + # df @ upB.t() and de @ gateB.t() each feed both a weight grad and dX; compute once. + up_dB = df @ upB.t() + gate_dB = de @ gateB.t() + # Up projection LoRA weights # d_upA = X.t() @ (df @ upB.t()) # d_upB = (upA.t() @ X.t()) @ df # d_upA *= upS # d_upB *= upS - d_upA.addmm_(X.t(), df @ upB.t(), alpha = upS, beta = 0) + d_upA.addmm_(X.t(), up_dB, alpha = upS, beta = 0) d_upB.addmm_(upA.t() @ X.t(), df, alpha = upS, beta = 0) # Gate projection LoRA weights @@ -195,7 +199,7 @@ class LoRA_MLP(torch.autograd.Function): # d_gateB = (gateA.t() @ X.t()) @ de # d_gateA *= gateS # d_gateB *= gateS - d_gateA.addmm_(X.t(), de @ gateB.t(), alpha = gateS, beta = 0) + d_gateA.addmm_(X.t(), gate_dB, alpha = gateS, beta = 0) d_gateB.addmm_(gateA.t() @ X.t(), de, alpha = gateS, beta = 0) # dX = matmul_lora(df, upW.t(), upW_quant, upB, upA, upS) @@ -204,14 +208,14 @@ class LoRA_MLP(torch.autograd.Function): dX = torch.matmul(df, upW.t(), out = X if ctx.inplace else None) del upW # dX += df @ upB.to(dtype).t() @ (upS * upA.to(dtype).t()) - dX.addmm_(df @ upB.t(), upA.t(), alpha = upS) + dX.addmm_(up_dB, upA.t(), alpha = upS) gateW = fast_dequantize(gateW.t(), gateW_quant) # dX += de @ gateW.t() dX.addmm_(de, gateW.t()) del gateW # dX += de @ gateB.to(dtype).t() @ (gateS * gateA.to(dtype).t()) - dX.addmm_(de @ gateB.t(), gateA.t(), alpha = gateS) + dX.addmm_(gate_dB, gateA.t(), alpha = gateS) # gateW, gateW_quant, gateA, gateB, gateS, # upW, upW_quant, upA, upB, upS, @@ -494,12 +498,17 @@ class LoRA_QKV(torch.autograd.Function): d_VA = torch.empty_like(VA) d_VB = torch.empty_like(VB) + # d @ B.t() each feed both a weight grad and dX; compute once. + q_dB = dQ @ QB.t() + k_dB = dK @ KB.t() + v_dB = dV @ VB.t() + # Q Projection # d_QA = X.t() @ (dQ @ QB.t()) # d_QB = (QA.t() @ X.t()) @ dQ # d_QA *= QS # d_QB *= QS - d_QA.addmm_(X.t(), dQ @ QB.t(), alpha = QS, beta = 0) + d_QA.addmm_(X.t(), q_dB, alpha = QS, beta = 0) d_QB.addmm_(QA.t() @ X.t(), dQ, alpha = QS, beta = 0) # K Projection @@ -507,7 +516,7 @@ class LoRA_QKV(torch.autograd.Function): # d_KB = (KA.t() @ X.t()) @ dK # d_KA *= KS # d_KB *= KS - d_KA.addmm_(X.t(), dK @ KB.t(), alpha = KS, beta = 0) + d_KA.addmm_(X.t(), k_dB, alpha = KS, beta = 0) d_KB.addmm_(KA.t() @ X.t(), dK, alpha = KS, beta = 0) # V Projection @@ -515,7 +524,7 @@ class LoRA_QKV(torch.autograd.Function): # d_VB = (VA.t() @ X.t()) @ dV # d_VA *= VS # d_VB *= VS - d_VA.addmm_(X.t(), dV @ VB.t(), alpha = VS, beta = 0) + d_VA.addmm_(X.t(), v_dB, alpha = VS, beta = 0) d_VB.addmm_(VA.t() @ X.t(), dV, alpha = VS, beta = 0) # Combine derivatives to find dX @@ -524,7 +533,7 @@ class LoRA_QKV(torch.autograd.Function): dX = torch.matmul(dQ, QW.t(), out = X if ctx.inplace else None) del QW # dX += (dQ @ QB.to(dtype).t() @ (QS * QA.to(dtype).t())) - dX.addmm_(dQ @ QB.t(), QA.t(), alpha = QS) + dX.addmm_(q_dB, QA.t(), alpha = QS) # dK KW = fast_dequantize(KW.t(), KW_quant) @@ -532,7 +541,7 @@ class LoRA_QKV(torch.autograd.Function): dX.addmm_(dK, KW.t()) del KW # dX += dK @ KB.to(dtype).t() @ (KS * KA.to(dtype).t()) - dX.addmm_(dK @ KB.t(), KA.t(), alpha = KS) + dX.addmm_(k_dB, KA.t(), alpha = KS) # dV VW = fast_dequantize(VW.t(), VW_quant) @@ -540,7 +549,7 @@ class LoRA_QKV(torch.autograd.Function): dX.addmm_(dV, VW.t()) del VW # dX += dV @ VB.to(dtype).t() @ (VS * VA.to(dtype).t()) - dX.addmm_(dV @ VB.t(), VA.t(), alpha = VS) + dX.addmm_(v_dB, VA.t(), alpha = VS) # QW, QW_quant, QA, QB, QS, # KW, KW_quant, KA, KB, KS, @@ -667,13 +676,16 @@ class LoRA_W(torch.autograd.Function): d_A = torch.empty_like(A) d_B = torch.empty_like(B) + # dY @ B.t() feeds both the d_A weight grad and dX; compute once. + y_dB = dY @ B.t() + ### Weight projection LoRA weights # Weight projection # d_A = X.t() @ (dY @ B.t()) # d_B = (A.t() @ X.t()) @ dY # d_A *= S # d_B *= S - d_A.addmm_(X.t(), dY @ B.t(), alpha = S, beta = 0) + d_A.addmm_(X.t(), y_dB, alpha = S, beta = 0) d_B.addmm_(A.t() @ X.t(), dY, alpha = S, beta = 0) # Get derivative for dX @@ -681,7 +693,7 @@ class LoRA_W(torch.autograd.Function): dX = dY @ W.t() del W # dX += dY @ B.to(dtype).t() @ (S * A.to(dtype).t()) - dX.addmm_(dY @ B.t(), A.t(), alpha = S) + dX.addmm_(y_dB, A.t(), alpha = S) # W, W_quant, A, B, S dX = dX.view(batch, seq_len, hd)