Skip to content

Commit fbfa3db

Browse files
ClashLukeClashLuke
committed
ruff
1 parent 46eb010 commit fbfa3db

4 files changed

Lines changed: 106 additions & 43 deletions

File tree

‎benchmarks/bench_klsoap_pinv.py‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -64,7 +64,9 @@ def summarize(rows):
6464
buckets[(r["shape"], r["dtype"], r["k"], r["eps"])].append(r)
6565

6666
cols = [(stat, m) for stat in ("maxinv", "err") for m in METHODS]
67-
header = f"{'shape':<10} {'dtype':<5} {'k':>3} {'eps':>9} {'rank/d':>8} " + " ".join(f"{stat + '_' + m:<13}" for stat, m in cols)
67+
header = f"{'shape':<10} {'dtype':<5} {'k':>3} {'eps':>9} {'rank/d':>8} " + " ".join(
68+
f"{stat + '_' + m:<13}" for stat, m in cols
69+
)
6870
print(f"\n{header}\n{'-' * len(header)}")
6971
for (shape, dt, k, eps), items in sorted(buckets.items()):
7072
rank, d = items[0]["rank"], items[0]["d"]
@@ -80,9 +82,7 @@ def main():
8082

8183
rows = [
8284
run_case(shape, k, eps, dtype, seed)
83-
for shape, k, eps, dtype, seed in product(
84-
SHAPES, WARMUP_K, EPS_VALS, DTYPES, range(args.seeds)
85-
)
85+
for shape, k, eps, dtype, seed in product(SHAPES, WARMUP_K, EPS_VALS, DTYPES, range(args.seeds))
8686
]
8787
summarize(rows)
8888

‎benchmarks/bench_soap_variance_rotation.py‎

Lines changed: 1 addition & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -114,9 +114,7 @@ def main():
114114

115115
rows = [
116116
run_case(shape, theta, kind, dtype, seed)
117-
for shape, theta, kind, dtype, seed in product(
118-
SHAPES, ANGLES, V_KINDS, DTYPES, range(args.seeds)
119-
)
117+
for shape, theta, kind, dtype, seed in product(SHAPES, ANGLES, V_KINDS, DTYPES, range(args.seeds))
120118
]
121119
summarize(rows)
122120

‎heavyball/chainable.py‎

Lines changed: 92 additions & 33 deletions
Original file line numberDiff line numberDiff line change
@@ -490,9 +490,7 @@ def __call__(self, state, group, update, grad, param, *args, **kwargs):
490490
eccs = [getattr(v, "_ecc", None) for v in views]
491491
stacked_corr = None
492492
if eccs[0] is not None:
493-
stacked_corr = (
494-
eccs[0].correction[None] if n == 1 else torch.stack([e.correction for e in eccs], 0)
495-
)
493+
stacked_corr = eccs[0].correction[None] if n == 1 else torch.stack([e.correction for e in eccs], 0)
496494
slab_p._ecc = utils._ULPState(stacked_corr, eccs[0].smax)
497495

498496
bucket_state = states[indices[0]].setdefault(bucket_key, {})
@@ -1185,8 +1183,9 @@ def _init_soap(state, group, update, grad, param):
11851183
)
11861184

11871185

