[avgvi]

Summary

PSGDPro stores every argument without checking it. warmup_steps=0 is only noticed on the first step(), as a ZeroDivisionError from _get_precond_lr:

scheduled_lr = precond_lr / math.sqrt(1.0 + step / warmup_steps)

Out-of-range values for the other arguments are accepted silently: a negative lr, weight_decay or precond_lr, or a momentum or beta_lip of 1 or more. The scalar optimizers already reject the shared ones in _validate_common_hparams, and Shampoo checks lr (shampoo.py:346).

This PR adds the same kind of checks at the top of PSGDPro.__init__, with the same Invalid ... messages:

Argument	Rejected when
lr	< 0
weight_decay	< 0
momentum	not in [0, 1)
beta_lip	not in [0, 1)
precond_lr	< 0
min_precond_lr	< 0
warmup_steps	<= 0

Test

TestPSGDPro.test_invalid_arguments_raise in tests/test_psgd_pro.py is parameterized over the seven cases, in the same shape as test_invalid_arguments_raise in tests/test_shampoo.py.

On main:

$ python -c "..."   # PSGDPro([p], **kwargs) then .step()
PSGDPro([p], warmup_steps=0).step(): ZeroDivisionError: division by zero
PSGDPro([p], **{'lr': -0.01}): accepted
PSGDPro([p], **{'momentum': 1.0}): accepted
PSGDPro([p], **{'beta_lip': 1.0}): accepted

$ python tests/test_psgd_pro.py --device=cpu
FAIL: test_invalid_arguments_raise0 (kwargs={'lr': -0.01}, message='Invalid learning rate')
...
FAIL: test_invalid_arguments_raise6 (kwargs={'warmup_steps': 0}, message='Invalid warmup_steps')
Ran 8 tests
FAILED (failures=7)

With this change:

$ python tests/test_psgd_pro.py --device=cpu --seed=42
Ran 8 tests
OK

Verification

On CPU, torch 2.14.1, Python 3.12:

• python tests/test_psgd_pro.py --device=cpu with --seed=42 and --seed=3735928559: 8 OK
• tests/test_psgd_utils.py 31 OK, tests/test_procrustes_step.py 28 OK, tests/test_psgd_contractions.py 11 OK
• tests/convergence/psgd_convergence_test.py --device=cpu: 3 OK
• pre-commit run --files emerging_optimizers/psgd/psgd.py tests/test_psgd_pro.py: all hooks pass

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[greptile-apps]

Confidence Score: 5/5

[Low impact] The PR appears safe to merge.

Summary

PSGDPro now rejects invalid constructor arguments before training starts.

• PSGDPro rejects out-of-range settings when callers create it.

No actionable issues found. Tests were not run during this review.

_{Reviews (1) · Last reviewed commit: "Validate PSGDPro hyperparameters at cons..." · Reviewed by Greptile}

[skyw]

Thanks for the contribution.

We should probably unify optimizer argument validation at some point FWIW.

[skyw]

/ok to test e01e212