Implementing efficient orthogonalization with QDWH

[valentyn1boreiko]

Is your feature request related to a problem? Please describe.
QR-based Dynamically Weighted Halley (QDWH) can be used for orthogonalization and has cubic convergence (https://arxiv.org/abs/2505.21799). Would it be possible to add it as well as an alternative to Newton-Schulz?

Describe the solution you'd like
A clear and concise description of what you want to happen.

Describe alternatives you've considered
A clear and concise description of any alternative solutions or features you've considered.

Additional context
Add any other context or screenshots about the feature request here.

[snowmanwwg]

Thanks @valentyn1boreiko - is this something you want to contribute?

[skyw]

I've tested 16 steps Newton-Schulz with polar express coefficient sets. It gets very close to polar (by scipy).
Will need to see whether (or how difficult) a good GPU implementation of QDWH will be faster than ~50 matrix multiplies involved in 16 steps of NS.

[mkhona-nvidia]

QDWH will also let us implement PolarGrad by letting us compute the nuclear norm term.

[skyw]

implemented a version via SVD. It is not too slow, should be enough for evaluation purpose and motivate kernel work if turned out to be promising.

[timlautk]

PolarGrad

It might be a bit late, but as the author of PolarGrad, I would like to mention that we can also compute the nuclear norm term via Newton--Schulz or Polar Express. The nuclear norm term can be computed via torch.sum(grad, orth_grad) regardless of how orth_grad is computed; see the paper for why this is the formula.

QDWH has the benefit of handling ill-conditioned gradients (could be beneficial to the embedding and head layers although people don't really use Muon or PolarGrad for these layers); see the latest version of PolarGrad for detailed explanation on this.

[mkhona-nvidia]

@timlautk do you want to submit an MR to OrthogonalizedOptimizer and/or Muon in the codebase to optionally use polargrad's nuclear norm term? Or you can subclass OrthogonalizedOptimizer, overwrite the orthogonalize function and add in the nuclear norm term through orth_grad@grad so we can reuse newton-schulz

[timlautk]

@timlautk do you want to submit an MR to OrthogonalizedOptimizer and/or Muon in the codebase to optionally use polargrad's nuclear norm term? Or you can subclass OrthogonalizedOptimizer, overwrite the orthogonalize function and add in the nuclear norm term through orth_grad@grad so we can reuse newton-schulz

Yes, I am working on it. I think I will subclass and add a polargrad.py.

[skyw]

We decided to upstream to PyTorch, stay tuned. pytorch/pytorch#175636