1188-
def _apply_soap_preconditioner(group, update, Q, GG, *exp_avgs, use_kl: bool = False, eps=1e-8,
1189-
exp_avg_sq=None, heavy: bool = False):
1186+
def _apply_soap_preconditioner(
1187+
group, update, Q, GG, *exp_avgs, use_kl: bool = False, eps=1e-8, exp_avg_sq=None, heavy: bool = False
1188+
):
11901189
beta = utils.beta_debias(group["shampoo_beta"], group["step"])
11911190
max_dim, p1d = group["max_precond_dim"], group["precondition_1d"]
11921191
eas = exp_avg_sq or [None] * len(update)
@@ -1208,8 +1207,13 @@ def _apply_soap_preconditioner(group, update, Q, GG, *exp_avgs, use_kl: bool = F
12081207
def scale_by_soap(group, update, grad, param, exp_avg, exp_avg_sq, Q, GG):
12091208
grad_projected = [utils.project(utils.promote(u), q, False) for u, q in zip(update, Q)]
12101209
precond = utils.adam_(
1211-
exp_avg, exp_avg_sq, grad_projected,
1212-
utils.get_beta1(group), utils.get_beta2(group), group["step"] - 1, group["eps"],
1210+
exp_avg,
1211+
exp_avg_sq,
1212+
grad_projected,
1213+
utils.get_beta1(group),
1214+
utils.get_beta2(group),
1215+
group["step"] - 1,
1216+
group["eps"],
12131217
)
12141218
precond = [utils.project(p, q, True) for p, q in zip(precond, Q)]
12151219
_apply_soap_preconditioner(group, update, Q, GG, exp_avg)
@@ -1224,8 +1228,13 @@ def scale_by_soap(group, update, grad, param, exp_avg, exp_avg_sq, Q, GG):
12241228
def scale_by_kl_soap(group, update, grad, param, exp_avg, exp_avg_sq, Q, GG):
12251229
grad_projected = [utils.project(utils.promote(u), q, False) for u, q in zip(update, Q)]
12261230
precond = utils.adam_(
1227-
exp_avg, exp_avg_sq, grad_projected,
1228-
utils.get_beta1(group), utils.get_beta2(group), group["step"] - 1, group["eps"],
1231+
exp_avg,
1232+
exp_avg_sq,
1233+
grad_projected,
1234+
utils.get_beta1(group),
1235+
utils.get_beta2(group),
1236+
group["step"] - 1,
1237+
group["eps"],
12291238
)
12301239
precond = [utils.project(p, q, True) for p, q in zip(precond, Q)]
12311240
_apply_soap_preconditioner(group, update, Q, GG, exp_avg, use_kl=True, eps=group["eps"])
@@ -1259,9 +1268,18 @@ def scale_by_kl_shampoo(group, update, grad, param, exp_avg, Q, GG):
12591268
def scale_by_soap_nadam(group, update, grad, param, exp_avg, exp_avg_sq, mu_product, Q, GG):
12601269
grad_projected = [utils.project(utils.promote(u), q, False) for u, q in zip(update, Q)]
12611270
precond = utils.nadam_(
1262-
grad_projected, exp_avg, exp_avg_sq, mu_product, grad_projected,
1263-
utils.get_beta1(group), utils.get_beta2(group), group["step"] - 1,
1264-
group["momentum_decay"], group["eps"], 0.0, False,
1271+
grad_projected,
1272+
exp_avg,
1273+
exp_avg_sq,
1274+
mu_product,
1275+
grad_projected,
1276+
utils.get_beta1(group),
1277+
utils.get_beta2(group),
1278+
group["step"] - 1,
1279+
group["momentum_decay"],
1280+
group["eps"],
1281+
0.0,
1282+
False,
12651283
)
12661284
precond = [utils.project(p, q, True) for p, q in zip(precond, Q)]
12671285
_apply_soap_preconditioner(group, update, Q, GG, exp_avg)
@@ -1276,8 +1294,12 @@ def scale_by_soap_nadam(group, update, grad, param, exp_avg, exp_avg_sq, mu_prod
12761294
def scale_by_soap_laprop(group, update, grad, param, exp_avg, exp_avg_sq, Q, GG):
12771295
grad_projected = [utils.project(utils.promote(u), q, False) for u, q in zip(update, Q)]
12781296
precond = utils.laprop_(
1279-
exp_avg, exp_avg_sq, grad_projected,
1280-
utils.get_beta1(group), utils.get_beta2(group), group["step"] - 1,
1297+
exp_avg,
1298+
exp_avg_sq,
1299+
grad_projected,
1300+
utils.get_beta1(group),
1301+
utils.get_beta2(group),
1302+
group["step"] - 1,
12811303
)
12821304
precond = [utils.project(p, q, True) for p, q in zip(precond, Q)]
12831305
_apply_soap_preconditioner(group, update, Q, GG, exp_avg)
@@ -1292,9 +1314,16 @@ def scale_by_soap_laprop(group, update, grad, param, exp_avg, exp_avg_sq, Q, GG)
12921314
def scale_by_soap_ademamix(group, update, grad, param, exp_avg_fast, exp_avg_slow, exp_avg_sq, Q, GG):
12931315
grad_projected = [utils.project(utils.promote(u), q, False) for u, q in zip(update, Q)]
12941316
precond = utils.ademamix_(
1295-
exp_avg_fast, exp_avg_slow, exp_avg_sq, grad_projected,
1296-
group["betas"], group["step"] - 1, group["eps"], group["alpha"],
1297-
group.get("beta3_warmup"), group.get("alpha_warmup"),
1317+
exp_avg_fast,
1318+
exp_avg_slow,
1319+
exp_avg_sq,
1320+
grad_projected,
1321+
group["betas"],
1322+
group["step"] - 1,
1323+
group["eps"],
1324+
group["alpha"],
1325+
group.get("beta3_warmup"),
1326+
group.get("alpha_warmup"),
12981327
)
12991328
precond = [utils.project(p, q, True) for p, q in zip(precond, Q)]
13001329
_apply_soap_preconditioner(group, update, Q, GG, exp_avg_slow, exp_avg_fast)
@@ -1309,8 +1338,13 @@ def scale_by_soap_ademamix(group, update, grad, param, exp_avg_fast, exp_avg_slo
13091338
def scale_by_heavy_soap(group, update, grad, param, exp_avg, exp_avg_sq, Q, GG):
13101339
grad_projected = [utils.project(utils.promote(u), q, False) for u, q in zip(update, Q)]
13111340
precond = utils.adam_(
1312-
exp_avg, exp_avg_sq, grad_projected,
1313-
utils.get_beta1(group), utils.get_beta2(group), group["step"] - 1, group["eps"],
1341+
exp_avg,
1342+
exp_avg_sq,
1343+
grad_projected,
1344+
utils.get_beta1(group),
1345+
utils.get_beta2(group),
1346+
group["step"] - 1,
1347+
group["eps"],
13141348
)
13151349
precond = [utils.project(p, q, True) for p, q in zip(precond, Q)]
13161350
_apply_soap_preconditioner(group, update, Q, GG, exp_avg, exp_avg_sq=exp_avg_sq, heavy=True)
@@ -1325,12 +1359,18 @@ def scale_by_heavy_soap(group, update, grad, param, exp_avg, exp_avg_sq, Q, GG):
13251359
def scale_by_heavy_kl_soap(group, update, grad, param, exp_avg, exp_avg_sq, Q, GG):
13261360
grad_projected = [utils.project(utils.promote(u), q, False) for u, q in zip(update, Q)]
13271361
precond = utils.adam_(
1328-
exp_avg, exp_avg_sq, grad_projected,
1329-
utils.get_beta1(group), utils.get_beta2(group), group["step"] - 1, group["eps"],
1362+
exp_avg,
1363+
exp_avg_sq,
1364+
grad_projected,
1365+
utils.get_beta1(group),
1366+
utils.get_beta2(group),
1367+
group["step"] - 1,
1368+
group["eps"],
13301369
)
13311370
precond = [utils.project(p, q, True) for p, q in zip(precond, Q)]
1332-
_apply_soap_preconditioner(group, update, Q, GG, exp_avg, use_kl=True, eps=group["eps"],
1333-
exp_avg_sq=exp_avg_sq, heavy=True)
1371+
_apply_soap_preconditioner(
1372+
group, update, Q, GG, exp_avg, use_kl=True, eps=group["eps"], exp_avg_sq=exp_avg_sq, heavy=True
1373+
)
13341374
return precond
13351375

13361376

@@ -1361,9 +1401,18 @@ def scale_by_heavy_kl_shampoo(group, update, grad, param, exp_avg, Q, GG):
13611401
def scale_by_heavy_soap_nadam(group, update, grad, param, exp_avg, exp_avg_sq, mu_product, Q, GG):
13621402
grad_projected = [utils.project(utils.promote(u), q, False) for u, q in zip(update, Q)]
13631403
precond = utils.nadam_(
1364-
grad_projected, exp_avg, exp_avg_sq, mu_product, grad_projected,
1365-
utils.get_beta1(group), utils.get_beta2(group), group["step"] - 1,
1366-
group["momentum_decay"], group["eps"], 0.0, False,
1404+
grad_projected,
1405+
exp_avg,
1406+
exp_avg_sq,
1407+
mu_product,
1408+
grad_projected,
1409+
utils.get_beta1(group),
1410+
utils.get_beta2(group),
1411+
group["step"] - 1,
1412+
group["momentum_decay"],
1413+
group["eps"],
1414+
0.0,
1415+
False,
13671416
)
13681417
precond = [utils.project(p, q, True) for p, q in zip(precond, Q)]
13691418
_apply_soap_preconditioner(group, update, Q, GG, exp_avg, exp_avg_sq=exp_avg_sq, heavy=True)
@@ -1378,8 +1427,12 @@ def scale_by_heavy_soap_nadam(group, update, grad, param, exp_avg, exp_avg_sq, m
13781427
def scale_by_heavy_soap_laprop(group, update, grad, param, exp_avg, exp_avg_sq, Q, GG):
13791428
grad_projected = [utils.project(utils.promote(u), q, False) for u, q in zip(update, Q)]
13801429
precond = utils.laprop_(
1381-
exp_avg, exp_avg_sq, grad_projected,
1382-
utils.get_beta1(group), utils.get_beta2(group), group["step"] - 1,
1430+
exp_avg,
1431+
exp_avg_sq,
1432+
grad_projected,
1433+
utils.get_beta1(group),
1434+
utils.get_beta2(group),
1435+
group["step"] - 1,
13831436
)
13841437
precond = [utils.project(p, q, True) for p, q in zip(precond, Q)]
13851438
_apply_soap_preconditioner(group, update, Q, GG, exp_avg, exp_avg_sq=exp_avg_sq, heavy=True)
@@ -1394,13 +1447,19 @@ def scale_by_heavy_soap_laprop(group, update, grad, param, exp_avg, exp_avg_sq,
13941447
def scale_by_heavy_soap_ademamix(group, update, grad, param, exp_avg_fast, exp_avg_slow, exp_avg_sq, Q, GG):
13951448
grad_projected = [utils.project(utils.promote(u), q, False) for u, q in zip(update, Q)]
13961449
precond = utils.ademamix_(
1397-
exp_avg_fast, exp_avg_slow, exp_avg_sq, grad_projected,
1398-
group["betas"], group["step"] - 1, group["eps"], group["alpha"],
1399-
group.get("beta3_warmup"), group.get("alpha_warmup"),
1450+
exp_avg_fast,
1451+
exp_avg_slow,
1452+
exp_avg_sq,
1453+
grad_projected,
1454+
group["betas"],
1455+
group["step"] - 1,
1456+
group["eps"],
1457+
group["alpha"],
1458+
group.get("beta3_warmup"),
1459+
group.get("alpha_warmup"),
14001460
)
14011461
precond = [utils.project(p, q, True) for p, q in zip(precond, Q)]
1402-
_apply_soap_preconditioner(group, update, Q, GG, exp_avg_slow, exp_avg_fast,
1403-
exp_avg_sq=exp_avg_sq, heavy=True)
1462+
_apply_soap_preconditioner(group, update, Q, GG, exp_avg_slow, exp_avg_fast, exp_avg_sq=exp_avg_sq, heavy=True)
14041463
return precond
14051464

14061465

‎heavyball/utils.py‎

Lines changed: 9 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -886,8 +886,11 @@ def get_orthogonal_matrix_QR(
886886
copy_stochastic_(r, compiled_einsum(subs, promote(r), *Q_kept, *Qn_kept))
887887

888888
if heavy and exp_avg_sq is not None:
889-
Rsq = [compiled_einsum("...ji,...jk->...ik", promote(qo.data), promote(qn)).square()
890-
for qo, qn in zip(Q, new_qs) if qo is not None]
889+
Rsq = [
890+
compiled_einsum("...ji,...jk->...ik", promote(qo.data), promote(qn)).square()
891+
for qo, qn in zip(Q, new_qs)
892+
if qo is not None
893+
]
891894
sq_terms = ",".join([f"...{i}{o}" for q, i, o in zip(Q, in_str, out_str) if q is not None])
892895
out_sq = "".join([o if o in sq_terms else i for i, o in zip(in_str, out_str)])
893896
subs = f"...{in_str},{sq_terms}->...{out_sq}"
@@ -3162,7 +3165,10 @@ def psgd_pro_update_precond(
31623165
q_ = q_ - (covariance_PP @ q_ - target_energy * q_) / ell_b * precond_lr
31633166

31643167
R = (q_.mT - q_).contiguous()
3165-
R = R / (max_singular_value(R, power_iter=power_iter).unsqueeze(-1).unsqueeze(-1) + torch.finfo(R.dtype).smallest_normal)
3168+
R = R / (
3169+
max_singular_value(R, power_iter=power_iter).unsqueeze(-1).unsqueeze(-1)
3170+
+ torch.finfo(R.dtype).smallest_normal
3171+
)
31663172
RQ = R @ q_
31673173
RRQ = R @ RQ
31683174
c1 = RQ.diagonal(dim1=-2, dim2=-1).sum(dim=-1)

0 commit comments

Comments
 (0)