Balance LayerWise optimizer shards by Newton-Schulz cost, not parameter size

[deepakn94]

Problem

LayerWiseDistributedOptimizer assigns whole Muon matrices to dp_size shards, packing them largest-first onto the least-loaded shard. It measures "loaded" by numel, which breaks under GTP: a GTP-sharded weight stores only its local shard, but Newton-Schulz runs on the full all-gathered matrix, so its real cost is gtp_remat_size times what its numel suggests.

Ordering by numel therefore starts with the cheapest work and leaves the most expensive matrices to fill in wherever there is room, which piles them onto shards that are already busy. Every rank waits for the slowest, so the Muon step is gated on whichever shard collected them.

Changes

1. _ns_compute_cost() estimates Newton-Schulz cost from the full post-all-gather shape (GTP shards along dim 0): max(M, N) * min(M, N)^2, the dominant term in the orthogonalization.
2. Order and assign by that cost instead of numel.
3. Track shard_compute_loads persistently across buckets, so expensive matrices spread over the whole buffer rather than clustering inside each bucket. shard_cursors still resets per bucket, which preserves the memory layout. A per-bucket numel cap (_NUMEL_EPSILON = 0.3) bounds the extra padding this can introduce, and assignment falls back to the previous least-numel rule when no shard satisfies the cap.
4. Correct the bucket-cutter comment from Treat LayerWise bucket_size as a soft minimum and fill would-be padding with real parameters #5415. See Notes.

Results

Measured end-to-end on Ultra-half (54-layer hybrid Mamba-MoE, 1024 GB200 GPUs, Muon + GTP=64): 2080 -> 1707 ms/iter, a 19% throughput improvement.

Tests

tests/unit_tests/distributed/test_layer_wise_param_layout.py gains TestComputeBalancedLayout, built so numel and cost disagree: three GTP-sharded matrices whose Newton-Schulz cost is 16x the largest dense matrix but whose numel is the smallest in the set. Numel-ordered placement puts all three on one shard for a 3.77x compute imbalance; cost-ordered placement spreads them across three shards for 1.33x. The test asserts both the spread and the imbalance bound, so it fails against the previous heuristic.

The pre-existing cases in TestSizeMatchingLayout all use 1-D params, where _ns_compute_cost falls back to nelement() and cost equals numel. They therefore pass identically under both heuristics and cannot detect this change; a class docstring now says so.

Notes

_emit_bucket sorts a chunk before packing it, while the incremental estimate in the bucket cutter walks params in backprop order. #5415's comment claimed the two mirror each other. They do not, and after this PR they also differ in what they measure. Equal-sized params make both differences vanish; mixed sizes can send params to different shards, so absorbing occasionally enlarges a bucket instead of filling it.

@kunlunl raised this in review of #5415 with a counterexample, now pinned as test_mixed_sizes_can_absorb_into_larger_bucket: dp_size=2, backprop-order numels [192, 192, 256, 128, 128, 192], bucket_size=448 emits 1408 elements where closing at the threshold emits 1280. Documented rather than fixed, since an exact estimate would mean re-sorting the chunk on every param, and the target case is equal-sized expert matrices.

This helps without GTP too, though less dramatically. Newton-Schulz cost is numel * min(M, N), so numel misvalues a matrix by its aspect ratio: across this model's shapes that factor ranges from 512 to 10240, meaning a tall-thin matrix and a square one of the same byte count do very different amounts of work. Balancing numel equalises bytes per shard; balancing cost equalises the Newton-Schulz FLOPs that actually gate the step. Without GTP the two orderings largely agree and the gain is small, but it grows with dp_size, where there is less slack to absorb a misranked matrix, and it is never worse.

The shared-embedding branch is untouched.

🤖 Generated with Claude Code

[copy-pr-bot]

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

[claude]

LGTM.

Traced the new packing rule against every pre-existing TestSizeMatchingLayout case: the 30% numel cap does force the empty-candidates fallback in a few of them, but since the fallback reduces to the old least-numel rule the placements and asserted buffer sizes are unchanged. The candidate filter also mirrors the real placement arithmetic (pad_param_start(cursor) + numel), so an admitted shard genuinely stays under the cap.

The claim that equal-sized params make the _absorbs estimate exact holds: with equal costs, persistent shard_compute_loads only rotates which shard starts a bucket, so the multiset of shard_cursors — and hence max(shard_cursors) and the bucket size — is unchanged. The cost/imbalance numbers in the new tests and docstrings check out (268,435,456 vs 16,777,216; 1.33x; 1024 + 384 = 1408 vs 768 + 512 = 1280), and test_gtp_params_land_on_different_shards_in_each_bucket really does depend on the loads persisting, so it would fail if they were reset per bucket.

No new global process-group reads in megatron/core.

Nit: the _ns_compute_cost docstring says "for a 2D parameter", but the first branch is the non-2D nelement() fallback that the 1-D cases in TestSizeMatchingLayout rely on — "for a parameter" would match the body.

(Review by inspection; I did not execute the suite.)

[deepakn94]

/ok to test 4c62a5c

[svcnvidia-nemo-ci]

🔄 Merge queue validation started!

You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/31605644484