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

Update fast_lora.py

Browse source
This commit is contained in:
Daniel Han-Chen 2024-03-17 22:46:17 +11:00
commit 4b9aaee6b1
1 changed files with 2 additions and 0 deletions
Download patch file Download diff file Expand all files Collapse all files

2
unsloth/kernels/fast_lora.py
View file

@ -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)