Repository navigation
Embeddings / output parameter sharding with LayerWiseOptimizer (dist_muon) #2163
Description
Is your feature request related to a problem? Please describe.
The current implementation of LayerWiseOptimizer shards by params after sorting them by numel. In typical settings where the embedding and output layers are very significantly larger than any other params, it means that the last two ranks will be respectively storing the entirety of the embedding/output layers, on top of their previous assigned layers. In many cases, especially with high PP, the memory overhead will be so large than it will negate most of the benefit of distributing the optimizer states, because two ranks will have to store a large fraction of the optimizer states for their PP group.
Describe the solution you'd like
The ideal solution would be to have Muon params sharded by params as is currently the case, but keep Adam parameters distributed "the old way", e.g., dim-0 sharding. Though likely optimal, I understand that the current structure of Megatron optimizers make it very difficult. Nonetheless, I think it is primordial to find a way to shard the gigantic embedding/output states if we want to truly benefit from optim state distribution, especially since they're not affected by Muon computation and so there is no real reason to keep the matrix complete.
Describe alternatives you've considered
Perhaps a more accessible short term improvement would be to at least process params in decreasing number of numels, and instead of assigning to ranks in order, always assign to the rank with the least total numel currently assigned. At least the ranks storing embeddings and output would not be storing any other param on top of it. However, this would also cause issues, because it would lead to some ranks having a lot of small params, and others having a few large params, and Muon computation cost is not proportional to the number of parameters and would thus be imbalanced. I think there is no 'easy fix' for this and some form of heterogeneous sharding will be necessary at some point.
Activity
FDecaYed commentedon Nov 11, 2025
Both suggestion were considered during initial development:
Our bottom line is:
when user switching from distopt+adamw to dist_muon, they should not see OOM in case that were working for them, and they should not see slowdown on throughput.
beyond this, we will prioritize features over optimizations for now(muonclip, low precision, fsdp etc)
Contribution also welcomed! :)
charlesGE commentedon Nov 18, 2025
Thank you for the detailed response. I understand that the hybrid distribution of optimizer states (layerwise for Muon params and dim0 sharding or similar for Adam params or even just for embeddings/unembeddings) seems very difficult given the current logic in Megatron.
And I also totally get that crucial features like muonclip have a higher priority than low level optimization!
In this case though, I would argue that it might be a little bit more than an optimization, and in particular that it will likely be an issue regarding
For instance for Deepseek-V3, with vocab size ~130k and hdim 7168, the embedding layer has 930M params. With FP32 master weights + FP32 1st moments + FP32 2nd moments, the weight of the optim state (for the embedding layer alone) will be ~11Gb, if I'm not mistaken (and nearly 8Gb with bf16 1st and 2nd moments). For higher vocab size, the overhead can easily be over 20Gb. This is negligible once sharded along DP, but quite massive when stored on a single GPU, like with dist_muon. It scales only linearly with hdim, so even for much smaller models it will remain a sizeable overhead.
In our current settings, the switch would immediately break our configs. Removing second moment for Muon params is not even close to compensating the >20Gb overhead. We can always 'fix it' by scraping 20Gb from the first and last Pipeline stages to make things lighter, e.g. by adapting the PP layout, but obviously there will be a price to pay in model size and/or throughput, and having 20Gb of "dead" memory on a large number of GPU is a bit problematic.
I will try to see on my end if I can find a reasonable trick to shard the embeddings until a more long-term solution can be found.
FDecaYed commentedon Nov 19, 2025
I see. Thanks for the suggestion, I'll discuss with the team and see what are the options
lk137095576 commentedon Dec 3, 2025
@FDecaYed same question，I find rank 1（yes, rank1 not rank0） get highest memory， in my case ,10GB more than others in same node.
Connor-XY commentedon Jun 5, 2026
Hello @charlesGE and @lk137095576, thanks for the detailed memory analysis and the concrete large-vocab examples.
I agree this is more than a small optimization knob: with very large embedding/output matrices, LayerWiseOptimizer/dist_muon can concentrate Adam-state memory on one or two ranks, so a config that fit with distributed Adam can still OOM or lose most of the expected memory benefit after switching to dist_muon. The earlier response captured the initial tradeoff, but the open design question is whether we need a hybrid sharding mode for Adam-style params, or at least a better load-balancing mode for very large embedding/output params.
I am routing this to the optimizer/dist_muon owner so the backlog item has a current maintainer path.