`megatron.core.datasets` SFTDataset and `get_batch*` functions refactor - #3542

Closed

asolergi-nv wants to merge 127 commits into

NVIDIA:mainfrom

asolergi-nv:clean_thd

Commits on Feb 4, 2026

Reapply "Add MTP support for hybrid models (#2363)"

sancha
authored and
arendu
committed
b19565a
Reapply "Fix two minor bugs in MTP implementation for hybrid models (#3194)"

sancha
authored and
arendu
committed
f1401f1
Handle one more missed reference

sancha
authored and
arendu
committed
86eacb7
identity template and tool role mask

Soumye Singhal
authored and
arendu
committed
1fa37a6
revert assert

Soumye Singhal
authored and
arendu
committed
1bbdf4f
Skip empty sequences and chunks in MTP tensor roll

rkarimimahab
authored and
arendu
committed
3c0ca73
add option for SFTTokenizer to build_tokenizers
arendu
committed
8992c41
resolve conflicts between args.sft and args.hybrid_context_parallel
arendu
committed
cf3d0ea

Commits on Feb 5, 2026

merged main
arendu
committed
f3835d3
do not check for and add any EOD
arendu
committed
8d5e58c

Commits on Feb 6, 2026

Merge branch 'main' into adithyare/sft-ultra-v3-feb2026
arendu
authored
ec97314

Commits on Feb 13, 2026

Fix nan loss caused by zero token in MTP

BestJuly
authored and
arendu
committed
9ed15ae
Merge branch 'adithyare/sft-ultra-v3-feb2026' of https://github.com/arendu/Megatron-LM into adithyare/sft-ultra-v3-feb2026
arendu
committed
38dc023

Commits on Feb 18, 2026

rebased main
arendu
committed
fa016fd
identity template and tool role mask

Soumye Singhal
authored and
arendu
committed
e039525
revert assert

Soumye Singhal
authored and
arendu
committed
ae4717d
Skip empty sequences and chunks in MTP tensor roll

rkarimimahab
authored and
arendu
committed
44c5c24
add option for SFTTokenizer to build_tokenizers
arendu
committed
8ed4905
resolve conflicts between args.sft and args.hybrid_context_parallel
arendu
committed
30f1f32
do not check for and add any EOD
arendu
committed
72a2740
removed dup identity template
arendu
committed
f82168a
added from #2363
arendu
committed
68514e8
use core sft_tokenizer
arendu
committed
d10fa92
fmt
arendu
committed
edd864c
Merge remote-tracking branch 'github-upstream/main' into adithyare/sft-ultra-v3-feb2026
arendu
committed
7bd7bfd

Commits on Feb 19, 2026

Merge branch adithyare/sft-ultra-v3-feb2026 from arendu/Megatron-LM

asolergi-nv
and
cursoragent
committed
dc371e5
SFT: tokenizer, dataset, utils and sft_mamba updates (clean_thd)

asolergi-nv
and
cursoragent
committed
7b44bc1

Commits on Feb 23, 2026

add sft dataset to core datasets
asolergi-nv
committed
203fd61
Merge branch 'NVIDIA:main' into clean_thd
asolergi-nv
authored
20c1841
create NullSFTTokenizer
asolergi-nv
committed
38da580
Add SFT dataset test
asolergi-nv
committed
a4025c4
Merge branch 'NVIDIA:main' into clean_thd
asolergi-nv
authored
cc54478
patch cu_seqlens
asolergi-nv
committed
a4ae811
add sft mamba script and refactor cp sharding logic
asolergi-nv
committed
a287001

Commits on Feb 24, 2026

Merge branch 'NVIDIA:main' into clean_thd
asolergi-nv
authored
4d47c63
remove logging - revert b4 merge
asolergi-nv
committed
f2ac856
create cu seqlens + remove some logging we should reintroduce b4 merge
asolergi-nv
committed
bef1a64

Commits on Feb 25, 2026

small refactor to mtp_on_this_rank
asolergi-nv
committed
599a36d
Merge branch 'NVIDIA:main' into clean_thd
asolergi-nv
authored
9211b08
first step get batch on this tp rank
asolergi-nv
committed
af6385e
Merge origin/clean_thd (resolve sft_mamba conflict, keep local get_batch_on_this_tp_rank)
asolergi-nv
committed
e1669df
Merge branch 'clean_thd' of https://github.com/asolergi-nv/Megatron-LM into clean_thd
asolergi-nv
committed
3011ca9
nit
asolergi-nv
committed
538bd5f
add pad token id for tests
asolergi-nv
committed
bdbf0f1
tests passing for sft & pretrain path
asolergi-nv
committed
fc40a37
return batch elements alphabetically
asolergi-nv
committed
58f8758
Create position_ids in SFTDataset
asolergi-nv
committed
3de705b

Commits on Mar 3, 2026

Commits on Mar 5, 2026

add sft to preprocess data script
asolergi-nv
committed
cc626c7
add completitions only training and update test
asolergi-nv
committed
c29cc9b
Merge branch 'NVIDIA:main' into clean_thd
asolergi-nv
authored
ed5feb6
mock position ids fix
asolergi-nv
committed
0d41a42
cuseqlens in int32
asolergi-nv
committed
ccb3704
REVERT position ids and autotokenizer fix. Also fix for flash backend
asolergi-nv
committed
3075013
properly create position ids after padding & truncation like loss mask
asolergi-nv
committed
1a28a58

