add --overlap-param-gather support for layer-wise optimizer. lots of unit tests. - #3524

Merged

deepakn94 merged 34 commits into

NVIDIA:mainfrom

mchrzanowski:overlap-param-gather-muon-rebased

Mar 5, 2026

Commits on Feb 18, 2026

add --overlap-param-gather support for layer-wise optimizer (muon)

mchrzanowski
and
claude
committed
e317e8d
Fix NaN in overlap-param-gather for layer-wise optimizer (Muon)

mchrzanowski
and
claude
committed
d8189e4
Add unit tests for overlap-param-gather in layer-wise optimizer

mchrzanowski
and
claude
committed
2df43c9
Remove use_layer_wise_optimizer from DDP config

mchrzanowski
and
claude
committed
fa0ceb9
Add comments explaining overlap_param_gather replacing use_layer_wise_optimizer

mchrzanowski
and
claude
committed
bbed683

Commits on Feb 20, 2026

Merge branch 'main' into overlap-param-gather-muon-rebased
mchrzanowski
authored
97fba8b
Run autoformat (black, isort) on changed files

mchrzanowski
and
claude
committed
46bc317

Commits on Feb 21, 2026

Merge branch 'main' into overlap-param-gather-muon-rebased
mchrzanowski
authored
7027091

Commits on Feb 22, 2026

Commits on Feb 23, 2026

Autoformat: black formatting fixes

mchrzanowski
and
claude
committed
211a0ca
Remove assertions blocking overlap_grad_reduce and overlap_param_gather with muon optimizer

mchrzanowski
and
claude
committed
c7ee958
Fix dtype mismatch in layer-wise param gather broadcasts causing NCCL deadlock

mchrzanowski
and
claude
committed
cbed167
Fix timing-dependent NCCL deadlock in layer-wise param gather by waiting only on last broadcast handle

mchrzanowski
and
claude
committed
7726a35
Merge branch 'main' into overlap-param-gather-muon-rebased
mchrzanowski
authored
1b0f0b6
Add missing regression tests for layer-wise optimizer overlap param gather

mchrzanowski
and
claude
committed
5253153
Autoformat: black formatting fixes

mchrzanowski
and
claude
committed
189428b
Switch layer-wise optimizer tests from adam to dist_muon

mchrzanowski
and
claude
committed
d07bb57
Add back overlap assertions for muon (but not dist_muon)

mchrzanowski
and
claude
committed
3a0afd1
Remove redundant assert in finish_param_sync

mchrzanowski
and
claude
committed
97b9916

Commits on Feb 25, 2026

Commits on Feb 26, 2026

Merge branch 'main' into overlap-param-gather-muon-rebased
mchrzanowski
authored
87df1dd

Commits on Feb 28, 2026

Skip --overlap-grad-reduce requirement for layer-wise optimizers

mchrzanowski
and
claude
committed
a566a7b

Commits on Mar 2, 2026

Rename lw -> layerwise, improve docstring, clear handles in wait()

mchrzanowski
and
claude
committed
f30914e

Commits on Mar 3, 2026

Fix tests to use renamed layerwise_ attribute prefix

mchrzanowski
and
claude
committed
23c4516

Commits on Mar 4, 2026

Re-enable --overlap-grad-reduce assertion for --overlap-param-gather

mchrzanowski
and
claude
committed
d77abd5
Merge branch 'main' into overlap-param-gather-muon-rebased
deepakn94
authored
079e2a1
Run autoformat.sh

mchrzanowski
and
claude
committed
b22ecd3
Merge branch 'main' into overlap-param-gather-muon-rebased
mchrzanowski
authored
64dd7ea
Add missing docstring to fix lint error

mchrzanowski
and
claude
committed
62b5ae2