From 08a985d28abcb4d7186b1000ff15c287bc4cc7d0 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Thu, 9 Jul 2026 06:07:43 +0000 Subject: [PATCH] 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