Treat LayerWise bucket_size as a soft minimum and fill would-be padding with real parameters

[deepakn94]

What

LayerWiseDistributedOptimizer assigns whole parameter matrices to dp_size shards, and a bucket costs the tallest shard times dp_size. It closes a bucket as soon as the params it has accumulated cross a size threshold.

When that threshold lands mid-row, the bucket is cut with a partial shard row and the remainder opens a second bucket that pads out again. The split itself manufactures padding that neither bucket would have had.

Change

Treat the threshold as a soft minimum. Once it is met, keep absorbing params for as long as each one fits into padding the bucket already has, and close the bucket only when the next param would make it taller.

Effect

128 equal expert matrices (hidden x moe_ffn = 2688 x 1856), dp_size = 32, bucket_size = ~100 matrices:

	buckets	slots used	overhead
before	100 + 28	160	+25%
after	128	128	0%

Padding inherent to the shape is untouched: 100 matrices over 32 shards still occupy 128 slots either way, because the last row is genuinely partial.

Notes

• Single-bucket buffers (totals below the threshold) are unaffected.
• The single-dominating-param case keeps its behavior; PADDING_FLOOR still governs the soft minimum.
• Membership uses an incremental greedy estimate of _emit_bucket's packing. The final size-sorted packing is never larger than the membership decision assumed. Correction: that is not true. _emit_bucket sorts the chunk before packing it while the estimate walks params in backprop order, so for mixed sizes the two reach different shard loads and absorbing can enlarge the bucket. @kunlunl found a counterexample in review: dp_size=2, bucket_size=448, backprop-order numels [192, 192, 256, 128, 128, 192] emits 1408 elements where closing at the threshold emits 1280. Corrected comment and a regression test are in Balance LayerWise optimizer shards by Newton-Schulz cost, not parameter size #6379. Equal-sized params, the target case, are unaffected because sorting is then a no-op.
• Unit tests added in tests/unit_tests/distributed/test_layer_wise_param_layout.py; both assertions fail against the previous cut-on-threshold logic.

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

[deepakn94]

/ok to test 3d8dd18

[deepakn94]

/ok to test 94846ca

[erhoo82]

LGTM

[kunlunl]

One small perf-related thought, probably not relevant to the target case with equal-sized expert matrices: _place() tracks an online greedy assignment in backprop order, while _emit_bucket() eventually re-sorts the chunk by numel before running LPT. For mixed-size parameters, those two layouts are not always identical, so the simulated padding can occasionally differ from the final packing.
I found a synthetic aligned example with dp_size=2, bucket_size=448, and backprop-order numels [192, 192, 256, 128, 128, 192], where the new layout uses 1408 elements versus 1280 with the previous splitting behavior.
This looks perf-only rather than a correctness concern, and the pattern may not occur in realistic model layouts. It may just be worth softening the “never larger” wording, or adding a small mixed-size test to document the expected behavior. Not blocking from my side.

[svcnvidia-nemo-ci]

🔄 Merge queue validation started!

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

[deepakn94]

@kunlunl thanks, you are right and I reproduced your case exactly: dp_size=2, bucket_size=448, backprop-order numels [192, 192, 256, 128, 128, 192] gives 1408 elements against 1280 for the pre-absorb behavior.

The mechanism is the one you identified. _emit_bucket sorts the chunk before packing it, while the incremental estimate places params in backprop order, so the two reach different shard loads:

• estimate (backprop order): 192, 192, 256, 128, 128 -> shard loads [448, 448], max 448
• actual (sorted): 256, 192, 192, 128, 128 -> shard loads [512, 384], max 512

_absorbs compares against its own max of 448, admits the second 128, and the sorted packing then stacks both 128s onto one shard and grows the bucket. Equal-sized params make sorting a no-op, which is why the target case is unaffected.

Since #5415 is already merged, I have addressed both of your suggestions in #6379, which touches the same comment:

• The "never larger" claim is gone. The comment now explains why the estimate can under-predict and states that the emitted buffer is occasionally larger than closing the bucket early would have produced.
• Your counterexample is now a test, test_mixed_sizes_can_absorb_into_a_larger_bucket, asserting the 1408 figure so the behavior is pinned rather than assumed.

Left as documented behavior rather than fixed, per your read that it is perf-only and unlikely in realistic layouts. Making the estimate exact would mean sorting the chunk on every param, which is not worth it for a case that does not arise with equal-sized expert matrices. Happy to revisit if it shows up in a real model.