feat(moe): Support placing MTP layers into standalone stages - #2136

Merged

ericharper merged 22 commits into

NVIDIA:mainfrom

BestJuly:lit/mtp_layer_standalone_main

Jan 14, 2026

Commits on Nov 5, 2025

Squash following code changes (commits):

shifangx
authored and
BestJuly
committed
48792b7
Fix error when saving checkpoints
BestJuly
committed
259cd21
fix hang issue with torch.distributed.batch_isend_irecv

shifangx
authored and
BestJuly
committed
d515cb0
add conditions to communicate p2p shapes, will switch to calculate shapes later
BestJuly
committed
ebee833
make it compatible with 1f1b
BestJuly
committed
813985d
support mtp standalone in 1f1b

Wohox
authored and
BestJuly
committed
cebb164
fix duplicated final layer norm in mtp combine

Wohox
authored and
BestJuly
committed
9b61229
formatting
BestJuly
committed
e650d5a
update mtp offset in 1f1b callcables

Wohox
authored and
BestJuly
committed
e0dae54
fix 1f1b overlap ut & remove final layernorm in postprocess

Wohox
authored and
BestJuly
committed
db1435f
fix ut failure: no mtp_standalone attribute
BestJuly
committed
29e7ea0
add UT, fix pipeline failure and add document
BestJuly
committed
8947471
Update copyright
BestJuly
committed
41e0bee
format
BestJuly
committed
c1022a7

Commits on Nov 14, 2025

remove abundant argument in BlendedMegatronDatasetBuilder
BestJuly
committed
319390b
Merge branch 'main' into lit/mtp_layer_standalone_main
BestJuly
committed
95031fc

Commits on Nov 18, 2025

Merge branch 'main' into lit/mtp_layer_standalone_main
BestJuly
authored
e3116bb

Commits on Nov 25, 2025

Merge branch 'main' into lit/mtp_layer_standalone_main
BestJuly
committed
5d3f862

Commits on Dec 19, 2025

fix missing layernorm

Wohox
authored and
BestJuly
committed
a3fa990

Commits on Jan 6, 2026

Merge branch 'main' into lit/mtp_layer_standalone_main
BestJuly
committed
ae43919
Address review comments
BestJuly
committed
50d34a3

Commits on Jan 7, 2026

Merge branch 'main' into lit/mtp_layer_standalone_main
BestJuly
authored
69bf9bf