Skip to content

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
Merged

Balance LayerWise optimizer shards by Newton-Schulz cost, not parameter size#6379
deepakn94 merged 6 commits into
NVIDIA:mainfrom
deepakn94:dnarayanan/balance_layerwise_lpt

Widen the _ns_compute_cost docstring to cover non-2D params

4c62a5c
Select commit
Loading
Failed to load commit list.
DCO / DCO succeeded Aug 11, 2026 in 1s

DCO

All commits are signed off!