From cee52de54411fdf63581bfb98dd0509276c7ec2d Mon Sep 17 00:00:00 2001 From: Daniel Han-Chen Date: Sun, 17 Mar 2024 22:46:17 +1100 Subject: [PATCH] Update fast_lora.py --- unsloth/kernels/fast_lora.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/unsloth/kernels/fast_lora.py b/unsloth/kernels/fast_lora.py index 6568bba681..f1524653a5 100644 --- a/unsloth/kernels/fast_lora.py +++ b/unsloth/kernels/fast_lora.py @@ -87,6 +87,7 @@ class LoRA_MLP(torch.autograd.Function): downW, downW_quant, downA, downB, downS, _forward_function, _backward_function,): dtype = X.dtype + print("Forward ", dtype) e = matmul_lora(X, gateW, gateW_quant, gateA, gateB, gateS) g = matmul_lora(X, upW, upW_quant, upA, upB, upS) @@ -122,6 +123,7 @@ class LoRA_MLP(torch.autograd.Function): e = e .view(-1, e .shape[-1]) g = g .view(-1, g .shape[-1]) dtype = X.dtype + print("Backward ", dtype) DW = matmul_lora(dY, downW.t(), downW_quant, downB, downA, downS) DW, e, g = _backward_function(DW, e, g)