Feature Request: Preconditioned Matrix Optimizers

**Is your feature request related to a problem? Please describe.**
Preconditioned SGD is a powerful family of optimizers, including: FOOF, KFAC, new optimizers like Newton-Muon (which can be seen as Muon msign( ) over FOOF). Recently, a number of prominent names in optimization theory have called for PSGD to be scaled as well.

Tag @NVIDIA/mcore-oncall
to get oncall's attention to this issue.

**Describe the solution you'd like**

PSGD usually suffers from the problem of often having to communicate activation statistics: FOOF uses the input feature gram, KFAC uses both input feature gram and the gradient gram matrix. However, we can show that this is avoidable by calculating it with the `w_grad` inside TransformerEngine: `w_grad = dY.T @ X`, `feature_gram = X.T @ X` and `grad_gram = dY.T @ dY`.

In fact, some more obscure optimizers like the simple variant of Local Loss propagation (Locoprop-S) can be expressed as preconditioned gradient descent with these building blocks as well.

It is also straightforward, then, to calculate MuP parameterization for these systems, unlike AdamW, where the weight decay is a problem.

For FSDP, instead of the full preconditioning matrices, we may show that a diagonal/block-diagonal approximation often suffices, and imo can be cleanly implemented after we have FSDP support for Muon (#3179).

I have a reference implementation up at: https://github.com/plugyawn/Megaprop.

Megaprop pins TransformerEngine, Megatron-LM, as well as Emerging-Optimizers. It has been verified at 2/4/8 GPUs so far, but unfortunately I don't have the resources to check further.

@sbhavani not sure who the on-call is, so tagging since we had a small discussion about this.
Same: @mkhona-nvidia.