feat: colocated Megatron reshard - #3490

Merged

terrykong merged 25 commits into

NVIDIA-NeMo:mainfrom

tdene:tde/refit_colocated

Aug 18, 2026

Commits on Aug 17, 2026

feat: colocated Megatron reshard
tdene
committed
caafe62
colocated-reshard nightly (nanov3)
tdene
committed
e78254c
colocated reshard functional test
tdene
committed
f7d1ef4
Fix issues uncovered by tests
tdene
committed
84412d3
Single source of truth for config merge
tdene
committed
6e650b5
lint
tdene
committed
4a2dfaf
Unify Megatron refit behind a WeightSynchronizer
tdene
committed
fecd2e3
Add offload_optimizer_for_refit config
tdene
committed
75b1c81
Adopt setup_timing_metrics in the synchronizer init block

tdene
and
claude
committed
6d0a720
lint

tdene
and
claude
committed
356d25f
Fix the non-colocated refit test to dispatch via the synchronizer

tdene
and
claude
committed
ea559a4
Fix inference-optimized incompatibility
tdene
committed
9ffe449
Properly fix inference config mismatch
tdene
committed
844e964
Use a QKV-bias-free model in the reshard test
tdene
committed
0cd9e5e
Add refit test that doesn't need TMS
tdene
committed
7cf972a
Minimize the 4n4g reshard recipe config
tdene
committed
6fa85ca
Fix segment-size on nightlies
tdene
committed
c7d9592
Address reviewer comments
tdene
committed
7695ba5
Address review-pr-team comments
tdene
committed
5d7f2c4
Import torch_memory_saver correctly
tdene
committed
4c3e03c
lint rebase conflict
tdene
committed
3ad74f7
Address reviewer comments
tdene
committed
3ca38b7
Address reviewer comment
tdene
committed
6f63120
Fix default value of argument
tdene
committed
4a21615
Fix extra issues
tdene
committed
5f1f294