diff --git a/unsloth/kernels/fast_lora.py b/unsloth/kernels/fast_lora.py index f5d85e8088..b1206e4a2c 100644 --- a/unsloth/kernels/fast_lora.py +++ b/unsloth/kernels/fast_lora.py @@ -95,6 +95,13 @@ 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 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) + ctx.custom_saved_tensors = ( gateW, gateW_quant, @@ -134,6 +141,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), @@ -172,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 @@ -185,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) @@ -194,20 +208,23 @@ 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, # 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 +421,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 +470,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), @@ -468,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 @@ -481,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 @@ -489,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 @@ -498,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) @@ -506,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) @@ -514,13 +549,16 @@ 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, # 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 +642,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 +665,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) @@ -630,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 @@ -644,10 +693,13 @@ 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 - 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..5bc5ab9a74 100644 --- a/unsloth/kernels/utils.py +++ b/unsloth/kernels/utils.py @@ -1091,11 +1091,21 @@ def matmul_lora( W = W.dequantize() else: W = W.contiguous() + # custom_fwd disables autocast, so reconcile the activation dtype to the + # 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 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