Add MTP support for hybrid models - #2363

Merged

deepakn94 merged 55 commits into

NVIDIA:mainfrom

rkarimimahab:rkarimimahab/mtp

Feb 1, 2026

Commits on Dec 18, 2025

added mtp support

rkarimimahab
authored and
Rabeeh Mahabadi
committed
16e4ed6
addressed comments

rkarimimahab
authored and
Rabeeh Mahabadi
committed
3db9741
fixed the duplicate

rkarimimahab
authored and
Rabeeh Mahabadi
committed
aeecff9
added repeated mtp layers

rkarimimahab
authored and
Rabeeh Mahabadi
committed
9b4e2b5
adds mtp-num-per-layer to fix a bug

rkarimimahab
authored and
Rabeeh Mahabadi
committed
225c2dc
added the assert

rkarimimahab
authored and
Rabeeh Mahabadi
committed
61eb031
added the assert for vp-stage of MTP layers

rkarimimahab
authored and
Rabeeh Mahabadi
committed
e8304aa
added the LBL fixes
Rabeeh Mahabadi
committed
7398d87
added some fixes
Rabeeh Mahabadi
committed
47063da
fixed more issues
Rabeeh Mahabadi
committed
4896a80

Commits on Dec 19, 2025

updated
Rabeeh Mahabadi
committed
d5a8dd2

Commits on Dec 21, 2025

added mtp support

rkarimimahab
authored and
deepakn94
committed
6f5e06a
addressed comments

rkarimimahab
authored and
deepakn94
committed
3b057ca
fixed the duplicate

rkarimimahab
authored and
deepakn94
committed
7649f6e
added repeated mtp layers

rkarimimahab
authored and
deepakn94
committed
07d7e6c
adds mtp-num-per-layer to fix a bug

rkarimimahab
authored and
deepakn94
committed
bd1ac40
added the assert

rkarimimahab
authored and
deepakn94
committed
cbd19f1
added the assert for vp-stage of MTP layers

rkarimimahab
authored and
deepakn94
committed
d1ba228
added the LBL fixes

Rabeeh Mahabadi
authored and
deepakn94
committed
4d0cd4d
added some fixes

Rabeeh Mahabadi
authored and
deepakn94
committed
305ba71
fixed more issues

Rabeeh Mahabadi
authored and
deepakn94
committed
c60715a
updated

Rabeeh Mahabadi
authored and
deepakn94
committed
cb77b0e

Commits on Dec 22, 2025

fixed some lintingand small bugs
Rabeeh Mahabadi
committed
d732548
solved merge conflicts
Rabeeh Mahabadi
committed
c66109d
fixed linting errors
Rabeeh Mahabadi
committed
9d07525

Commits on Jan 13, 2026

Commits on Jan 14, 2026

Remove mamba-model's re-preprocessing
sancha
committed
82fab87
MLPLayer inherits from TransformerLayer
sancha
committed
7035cd8
gpt model to not support layer patterns
sancha
committed
af53f49
More thorough backward compatibility warnings for the args
sancha
committed
3611f2a
Merge branch 'main' into rkarimimahab/mtp
sancha
committed
d4f6a72
Remove addition of cp_group_size to mtp loss
sancha
committed
d81b58a

Commits on Jan 16, 2026

Refactor arg parsing to arguments.py
sancha
committed
a8a8eae

Commits on Jan 20, 2026

Handle mtp_standalone case, TELayerNormColumnParallel bug ..
sancha
committed
590315f

Commits on Jan 21, 2026

Address review comments
sancha
committed
0520329

Commits on Jan 22, 2026

standalone broken because loss not on non-mtp stage
sancha
committed
2bfa108
Merge remote-tracking branch 'gh/main' into rkarimimahab/mtp
sancha
committed
22de055
Switch to using mamba-stack for MTP inner layers
sancha
committed
ae7b12a

Commits on Jan 23, 2026

Merge remote-tracking branch 'gh/main' into rkarimimahab/mtp
sancha
committed
76c2443
Address review comments
sancha
committed
6683247

Commits on Jan 26, 2026

Address review comments
sancha
committed
7bfa95f

Commits on Jan 28, 2026

Merge branch 'gh/main' into rkarimimahab/mtp
sancha
committed
c65cbfa

Commits on Jan 31, 2026

Merge remote-tracking branch 'gh/main' into rkarimimahab/mtp
sancha
committed
d0773c5
Merge branch 'main' into rkarimimahab/mtp
deepakn94
authored
d355d67
Fix linting errors
sancha
committed
ea3c2bd
isort then black
sancha
committed
19a56e5
Fix pylint errors
sancha
committed
e50409b
bigfix: test failures
sancha
committed
906e513

Commits on Feb 1, 2026

bugfix: fixing standalone mtp test failures
sancha
committed
4bfc5ec
bugfix: expected value error strings
sancha
committed
9cc1668