PSGD-Kron's helper functions

[mkhona-nvidia]

Building the helper functions for the PSGD Kron optimizer By Xilin-Li.

Includes subspace-iteration based spectral lower bound calculation, an online procrustes solver and functions to apply the kronecker-factored upper triangular factored preconditioner.

[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.

[evanatyourservice]

@mkhona-nvidia This is really awesome work! The functions are looking accurate.

I'll mention these for down the road:

• LR: psgd's default LR should be /5.0 compared to adam's (e.g. if we consider adam's default to be .001, psgd's default is then .0002)
• momentum behavior: in Xilin's code, the default I believe is no momentum and not passing momentum into the preconditioner update. For modern deep learning, be sure to have the default as having momentum, passing it into the preconditioner update, and being what is whitened for the final update. To spell it out the pseudocode flow is:
  1) do momentum ema
  2) Pg = apply_qs(Qs, momentum_with_damping_noise)
  3) do Q update
  4) Pg_out = apply_qs(Qs_new, clean_momentum)
• betas: i've found momentum beta of 0.95 and betaL of 0.95 to work very well but xilin's defaults are 0.9 for both of these so you can probably choose whichever you'd like

[mkhona-nvidia]

/ok to test b066ed1

[mkhona-nvidia]

@mkhona-nvidia This is really awesome work! The functions are looking accurate.

I'll mention these for down the road:

• LR: psgd's default LR should be /5.0 compared to adam's (e.g. if we consider adam's default to be .001, psgd's default is then .0002)
• momentum behavior: in Xilin's code, the default I believe is no momentum and not passing momentum into the preconditioner update. For modern deep learning, be sure to have the default as having momentum, passing it into the preconditioner update, and being what is whitened for the final update. To spell it out the pseudocode flow is:
  1. do momentum ema
  2. Pg = apply_qs(Qs, momentum_with_damping_noise)
  3. do Q update
  4. Pg_out = apply_qs(Qs_new, clean_momentum)
• betas: i've found momentum beta of 0.95 and betaL of 0.95 to work very well but xilin's defaults are 0.9 for both of these so you can probably choose whichever you'd like

Sounds good; Once this PR is merged, I will have another PR that actually uses the building blocks of this PR to make PSGD-Kron-Pro with this flow and defaults

[mkhona-nvidia]

/ok to test 2062d44

[mkhona-nvidia]

/ok to test 0788dc9

[mkhona-nvidia]

/ok to test 2fe3ce3

[mkhona-nvidia]

/ok to test 5006bb9