add muon code #168
add muon code #168
Conversation
| if isinstance(g, DTensor): | ||
| g, meta = to_local(g, keep_sharded=False) |
There was a problem hiding this comment.
This will result in every rank doing the same orthogonalization computation?
| meta = None | ||
| if isinstance(g, DTensor): | ||
| g, meta = to_local(g, keep_sharded=False) | ||
| # gives NaNs when done with Dtensor, instead of throwing a typical op not supported error, quite sneaky |
There was a problem hiding this comment.
The code is mostly copy paste from this repo btw : https://github.com/ethansmith2000/fsdp_optimizers
How do you think it should be handled?
There was a problem hiding this comment.
Personally, if the optimizer code is meant to run locally per rank without any collectives and you own the optimizer implementation, then converting DTensors to local torch.Tensors for all of the computation seems fine (and will have slightly lower eager-mode overhead due to avoiding DTensor.__torch_dispatch__).
The main value add of DTensor in that case is providing the sharding info/metadata on the tensor, which could be useful in the state dict for example. For that, you would still want the optimizer states to be saved as DTensors.
There was a problem hiding this comment.
I see makes sense. Thanks !
No description provided.