PSGD-Kron-Pro(crustes) optimizer implementation

[mkhona-nvidia]

This builds on previous PRs for PSGD's helper functions to make the PSGD-Kron-Pro optimizer

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

[mkhona-nvidia]

/ok to test 780e3d7

[skyw]

Looks good overall. Comments are mostly style, although must fix.

[mkhona-nvidia]

@https://github.com/lixilinx does this look good to you?

[lixilinx]

Thanks, @mkhona-nvidia for the PSGD code. It looks good and well organized to me!

I once verified the correctness of psgd_kron_contractions by comparison with einsum.

In the norm_lower_bound_spd, we will set the default subspace dim to 32 for float32 (based on my test).

[mkhona-nvidia]

Thanks, @mkhona-nvidia for the PSGD code. It looks good and well organized to me!

I once verified the correctness of psgd_kron_contractions by comparison with einsum.

In the norm_lower_bound_spd, we will set the default subspace dim to 32 for float32 (based on my test).

The momentum dampening has also been changed to:

Dampened momentum calculation:

dampened_momentum = exp_avg + (
    damping_noise_scale + torch.finfo(exp_avg.dtype).eps * exp_avg.abs()
) * torch.randn_like(exp_avg)

with a damping_noise_scale default of 1e-9

[mkhona-nvidia]

/ok to test f9f12bd

[mkhona-nvidia]

/ok to test 54220e2