Commits on Mar 6, 2026

add hybrid cp inputs to get batch method, update test, remove old cp shard code and add hybrid cp objects to get inputs on tp rank and cp rank
asolergi-nv
committed
df867ab

Commits on Mar 11, 2026

add docstring to utils funcs
asolergi-nv
committed
4974826

Commits on Mar 12, 2026

Commits on Mar 13, 2026

Merge branch 'NVIDIA:main' into clean_thd
asolergi-nv
authored
e664a03
improve SFTDataset tests
asolergi-nv
committed
200d3e5

Commits on Mar 16, 2026

Commits on Mar 17, 2026

Merge remote-tracking branch 'upstream/main' into clean_thd
asolergi-nv
committed
f00322f

Commits on Mar 18, 2026

Merge branch 'NVIDIA:main' into clean_thd
asolergi-nv
authored
ffdbb99
fix conflicts
asolergi-nv
committed
b0a76ba
Merge branch 'NVIDIA:main' into clean_thd
asolergi-nv
authored
b51203e
fix missing tool_response_end_tokens and remove adding special tokens in SFTTokenizer
asolergi-nv
committed
363b288
remove megatron training SFTDataset and move IGNORE_INDEX
asolergi-nv
committed
6dfbc23
redo build tokenizer
asolergi-nv
committed
ee02362

Commits on Mar 19, 2026

fix think tokens tokenize and add knob for add_special_tokens
asolergi-nv
committed
c523331
Merge branch 'NVIDIA:main' into clean_thd
asolergi-nv
authored
8fdb933
cleaning
asolergi-nv
committed
949cf73
autoformat linting
asolergi-nv
committed
ae7f1f2
update pretrain mamba and gpt entry points
asolergi-nv
committed
4999b40
Merge branch 'NVIDIA:main' into clean_thd
asolergi-nv
authored
4ad9e80
add copyright headers
asolergi-nv
committed
9f0e3a6
fix ci install
asolergi-nv
committed
e449719
fix ci
asolergi-nv
committed
01e0a5f
Merge branch 'main' into clean_thd
asolergi-nv
authored
01481d3
More CI fixes
asolergi-nv
committed
f1a596e
Merge branch 'main' into clean_thd
asolergi-nv
authored
c832039
lint
asolergi-nv
committed
12ec5e6
remove null sft tokenizer
asolergi-nv
committed
c766445

Commits on Mar 20, 2026

fix more ci errors
asolergi-nv
committed
f86ea8a
test
asolergi-nv
committed
45b402d
Merge branch 'main' into clean_thd
asolergi-nv
authored
be6f0c8
lint
asolergi-nv
committed
4dddd94
ci
asolergi-nv
committed
a61dd23
Merge branch 'main' into clean_thd
asolergi-nv
authored
7740a2c
fix hybridcp cp sharding
asolergi-nv
committed
cb5e4fb
Merge branch 'main' into clean_thd
asolergi-nv
authored
5544ee2
What a bug
asolergi-nv
committed
e785d9b
nits
asolergi-nv
committed
b9fb498

Commits on Mar 23, 2026

Merge branch 'main' into clean_thd
asolergi-nv
authored
6302c75

Commits on Mar 24, 2026

Merge branch 'main' into clean_thd
asolergi-nv
authored
f80b452
nits before review
asolergi-nv
committed
3564e8d
doc nit and torch seeding
asolergi-nv
committed
c864503
work with numpy arrays directly
asolergi-nv
committed
4c87a7c
delete get_batch_on_this_tp_rank from megatron training utils
asolergi-nv
committed
74c9b6c
force cp > 1 for hybrid cp tests
asolergi-nv
committed
39d1623
Apply suggestions from code review

asolergi-nv
and
claude[bot]
authored
24adf0d
Merge branch 'main' into clean_thd
asolergi-nv
authored
33d9372
lint
asolergi-nv
committed
e969aa1
nits
asolergi-nv
committed
ac0cf2a
Merge branch 'main' into clean_thd
asolergi-nv
authored
3e1ddf7
lint
asolergi-nv
committed
754aa2d
fix cp calls
asolergi-nv
committed
957911b
Merge branch 'main' into clean_thd
asolergi-nv
authored
d6c9b29

Commits on Mar 25, 2026

Merge branch 'main' into clean_thd
asolergi-nv
authored
e450e34
nit
asolergi-nv
committed
f07ae4b

Commits on Mar 29, 2026

Merge branch 'main' into clean_thd
asolergi-nv
authored
b6e48b1

Commits on Mar 30, 2026

nit
asolergi-nv
committed
dd16ebe
Merge branch 'main' into clean_thd
asolergi-nv
authored
c093244
more hope
asolergi-nv
committed
2b8e74f
Merge branch 'main' into clean_thd
asolergi-nv
authored
807ee2e

Commits on Apr 1, 2026

Merge branch 'main' into clean_thd
asolergi-nv
authored
8b7126b