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

Merged

deepakn94 merged 6 commits into

NVIDIA:mainfrom

deepakn94:dnarayanan/balance_layerwise_lpt

Aug 12, 2026

Commits on Aug 11, 2026

Balance LayerWise LPT by Newton-Schulz compute cost, not shard numel

deepakn94
and
claude
committed
3d32382
Reconcile the bucket-cutter estimate with compute-balanced placement

deepakn94
and
claude
committed
f9e5314
Note why the 1-D layout tests cannot detect the placement change

deepakn94
and
claude
committed
d1e7b64
Document that absorbing can enlarge a bucket for mixed-size params

deepakn94
and
claude
committed
e64fa9a
Update the layout docstring and cover cross-bucket load persistence

deepakn94
and
claude
committed
8d204a8
Widen the _ns_compute_cost docstring to cover non-2D params

deepakn94
and
claude
committed
4c62a5c