Route non-Muon params through DistributedOptimizer - #4771

Merged

deepakn94 merged 10 commits into

NVIDIA:mainfrom

deepakn94:dnarayanan/distopt_for_non_muon_params

May 21, 2026

Commits on May 19, 2026

Route non-Muon params through DistributedOptimizer

deepakn94
and
claude
committed
875fc12
Sync only owned bucket groups in LayerWise/DistOpt step

deepakn94
and
claude
committed
dba39e6
Use LPT bin-packing for LayerWise shard-aligned bucket layout

deepakn94
and
claude
committed
9e0df66
Populate dist-optim sharding metadata in layer-wise mode

deepakn94
and
claude
committed
fa85d52
Address review feedback

deepakn94
and
claude
committed
3187acd
Grow LayerWise bucket to amortize per-shard padding at high dp

deepakn94
and
claude
committed
2412967
Assert overlap_param_gather_with_optimizer_step is off in layer-wise mode

deepakn94
and
claude
committed
f21a3ca
Rename layer-wise tag to ``is_managed_by_layer_wise_optimizer``

deepakn94
and
claude
committed
8f622f1

Commits on May 20, 2026