Scion optimizer

[mkhona-nvidia]

Added Scion.

The main change is the parametrization to allow Franke-Wolfe. We now DO NOT use weight decay and instead use step size (i.e. the learning rate) and spectral radius, with the unit_rms choice for width scaling

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

[valentyn1boreiko]

Amazing, thanks @mkhona-nvidia! Could you please also

• add 5D parallelism support (which includes LayerWiseDistributedOptimizer wrapper and TensorParallel similar to what TensorParallelMuon in this commit is doing);
• support 1D and 2D tensors in Scion with the automatic mapping to respective LMOs similar to how it is done here with class Auto and different norms in norms_dict;
• allow to import Scion from megatron/core/optimizer/init.py as an alternative Megatron optimizer and allow for passing layer-wise different radii to tune depending on parameter groups (separately for the router, other hidden layers, embedding, output layers, and 1D tensors for example)?

[skyw]

Some minor changes needed. otherwise LGTM

[skyw]

/ok to test 5521b58