M4 + Dist Checkpoint: Replace global parallel state with explicit group parameters - #2053

Merged

ko3n1g merged 124 commits into

NVIDIA:mainfrom

dimapihtar:yuya/m4_dist_ckpt

Dec 16, 2025

Commits on Jul 28, 2025

update get_rng_state to take in comm pgs
yaoyu-33
committed
fd8bdc1
update save and load checkpoint
yaoyu-33
committed
a7477e3
update checkpointing utils for m4
yaoyu-33
committed
1f4a19b
update `make_sharded_tensor_for_checkpoint` and `make_sharded_tensor_for_checkpoint` for m4 entrypoint
yaoyu-33
committed
bbff062
update fallback
yaoyu-33
committed
8fba285
update `make_sharded_tensor_for_checkpoint` function calls
yaoyu-33
committed
9bad6eb
update `make_tp_sharded_tensor_for_checkpoint` for m4
yaoyu-33
committed
4a0f021
chore: Format files
Mcore Bot
committed
794c0fc
add dp_cp_group to metadata of ckpt
yaoyu-33
committed
568da27
lint
yaoyu-33
committed
029a929
Merge remote-tracking branch 'origin/yuya/m4_dist_ckpt' into yuya/m4_dist_ckpt
yaoyu-33
committed
790e56d
lint
yaoyu-33
committed
9c45be8
chore: Format files
Mcore Bot
committed
1ecf614
bug fix
yaoyu-33
committed
bee469a
Merge remote-tracking branch 'origin/yuya/m4_dist_ckpt' into yuya/m4_dist_ckpt
yaoyu-33
committed
b67586c
chore: Format files
Mcore Bot
committed
3a9c8e6

Commits on Jul 29, 2025

fallback fix
yaoyu-33
committed
81d203d
chore: Format files
Mcore Bot
committed
f893946
bug fix
yaoyu-33
committed
85c59f4

Commits on Jul 30, 2025

bug fix
yaoyu-33
committed
7b59837

Commits on Jul 31, 2025

Merge branch 'refs/heads/main' into yuya/m4_dist_ckpt
yaoyu-33
committed
7f9cce8

Commits on Aug 3, 2025

metadata guard
yaoyu-33
committed
ce6ff71
reformats
yaoyu-33
committed
db5409e
chore: Format files
Mcore Bot
committed
d4dbe48
try remove expert_dist_ckpt_decorator
yaoyu-33
committed
6508c72
lint
yaoyu-33
committed
aefc4b6
Merge remote-tracking branch 'origin/yuya/m4_dist_ckpt' into yuya/m4_dist_ckpt
yaoyu-33
committed
87c2b13

Commits on Aug 4, 2025

condition fix
yaoyu-33
committed
b197a28

Commits on Aug 5, 2025

Merge branch 'main' into yuya/m4_dist_ckpt
yaoyu-33
committed
3bc6dea

Commits on Aug 6, 2025

lint
yaoyu-33
committed
7b1b555

Commits on Aug 7, 2025

Merge branch 'main' into yuya/m4_dist_ckpt
yaoyu-33
committed
9dcff5e
temp update test
yaoyu-33
committed
86f4292
try bug fix
yaoyu-33
committed
4d9a89e
revert and some fix
yaoyu-33
committed
02248d9
bug fix
yaoyu-33
committed
7dfb88a
lint
yaoyu-33
committed
787c800
Revert "temp update test"
yaoyu-33
committed
c4157b3
bug fix for missing default values
yaoyu-33
committed
4f100de
remove not used
yaoyu-33
committed
5b42ca3

Commits on Aug 13, 2025

Merge branch 'main' into yuya/m4_dist_ckpt
yaoyu-33
committed
4395f4d
chore: Format files
Mcore Bot
committed
9bcf4cb
bug fix for `sharded_state_dict_default`
yaoyu-33
committed
ecbbbc7
Merge remote-tracking branch 'origin/yuya/m4_dist_ckpt' into yuya/m4_dist_ckpt
yaoyu-33
committed
d0cc504

Commits on Sep 25, 2025

Merge branch 'refs/heads/main' into yuya/m4_dist_ckpt
yaoyu-33
committed
0a8446d
chore: Format files
Mcore Bot
committed
87e231e
main branch change
yaoyu-33
committed
bdeb0ed
Merge branch 'refs/heads/main' into yuya/m4_dist_ckpt
yaoyu-33
committed
c9ff988
interface fixes for checkpoint with m4
yaoyu-33
committed
9c75cb5
chore: Format files
Mcore Bot
committed
c9407dc
revert transformer_engine.py changes
yaoyu-33
committed
48be193
bug fixes
yaoyu-33
committed
3c30a13

Commits on Sep 26, 2025

lint
yaoyu-33
committed
5a671e9

Commits on Sep 29, 2025

update to use logger to pass lint
yaoyu-33
committed
6360922
chore: Format files
Mcore Bot
committed
39dcac3

