[RFC] distributed scion/muon

[rakkit]

The work is in progress #1630; there are two general questions from our side

1. Do we wanna bring semi-orthogonal init rather than using the default one? https://arxiv.org/abs/2310.17813 https://arxiv.org/abs/2405.14813
2. Do we wanna to have the log of training dynamics from the optimizer's side?

[rakkit]

The optimizer itself is mostly ready. I am not sure if we want to add a branch to handle the case of EP shard(1), which needs lots of all2all for experts * layers. @tianyu-l

[tianyu-l]

Thanks for the interesting work! It's amazing that you're able to integrate such cutting-edge features with torchtitan.

Like I comment in the PR, I wonder if there's a chance you could isolate the work so that (1) the core part can be merged into pytorch core, and then (2) the integration can happen in torchtitan?

[rakkit]

Hi @tianyu-l thanks. I will fix the optimizer build parts later, make the optimizer init as simple as possible, and no need to set lots of configs. The current code is more or less "hard-coded" to torchtitan's parallel-dims.

e.g. we need to pass parallel_dims to optimizer:

        self.world_mesh = parallel_dims.world_mesh
        self.fsdp_enabled = parallel_dims.fsdp_enabled
        self.expert_enabled = parallel_dims.ep_enabled
        self.dp_replicate_enabled = parallel_dims.dp_replicate_enabled
        self.tp_enabled = parallel_dims.tp_enabled

I am not sure how to deal with this properly outside the torchtitan

[fegin]

There seems to be an ongoing effort to make muon to PyTorch core. If this is the case, we should sync with @janeyx99 to see how to enable muon.

[rakkit]

clarification of the difference between this proposal and Muon.
"Muon" explicitly uses adamw for embedding and head's weights in LLM; one essentially needs two optimizers there.
Modular Norm or Scion allows us to properly handle embedding and heads' parameters in a universal theory, which traces back to spectral condition.

Ideally, we can let this optimizer mainly focus on scion, while leaving room for users who explicitly want to place embeddings, heads, or other parameters into AdamW. From my previous experience, fusing adamw into scion does not bring significant speedup, unless a separate CUDA stream is opened for the fused adamw update — and even then, the benefits are negligible.

[ofivite]

Adding to @rakkit comment, as we show in our upcoming paper (soon on arxiv) there is a great benefit of tracking norms and using the Scion framework for scaling up -- while Muon due to usage of Adam for input/output lacks this principled approach.

At the same time I agree that people still might want to use Muon for one reason or another and it makes sense to not reinvent the wheel and see if we can reuse some components from mainstream Muon implementation for this PR. Yet, I'd still view this PR as a broader contribution, due to the distributed support + modularity of specify how each layer should be trained + nativeness to torchtitan.