| import torch |
| import torch.nn.functional as F |
| import numpy as np |
|
|
| def _get_grad(g): |
| if isinstance(g, tuple): |
| return g[0] |
| return g |
|
|
| def _grads_dot(ga, gb): |
| dot = 0.0 |
| for a, b in zip(ga, gb): |
| a = _get_grad(a) |
| b = _get_grad(b) |
| if a is not None and b is not None: |
| dot = dot + (a.flatten() @ b.flatten()) |
| return dot |
|
|
| def compute_dpo_loss(beta, delta, sign): |
| return -F.logsigmoid(beta * sign * delta).mean() |
|
|
| def _solve_lambda_k2(b1, b2, H11, H12, H22, s): |
| d = (b2 - b1) + (H11 - 2 * H12 + H22) * 0.5 |
| eps = 1e-10 |
| if abs(d) > eps: |
| lam = 0.5 * (1.0 + (b1 - b2) / d) |
| lam = max(0.0, min(1.0, lam)) |
| else: |
| lam = 0.5 |
| return lam, {"case": "closed_form"} |
|
|
| def _build_update(g1, g2, w1, w2, p1, p2, s_val): |
| g1_0 = _get_grad(g1[0]) |
| g2_0 = _get_grad(g2[0]) |
| dev = g1_0.device if g1_0 is not None else g2_0.device |
| dtype = g1_0.dtype if g1_0 is not None else g2_0.dtype |
|
|
| p1_t = torch.tensor(p1, device=dev, dtype=dtype) |
| p2_t = torch.tensor(p2, device=dev, dtype=dtype) |
|
|
| gp_norm_sq = (p1_t * p1_t) * _grads_dot(g1, g1) \ |
| + (2.0 * p1_t * p2_t) * _grads_dot(g1, g2) \ |
| + (p2_t * p2_t) * _grads_dot(g2, g2) |
| gp_norm = torch.sqrt(torch.clamp(gp_norm_sq, min=0.0)) |
|
|
| if float(gp_norm) > 0.0: |
| scale_p = float(s_val / gp_norm) |
| a1 = w1 + scale_p * p1 |
| a2 = w2 + scale_p * p2 |
| else: |
| a1, a2 = w1, w2 |
|
|
| result = [] |
| for g1_i, g2_i in zip(g1, g2): |
| g1_i = _get_grad(g1_i) |
| g2_i = _get_grad(g2_i) |
| if g1_i is None and g2_i is None: |
| result.append(None) |
| elif g1_i is None: |
| result.append(g2_i.detach() * a2) |
| elif g2_i is None: |
| result.append(g1_i.detach() * a1) |
| else: |
| result.append(g1_i.detach() * a1 + g2_i.detach() * a2) |
| return result |
|
|
| def cagrad_update_k2(g1, g2, w1, w2, c=0.4): |
| H11 = _grads_dot(g1, g1) |
| H12 = _grads_dot(g1, g2) |
| H22 = _grads_dot(g2, g2) |
| g0_norm_sq = (w1 * w1) * H11 + (2.0 * w1 * w2) * H12 + (w2 * w2) * H22 |
| g0_norm = torch.sqrt(torch.clamp(g0_norm_sq, min=0.0)) |
| if float(g0_norm) == 0.0 or c == 0.0: |
| return [(_get_grad(g1_i) * w1 + _get_grad(g2_i) * w2) if _get_grad(g1_i) is not None and _get_grad(g2_i) is not None else (_get_grad(g1_i) * w1 if _get_grad(g1_i) is not None else _get_grad(g2_i) * w2) for g1_i, g2_i in zip(g1, g2)] |
|
|
| b1 = (w1 * H11) + (w2 * H12) |
| b2 = (w1 * H12) + (w2 * H22) |
| s_val = float(g0_norm) * c |
|
|
| lam, _ = _solve_lambda_k2(b1, b2, H11, H12, H22, s_val) |
| p1 = float(lam) |
| p2 = 1.0 - float(lam) |
|
|
| return _build_update(g1, g2, w1, w2, p1, p2, s_val) |
|
|
| def raco_update_k2(g1, g2, w1, w2, c=0.4): |
| H11 = _grads_dot(g1, g1) |
| H12 = _grads_dot(g1, g2) |
| H22 = _grads_dot(g2, g2) |
| g0_norm_sq = (w1 * w1) * H11 + (2.0 * w1 * w2) * H12 + (w2 * w2) * H22 |
| g0_norm = torch.sqrt(torch.clamp(g0_norm_sq, min=0.0)) |
| if float(g0_norm) == 0.0 or c == 0.0: |
| return [(_get_grad(g1_i) * w1 + _get_grad(g2_i) * w2) if _get_grad(g1_i) is not None and _get_grad(g2_i) is not None else (_get_grad(g1_i) * w1 if _get_grad(g1_i) is not None else _get_grad(g2_i) * w2) for g1_i, g2_i in zip(g1, g2)] |
|
|
| b1 = (w1 * H11) + (w2 * H12) |
| b2 = (w1 * H12) + (w2 * H22) |
| s_val = float(g0_norm) * c |
|
|
| lam, _ = _solve_lambda_k2(b1, b2, H11, H12, H22, s_val) |
| p1_raw = float(lam) |
| p2_raw = 1.0 - float(lam) |
|
|
| p1 = min(p1_raw, w1) |
| p2 = min(p2_raw, w2) |
|
|
| return _build_update(g1, g2, w1, w2, p1, p2, s_val) |
|
|
