LoRA
This commit is contained in:
Daniel Han
parent
09be705d7d
commit
3fa6ac944f
2 changed files with 94 additions and 46 deletions
|
|
@ -98,9 +98,6 @@ class LoRA_MLP(torch.autograd.Function):
|
|||
gateA, gateB, upA, upB, downA, downB, \
|
||||
X, e, g = ctx.saved_tensors
|
||||
|
||||
gateA, gateB, upA, upB, downA, downB = \
|
||||
gateA.t(), gateB.t(), upA.t(), upB.t(), downA.t(), downB.t()
|
||||
|
||||
batch, seq_len, hd = X.shape
|
||||
dY = dY.view(-1, dY.shape[-1])
|
||||
X = X .view(-1, X .shape[-1])
|
||||
|
|
@ -108,39 +105,61 @@ class LoRA_MLP(torch.autograd.Function):
|
|||
g = g .view(-1, g .shape[-1])
|
||||
dtype = X.dtype
|
||||
|
||||
gateA, gateB, upA, upB, downA, downB = \
|
||||
gateA.to(dtype), gateB.to(dtype), upA.to(dtype), upB.to(dtype), downA.to(dtype), downB.to(dtype)
|
||||
|
||||
gateA, gateB, upA, upB, downA, downB = \
|
||||
gateA.t(), gateB.t(), upA.t(), upB.t(), downA.t(), downB.t()
|
||||
|
||||
DW = matmul_lora(dY, downW.t(), downW_quant, downB, downA, downS)
|
||||
DW, e, g = _backward_function(DW, e, g)
|
||||
h, df, de = DW, e, g
|
||||
|
||||
d_downA = torch.empty_like(downA)
|
||||
d_downB = torch.empty_like(downB)
|
||||
d_gateA = torch.empty_like(gateA)
|
||||
d_gateB = torch.empty_like(gateB)
|
||||
d_upA = torch.empty_like(upA)
|
||||
d_upB = torch.empty_like(upB)
|
||||
|
||||
# Down projection LoRA weights
|
||||
d_downA = h.t() @ (dY @ downB.t())
|
||||
d_downB = (downA.t() @ h.t()) @ dY
|
||||
d_downA *= downS
|
||||
d_downB *= downS
|
||||
# d_downA = h.t() @ (dY @ downB.t())
|
||||
# d_downB = (downA.t() @ h.t()) @ dY
|
||||
# d_downA *= downS
|
||||
# d_downB *= downS
|
||||
d_downA.addmm_(h.t(), dY @ downB.t(), alpha = downS, beta = 0)
|
||||
d_downB.addmm_(downA.t() @ h.t(), dY, alpha = downS, beta = 0)
|
||||
|
||||
# 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 = 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_upB.addmm_(upA.t() @ X.t(), df, alpha = upS, beta = 0)
|
||||
|
||||
# Gate projection LoRA weights
|
||||
d_gateA = X.t() @ (de @ gateB.t())
|
||||
d_gateB = (gateA.t() @ X.t()) @ de
|
||||
d_gateA *= gateS
|
||||
d_gateB *= gateS
|
||||
# d_gateA = X.t() @ (de @ gateB.t())
|
||||
# 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_gateB.addmm_(gateA.t() @ X.t(), de, alpha = gateS, beta = 0)
|
||||
|
||||
# dX = matmul_lora(df, upW.t(), upW_quant, upB, upA, upS)
|
||||
# dX += matmul_lora(de, gateW.t(), gateW_quant, gateB, gateA, gateS)
|
||||
upW = fast_dequantize(upW.t(), upW_quant)
|
||||
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 += df @ upB.to(dtype).t() @ (upS * upA.to(dtype).t())
|
||||
dX.addmm_(df @ upB.t(), upA.t(), alpha = upS)
|
||||
|
||||
gateW = fast_dequantize(gateW.t(), gateW_quant)
|
||||
dX += de @ gateW.t()
|
||||
# dX += de @ gateW.t()
|
||||
dX.addmm_(de, gateW.t())
|
||||
del gateW
|
||||
dX += de @ gateB.to(dtype).t() @ (gateS * gateA.to(dtype).t())
|
||||
# dX += de @ gateB.to(dtype).t() @ (gateS * gateA.to(dtype).t())
|
||||
dX.addmm_(de @ gateB.t(), gateA.t(), alpha = gateS)
|
||||
|
||||
# gateW, gateW_quant, gateA, gateB, gateS,
|
||||
# upW, upW_quant, upA, upB, upS,
|
||||
|
|
@ -258,9 +277,6 @@ class LoRA_QKV(torch.autograd.Function):
|
|||
ctx.custom_saved_tensors
|
||||
X, QA, QB, KA, KB, VA, VB, = ctx.saved_tensors
|
||||
|
||||
QA, QB, KA, KB, VA, VB = \
|
||||
QA.t(), QB.t(), KA.t(), KB.t(), VA.t(), VB.t()
|
||||
|
||||
batch, seq_len, hd = X.shape
|
||||
dQ = dQ.view(-1, dQ.shape[-1])
|
||||
dK = dK.reshape(-1, dK.shape[-1]) # view doesn't work on K.T
|
||||
|
|
@ -268,45 +284,68 @@ class LoRA_QKV(torch.autograd.Function):
|
|||
X = X .view(-1, X .shape[-1])
|
||||
dtype = X.dtype
|
||||
|
||||
QA, QB, KA, KB, VA, VB = \
|
||||
QA.to(dtype), QB.to(dtype), KA.to(dtype), KB.to(dtype), VA.to(dtype), VB.to(dtype)
|
||||
|
||||
QA, QB, KA, KB, VA, VB = \
|
||||
QA.t(), QB.t(), KA.t(), KB.t(), VA.t(), VB.t()
|
||||
|
||||
### Weight projection LoRA weights
|
||||
# See our blogpost for more details.
|
||||
d_QA = torch.empty_like(QA)
|
||||
d_QB = torch.empty_like(QB)
|
||||
d_KA = torch.empty_like(KA)
|
||||
d_KB = torch.empty_like(KB)
|
||||
d_VA = torch.empty_like(VA)
|
||||
d_VB = torch.empty_like(VB)
|
||||
|
||||
# Q Projection
|
||||
d_QA = X.t() @ (dQ @ QB.t())
|
||||
d_QB = (QA.t() @ X.t()) @ dQ
|
||||
d_QA *= QS
|
||||
d_QB *= QS
|
||||
# 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_QB.addmm_(QA.t() @ X.t(), dQ, alpha = QS, beta = 0)
|
||||
|
||||
# K Projection
|
||||
d_KA = X.t() @ (dK @ KB.t())
|
||||
d_KB = (KA.t() @ X.t()) @ dK
|
||||
d_KA *= KS
|
||||
d_KB *= KS
|
||||
# d_KA = X.t() @ (dK @ KB.t())
|
||||
# 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_KB.addmm_(KA.t() @ X.t(), dK, alpha = KS, beta = 0)
|
||||
|
||||
# V Projection
|
||||
d_VA = X.t() @ (dV @ VB.t())
|
||||
d_VB = (VA.t() @ X.t()) @ dV
|
||||
d_VA *= VS
|
||||
d_VB *= VS
|
||||
# d_VA = X.t() @ (dV @ VB.t())
|
||||
# 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_VB.addmm_(VA.t() @ X.t(), dV, alpha = VS, beta = 0)
|
||||
|
||||
# Combine derivatives to find dX
|
||||
# dQ
|
||||
QW = fast_dequantize(QW.t(), QW_quant)
|
||||
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 += (dQ @ QB.to(dtype).t() @ (QS * QA.to(dtype).t()))
|
||||
dX.addmm_(dQ @ QB.t(), QA.t(), alpha = QS)
|
||||
|
||||
# dK
|
||||
KW = fast_dequantize(KW.t(), KW_quant)
|
||||
dX += dK @ KW.t()
|
||||
# dX += dK @ KW.t()
|
||||
dX.addmm_(dK, KW.t())
|
||||
del KW
|
||||
dX += dK @ KB.to(dtype).t() @ (KS * KA.to(dtype).t())
|
||||
# dX += dK @ KB.to(dtype).t() @ (KS * KA.to(dtype).t())
|
||||
dX.addmm_(dK @ KB.t(), KA.t(), alpha = KS)
|
||||
|
||||
# dV
|
||||
VW = fast_dequantize(VW.t(), VW_quant)
|
||||
dX += dV @ VW.t()
|
||||
# dX += dV @ VW.t()
|
||||
dX.addmm_(dV, VW.t())
|
||||
del VW
|
||||
dX += dV @ VB.to(dtype).t() @ (VS * VA.to(dtype).t())
|
||||
# dX += dV @ VB.to(dtype).t() @ (VS * VA.to(dtype).t())
|
||||
dX.addmm_(dV @ VB.t(), VA.t(), alpha = VS)
|
||||
|
||||
# QW, QW_quant, QA, QB, QS,
|
||||
# KW, KW_quant, KA, KB, KS,
|
||||
|
|
@ -378,25 +417,33 @@ class LoRA_W(torch.autograd.Function):
|
|||
W, W_quant, S = ctx.custom_saved_tensors
|
||||
A, B, X = ctx.saved_tensors
|
||||
|
||||
A, B = A.t(), B.t()
|
||||
|
||||
batch, seq_len, hd = X.shape
|
||||
dY = dY.reshape(-1, dY.shape[-1]) # Must be reshape
|
||||
X = X .reshape(-1, X .shape[-1]) # Must be reshape
|
||||
dtype = X.dtype
|
||||
|
||||
A, B = A.to(dtype), B.to(dtype)
|
||||
|
||||
A, B = A.t(), B.t()
|
||||
|
||||
d_A = torch.empty_like(A)
|
||||
d_B = torch.empty_like(B)
|
||||
|
||||
### 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 = 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_B.addmm_(A.t() @ X.t(), dY, alpha = S, beta = 0)
|
||||
|
||||
# Get derivative for dX
|
||||
W = fast_dequantize(W.t(), W_quant)
|
||||
dX = dY @ W.t()
|
||||
del W
|
||||
dX += dY @ B.to(dtype).t() @ (S * A.to(dtype).t())
|
||||
# dX += dY @ B.to(dtype).t() @ (S * A.to(dtype).t())
|
||||
dX.addmm_(dY @ B.t(), A.t(), alpha = S)
|
||||
|
||||
# W, W_quant, A, B, S
|
||||
return dX.view(batch, seq_len, hd), \
|
||||
|
|
|
|||
|
|
@ -465,7 +465,8 @@ def matmul_lora(X, W, W_quant, A, B, s, out = None):
|
|||
if A is not None:
|
||||
# LoRA is enabled
|
||||
A, B = A.t(), B.t()
|
||||
out = torch_addmm(X @ A.to(dtype), B.to(dtype), alpha = s, beta = 1.0, out = out)
|
||||
XA = torch_matmul(X, A.to(dtype))
|
||||
out.addmm_(XA, B.to(dtype), alpha = s)
|
||||
# out += (X @ A.to(dtype)) @ (s * B.to(dtype))
|
||||
pass
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue