diff options
| -rw-r--r-- | ep_run/casc_eq_train.py | 4 | ||||
| -rw-r--r-- | ep_run/muon.py | 28 | ||||
| -rw-r--r-- | ep_run/psgd_vendor.py | 1568 |
3 files changed, 1597 insertions, 3 deletions
diff --git a/ep_run/casc_eq_train.py b/ep_run/casc_eq_train.py index d19504d..db6d71a 100644 --- a/ep_run/casc_eq_train.py +++ b/ep_run/casc_eq_train.py @@ -22,7 +22,7 @@ ap.add_argument('--wandb_run', default='') ap.add_argument('--kmax', type=int, default=8) # adaptive fb rounds cap ap.add_argument('--noguard', action='store_true') # diagnosis: skip only non-finite grads ap.add_argument('--untie', action='store_true') # separate readout matrix (untied from tok) -ap.add_argument('--opt', choices=['adamw', 'muon', 'sgdm', 'lion', 'olion', 'adafactor', 'signline', 'conslion', 'ditherlion', 'cautlion', 'olionns', 'olionk1', 'olionk2', 'olionk3'], default='adamw') +ap.add_argument('--opt', choices=['adamw', 'muon', 'sgdm', 'lion', 'olion', 'adafactor', 'signline', 'conslion', 'ditherlion', 'cautlion', 'olionns', 'olionk1', 'olionk2', 'olionk3', 'psgdquad'], default='adamw') ap.add_argument('--muon_lr', type=float, default=0.02) ap.add_argument('--tok_init', type=float, default=0.0) # >0: init tok/pos with this std (GPT-standard 0.02) ap.add_argument('--compile', action='store_true') # torch.compile each block (free speed where supported) @@ -345,7 +345,7 @@ if args.bf16: if args.untie: with torch.no_grad(): W_out.data = W_out.data.to(torch.bfloat16) print('[bf16] model cast to bfloat16 (E-accum + sigma stay fp32)', flush=True) -if args.opt in ('sgdm', 'lion', 'olion', 'adafactor', 'signline', 'conslion', 'ditherlion', 'cautlion', 'olionns', 'olionk1', 'olionk2', 'olionk3'): +if args.opt in ('sgdm', 'lion', 'olion', 'adafactor', 'signline', 'conslion', 'ditherlion', 'cautlion', 'olionns', 'olionk1', 'olionk2', 'olionk3', 'psgdquad'): from muon import build_alt opt, sched = build_alt(args.opt, blocks, all_params, args.lr, args.warmup, total_steps=(args.steps if args.cosine else 0), lr_min_ratio=args.lr_min_ratio, diff --git a/ep_run/muon.py b/ep_run/muon.py index 7bad679..86c81d3 100644 --- a/ep_run/muon.py +++ b/ep_run/muon.py @@ -157,6 +157,29 @@ class Adafactor2D(torch.optim.Optimizer): p.add_(u, alpha=-g_['lr']) +class PSGDQuadWrap(torch.optim.Optimizer): + """Xi-Lin Li's KronWhiten (dQ='QUAD') as a proper torch Optimizer so LambdaLR accepts it. + Our EP flow pre-populates p.grad; the closure hands PSGD a synthetic scalar whose autograd + gradient equals the stored p.grad (loss = sum <p, p.grad_detached>). Screening-tier: the + inner preconditioner state is NOT checkpointed (base state_dict covers param_groups only).""" + def __init__(self, params, lr=1e-3, momentum=0.95): + params = list(params) + super().__init__(params, dict(lr=lr)) + from psgd_vendor import KronWhiten + self._flat = [p for g_ in self.param_groups for p in g_['params']] + self.inner = KronWhiten(self._flat, preconditioner_init_scale=1.0, + lr_params=lr, lr_preconditioner=0.1, momentum=momentum, + whiten_grad=True, dQ="QUAD") + + @torch.no_grad() + def step(self, closure=None): + self.inner.lr_params = self.param_groups[0]['lr'] + flat = self._flat + def _closure(): + return sum((p * p.grad.detach()).sum() for p in flat if p.grad is not None) + self.inner.step(_closure) + + def build_alt(opt_name, blocks, other_params, lr, warmup, total_steps=0, lr_min_ratio=0.1, lr_matrix=None, wd=0.0): """Screening-tier builder for the optimizer price list: OPT on block matrices + AdamW on the @@ -181,7 +204,8 @@ def build_alt(opt_name, blocks, other_params, lr, warmup, total_steps=0, lr_min_ 'olionns': lambda: OLion(mats, lr=lm, wd=wd, ns_steps=0), 'olionk1': lambda: OLion(mats, lr=lm, wd=wd, ns_steps=1), 'olionk2': lambda: OLion(mats, lr=lm, wd=wd, ns_steps=2), - 'olionk3': lambda: OLion(mats, lr=lm, wd=wd, ns_steps=3)}[opt_name]() + 'olionk3': lambda: OLion(mats, lr=lm, wd=wd, ns_steps=3), + 'psgdquad': lambda: PSGDQuadWrap(mats, lr=lm)}[opt_name]() oa = torch.optim.AdamW(rest, lr=lr, weight_decay=1e-4) opts = [om, oa] if total_steps > 0: @@ -296,3 +320,5 @@ class DitherLion(torch.optim.Optimizer): if g_['wd'] > 0: p.mul_(1 - g_['lr'] * g_['wd']) p.add_((u + d).sign_(), alpha=-g_['lr']) m.mul_(b2).add_(p.grad, alpha=1 - b2) + + diff --git a/ep_run/psgd_vendor.py b/ep_run/psgd_vendor.py new file mode 100644 index 0000000..1155b87 --- /dev/null +++ b/ep_run/psgd_vendor.py @@ -0,0 +1,1568 @@ +""" +The new PSGD-Kron Newton/Whitening preconditioners support five kinds of local coordinates for updating Q: + + QUAD): It's a specific form for updating Q to ensure that Q > 0 (thus Q is symmetric/Hermitian). + It still is numerically stable even if round-off errors break the SPD property of Q. + + QEQ): dQ = Q * mathcal{E} * Q + This leads to another simple way for updating Q (Q is in the general linear group). + + Q0.5EQ1.5/Q0p5EQ1p5): dQ = Q^0.5 * mathcal{E} * Q^1.5 + The default and recommended choice for fitting Q. + An online orthogonal Procrustes problem solver is used to keep Q approximately SPD (no need to be exactly SPD). + + EQ): dQ = mathcal{E} * Q + This choice recovers the old PSGD way for updating Q in Lie groups (Q is triangular). + Its main drawback is that triangular solvers are required for updating Q. + + QEP): dQ = Q * mathcal{E} * P + This last choice works very well if it does. Q is in the general linear group. + But, one drawback is that Q might get stuck around ill-conditioned matrices (not strongly convex). + +Both the QUAD and Q0.5EQ1.5 methods can be used to update P directly with little changes. +We call them QUAD4P and PRO4P (PRO is from Procrustes), respectively. +PRO4P still is a competitive and recommended choice for optimization with single precisions. + +The PSGD-LRA Newton/Whitening preconditioners still adopt local coordinate dQ = mathcal{E} * Q, +and needs a small linear solver to update the preconditioner. + +We also keep the PSGD dense matrix Newton-type preconditioner here to illustrate the math. +It supports all the five methods for updating Q, +and can be a good alternative to the BFGS like quasi-Newton optimizers as no line search is required. + +Please refer to + https://github.com/lixilinx/psgd_torch/blob/master/wrapped_as_torch_optimizer_for_ddp.py + https://github.com/lixilinx/psgd_torch/blob/master/wrapped_as_torch_optimizer_for_dtensor.py +for torch.optim optimizer wrappings for DDP, FSDP, FP, etc. trainings and typical settings. + +Xi-Lin Li, lixilinx@gmail.com; last updated in Oct., 2025. +Main refs: https://arxiv.org/abs/1512.04202; https://arxiv.org/abs/2402.11858. +""" + +import opt_einsum +import torch + + +def norm_lower_bound_spd(A, k=32, half_iters=2): + """ + Returns a cheap lower bound for the spectral norm of a symmetric positive definite matrix A, where, + k: the dim of subspace, suggesting 128 for bfloat16 and 32 for float32 (tested on my laptop 4070 GPU); + half_iters: half of the number of subspace iterations, suggesting 2. + A rough norm estimation with bfloat16 is good enough, and we don't orthonormalize the subspace vectors. + + The initial noise space V is rotated such that its centroid aligns with the largest row of A. + Hence, each row of V and the largest row of A has an angle about acos(1/sqrt(k)) when k << dim(A). + This feature makes the subspace iteration more robust for large matrices with very low rank. + A simplified branchless approximate implementation is provided here. + """ + smallest_normal = torch.finfo(A.dtype).smallest_normal + normalizing_factor = A.diagonal().real.amax() + smallest_normal + A = A / normalizing_factor # (complex tensor) / (subnormal number) could produce inf or nan unexpectedly + j = torch.argmax(torch.linalg.vector_norm(A, dim=1)) + V = torch.randn(k, A.shape[1], dtype=A.dtype, device=A.device) + V = A[j] + torch.sgn(torch.sum(A[j] * V.conj(), dim=1, keepdim=True)) * V # torch.sign for real + for _ in range(half_iters): + V = V @ A + V /= torch.linalg.vector_norm(V, dim=1, keepdim=True) + smallest_normal + V = V @ A + return normalizing_factor * torch.amax(torch.linalg.vector_norm(V, dim=1)) + + +def norm_lower_bound_skh(A, k=32, half_iters=2): + """ + Returns a cheap lower bound for the spectral norm of a skew-Hermitian matrix A, + k: the dim of subspace, suggesting 128 for bfloat16 and 32 for float32 (tested on my laptop 4070 GPU); + half_iters: half of the number of subspace iterations, suggesting 2. + A rough norm estimation with bfloat16 is good enough, and we don't orthonormalize the subspace vectors. + + The initial noise space V is rotated such that its centroid aligns with the largest row of A. + Hence, each row of V and the largest row of A has an angle about acos(1/sqrt(k)) when k << dim(A). + This feature makes the subspace iteration more robust for large matrices with very low rank. + A simplified branchless approximate implementation is provided here. + """ + smallest_normal = torch.finfo(A.dtype).smallest_normal + normalizing_factor = A.abs().amax() + smallest_normal + A = A / normalizing_factor # (complex tensor) / (subnormal number) could produce inf or nan unexpectedly + j = torch.argmax(torch.linalg.vector_norm(A, dim=1)) + V = torch.randn(k, A.shape[1], dtype=A.dtype, device=A.device) + V = A[j] + torch.sgn(torch.sum(A[j] * V.conj(), dim=1, keepdim=True)) * V # torch.sign for real + for _ in range(half_iters): + V = V @ A + V /= torch.linalg.vector_norm(V, dim=1, keepdim=True) + smallest_normal + V = V @ A + return normalizing_factor * torch.amax(torch.linalg.vector_norm(V, dim=1)) + + +def lift2single(x): + # lift half or lower precision to single precision; leave single precision unchanged + return x.to(torch.float32) if torch.finfo(x.dtype).eps > 1e-6 else x + + +def procrustes_step2(Q, max_step_size=1/8): + """ + A in-place (update Q directly) online solver for the orthogonal Procrustes problem, + min_U || U Q - I ||_F, s.t. U^H U = I + by rotating Q as exp(a R) Q, where R = Q^H - Q is the generator and ||a R|| < 1. + + We expand U = exp(a R) to its 2nd term as + U ~ I + aR + (aR)^2/2 + and the truncation error ||U^H U - I|| is upper bounded as ||a R||^4/4. + Set max_step_size <= 1/4 such that the truncation error <= (1/4)^4/4 < 1e-3. + + Note that U(n) is connected and such rotations can make almost any complex Q SPD except for convergence to saddle points. + However, O(n) is not connected. Hence, such SO(n) rotations can only make real Q with det(Q) > 0 SPD. + + We have simplified the original implementation. The one branch here is necessary for line search. + """ + R = Q.H - Q + R /= norm_lower_bound_skh(R) + torch.finfo(R.dtype).smallest_normal # normalize R as typically it's too small + RQ = R @ Q + RRQ = R @ RQ + tr_RQ = RQ.diagonal().real.sum() # tr_RQ >=0 by theory; torch.trace not implemented for CPU bfloat16, so using sum(diag(.)) here + tr_RRQ = RRQ.diagonal().real.sum() # line search is needed if tr_RRQ < 0 + a = torch.where(tr_RRQ < 0, torch.clamp(-tr_RQ / tr_RRQ, max=max_step_size), max_step_size) + Q.add_(a * (RQ + 0.5 * a * RRQ)) + + +def procrustes_step3(Q, max_step_size=1/3): + """ + A in-place (update Q directly) online solver for the orthogonal Procrustes problem, + min_U || U Q - I ||_F, s.t. U^H U = I + by rotating Q as exp(a R) Q, where R = Q^H - Q is the generator and ||a R|| < 1. + + We expand U = exp(a R) to its 3rd term as (not the same as the Taylor series of U) + U ~ I + aR + (aR)^2/2 + (aR)^3/8 + and the truncation error ||U^H U - I|| is upper bounded as ||a R||^6/64. + Set max_step_size <= 5/8 such that the truncation error <= (5/8)^6/64 < 1e-3. + + Note that U(n) is connected and such rotations can make almost any complex Q SPD except for convergence to saddle points. + However, O(n) is not connected. Hence, such SO(n) rotations can only make real Q with det(Q) > 0 SPD. + """ + R = Q.H - Q + R /= norm_lower_bound_skh(R) + torch.finfo(R.dtype).smallest_normal # normalize R as typically it's too small + RQ = R @ Q + RRQ = R @ RQ + RRRQ = R @ RRQ + tr_RQ = RQ.diagonal().real.sum() # tr_RQ >=0 by theory; torch.trace not implemented for CPU bfloat16, so using sum(diag(.)) here + tr_RRQ = RRQ.diagonal().real.sum() + tr_RRRQ = RRRQ.diagonal().real.sum() # tr_RRRQ <=0 + if tr_RQ > 0 and tr_RRRQ < 0: # otherwise, Q^T = Q up to machine precision + # optimal a is the larger root of tr_RQ + 2 * a * tr_RRQ / 2 + 3 * a^2 * tr_RRRQ / 8 = 0 + if torch.finfo(tr_RQ.dtype).eps > 1e-6: # half precision is not accurate enough when tr_RRQ < 0 + tr_RQ, tr_RRQ, tr_RRRQ = tr_RQ.to(torch.float32), tr_RRQ.to(torch.float32), tr_RRRQ.to(torch.float32) + a = (-tr_RRQ - torch.sqrt(tr_RRQ*tr_RRQ - 1.5*tr_RQ*tr_RRRQ)) / (0.75*tr_RRRQ) + a = torch.clamp(a, max=max_step_size) + Q.add_(a * (RQ + 0.5 * a * (RRQ + 0.25 * a * RRRQ))) + + +############# Begin of PSGD Kronecker product preconditioners ############# + + +def init_kron(t, Scale=1.0, max_size=float("inf"), max_skew=1.0, dQ="Q0.5EQ1.5"): + """ + For a scalar or tensor t, we initialize its states (preconditioner Q and Lipschitz smoothness constant L), + and reusable contraction expressions for updating Q and preconditioning gradient. + + 1, The preconditioner Q is initialized to + Q = Scale * I = Scale * kron(eye(t.shape[0]), eye(t.shape[1]), ...) + where the eye(.) may be replaced with diag(ones(.)) if that dim is too large, determined by max_size and max_skew. + + The Lipschitz smoothness constant L for Q is initialized to zero. + + 2, A series of einsum contract expressions. The following subscript examples are for a 5th order tensor. + 2.1, exprP is the expression for applying the Preconditioner on the gradient, e.g., + 'aA,bB,cC,dD,eE,aα,bβ,cγ,dδ,eε,αβγδε->ABCDE' + 2.2, the i-th expression of exprGs is for the contraction of two tensors that only keeps the i-th dim, e.g., + 'abCde,abγde->Cγ' + for i=2. It's useful for Gradient calculation. + 2.3, exprA is the expression for applying All the factors of Q on a tensor, e.g., + 'aA,bB,cC,dD,eE,ABCDE->abcde' + 2.4, the i-th expression of exprQs is the expression for applying the i-th factor of Q on a tensor, e.g., + 'Cγ,abγde->abCde' + for i=2. + + Please check https://drive.google.com/file/d/1CEEq7A3_l8EcPEDa_sYtqr5aMLVeZWL7/view?usp=drive_link for notations and derivations. + """ + if dQ in {"QUAD4P", "PRO4P"}: # the only two cases that we fit P directly; so square Scale + Scale = Scale ** 2 + shape = t.shape + if len(shape)==0: # scalar + Q = [Scale * torch.ones_like(t),] + L = [lift2single(torch.zeros_like(t.real)),] + exprA = opt_einsum.contract_expression(",->", Q[0].shape, t.shape) + exprP = opt_einsum.contract_expression(",,->", Q[0].shape, Q[0].shape, t.shape) + exprGs = [opt_einsum.contract_expression(",->", t.shape, t.shape),] + exprQs = [opt_einsum.contract_expression(",->", Q[0].shape, t.shape),] + else: # tensor + if len(shape) > 26: + raise ValueError(f"Got tensor with dim {len(t.shape)}; einsum runs out of letters; replace 26 with larger numbers.") + + scale = Scale ** (1/len(shape)) + + Q, L = [], [] + exprGs, exprQs = [], [] + piece1A, piece2A, piece3A = [], "", "" # used for getting the subscripts for exprA + piece1P, piece2P, piece3P, piece4P = [], [], "", "" # used for getting the subscripts for exprP + for i, size in enumerate(shape): + L.append(lift2single(torch.zeros([], dtype=t.real.dtype, device=t.device))) + if size <= 1 or size > max_size or size**2 > max_skew * t.numel(): + # use diagonal matrix as preconditioner for this dim + Q.append(scale * torch.ones(size, dtype=t.dtype, device=t.device)) + + piece1A.append(opt_einsum.get_symbol(i)) + piece2A = piece2A + opt_einsum.get_symbol(i) + piece3A = piece3A + opt_einsum.get_symbol(i) + + piece1P.append(opt_einsum.get_symbol(i + 26)) + piece2P.append(opt_einsum.get_symbol(i + 26)) + piece3P = piece3P + opt_einsum.get_symbol(i + 26) + piece4P = piece4P + opt_einsum.get_symbol(i + 26) + + piece1 = "".join([opt_einsum.get_symbol(i+26) if j==i else opt_einsum.get_symbol(j) for j in range(len(shape))]) + subscripts = piece1 + "," + piece1 + "->" + opt_einsum.get_symbol(i+26) + exprGs.append(opt_einsum.contract_expression(subscripts, t.shape, t.shape)) + + subscripts = opt_einsum.get_symbol(i+26) + "," + piece1 + "->" + piece1 + exprQs.append(opt_einsum.contract_expression(subscripts, Q[-1].shape, t.shape)) + else: # use matrix preconditioner for this dim + Q.append(scale * torch.eye(size, dtype=t.dtype, device=t.device)) + + piece1A.append(opt_einsum.get_symbol(i) + opt_einsum.get_symbol(i + 26)) + piece2A = piece2A + opt_einsum.get_symbol(i + 26) + piece3A = piece3A + opt_einsum.get_symbol(i) + + a, b, c = opt_einsum.get_symbol(i), opt_einsum.get_symbol(i + 26), opt_einsum.get_symbol(i + 805) + piece1P.append(a + b) + piece2P.append(a + c) + piece3P = piece3P + c + piece4P = piece4P + b + + piece1 = "".join([opt_einsum.get_symbol(i+26) if j==i else opt_einsum.get_symbol(j) for j in range(len(shape))]) + piece2 = "".join([opt_einsum.get_symbol(i+805) if j==i else opt_einsum.get_symbol(j) for j in range(len(shape))]) + subscripts = piece1 + "," + piece2 + "->" + opt_einsum.get_symbol(i+26) + opt_einsum.get_symbol(i+805) + exprGs.append(opt_einsum.contract_expression(subscripts, t.shape, t.shape)) + + subscripts = opt_einsum.get_symbol(i+26) + opt_einsum.get_symbol(i+805) + "," + piece2 + "->" + piece1 + exprQs.append(opt_einsum.contract_expression(subscripts, Q[-1].shape, t.shape)) + + subscripts = ",".join(piece1A) + "," + piece2A + "->" + piece3A + exprA = opt_einsum.contract_expression(subscripts, *[q.shape for q in Q], t.shape) + + subscripts = ",".join(piece1P) + "," + ",".join(piece2P) + "," + piece3P + "->" + piece4P + exprP = opt_einsum.contract_expression(subscripts, *[q.shape for q in Q], *[q.shape for q in Q], t.shape) + + exprGs, exprQs = tuple(exprGs), tuple(exprQs) + if dQ == "QEP": + return [[Q, L], (exprP, exprGs, exprQs)] + elif dQ == "EQ": + return [[Q, L], (exprP, exprGs, exprA)] + elif dQ in {"QEQ", "QUAD", "Q0p5EQ1p5", "Q0.5EQ1.5"}: + return [[Q, L], (exprP, exprGs)] + else: # the only two cases that we fit P directly; dQ actually is dP + assert dQ in {"QUAD4P", "PRO4P"}, "Invalid choice for dQ" + return [[Q, L], (exprA, exprGs)] + + +def balance_kron_precond(Q): + """ + In place balancing the dynamic ranges of the factors of Q to avoid over/under-flow. + """ + order = len(Q) # order of tensor or the number of factors in Q + if order>1: + norms = [torch.max(torch.abs(q)) for q in Q] + gmean = torch.prod(torch.stack(norms))**(1/order) # geometric mean + for i, q in enumerate(Q): + q.mul_(gmean/norms[i]) + + +def update_precond_kron_eq(QL, exprs, V, Hvp, lr=0.1, betaL=0.9): + """ + The raw function for updating the Kron preconditioner Q and Lipschitz smoothness constant L with pair (V, Hvp), + where Q is update as dQ = E*Q, + the pair (V, Hvp) can be (vector, hess-vector-prod) or (randn, gradient/momentum). + The damping logic is not included here. + """ + Q, L = QL + _, exprGs, exprA = exprs + + def solve_triangular_right(B, A): + # return B @ inv(A) + if B.dim()>1: + return torch.linalg.solve_triangular(lift2single(A), lift2single(B), upper=True, left=False).to(B.dtype) + else: # torch.linalg.solve_triangular complains if B.dim() < 2. So insert None. + return (torch.linalg.solve_triangular(lift2single(A), lift2single(B[None,:]), upper=True, left=False)[0]).to(B.dtype) + + A = exprA(*Q, Hvp) + + order = V.dim() + p = list(range(order)) + conjB = torch.permute(V.conj(), p[1:] + p[:1]) # permute dims like [0,1,2,3,4] -> [1,2,3,4,0] + for i, q in enumerate(Q): + conjB = conjB/q if q.dim()<2 else solve_triangular_right(conjB, q) + if i < order - 1: # transpose dims like [1,2,3,4,0]->[0,2,3,4,1]->[0,1,3,4,2]->[0,1,2,4,3]->[0,1,2,3,4] + conjB = torch.transpose(conjB, i, order - 1) + + for i, q in enumerate(Q): + term1 = exprGs[i](A, A.conj()) + term2 = exprGs[i](conjB.conj(), conjB) + + if q.dim() < 2: # q is a diagonal matrix or scalar preconditioner + ell = torch.max(torch.real(term1 + term2)) + L[i].copy_(torch.max(betaL*L[i] + (1 - betaL)*ell, ell)) + q.sub_(lr/L[i] * (term1 - term2) * q) # q.mul_(1 - lr/L[i] * (term1 - term2)): larger roundoff errors + else: # q is a matrix preconditioner + ell = norm_lower_bound_spd(term1 + term2) + L[i].copy_(torch.max(betaL*L[i] + (1 - betaL)*ell, ell)) + q.sub_(lr/L[i] * torch.triu(term1 - term2) @ q) + + if torch.rand([]) < 0.01: # balance factors of Q + balance_kron_precond(Q) + + +def precond_grad_kron(QL, exprs, G): + """ + Precondition gradient G with Kron preconditioner Q. + """ + Q, exprP = QL[0], exprs[0] + return exprP(*[q.conj() for q in Q], *Q, G) + + +def update_precond_kron_whiten_eq(QL, exprs, G, lr=0.1, betaL=0.9, damping=1e-9): + """ + Update the Kron preconditioner Q as dQ = E*Q. + """ + V = torch.randn_like(G) + damping = damping + torch.finfo(G.dtype).eps * G.abs() + update_precond_kron_eq(QL, exprs, V, G + damping*V, lr=lr, betaL=betaL) + + +def update_precond_kron_whiten_qep(QL, exprs, G, lr=0.1, betaL=0.9, damping=1e-9): + """ + Update the Kron preconditioner Q as dQ = Q*E*P. + """ + Q, L = QL + exprP, exprGs, exprQs = exprs + + # balancing is not optional as L for each factor is not scaling invariant + balance_kron_precond(Q) + + total_numel = G.numel() + damping = damping + torch.finfo(G.dtype).eps * G.abs() + Pg = exprP(*[q.conj() for q in Q], *Q, G + damping*torch.randn_like(G)) + for i, q in enumerate(Q): + QPg = exprQs[i](q, Pg) + term1 = exprGs[i](QPg, QPg.conj()) + if q.dim() < 2: # diagonal or scalar Q + term2 = total_numel/q.numel() * q * q.conj() + ell = torch.max(torch.real(term1 + term2)) + L[i].copy_(torch.max(betaL*L[i] + (1 - betaL)*ell, ell)) + q.mul_(1 - lr/L[i] * (term1 - term2)) + else: # matrix Q + term2 = total_numel/q.shape[0] * q @ q.H + ell = norm_lower_bound_spd(term1 + term2) + L[i].copy_(torch.max(betaL*L[i] + (1 - betaL)*ell, ell)) + q.sub_(lr/L[i] * (term1 - term2) @ q) + + +def update_precond_kron_whiten_qeq(QL, exprs, G, lr=0.1, betaL=0.9, damping=1e-9): + """ + Update the Kron preconditioner Q as dQ = Q*E*Q. + """ + Q, L = QL + exprP, exprGs = exprs + + total_numel = G.numel() + damping = damping + torch.finfo(G.dtype).eps * G.abs() + Pg = exprP(*[q.conj() for q in Q], *Q, G + damping*torch.randn_like(G)) + for i, q in enumerate(Q): + term1 = exprGs[i](Pg, Pg.conj()) + if q.dim() < 2: # diagonal or scalar Q + term2 = total_numel/q.numel() # times I + ell = torch.max(torch.real(term1)) + term2 + L[i].copy_(torch.max(betaL*L[i] + (1 - betaL)*ell, ell)) + q.mul_(1 - lr/L[i] * (term1 - term2)) + else: # matrix Q + term2 = total_numel/q.shape[0] # times I + ell = norm_lower_bound_spd(term1) + term2 + L[i].copy_(torch.max(betaL*L[i] + (1 - betaL)*ell, ell)) + q.sub_(lr/L[i] * (q @ term1 - q * term2)) + + if torch.rand([]) < 0.01: # balance factors of Q + balance_kron_precond(Q) + + +def update_precond_kron_whiten_q0p5eq1p5(QL, exprs, G, lr=0.1, betaL=0.9, damping=1e-9): + """ + Update the Kron preconditioner Q as dQ = Q^0.5 * E * Q^1.5. + """ + Q, L = QL + exprP, exprGs = exprs + + total_numel = G.numel() + damping = damping + torch.finfo(G.dtype).eps * G.abs() + Pg = exprP(*[q.conj() for q in Q], *Q, G + damping*torch.randn_like(G)) + for i, q in enumerate(Q): + term1 = exprGs[i](Pg, Pg.conj()) + if q.dim() < 2: # diagonal or scalar Q + term2 = total_numel/q.numel() # times I + ell = torch.max(torch.real(term1)) + term2 + L[i].copy_(torch.max(betaL*L[i] + (1 - betaL)*ell, ell)) + q.mul_(1 - lr/L[i] * (term1 - term2)) + else: # matrix Q + term2 = total_numel/q.shape[0] # times I + ell = norm_lower_bound_spd(term1) + term2 + L[i].copy_(torch.max(betaL*L[i] + (1 - betaL)*ell, ell)) + q.sub_(lr/L[i] * (term1 @ q - term2 * q)) + procrustes_step2(q) + + if torch.rand([]) < 0.01: # balance factors of Q + balance_kron_precond(Q) + + +def update_precond_kron_whiten_pro4p(QL, exprs, G, lr=0.1, betaL=0.9, damping=1e-9): + """ + Update the Kron preconditioner P as dP = P^0.5 * E * P. + Almost the same as update_precond_kron_whiten_q0p5eq1p5. But the Q here actually is P. + Unlike fitting Q, fitting P directly is more sensitive to numerical round-off errors. + """ + Q, L = QL + exprA, exprGs = exprs + + total_numel = G.numel() + damping = damping + torch.finfo(G.dtype).eps * G.abs() + Pg = exprA(*Q, G + damping*torch.randn_like(G)) # Q actually is P; so just applying all its factors once. + for i, q in enumerate(Q): + term1 = exprGs[i](Pg, Pg.conj()) + if q.dim() < 2: # diagonal or scalar Q + term2 = total_numel/q.numel() # times I + ell = torch.max(torch.real(term1)) + term2 + L[i].copy_(torch.max(betaL*L[i] + (1 - betaL)*ell, ell)) + q.mul_(1 - lr/L[i] * (term1 - term2)) + else: # matrix Q + term2 = total_numel/q.shape[0] # times I + ell = norm_lower_bound_spd(term1) + term2 + L[i].copy_(torch.max(betaL*L[i] + (1 - betaL)*ell, ell)) + q.sub_(lr/L[i] * (term1 @ q - term2 * q)) + for _ in range(10): + procrustes_step3(q) + if (q.H - q).abs().amax() < 0.001 * q.abs().amax(): + break # q is almost Hermitian + + if torch.rand([]) < 0.01: # balance factors of P + balance_kron_precond(Q) + + +def update_precond_kron_whiten_quad(QL, exprs, G, lr=0.1, betaL=0.9, damping=1e-9): + """ + Update the Kron preconditioner Q with a quadratic form. + """ + Q, L = QL + exprP, exprGs = exprs + + total_numel = G.numel() + damping = damping + torch.finfo(G.dtype).eps * G.abs() + Pg = exprP(*[q.conj() for q in Q], *Q, G + damping*torch.randn_like(G)) + for i, q in enumerate(Q): + term1 = exprGs[i](Pg, Pg.conj()) + if q.dim() < 2: # diagonal or scalar Q + term2 = total_numel/q.numel() # times I + ell = torch.max(torch.real(term1)) + term2 + L[i].copy_(torch.max(betaL*L[i] + (1 - betaL)*ell, ell)) + gain = 1 - lr/2/L[i] * (term1 - term2) + q.mul_(gain * gain) + else: # matrix Q + term2 = total_numel/q.shape[0] # times I + ell = norm_lower_bound_spd(term1) + term2 + L[i].copy_(torch.max(betaL*L[i] + (1 - betaL)*ell, ell)) + p = q - lr/2/L[i] * (term1 @ q - term2 * q) + p = p - lr/2/L[i] * (p @ term1 - p * term2) + q.copy_((p + p.H)/2) # p must be symmetric/hermitian + + if torch.rand([]) < 0.01: # balance factors of Q + balance_kron_precond(Q) + + +def update_precond_kron_whiten_quad4p(QL, exprs, G, lr=0.1, betaL=0.9, damping=1e-9): + """ + Almost the same as function update_precond_kron_whiten_quad except that fitting P directly (Q here actually is P). + Vulnerable to numerical errors as the round-off errors could break the SPD property of P. + """ + Q, L = QL + exprA, exprGs = exprs + + total_numel = G.numel() + damping = damping + torch.finfo(G.dtype).eps * G.abs() + Pg = exprA(*Q, G + damping*torch.randn_like(G)) # Q actually is P; so just applying all its factors once. + for i, q in enumerate(Q): + term1 = exprGs[i](Pg, Pg.conj()) + if q.dim() < 2: # diagonal or scalar Q + term2 = total_numel/q.numel() # times I + ell = torch.max(torch.real(term1)) + term2 + L[i].copy_(torch.max(betaL*L[i] + (1 - betaL)*ell, ell)) + gain = 1 - lr/L[i] * (term1 - term2) + q.mul_(gain * gain) + else: # matrix Q + term2 = total_numel/q.shape[0] # times I + ell = norm_lower_bound_spd(term1) + term2 + L[i].copy_(torch.max(betaL*L[i] + (1 - betaL)*ell, ell)) + p = q - lr/L[i] * (term1 @ q - term2 * q) + p = p - lr/L[i] * (p @ term1 - p * term2) + q.copy_((p + p.H)/2) # p must be symmetric/hermitian + + if torch.rand([]) < 0.01: # balance factors of Q + balance_kron_precond(Q) + + +class KronWhiten: + """ + Implements the PSGD optimizer with the Kronecker product gradient/momentum whitening preconditioner. + Most of the time, the hyperparameter name says it all. Here are some comments on a few key hyperparameters. + + 1, preconditioner_max_size and preconditioner_max_skew. These two together control the complexity of the preconditioners. + For example, we are to precondition a 2D gradient with shape 10 x 50. + With preconditioner_max_size 20, we use a dense preconditioner for the first dim since 10 <= 20 and diagonal preconditioner for the second dim since 50 > 20. + With preconditioner_max_skew 1.5, we use a dense preconditioner for the first dim since 10/50 <= 1.5 and diagonal preconditioner for the second dim since 50/10 > 1.5. + + 2, grad_clip_max_amps, betaL and damping. These three together help to stabilize the training. + PSGD here tries to normalize the gradients to unit amplitude. This can be problematic when gradients approach zeros. + The most effective way is to clip the preconditioned gradients if their average/element-wise amplitudes exceed grad_clip_max_amps[0]/[1], respectively. + Another way is to damp and upper bound the fitted preconditioner such that P < eye/damping. + For extremely sparse gradients, increasing betaL (say to 0.999) helps a lot, where betaL is the EMA factor for the L-smoothness constant (wrt Q) estimation. + + 3, Lastly, dQ is for the selection of geometry for preconditioner update. + The two recommended choices are dQ = Q0.5EQ1.5 and dP = P0.5EP (online Newton-Schulz iterations). + Q is initialized to preconditioner_init_scale * eye. Boolean setting whiten_grad decides to whiten whether the gradient or momentum. + Always good to check https://arxiv.org/abs/2402.11858 for math details. + """ + def __init__(self, params_with_grad, + preconditioner_max_size=float("inf"), preconditioner_max_skew=1.0, preconditioner_init_scale:float|None=None, + lr_params=0.001, lr_preconditioner=0.1, betaL=0.9, damping=1e-9, momentum=0.0, grad_clip_max_amps=(2.0, 10.0), + preconditioner_update_probability=1.0, update_preconditioner_first=True, whiten_grad=True, dQ="Q0.5EQ1.5"): + # mutable members + self.lr_params = lr_params + self.lr_preconditioner = lr_preconditioner + self.betaL = betaL # beta for the Lipschitz smoothness constant estimation; set to a large value for sparse gradients + self.damping = damping # to damp and upper bound the preconditioner such that P < eye/damping + self.momentum = momentum if (0<momentum<1) else 0.0 + self.grad_clip_max_amps = grad_clip_max_amps # clip grad with thresholds (max average amplitude, max element-wise amplitude) + self.preconditioner_update_probability = preconditioner_update_probability + self.update_preconditioner_first = update_preconditioner_first # True for biased update; False for unbiased update. + # protected members + self._preconditioner_max_size = preconditioner_max_size + self._preconditioner_max_skew = preconditioner_max_skew + params_with_grad = [params_with_grad,] if isinstance(params_with_grad, torch.Tensor) else params_with_grad + self._params_with_grad = [param for param in params_with_grad if param.requires_grad] # double check requires_grad flag + if preconditioner_init_scale is None: + self._QLs_exprs = None # initialize on the fly + print("FYI: Will set the preconditioner initial scale on the fly. Recommend to set it manually.") + else: + self._QLs_exprs = [init_kron(p.squeeze(), preconditioner_init_scale, preconditioner_max_size, preconditioner_max_skew, dQ) for p in self._params_with_grad] + self._ms, self._counter_m = None, 0 # momentum buffers and counter + self._whiten_grad = whiten_grad # set to False to whiten momentum. + if not whiten_grad: + assert self.momentum > 0, "Cannot whiten momentum if the momentum setting is invalid." + print(f"Recommend reducing the lr_params for gradient whitening by a factor of {((1 + self.momentum)/(1 - self.momentum))**0.5} for this momentum whitening setting.") + self._dQ = dQ + if dQ in {"QUAD4P", "PRO4P"}: # the only two cases that we fit P directly + if max([torch.finfo(p.dtype).eps for p in self._params_with_grad]) > 1e-6: + print("Fitting P directly with half precision is risky.") + if dQ == "QUAD4P": + self._update_precond = update_precond_kron_whiten_quad4p + else: # dP = P^0.5 * E * P + self._update_precond = update_precond_kron_whiten_pro4p + self._precond_grad = lambda QL, exprs, G: exprs[0](*QL[0], G) # it's exprA(*Q, G) + else: + self._precond_grad = precond_grad_kron + if dQ == "QEP": + self._update_precond = update_precond_kron_whiten_qep + elif dQ == "EQ": + self._update_precond = update_precond_kron_whiten_eq + elif dQ == "QEQ": + self._update_precond = update_precond_kron_whiten_qeq + elif dQ == "QUAD": + self._update_precond = update_precond_kron_whiten_quad + else: + assert dQ in {"Q0.5EQ1.5", "Q0p5EQ1p5"}, "Invalid choice for dQ" + self._update_precond = update_precond_kron_whiten_q0p5eq1p5 + + + @torch.no_grad() + def step(self, closure): + """ + Performs one step of PSGD with the Kronecker product gradient/momentum whitening preconditioner. + """ + with torch.enable_grad(): + closure_returns = closure() + loss = closure_returns if isinstance(closure_returns, torch.Tensor) else closure_returns[0] + grads = [g.squeeze() for g in torch.autograd.grad(loss, self._params_with_grad)] + + if self._QLs_exprs is None: + scale = max([torch.mean((torch.abs(g))**4) for g in grads]) + scale = (scale + self.damping**4)**(-1/8) + self._QLs_exprs = [init_kron(g, scale, self._preconditioner_max_size, self._preconditioner_max_skew, self._dQ) for g in grads] + + if self.momentum > 0: + beta = min(self._counter_m/(1 + self._counter_m), self.momentum) + self._counter_m += 1 + if self._ms is None: + self._ms = [torch.zeros_like(g) for g in grads] + + for (m, g) in zip(self._ms, grads): + m.mul_(beta).add_(g, alpha=1 - beta) + else: + self._ms, self._counter_m = None, 0 + + if torch.rand([]) < self.preconditioner_update_probability: + update_preconditioner_first, update_preconditioner_last = self.update_preconditioner_first, not self.update_preconditioner_first + else: + update_preconditioner_first, update_preconditioner_last = False, False + + if update_preconditioner_first: # update Q + if self._whiten_grad: # Q whitens gradient + for (QL_exprs, g) in zip(self._QLs_exprs, grads): + self._update_precond(*QL_exprs, g, lr=self.lr_preconditioner, betaL=self.betaL, damping=self.damping) + else: # Q whitens momentum + for (QL_exprs, m) in zip(self._QLs_exprs, self._ms): + self._update_precond(*QL_exprs, m, lr=self.lr_preconditioner, betaL=self.betaL, damping=self.damping) + + if self.momentum > 0: # precondition momentum + pre_grads = [self._precond_grad(*QL_exprs, m) for (QL_exprs, m) in zip(self._QLs_exprs, self._ms)] + else: # precondition gradient + pre_grads = [self._precond_grad(*QL_exprs, g) for (QL_exprs, g) in zip(self._QLs_exprs, grads)] + + if update_preconditioner_last: # update Q + if self._whiten_grad: # Q whitens gradient + for (QL_exprs, g) in zip(self._QLs_exprs, grads): + self._update_precond(*QL_exprs, g, lr=self.lr_preconditioner, betaL=self.betaL, damping=self.damping) + else: # Q whitens momentum + for (QL_exprs, m) in zip(self._QLs_exprs, self._ms): + self._update_precond(*QL_exprs, m, lr=self.lr_preconditioner, betaL=self.betaL, damping=self.damping) + + # Update the parameters after clipping the preconditioned gradient per tensor + max_avg_amp, max_element_amp = self.grad_clip_max_amps + for param, g in zip(self._params_with_grad, pre_grads): + avg_amp = torch.sqrt(torch.real(torch.mean(g*g.conj()))) + if avg_amp > max_avg_amp: + g *= max_avg_amp/avg_amp + if torch.is_complex(g): + g /= torch.clamp(torch.abs(g)/max_element_amp, min=1.0) + else: + g.clamp_(min=-max_element_amp, max=max_element_amp) + param.subtract_(g.view_as(param), alpha=self.lr_params) + + # return whatever closure returns + return closure_returns + + +def update_precond_kron_newton_eq(QL, exprs, V, Hvp, lr=0.1, betaL=0.9, damping=1e-9): + """ + Update the Kron Newton-type preconditioner Q as dQ = E*Q with a pair of vector and hvp, (V, Hvp). + """ + damping = damping + torch.finfo(Hvp.dtype).eps * Hvp.abs() + update_precond_kron_eq(QL, exprs, V, Hvp + damping*torch.randn_like(Hvp), lr=lr, betaL=betaL) + + +def update_precond_kron_newton_qep(QL, exprs, V, Hvp, lr=0.1, betaL=0.9, damping=1e-9): + """ + Update the Kron Newton-type preconditioner Q as dQ = Q*E*P with a pair of vector and hvp, (V, Hvp). + """ + Q, L = QL + exprP, exprGs, exprQs = exprs + + # balancing is not optional as L for each factor is not scaling invariant + balance_kron_precond(Q) + damping = damping + torch.finfo(Hvp.dtype).eps * Hvp.abs() + Ph = exprP(*[q.conj() for q in Q], *Q, Hvp + damping*torch.randn_like(Hvp)) + + for i, q in enumerate(Q): + QPh = exprQs[i](q, Ph) + Qv = exprQs[i](q, V) + term1 = exprGs[i](QPh, QPh.conj()) + term2 = exprGs[i](Qv, Qv.conj()) + if q.dim() < 2: # diagonal or scalar Q + ell = torch.max(torch.real(term1 + term2)) + L[i].copy_(torch.max(betaL*L[i] + (1 - betaL)*ell, ell)) + q.mul_(1 - lr/L[i] * (term1 - term2)) + else: # matrix Q + ell = norm_lower_bound_spd(term1 + term2) + L[i].copy_(torch.max(betaL*L[i] + (1 - betaL)*ell, ell)) + q.sub_(lr/L[i] * (term1 - term2) @ q) + + +def update_precond_kron_newton_qeq(QL, exprs, V, Hvp, lr=0.1, betaL=0.9, damping=1e-9): + """ + Update the Kron Newton-type preconditioner Q as dQ = Q*E*Q with a pair of vector and hvp, (V, Hvp). + """ + Q, L = QL + exprP, exprGs = exprs + damping = damping + torch.finfo(Hvp.dtype).eps * Hvp.abs() + Ph = exprP(*[q.conj() for q in Q], *Q, Hvp + damping*torch.randn_like(Hvp)) + + for i, q in enumerate(Q): + term1 = exprGs[i](Ph, Ph.conj()) + term2 = exprGs[i](V, V.conj()) + if q.dim() < 2: # diagonal or scalar Q + ell = torch.max(torch.real(term1 + term2)) + L[i].copy_(torch.max(betaL*L[i] + (1 - betaL)*ell, ell)) + q.mul_(1 - lr/L[i] * (term1 - term2)) + else: # matrix Q + ell = norm_lower_bound_spd(term1 + term2) + L[i].copy_(torch.max(betaL*L[i] + (1 - betaL)*ell, ell)) + q.sub_(lr/L[i] * q @ (term1 - term2)) + + if torch.rand([]) < 0.01: # balance factors of Q + balance_kron_precond(Q) + + +def update_precond_kron_newton_q0p5eq1p5(QL, exprs, V, Hvp, lr=0.1, betaL=0.9, damping=1e-9): + """ + Update the Kron Newton-type preconditioner Q as dQ = Q^0.5 * E * Q^1.5 with a pair of vector and hvp, (V, Hvp). + """ + Q, L = QL + exprP, exprGs = exprs + damping = damping + torch.finfo(Hvp.dtype).eps * Hvp.abs() + Ph = exprP(*[q.conj() for q in Q], *Q, Hvp + damping*torch.randn_like(Hvp)) + + for i, q in enumerate(Q): + term1 = exprGs[i](Ph, Ph.conj()) + term2 = exprGs[i](V, V.conj()) + if q.dim() < 2: # diagonal or scalar Q + ell = torch.max(torch.real(term1 + term2)) + L[i].copy_(torch.max(betaL*L[i] + (1 - betaL)*ell, ell)) + q.mul_(1 - lr/L[i] * (term1 - term2)) + else: # matrix Q + ell = norm_lower_bound_spd(term1 + term2) + L[i].copy_(torch.max(betaL*L[i] + (1 - betaL)*ell, ell)) + q.sub_(lr/L[i] * (term1 - term2) @ q) + procrustes_step2(q) + + if torch.rand([]) < 0.01: # balance factors of Q + balance_kron_precond(Q) + + +def update_precond_kron_newton_pro4p(QL, exprs, V, Hvp, lr=0.1, betaL=0.9, damping=1e-9): + """ + Update the Kron Newton-type preconditioner P as dP = P^0.5 * E * P with a pair of vector and hvp, (V, Hvp). + It is almost the same as update_precond_kron_newton_q0p5eq1p5. But we fit P directly (Q actually is P here). + """ + Q, L = QL + exprA, exprGs = exprs + damping = damping + torch.finfo(Hvp.dtype).eps * Hvp.abs() + Ph = exprA(*Q, Hvp + damping*torch.randn_like(Hvp)) # Q actually is P; so only need to apply its factors once. + + for i, q in enumerate(Q): + term1 = exprGs[i](Ph, Ph.conj()) + term2 = exprGs[i](V, V.conj()) + if q.dim() < 2: # diagonal or scalar Q + ell = torch.max(torch.real(term1 + term2)) + L[i].copy_(torch.max(betaL*L[i] + (1 - betaL)*ell, ell)) + q.mul_(1 - lr/L[i] * (term1 - term2)) + else: # matrix Q + ell = norm_lower_bound_spd(term1 + term2) + L[i].copy_(torch.max(betaL*L[i] + (1 - betaL)*ell, ell)) + q.sub_(lr/L[i] * (term1 - term2) @ q) + for _ in range(10): + procrustes_step3(q) + if (q.H - q).abs().amax() < 0.001 * q.abs().amax(): + break + + if torch.rand([]) < 0.01: # balance factors of Q + balance_kron_precond(Q) + + +def update_precond_kron_newton_quad(QL, exprs, V, Hvp, lr=0.1, betaL=0.9, damping=1e-9): + """ + Update the Kron Newton-type preconditioner Q with a quadratic form for dQ and pair of vector and hvp, (V, Hvp). + """ + Q, L = QL + exprP, exprGs = exprs + damping = damping + torch.finfo(Hvp.dtype).eps * Hvp.abs() + Ph = exprP(*[q.conj() for q in Q], *Q, Hvp + damping*torch.randn_like(Hvp)) + + for i, q in enumerate(Q): + term1 = exprGs[i](Ph, Ph.conj()) + term2 = exprGs[i](V, V.conj()) + if q.dim() < 2: # diagonal or scalar Q + ell = torch.max(torch.real(term1 + term2)) + L[i].copy_(torch.max(betaL*L[i] + (1 - betaL)*ell, ell)) + gain = 1 - lr/2/L[i] * (term1 - term2) + q.mul_(gain * gain) + else: # matrix Q + ell = norm_lower_bound_spd(term1 + term2) + L[i].copy_(torch.max(betaL*L[i] + (1 - betaL)*ell, ell)) + err = lr/2/L[i] * (term1 - term2) + p = q - err @ q # p = q - lr/L[i]/2 * (term1 - term2) @ q + p = p - p @ err # p = p - lr/L[i]/2 * p @ (term1 - term2) + q.copy_((p + p.H)/2) # p must be symmetric or hermitian + + if torch.rand([]) < 0.01: # balance factors of Q + balance_kron_precond(Q) + + +def update_precond_kron_newton_quad4p(QL, exprs, V, Hvp, lr=0.1, betaL=0.9, damping=1e-9): + """ + Almost the same as function update_precond_kron_newton_quad except that we fit P directly (Q here actually is P). + It's vulnerable to numerical errors as the round-off errors could break the SPD property of P. + """ + Q, L = QL + exprA, exprGs = exprs + damping = damping + torch.finfo(Hvp.dtype).eps * Hvp.abs() + Ph = exprA(*Q, Hvp + damping*torch.randn_like(Hvp)) # Q actually is P; so only need to apply its factors once. + + for i, q in enumerate(Q): + term1 = exprGs[i](Ph, Ph.conj()) + term2 = exprGs[i](V, V.conj()) + if q.dim() < 2: # diagonal or scalar Q + ell = torch.max(torch.real(term1 + term2)) + L[i].copy_(torch.max(betaL*L[i] + (1 - betaL)*ell, ell)) + gain = 1 - lr/L[i] * (term1 - term2) + q.mul_(gain * gain) + else: # matrix Q + ell = norm_lower_bound_spd(term1 + term2) + L[i].copy_(torch.max(betaL*L[i] + (1 - betaL)*ell, ell)) + err = lr/L[i] * (term1 - term2) + p = q - err @ q # p = q - lr/L[i] * (term1 - term2) @ q + p = p - p @ err # p = p - lr/L[i] * p @ (term1 - term2) + q.copy_((p + p.H)/2) # p must be symmetric or hermitian + + if torch.rand([]) < 0.01: # balance factors of Q + balance_kron_precond(Q) + + +class KronNewton: + """ + Implements the Kronecker product Newton-type preconditioner as a class. + Most of the time, the hyperparameter name says it all. Here are some comments on a few key parameters. + + 1, preconditioner_max_size and preconditioner_max_skew. These two together control the complexity of the preconditioners. + For example, we are to precondition a 2D gradient with shape 10 x 50. + With preconditioner_max_size 20, we use a dense preconditioner for the first dim since 10 <= 20 and diagonal preconditioner for the second dim since 50 > 20. + With preconditioner_max_skew 1.5, we use a dense preconditioner for the first dim since 10/50 <= 1.5 and diagonal preconditioner for the second dim since 50/10 > 1.5. + + 2, grad_clip_max_norm, betaL and damping. These three together help to stabilize the training. + The grad_clip_max_norm is used to clip the preconditioned gradient to stabilize the optimization as in the classic trust region method. + Setting damping is used to damp and upper bound the fitted preconditioner such that P < eye/damping. + For extremely sparse Hess-vector-prod, a large betaL (say 0.999) helps a lot, where betaL is the EMA factor for the L-smoothness constant (wrt Q) estimation. + + 3, exact_hessian_vector_product. + By setting this flag to False, the finite difference method will be used for Hvp approximation. + Be cautious with the finite difference method (possible numerical issues; the closure must behave like a stateless function). + + 4, Lastly, dQ is for the selection of geometry for preconditioner update. + The two recommended choices are dQ = Q0.5EQ1.5 and dP = P0.5EP (online Newton-Schulz iterations). + Both lr_params and lr_preconditioner are normalized learning rates. + Q is initialized to preconditioner_init_scale * eye. + Always good to check https://arxiv.org/abs/2402.11858 for math details. + """ + def __init__(self, params_with_grad, preconditioner_max_size=float("inf"), preconditioner_max_skew=1.0, preconditioner_init_scale:float|None=None, + lr_params=0.01, lr_preconditioner=0.1, betaL=0.9, damping=1e-9, momentum=0.0, + grad_clip_max_norm=float("inf"), preconditioner_update_probability=1.0, + exact_hessian_vector_product=True, dQ="Q0.5EQ1.5"): + # mutable members + self.lr_params = lr_params + self.lr_preconditioner = lr_preconditioner + self.betaL = betaL # beta for Lipschitz smoothness constant estimation; set to a large value for sparse Hvp + self.damping = damping # used to damp and upper bound P as P < eye/damping + self.momentum = momentum if (0<momentum<1) else 0.0 + self.grad_clip_max_norm = grad_clip_max_norm + self.preconditioner_update_probability = preconditioner_update_probability + # protected members + self._preconditioner_max_size = preconditioner_max_size + self._preconditioner_max_skew = preconditioner_max_skew + params_with_grad = [params_with_grad,] if isinstance(params_with_grad, torch.Tensor) else params_with_grad + self._params_with_grad = [param for param in params_with_grad if param.requires_grad] # double check requires_grad flag + eps = max([torch.finfo(p.dtype).eps for p in self._params_with_grad]) + self._delta_param_scale = eps ** 0.5 + if preconditioner_init_scale is None: + self._QLs_exprs = None # initialize on the fly + print("FYI: Will set the preconditioner initial scale on the fly. Recommend to set it manually.") + else: + self._QLs_exprs = [init_kron(p.squeeze(), preconditioner_init_scale, preconditioner_max_size, preconditioner_max_skew, dQ) for p in self._params_with_grad] + self._ms, self._counter_m = None, 0 # momentum buffers and counter + self._exact_hessian_vector_product = exact_hessian_vector_product + if not exact_hessian_vector_product: + print("FYI: Approximate Hvp with finite-difference method. Make sure that: 1) the closure behaves like a stateless function; 2) delta param scale is proper.") + self._dQ = dQ + if dQ in {"QUAD4P", "PRO4P"}: # the only two cases that fits P directly and dQ actually is dP + if eps > 1e-6: + print("Fitting P directly with half precision is risky.") + if dQ == "QUAD4P": + self._update_precond = update_precond_kron_newton_quad4p + else: # dP = P^0.5 * E * P + self._update_precond = update_precond_kron_newton_pro4p + self._precond_grad = lambda QL, exprs, G: exprs[0](*QL[0], G) # it's exprA(*Q, G) + else: + self._precond_grad = precond_grad_kron + if dQ == "QUAD": + self._update_precond = update_precond_kron_newton_quad + elif dQ == "QEP": + self._update_precond = update_precond_kron_newton_qep + elif dQ == "EQ": + self._update_precond = update_precond_kron_newton_eq + elif dQ == "QEQ": + self._update_precond = update_precond_kron_newton_qeq + else: + assert dQ in {"Q0.5EQ1.5", "Q0p5EQ1p5"}, "Invalid choice for dQ" + self._update_precond = update_precond_kron_newton_q0p5eq1p5 + + + @torch.no_grad() + def step(self, closure): + """ + Performs one step of PSGD with the Kronecker product Newton-type preconditioner. + """ + if (torch.rand([]) < self.preconditioner_update_probability) or (self._QLs_exprs is None): + # evaluates gradients, Hessian-vector product, and updates the preconditioner + if self._exact_hessian_vector_product: + with torch.enable_grad(): + closure_returns = closure() + loss = closure_returns if isinstance(closure_returns, torch.Tensor) else closure_returns[0] + grads = torch.autograd.grad(loss, self._params_with_grad, create_graph=True) + vs = [torch.randn_like(p) for p in self._params_with_grad] + Hvs = torch.autograd.grad(grads, self._params_with_grad, vs) # this line also works for complex matrices + else: # approximate the Hessian-vector product via finite-difference formulae. Use it with cautions. + with torch.enable_grad(): + closure_returns = closure() + loss = closure_returns if isinstance(closure_returns, torch.Tensor) else closure_returns[0] + grads = torch.autograd.grad(loss, self._params_with_grad) + + vs = [torch.randn_like(p) for p in self._params_with_grad] + for (p, v) in zip(self._params_with_grad, vs): # add perturbation + p.add_(v, alpha=self._delta_param_scale) + with torch.enable_grad(): + perturbed_returns = closure() + perturbed_loss = perturbed_returns if isinstance(perturbed_returns, torch.Tensor) else perturbed_returns[0] + perturbed_grads = torch.autograd.grad(perturbed_loss, self._params_with_grad) + Hvs = [(perturbed_g - g)/self._delta_param_scale for (perturbed_g, g) in zip(perturbed_grads, grads)] + for (p, v) in zip(self._params_with_grad, vs): # remove the perturbation + p.sub_(v, alpha=self._delta_param_scale) + + if self._QLs_exprs is None: # initialize QLs on the fly if it is None + scale = (sum([torch.sum(torch.abs(v)**2) for v in vs])/sum([v.numel() for v in vs])) ** (1/4) # (mean(|v|^2))^(1/4) + scale = scale * (max([torch.mean((torch.abs(h))**4) for h in Hvs]) + self.damping**4) ** (-1/8) # (mean(|v|^2))^(1/4) * (mean(|h|^4))^(-1/8) + self._QLs_exprs = [init_kron(h.squeeze(), scale, self._preconditioner_max_size, self._preconditioner_max_skew, self._dQ) for h in Hvs] + # update preconditioner + for (QL_exprs, v, h) in zip(self._QLs_exprs, vs, Hvs): + self._update_precond(*QL_exprs, v.squeeze(), h.squeeze(), lr=self.lr_preconditioner, betaL=self.betaL, damping=self.damping) + else: # only evaluate the gradients + with torch.enable_grad(): + closure_returns = closure() + loss = closure_returns if isinstance(closure_returns, torch.Tensor) else closure_returns[0] + grads = torch.autograd.grad(loss, self._params_with_grad) + + grads = [g.squeeze() for g in grads] + if self.momentum > 0: # precondition the momentum + beta = min(self._counter_m/(1 + self._counter_m), self.momentum) + self._counter_m += 1 + if self._ms is None: + self._ms = [torch.zeros_like(g) for g in grads] + + for (m, g) in zip(self._ms, grads): + m.mul_(beta).add_(g, alpha=1 - beta) + pre_grads = [self._precond_grad(*QL_exprs, m) for (QL_exprs, m) in zip(self._QLs_exprs, self._ms)] + else: # precondition the gradient + self._ms, self._counter_m = None, 0 # clear the buffer and counter when momentum is set to zero + pre_grads = [self._precond_grad(*QL_exprs, g) for (QL_exprs, g) in zip(self._QLs_exprs, grads)] + + lr = self.lr_params + if self.grad_clip_max_norm < float("inf"): + grad_norm = torch.sqrt(torch.real(sum([torch.sum(g*g.conj()) for g in pre_grads]))) + if grad_norm > self.grad_clip_max_norm: + lr = lr * self.grad_clip_max_norm / grad_norm + + # Update the parameters + for (param, g) in zip(self._params_with_grad, pre_grads): + param.subtract_(lr*g.view_as(param)) + + # return whatever closure returns + return closure_returns + + +############# End of PSGD Kronecker product preconditioners ############# + + +############# Begin of PSGD LRA (low rank approximation) preconditioners ############# + + +def IpUVtmatvec(U, V, x): + """ + Returns (I + U*V')*x. All variables are either matrices or column vectors. + """ + return x + U.mm(V.t().mm(x)) + + +def update_precond_lra(UVd, Luvd, v, h, lr=0.1, betaL=0.9): + """ + The raw function for updating the LRA preconditioner Q = (I + U*V')*diag(d) with pair (v, h), + where h can be a Hvp associated with v, or a gradient/momentum independent of v. + State variables (U, V, d) and their Lipschitz smoothness constant estimates (Lu, Lv, Ld) are updated inplace. + Damping logic is not implemented here. + Note that U, V, d, v, and h all are either matrices or column vectors. + """ + U, V, d = UVd + Lu, Lv, Ld = Luvd + + # Approximately balancing U and V such that U^T U = V^T V (exact balancing needs three EVDs) + UtU, VtV = U.t() @ U, V.t() @ V + trUtU, trVtV = torch.sum(UtU.diagonal()), torch.sum(VtV.diagonal()) + rho = (trUtU/trVtV) ** (1/4) # will scale U and V as U <-- U/rho and V <-- V*rho + rho2 = rho * rho + E = 0.1 * (UtU/rho2 - VtV*rho2)/(trUtU/rho2 + trVtV*rho2) # errors after scaling U and V + E2 = 0.5 * E @ E # using this E2 term to make (I - E + E^2/2)(I + E + E^2/2) = (I + E^2/2)^2 - E^2 = I + E^4/4 + U.div_(rho) # scale U and V to have ||U||_F = ||V||_F + V.mul_(rho) + U.sub_(U @ (E - E2)) # rotate (as tr(E)=0) U and V to approach U^TU = V^TV + V.add_(V @ (E + E2)) + + Qh = IpUVtmatvec(U, V, d * h) + Ph = d*IpUVtmatvec(V, U, Qh) + + IpVtU = V.t().mm(U) + IpVtU.diagonal().add_(1) # avoid forming matrix I explicitly + invQtv = v/d + LU, pivots = torch.linalg.lu_factor(lift2single(IpVtU)) + invQtv = invQtv - V.mm(torch.linalg.lu_solve(LU, pivots, lift2single(U.t().mm(invQtv)), adjoint=True).to(V.dtype)) + invPv = invQtv - U.mm(torch.linalg.lu_solve(LU, pivots, lift2single(V.t().mm(invQtv))).to(U.dtype)) + invPv = invPv/d + + # update d + Phh, vinvPv = Ph*h, v*invPv + ell = torch.max(torch.abs(Phh)) + torch.max(torch.abs(vinvPv)) + Ld.copy_(torch.max(betaL*Ld + (1 - betaL)*ell, ell)) + d.sub_(lr/Ld*(Phh - vinvPv)*d) # d.mul_(1 - lr/Ld*(Phh - vinvPv)): larger roundoff errors, unstable with bfloat16 and lr<<1 + + a, b = Qh, invQtv + if torch.rand([]) < 0.5: # only update U + atV = a.t().mm(V) + btV = b.t().mm(V) + atVVt = atV.mm(V.t()) + btVVt = btV.mm(V.t()) + ell = (torch.linalg.vector_norm(a)*torch.linalg.vector_norm(atVVt) + + torch.linalg.vector_norm(b)*torch.linalg.vector_norm(btVVt)) + Lu.copy_(torch.max(betaL*Lu + (1 - betaL)*ell, ell)) + U.sub_(lr/Lu * ( a.mm(atV.mm(IpVtU)) - b.mm(btV.mm(IpVtU)) )) + else: # only update V + atU = a.t().mm(U) + btU = b.t().mm(U) + UUta = U.mm(atU.t()) + UUtb = U.mm(btU.t()) + ell = (torch.linalg.vector_norm(a)*torch.linalg.vector_norm(UUta) + + torch.linalg.vector_norm(b)*torch.linalg.vector_norm(UUtb)) + Lv.copy_(torch.max(betaL*Lv + (1 - betaL)*ell, ell)) + V.sub_(lr/Lv * ( (a + V.mm(atU.t())).mm(atU) - (b + V.mm(btU.t())).mm(btU) )) + + +def precond_grad_lra(UVd, g): + """ + Precondition gradient g with Q = (I + U*V')*diag(d). + All variables here are either matrices or column vectors. + """ + U, V, d = UVd + g = IpUVtmatvec(U, V, d * g) + g = d * IpUVtmatvec(V, U, g) + return g + + +def update_precond_lra_whiten(UVd, Luvd, g, lr=0.1, betaL=0.9, damping=1e-9): + """ + Update the LRA whiten preconditioner. + """ + v = torch.randn_like(g) + damping = damping + torch.finfo(g.dtype).eps * g.abs() + update_precond_lra(UVd, Luvd, v, g + damping*v, lr=lr, betaL=betaL) + + +class LRAWhiten: + """ + Implements the PSGD LRA gradient/momentum whitening preconditioner as a class. + Most of the time, the hyperparameter name says it all. Here are some comments on a few key parameters. + + 1, rank_of_approximation. + Preconditioner Q has a diagonal part and a low rank part, whose rank is decided by this setting. + Rank 0 reduces Q to a diagonal preconditioner. + + 2, grad_clip_max_amps, betaL and damping. These three together help to stabilize the training. + PSGD here tries to normalize the gradients to unit amplitude. This can be problematic when gradients approach zeros. + The most effective way is to clip the preconditioned gradients when their average/element-wise amplitudes exceed grad_clip_max_amps[0]/[1], respectively. + Another way is to damp and upper bound the fitted preconditioner as P < eye/damping. + For extremely sparse gradient, increasing betaL (say to 0.999) also helps a lot, where betaL is the EMA factor for the L-smoothness constant (wrt Q) estimation. + + 3, Lastly, Q is initialized to preconditioner_init_scale * eye. + Boolean setting whiten_grad decides to whiten whether the gradient or momentum. + Always good to check https://arxiv.org/abs/2402.11858 for math details. + """ + def __init__(self, params_with_grad, rank_of_approximation:int=10, preconditioner_init_scale:float|None=None, + lr_params=0.001, lr_preconditioner=0.1, betaL=0.9, damping=1e-9, momentum=0.0, grad_clip_max_amps=(2.0, 10.0), + preconditioner_update_probability=1.0, update_preconditioner_first=True, whiten_grad=True): + # mutable members + self.lr_params = lr_params + self.lr_preconditioner = lr_preconditioner + self.betaL = betaL # set to a large betaL for sparse gradients + self.damping = damping # to damp and upper bound P as P < eye/damping + self.momentum = momentum if (0<momentum<1) else 0.0 + self.grad_clip_max_amps = grad_clip_max_amps + self.preconditioner_update_probability = preconditioner_update_probability + self.update_preconditioner_first = update_preconditioner_first # True for biased update; False for unbiased update. + # protected members + params_with_grad = [params_with_grad,] if isinstance(params_with_grad, torch.Tensor) else params_with_grad + self._params_with_grad = [param for param in params_with_grad if param.requires_grad] # double check requires_grad flag + dtype, device = self._params_with_grad[0].dtype, self._params_with_grad[0].device + self._param_sizes = [torch.numel(param) for param in self._params_with_grad] + self._param_cumsizes = torch.cumsum(torch.tensor(self._param_sizes), 0) + num_params = self._param_cumsizes[-1] + assert 0 <= rank_of_approximation < num_params, "Rank r should be in range [0, number of total parameters)" + self._UVd = [] # saves U, V and d + self._UVd.append(torch.randn(num_params, rank_of_approximation, dtype=dtype, device=device)) # U + self._UVd[0] *= 0.1**0.5 / torch.linalg.vector_norm(self._UVd[0]) + self._UVd.append(torch.randn(num_params, rank_of_approximation, dtype=dtype, device=device)) # V + self._UVd[1] *= 0.1**0.5 / torch.linalg.vector_norm(self._UVd[1]) + if preconditioner_init_scale is None: + print("FYI: Will set the preconditioner initial scale on the fly. Recommend to set it manually.") + else: + self._UVd.append(torch.ones(num_params, 1, dtype=dtype, device=device) * preconditioner_init_scale) + self._Luvd = [lift2single(torch.zeros([], dtype=dtype, device=device)) for _ in range(3)] + self._m, self._counter_m = None, 0 # momentum buffer and counter + self._whiten_grad = whiten_grad + if (not whiten_grad): + assert self.momentum > 0, "Cannot whiten momentum if the momentum setting is invalid." + print(f"Recommend reducing the lr_params for gradient whitening by a factor of {((1 + self.momentum)/(1 - self.momentum))**0.5} for this momentum whitening setting.") + + + @torch.no_grad() + def step(self, closure): + """ + Performs one step of the PSGD LRA gradient/momentum whitening optimizer. + """ + with torch.enable_grad(): + closure_returns = closure() + loss = closure_returns if isinstance(closure_returns, torch.Tensor) else closure_returns[0] + grads = torch.autograd.grad(loss, self._params_with_grad) + + # cat grads into a single vector + grad = torch.cat([torch.reshape(g, [-1, 1]) for g in grads]) # column vector + + if len(self._UVd) < 3: # initialize d on the fly + self._UVd.append((torch.mean(grad**4) + self.damping**4)**(-1/8) * torch.ones_like(grad)) + + if self.momentum > 0: + beta = min(self._counter_m/(1 + self._counter_m), self.momentum) + self._counter_m += 1 + if self._m is None: + self._m = torch.zeros_like(grad) + + self._m.mul_(beta).add_(grad, alpha=1 - beta) + else: # clear the momentum buffer and counter when momentum is set to zero + self._m, self._counter_m = None, 0 + + if torch.rand([]) < self.preconditioner_update_probability: + update_preconditioner_first, update_preconditioner_last = self.update_preconditioner_first, not self.update_preconditioner_first + else: + update_preconditioner_first, update_preconditioner_last = False, False + + if update_preconditioner_first: # update preconditioner first + if self._whiten_grad: # whitens gradient + update_precond_lra_whiten(self._UVd, self._Luvd, grad, lr=self.lr_preconditioner, betaL=self.betaL, damping=self.damping) + else: # whitens momentum + update_precond_lra_whiten(self._UVd, self._Luvd, self._m, lr=self.lr_preconditioner, betaL=self.betaL, damping=self.damping) + + if self.momentum > 0: # precondition momentum + pre_grad = precond_grad_lra(self._UVd, self._m) + else: # precondition gradient + pre_grad = precond_grad_lra(self._UVd, grad) + + if update_preconditioner_last: # update preconditioner later + if self._whiten_grad: # whitens gradient + update_precond_lra_whiten(self._UVd, self._Luvd, grad, lr=self.lr_preconditioner, betaL=self.betaL, damping=self.damping) + else: # whitens momentum + update_precond_lra_whiten(self._UVd, self._Luvd, self._m, lr=self.lr_preconditioner, betaL=self.betaL, damping=self.damping) + + max_avg_amp, max_element_amp = self.grad_clip_max_amps + avg_amp = torch.sqrt(torch.mean(pre_grad * pre_grad)) + if avg_amp > max_avg_amp: + pre_grad *= max_avg_amp/avg_amp + pre_grad.clamp_(min=-max_element_amp, max=max_element_amp) + + # update the parameters + for (param, i, j) in zip(self._params_with_grad, self._param_sizes, self._param_cumsizes): + param.subtract_(pre_grad[j - i:j].view_as(param), alpha=self.lr_params) + + # return whatever closure returns + return closure_returns + + +def update_precond_lra_newton(UVd, Luvd, v, h, lr=0.1, betaL=0.9, damping=1e-9): + """ + Update the LRA Newton preconditioner. + """ + damping = damping + torch.finfo(h.dtype).eps * h.abs() + update_precond_lra(UVd, Luvd, v, h + damping*torch.randn_like(h), lr=lr, betaL=betaL) + + +class LRANewton: + """ + Implements the PSGD LRA Newton-type preconditioner as a class. + Most of the time, the hyperparameter name says it all. Here are some comments on a few key parameters. + + 1, rank_of_approximation. + Preconditioner Q has a diagonal part and a low rank part, whose rank is decided by this setting. + Rank 0 reduces Q to a diagonal preconditioner. + + 2, grad_clip_max_norm, betaL and damping. These three together help to stabilize the training. + The grad_clip_max_norm is used to clip the preconditioned gradient to stabilize the optimization as in the classic trust region method. + Setting damping is used to damp and upper bound the preconditioner as P < eye/damping. + For extremely sparse hess-vector-prods, a large betaL (say 0.999) helps a lot, where betaL is the EMA factor for the L-smoothness constant (wrt Q) estimation. + + 3, exact_hessian_vector_product. + By setting this flag to False, the finite difference method will be used for Hvp approximation. + Be cautious with the finite difference method (possible numerical issues; the closure must behave like a stateless function). + + 4, Lastly, Q is initialized to preconditioner_init_scale * eye. + Both lr_params and lr_preconditioner are normalized learning rates. + Always good to check https://arxiv.org/abs/2402.11858 for math details. + """ + def __init__(self, params_with_grad, rank_of_approximation:int=10, preconditioner_init_scale:float|None=None, + lr_params=0.01, lr_preconditioner=0.1, betaL=0.9, damping=1e-9, momentum=0.0, + grad_clip_max_norm=float("inf"), preconditioner_update_probability=1.0, + exact_hessian_vector_product=True): + # mutable members + self.lr_params = lr_params + self.lr_preconditioner = lr_preconditioner + self.betaL = betaL # set to a large betaL for sparse Hvp + self.damping = damping # to damp and upper bound the preconditioner as P < eye/damping + self.momentum = momentum if (0<momentum<1) else 0.0 + self.grad_clip_max_norm = grad_clip_max_norm + self.preconditioner_update_probability = preconditioner_update_probability + # protected members + params_with_grad = [params_with_grad,] if isinstance(params_with_grad, torch.Tensor) else params_with_grad + self._params_with_grad = [param for param in params_with_grad if param.requires_grad] # double check requires_grad flag + dtype, device = self._params_with_grad[0].dtype, self._params_with_grad[0].device + self._delta_param_scale = torch.finfo(dtype).eps**0.5 + self._param_sizes = [torch.numel(param) for param in self._params_with_grad] + self._param_cumsizes = torch.cumsum(torch.tensor(self._param_sizes), 0) + num_params = self._param_cumsizes[-1] + assert 0 <= rank_of_approximation < num_params, "Rank r should be in range [0, number of total parameters)" + self._UVd = [] # saves U, V and d + self._UVd.append(torch.randn(num_params, rank_of_approximation, dtype=dtype, device=device)) # U + self._UVd[0] *= 0.1**0.5 / torch.linalg.vector_norm(self._UVd[0]) + self._UVd.append(torch.randn(num_params, rank_of_approximation, dtype=dtype, device=device)) # V + self._UVd[1] *= 0.1**0.5 / torch.linalg.vector_norm(self._UVd[1]) + if preconditioner_init_scale is None: + print("FYI: Will set the preconditioner initial scale on the fly. Recommend to set it manually.") + else: + self._UVd.append(torch.ones(num_params, 1, dtype=dtype, device=device) * preconditioner_init_scale) + self._Luvd = [lift2single(torch.zeros([], dtype=dtype, device=device)) for _ in range(3)] + self._m, self._counter_m = None, 0 # momentum buffer and counter + self._exact_hessian_vector_product = exact_hessian_vector_product + if not exact_hessian_vector_product: + print("FYI: Approximate Hvp with finite-difference method. Make sure that: 1) the closure behaves like a stateless function; 2) delta param scale is proper.") + + + @torch.no_grad() + def step(self, closure): + """ + Performs one step of the PSGD LRA Newton optimizer. + """ + if (torch.rand([]) < self.preconditioner_update_probability) or (len(self._UVd) < 3): + # evaluates gradients, Hessian-vector product, and updates the preconditioner + if self._exact_hessian_vector_product: + with torch.enable_grad(): + closure_returns = closure() + loss = closure_returns if isinstance(closure_returns, torch.Tensor) else closure_returns[0] + grads = torch.autograd.grad(loss, self._params_with_grad, create_graph=True) + vs = [torch.randn_like(param) for param in self._params_with_grad] + Hvs = torch.autograd.grad(grads, self._params_with_grad, vs) + else: # approximate Hessian-vector product via finite-difference formulae. Use it with cautions. + with torch.enable_grad(): + closure_returns = closure() + loss = closure_returns if isinstance(closure_returns, torch.Tensor) else closure_returns[0] + grads = torch.autograd.grad(loss, self._params_with_grad) + + vs = [torch.randn_like(param) for param in self._params_with_grad] + for (param, v) in zip(self._params_with_grad, vs): + param.add_(v, alpha=self._delta_param_scale) + with torch.enable_grad(): + perturbed_returns = closure() + perturbed_loss = perturbed_returns if isinstance(perturbed_returns, torch.Tensor) else perturbed_returns[0] + perturbed_grads = torch.autograd.grad(perturbed_loss, self._params_with_grad) + Hvs = [(perturbed_g - g)/self._delta_param_scale for (perturbed_g, g) in zip(perturbed_grads, grads)] + for (param, v) in zip(self._params_with_grad, vs): + param.sub_(v, alpha=self._delta_param_scale) + + v = torch.cat([torch.reshape(v, [-1, 1]) for v in vs]) # column vector + h = torch.cat([torch.reshape(h, [-1, 1]) for h in Hvs]) # column vector + if len(self._UVd) < 3: # init d if it's not in the UVd list + self._UVd.append((torch.mean(v*v))**(1/4) * (torch.mean(h**4) + self.damping**4)**(-1/8) * torch.ones_like(v)) + + # update preconditioner + update_precond_lra_newton(self._UVd, self._Luvd, v, h, lr=self.lr_preconditioner, betaL=self.betaL, damping=self.damping) + else: # only evaluates the gradients + with torch.enable_grad(): + closure_returns = closure() + loss = closure_returns if isinstance(closure_returns, torch.Tensor) else closure_returns[0] + grads = torch.autograd.grad(loss, self._params_with_grad) + + # cat grads + grad = torch.cat([torch.reshape(g, [-1, 1]) for g in grads]) # column vector + + if self.momentum > 0: # precondition momentum + beta = min(self._counter_m/(1 + self._counter_m), self.momentum) + self._counter_m += 1 + if self._m is None: + self._m = torch.zeros_like(grad) + + self._m.mul_(beta).add_(grad, alpha=1 - beta) + pre_grad = precond_grad_lra(self._UVd, self._m) + else: # precondition gradient + self._m, self._counter_m = None, 0 # clear the buffer and counter when momentum is set to zero + pre_grad = precond_grad_lra(self._UVd, grad) + + lr = self.lr_params + if self.grad_clip_max_norm < float("inf"): + grad_norm = torch.linalg.vector_norm(pre_grad) + if grad_norm > self.grad_clip_max_norm: + lr = lr * self.grad_clip_max_norm / grad_norm + + # update the parameters + for (param, i, j) in zip(self._params_with_grad, self._param_sizes, self._param_cumsizes): + param.subtract_(lr * pre_grad[j - i:j].view_as(param)) + + # return whatever closure returns + return closure_returns + + +############# End of PSGD LRA preconditioners ############# + + +############# Begin of PSGD dense matrix Newton-type preconditioner ############# + + +def update_precond_dense_eq(Q, L, v, h, lr=0.1, betaL=0.9, damping=1e-9): + """ + Update dense matrix Newton-type preconditioner Q with local coordinate dQ = mathcal{E} * Q. + """ + damping = damping + torch.finfo(h.dtype).eps * h.abs() + a = Q.mm(h + damping*torch.randn_like(h)) + b = torch.linalg.solve_triangular(lift2single(Q.t()), lift2single(v), upper=False).to(v.dtype) + ell = torch.sum(a*a + b*b) + L.copy_(torch.max(betaL*L + (1 - betaL)*ell, ell)) + Q.sub_(lr/L * torch.triu(a.mm(a.t()) - b.mm(b.t())) @ Q) + + +def update_precond_dense_qep(Q, L, v, h, lr=0.1, betaL=0.9, damping=1e-9): + """ + Update dense matrix Newton-type preconditioner Q with local coordinate dQ = Q * mathcal{E} * P. + """ + damping = damping + torch.finfo(h.dtype).eps * h.abs() + a = Q @ (Q.T @ (Q @ (h + damping*torch.randn_like(h)))) + b = Q @ v + ell = torch.sum(a*a + b*b) + L.copy_(torch.max(betaL*L + (1 - betaL)*ell, ell)) + Q.sub_(lr/L * (a @ (a.T @ Q) - b @ (b.T @ Q))) + + +def update_precond_dense_qeq(Q, L, v, h, lr=0.1, betaL=0.9, damping=1e-9): + """ + Update dense matrix Newton-type preconditioner Q with local coordinate dQ = Q * mathcal{E} * Q. + """ + damping = damping + torch.finfo(h.dtype).eps * h.abs() + a = Q.T @ (Q @ (h + damping*torch.randn_like(h))) + ell = torch.sum(a*a + v*v) + L.copy_(torch.max(betaL*L + (1 - betaL)*ell, ell)) + Q.sub_(lr/L * ((Q @ a) @ a.T - (Q @ v) @ v.T)) + + +def update_precond_dense_q0p5eq1p5(Q, L, v, h, lr=0.1, betaL=0.9, damping=1e-9): + """ + Update dense matrix Newton-type preconditioner Q with local coordinate dQ = Q^0.5 * mathcal{E} * Q^1.5. + """ + damping = damping + torch.finfo(h.dtype).eps * h.abs() + a = Q.T @ (Q @ (h + damping*torch.randn_like(h))) + ell = torch.sum(a*a + v*v) + L.copy_(torch.max(betaL*L + (1 - betaL)*ell, ell)) + Q.sub_(lr/L * (a @ (a.T @ Q) - v @ (v.T @ Q))) + procrustes_step2(Q) + + +def update_precond_dense_pro4p(Q, L, v, h, lr=0.1, betaL=0.9, damping=1e-9): + """ + Update dense matrix Newton-type preconditioner P with local coordinate dP = P^0.5 * mathcal{E} * P. + """ + damping = damping + torch.finfo(h.dtype).eps * h.abs() + a = Q @ (h + damping*torch.randn_like(h)) # Q actually is P; so just apply it once. + ell = torch.sum(a*a + v*v) + L.copy_(torch.max(betaL*L + (1 - betaL)*ell, ell)) + Q.sub_(lr/L * (a @ (a.T @ Q) - v @ (v.T @ Q))) + for _ in range(10): + procrustes_step3(Q) + if (Q.T - Q).abs().amax() < 0.001 * Q.abs().amax(): + break + + +def update_precond_dense_quad(Q, L, v, h, lr=0.1, betaL=0.9, damping=1e-9): + """ + Update dense matrix Newton-type preconditioner Q with a quadratic form for dQ. + """ + damping = damping + torch.finfo(h.dtype).eps * h.abs() + a = Q @ (Q @ (h + damping*torch.randn_like(h))) # Q is symmetric here + ell = torch.sum(a*a + v*v) + L.copy_(torch.max(betaL*L + (1 - betaL)*ell, ell)) + p = Q - lr/2/L * (a @ (a.T @ Q) - v @ (v.T @ Q)) + p = p - lr/2/L * ((p @ a) @ a.T - (p @ v) @ v.T) + Q.copy_((p + p.T)/2) + + +def update_precond_dense_quad4p(Q, L, v, h, lr=0.1, betaL=0.9, damping=1e-9): + """ + Almost the same as update_precond_dense_quad. But it fits P directly. + """ + damping = damping + torch.finfo(h.dtype).eps * h.abs() + a = Q @ (h + damping*torch.randn_like(h)) # Q actually is P; so just apply it once. + ell = torch.sum(a*a + v*v) + L.copy_(torch.max(betaL*L + (1 - betaL)*ell, ell)) + p = Q - lr/L * (a @ (a.T @ Q) - v @ (v.T @ Q)) + p = p - lr/L * ((p @ a) @ a.T - (p @ v) @ v.T) + Q.copy_((p + p.T)/2) + + +class DenseNewton: + """ + Implements the PSGD dense matrix Newton-type preconditioner as a class. + Be extra cautious when using the finite difference method for Hvp approximation (the closure must behave like a stateless function). + It's mainly for illustrating how PSGD works due to its simplicity. + It's also a good alternative to the BFGS like quasi-Newton methods as no line search is required. + """ + def __init__(self, params_with_grad, preconditioner_init_scale:float|None=None, + lr_params=0.01, lr_preconditioner=0.1, betaL=0.9, damping=1e-9, momentum=0.0, + grad_clip_max_norm=float("inf"), preconditioner_update_probability=1.0, + exact_hessian_vector_product=True, dQ="Q0.5EQ1.5"): + # mutable members + self.lr_params = lr_params + self.lr_preconditioner = lr_preconditioner + self.betaL = betaL # set to a large betaL for sparse Hvp + self.damping = damping # to damp and upper bound the preconditioner as P < eye/damping + self.momentum = momentum if (0<momentum<1) else 0.0 + self.grad_clip_max_norm = grad_clip_max_norm + self.preconditioner_update_probability = preconditioner_update_probability + # protected members + params_with_grad = [params_with_grad,] if isinstance(params_with_grad, torch.Tensor) else params_with_grad + self._params_with_grad = [param for param in params_with_grad if param.requires_grad] # double check requires_grad flag + dtype, device = self._params_with_grad[0].dtype, self._params_with_grad[0].device + self._delta_param_scale = torch.finfo(dtype).eps ** 0.5 + self._param_sizes = [torch.numel(param) for param in self._params_with_grad] + self._param_cumsizes = torch.cumsum(torch.tensor(self._param_sizes), 0) + num_params = self._param_cumsizes[-1] + if preconditioner_init_scale is None: # initialize Q on the fly + self._Q = None + else: + if dQ in {"QUAD4P", "PRO4P"}: # Q and dQ actually are P and dP, respectively + preconditioner_init_scale *= preconditioner_init_scale + self._Q = torch.eye(num_params, dtype=dtype, device=device) * preconditioner_init_scale + self._L = lift2single(torch.zeros([], dtype=dtype, device=device)) # Lipschitz smoothness constant estimation for the psgd criterion + self._m, self._counter_m = None, 0 # buffer and counter for momentum + self._exact_hessian_vector_product = exact_hessian_vector_product + if not exact_hessian_vector_product: + print("FYI: Approximate Hvp with finite-difference method. Make sure that: 1) the closure behaves like a stateless function; 2) delta param scale is proper.") + self._dQ = dQ + if dQ in {"QUAD4P", "PRO4P"}: # the only two cases that we fit P directly + if torch.finfo(dtype).eps > 1e-6: + print("Fitting P directly with half precision is risky.") + if dQ == "QUAD4P": + self._update_precond = update_precond_dense_quad4p + else: # dP = P^0.5 * E * P + self._update_precond = update_precond_dense_pro4p + self._precond_grad = lambda Q, g: Q @ g + elif dQ == "QUAD": + self._update_precond = update_precond_dense_quad + self._precond_grad = lambda Q, g: Q @ (Q @ g) # Q is symmetric here; so Q^T = Q + else: + self._precond_grad = lambda Q, g: Q.T @ (Q @ g) + if dQ == "QEP": + self._update_precond = update_precond_dense_qep + elif dQ == "EQ": + self._update_precond = update_precond_dense_eq + elif dQ == "QEQ": + self._update_precond = update_precond_dense_qeq + else: + assert dQ in {"Q0p5EQ1p5", "Q0.5EQ1.5"}, "Invalid choice for dQ" + self._update_precond = update_precond_dense_q0p5eq1p5 + + + @torch.no_grad() + def step(self, closure): + """ + Performs one step of PSGD with the dense matrix Newton-type preconditioner. + """ + if (torch.rand([]) < self.preconditioner_update_probability) or (self._Q is None): + # evaluates gradients, Hessian-vector product, and updates the preconditioner + if self._exact_hessian_vector_product: # exact Hessian-vector product + with torch.enable_grad(): + closure_returns = closure() + loss = closure_returns if isinstance(closure_returns, torch.Tensor) else closure_returns[0] + grads = torch.autograd.grad(loss, self._params_with_grad, create_graph=True) + vs = [torch.randn_like(param) for param in self._params_with_grad] + Hvs = torch.autograd.grad(grads, self._params_with_grad, vs) + else: # approximate Hessian-vector product via finite-difference formulae. Use it with cautions. + with torch.enable_grad(): + closure_returns = closure() + loss = closure_returns if isinstance(closure_returns, torch.Tensor) else closure_returns[0] + grads = torch.autograd.grad(loss, self._params_with_grad) + + vs = [torch.randn_like(param) for param in self._params_with_grad] + for (param, v) in zip(self._params_with_grad, vs): + param.add_(v, alpha=self._delta_param_scale) + with torch.enable_grad(): + perturbed_returns = closure() + perturbed_loss = perturbed_returns if isinstance(perturbed_returns, torch.Tensor) else perturbed_returns[0] + perturbed_grads = torch.autograd.grad(perturbed_loss, self._params_with_grad) + Hvs = [(perturbed_g - g)/self._delta_param_scale for (perturbed_g, g) in zip(perturbed_grads, grads)] + for (param, v) in zip(self._params_with_grad, vs): + param.sub_(v, alpha=self._delta_param_scale) + + v = torch.cat([torch.reshape(v, [-1, 1]) for v in vs]) + h = torch.cat([torch.reshape(h, [-1, 1]) for h in Hvs]) + if self._Q is None: # initialize Q on the fly if it is None + scale = (torch.mean(v*v))**(1/4) * (torch.mean(h**4) + self.damping**4)**(-1/8) + if self._dQ in {"QUAD4P", "PRO4P"}: # Q actually is P in this case + scale *= scale + self._Q = torch.eye(len(v), dtype=v.dtype, device=v.device) * scale + + # update preconditioner + self._update_precond(self._Q, self._L, v, h, lr=self.lr_preconditioner, betaL=self.betaL, damping=self.damping) + else: # only evaluates the gradients + with torch.enable_grad(): + closure_returns = closure() + loss = closure_returns if isinstance(closure_returns, torch.Tensor) else closure_returns[0] + grads = torch.autograd.grad(loss, self._params_with_grad) + + # cat grads + grad = torch.cat([torch.reshape(g, [-1, 1]) for g in grads]) + + if self.momentum > 0: # precondition momentum + beta = min(self._counter_m/(1 + self._counter_m), self.momentum) + self._counter_m += 1 + if self._m is None: + self._m = torch.zeros_like(grad) + + self._m.mul_(beta).add_(grad, alpha=1 - beta) + pre_grad = self._precond_grad(self._Q, self._m) + else: + self._m, self._counter_m = None, 0 # clear the buffer and counter when momentum is set to zero + pre_grad = self._precond_grad(self._Q, grad) + + lr = self.lr_params + if self.grad_clip_max_norm < float("inf"): + grad_norm = torch.linalg.vector_norm(pre_grad) + if grad_norm > self.grad_clip_max_norm: + lr = lr * self.grad_clip_max_norm / grad_norm + + # update the parameters + for (param, i, j) in zip(self._params_with_grad, self._param_sizes, self._param_cumsizes): + param.subtract_(lr * pre_grad[j - i:j].view_as(param)) + + # return whatever closure returns + return closure_returns + + +############# End of PSGD dense matrix Newton-type preconditioner ############# + +""" end of psgd """ |