Commits on Sep 30, 2025

Guard for cases metadata is not provided
yaoyu-33
committed
1b3cf6e
chore: Format files
Mcore Bot
committed
d2258c0
Merge branch 'main' into yuya/m4_dist_ckpt
yaoyu-33
committed
ad61ac8
update to ensure_metadata_has_dp_cp_group util
yaoyu-33
committed
1da7644
Add debug assert
yaoyu-33
committed
321a68b
chore: Format files
Mcore Bot
committed
6021b37
lint
yaoyu-33
committed
aed7622
Merge remote-tracking branch 'origin/yuya/m4_dist_ckpt' into yuya/m4_dist_ckpt
yaoyu-33
committed
dfc7978

Commits on Oct 30, 2025

Merge branch 'main' into yuya/m4_dist_ckpt
dimapihtar
authored
e71a0b0
minor fixes
dimapihtar
committed
10a137f

Commits on Nov 3, 2025

fix te module issue
dimapihtar
committed
386f840
fix te module issue
dimapihtar
committed
ab71e71
add tp_group param
dimapihtar
committed
7438282
minor fix
dimapihtar
committed
d02bbf1
clear metadata
dimapihtar
committed
070fa5a

Commits on Nov 5, 2025

fix tp_group
dimapihtar
committed
ab24fa4
remove tp_group extra usage
dimapihtar
committed
a961335
add tp_group to state_dict
dimapihtar
committed
406dff3
add tp_group to state_dict
dimapihtar
committed
ca4d1bf
add tp_group to state_dict
dimapihtar
committed
d8845f8
add tp_group to state_dict
dimapihtar
committed
df58591
fix tp_group
dimapihtar
committed
778b01d
fix tp_group usage
dimapihtar
committed
465d803
fix typo
dimapihtar
committed
bcf420f
add tp_group param to state_dict
dimapihtar
committed
8cd778b
add tp_group param to state_dict
dimapihtar
committed
f19cb10
fix tp_group usage
dimapihtar
committed
bd3f9e4
fix metadata
dimapihtar
committed
787ec6a
fix unit tests
dimapihtar
committed
6e9af69
fix style
dimapihtar
committed
146d3c1
fix style
dimapihtar
committed
3a203d8
remove unused imports
dimapihtar
committed
8bf8220
add headers
dimapihtar
committed
ba4e32f
Merge branch 'main' into yuya/m4_dist_ckpt
dimapihtar
authored
b33362d
fix style
dimapihtar
committed
1311d81

Commits on Nov 6, 2025

pass tp_group param
dimapihtar
committed
f4e828a
fix tp_group usage
dimapihtar
committed
f5bfcee
fix tp_group usage
dimapihtar
committed
5e375a2
remove debug assertions
dimapihtar
committed
12cc898
fix style
dimapihtar
committed
9cbb410
pass tp_param to state_dict
dimapihtar
committed
a599cd1
add tp_group to state_dict
dimapihtar
committed
4e11170
pass tp_group param to sharded_state_dict
dimapihtar
committed
088702a
pass tp_group to sharded_state_dict
dimapihtar
committed
bcc30d7
pass tp_group to sharded_state_dict
dimapihtar
committed
d10b644
fix training/checkpointing
yaoyu-33
committed
17cf6c9
Merge remote-tracking branch 'dim/yuya/m4_dist_ckpt' into yuya/m4_dist_ckpt
yaoyu-33
committed
e8a1f84
pass metadata arg
dimapihtar
committed
93e6099
remove tp_group arg from sharded_state_dict
dimapihtar
committed
34c02b0
set self.tp_group
dimapihtar
committed
2900114

Commits on Nov 7, 2025

pass tp_group & pp_group to get_rng_state
dimapihtar
committed
59b9485
minor fix
dimapihtar
committed
415b371
remove extra tp_group arg
dimapihtar
committed
10af8c0
remove extra tp_group arg
dimapihtar
committed
e612b04
fix code style
dimapihtar
committed
3456575
fix code style
dimapihtar
committed
0fc5fa0
set te_group properly
dimapihtar
committed
d59a993

Commits on Nov 25, 2025

Merge branch 'main' into yuya/m4_dist_ckpt
dimapihtar
authored
8344474

Commits on Dec 4, 2025

Commits on Dec 5, 2025

Merge branch 'main' into yuya/m4_dist_ckpt
yaoyu-33
authored
272f126

Commits on Dec 9, 2025

Merge branch 'main' into yuya/m4_dist_ckpt
dimapihtar
authored
c60d4ca
fix code style
dimapihtar
committed
467bcb2

Commits on Dec 16, 2025

Merge branch 'main' into yuya/m4_dist_ckpt
dimapihtar
authored
c12f0bf
fix API backward comp. check
dimapihtar
committed
ee1d6b5
fix uni tests
dimapihtar
committed
49df433
add deprecation decorator
dimapihtar
committed
defaf33
fix code style
dimapihtar
committed
d37fdfb