heavyball 0.21.4__tar.gz → 0.21.5__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- {heavyball-0.21.4 → heavyball-0.21.5}/PKG-INFO +1 -1
- {heavyball-0.21.4 → heavyball-0.21.5}/heavyball/delayed_psgd.py +2 -2
- {heavyball-0.21.4 → heavyball-0.21.5}/heavyball.egg-info/PKG-INFO +1 -1
- {heavyball-0.21.4 → heavyball-0.21.5}/setup.py +1 -1
- {heavyball-0.21.4 → heavyball-0.21.5}/test/test_bf16_params.py +1 -1
- {heavyball-0.21.4 → heavyball-0.21.5}/LICENSE +0 -0
- {heavyball-0.21.4 → heavyball-0.21.5}/README.md +0 -0
- {heavyball-0.21.4 → heavyball-0.21.5}/heavyball/__init__.py +0 -0
- {heavyball-0.21.4 → heavyball-0.21.5}/heavyball/cached_delayed_psgd_kron.py +0 -0
- {heavyball-0.21.4 → heavyball-0.21.5}/heavyball/cached_psgd_kron.py +0 -0
- {heavyball-0.21.4 → heavyball-0.21.5}/heavyball/foreach_adamw.py +0 -0
- {heavyball-0.21.4 → heavyball-0.21.5}/heavyball/foreach_adopt.py +0 -0
- {heavyball-0.21.4 → heavyball-0.21.5}/heavyball/foreach_laprop.py +0 -0
- {heavyball-0.21.4 → heavyball-0.21.5}/heavyball/foreach_sfadamw.py +0 -0
- {heavyball-0.21.4 → heavyball-0.21.5}/heavyball/foreach_soap.py +0 -0
- {heavyball-0.21.4 → heavyball-0.21.5}/heavyball/p_adam.py +0 -0
- {heavyball-0.21.4 → heavyball-0.21.5}/heavyball/palm_foreach_sfadamw.py +0 -0
- {heavyball-0.21.4 → heavyball-0.21.5}/heavyball/palm_foreach_soap.py +0 -0
- {heavyball-0.21.4 → heavyball-0.21.5}/heavyball/precond_schedule_foreach_soap.py +0 -0
- {heavyball-0.21.4 → heavyball-0.21.5}/heavyball/precond_schedule_palm_foreach_soap.py +0 -0
- {heavyball-0.21.4 → heavyball-0.21.5}/heavyball/precond_schedule_sfpsoap.py +0 -0
- {heavyball-0.21.4 → heavyball-0.21.5}/heavyball/psgd_kron.py +0 -0
- {heavyball-0.21.4 → heavyball-0.21.5}/heavyball/pure_psgd.py +0 -0
- {heavyball-0.21.4 → heavyball-0.21.5}/heavyball/schedule_free_palm_foreach_soap.py +0 -0
- {heavyball-0.21.4 → heavyball-0.21.5}/heavyball/utils.py +0 -0
- {heavyball-0.21.4 → heavyball-0.21.5}/heavyball.egg-info/SOURCES.txt +0 -0
- {heavyball-0.21.4 → heavyball-0.21.5}/heavyball.egg-info/dependency_links.txt +0 -0
- {heavyball-0.21.4 → heavyball-0.21.5}/heavyball.egg-info/requires.txt +0 -0
- {heavyball-0.21.4 → heavyball-0.21.5}/heavyball.egg-info/top_level.txt +0 -0
- {heavyball-0.21.4 → heavyball-0.21.5}/setup.cfg +0 -0
- {heavyball-0.21.4 → heavyball-0.21.5}/test/test_bf16_q.py +0 -0
- {heavyball-0.21.4 → heavyball-0.21.5}/test/test_bf16_storage.py +0 -0
- {heavyball-0.21.4 → heavyball-0.21.5}/test/test_closure.py +0 -0
- {heavyball-0.21.4 → heavyball-0.21.5}/test/test_ema.py +0 -0
- {heavyball-0.21.4 → heavyball-0.21.5}/test/test_foreach.py +0 -0
- {heavyball-0.21.4 → heavyball-0.21.5}/test/test_memory.py +0 -0
- {heavyball-0.21.4 → heavyball-0.21.5}/test/test_merge.py +0 -0
- {heavyball-0.21.4 → heavyball-0.21.5}/test/test_no_grad.py +0 -0
- {heavyball-0.21.4 → heavyball-0.21.5}/test/test_psgd.py +0 -0
- {heavyball-0.21.4 → heavyball-0.21.5}/test/test_soap.py +0 -0
- {heavyball-0.21.4 → heavyball-0.21.5}/test/test_stochastic_updates.py +0 -0
@@ -11,8 +11,8 @@ from .utils import update_param_, warmup, psgd_precond_grad, init_Q_exprs, trust
|
|
11
11
|
split_p_and_g_in_group, triu_to_line, line_to_triu, promote
|
12
12
|
|
13
13
|
|
14
|
-
@torch.compile(mode='max-autotune-no-cudagraphs', fullgraph=True, dynamic=
|
15
|
-
def _compilable_psgd_precond_grad_(q, exprs, ea, p, lr,
|
14
|
+
@torch.compile(mode='max-autotune-no-cudagraphs', fullgraph=True, dynamic=False)
|
15
|
+
def _compilable_psgd_precond_grad_(q, exprs, ea, p, lr, weight_decay, clip_fn):
|
16
16
|
new = psgd_precond_grad(q, exprs, ea)
|
17
17
|
update_param_([p], clip_fn([new]), lr, weight_decay)
|
18
18
|
|
@@ -20,7 +20,7 @@ def get_memory():
|
|
20
20
|
|
21
21
|
@pytest.mark.parametrize("opt", heavyball.__all__)
|
22
22
|
@pytest.mark.parametrize("size,depth", [(256, 2)])
|
23
|
-
def test_foreach(opt, size, depth: int, iterations: int =
|
23
|
+
def test_foreach(opt, size, depth: int, iterations: int = 512, outer_iterations: int = 3):
|
24
24
|
set_torch()
|
25
25
|
if 'psgd' not in opt.lower() and 'padam' not in opt.lower():
|
26
26
|
raise pytest.skip('Only PSGD and PaLMPAdam are supported')
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|