Add new SFTDataset to megatron.core.datasets - #5017
Conversation
|
Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually. Contributors can view more details about this message here. |
Signed-off-by: Antoni-Joan Solergibert <asolergibert@nvidia.com>
|
/ok to test b1f4fc2 |
|
/ok to test efa5da6 |
|
/claude review |
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
| return tokens | ||
|
|
||
| def text_to_ids(self, text: str) -> List[int]: | ||
| def text_to_ids(self, text: str, add_special_tokens: bool = True) -> List[int]: |
There was a problem hiding this comment.
Bug: add_special_tokens is accepted but never used in the method body. When include_special_tokens=True (the default), self.tokenizer(text).input_ids always includes BOS/EOS tokens regardless of the caller's request.
SFTDatasetConfig.__post_init__ calls self.tokenizer.tokenize(..., add_special_tokens=False) to tokenize chat template delimiters. With a HuggingFace tokenizer, this silently returns token sequences that include BOS/EOS, so extract_segments will fail to match them in pretokenized conversations — producing incorrect (all-zeros) loss masks.
The same issue applies to tiktoken_tokenizer.py, bytelevel_tokenizer.py, and sentencepiece_tokenizer.py — all accept the parameter but ignore it.
Suggested fix for this file:
def text_to_ids(self, text: str, add_special_tokens: bool = True) -> List[int]:
"""Converts text to tokens ids."""
if self.include_special_tokens and add_special_tokens:
return self.tokenizer(text).input_ids
tokens = self.text_to_tokens(text)
ids = self.tokens_to_ids(tokens)
return ids
There was a problem hiding this comment.
Review — light pass
One bug found:
add_special_tokens silently ignored in most tokenizer backends — The abstract text_to_ids signature was updated to accept add_special_tokens, and SFTDatasetConfig.__post_init__ relies on tokenize(..., add_special_tokens=False) to get clean delimiter token sequences. However, only sft_tokenizer.py actually propagates the parameter. The HuggingFace, TikToken, SentencePiece, and ByteLevel tokenizer implementations accept the parameter but ignore it. For HuggingFace in particular (the most likely backend for real SFT workloads), when include_special_tokens=True (the default), BOS/EOS tokens will be included in the delimiter sequences, causing extract_segments to never match them — resulting in all-zeros loss masks and no training signal on assistant responses.
The unit tests pass because they use NullTokenizer, which has no special tokens.
See inline comment on huggingface_tokenizer.py for a suggested fix.
Sync PR 5017 (SFTDataset) with main (225 commits since 873678a). Conflicts resolved: - pretrain_hybrid.py (core_gpt_dataset_config_from_args): keep the PR's data_args dict + SFTDatasetConfig dispatch, and add main's new inter_document_masking key (mirrors pretrain_gpt.py). - megatron/elastification/pretrain_hybrid_flex.py: import get_batch_on_this_*_rank from megatron.core.utils (main relocated them there) and drop the stale megatron.training.datasets.sft_dataset import (SFTDataset moved to core.datasets). - megatron/training/datasets/sft_dataset.py: keep deleted (relocated to megatron/core/datasets/sft_dataset.py by this PR). Auto-merged files reconciled with main's refactors: get_batch_on_this_tp_rank's is_sft param -> has_cu_seqlens, new flatten_batch_for_packed_sequences step, and get_batch_on_this_cp_rank's use_per_sequence_balancing kwarg. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: asolergi-nv <asolergibert@nvidia.com>
45119d0 to
a25e4b0
Compare
Signed-off-by: Antoni-Joan Solergibert <asolergibert@nvidia.com>
|
/ok to test fe2c20e |
| return tokens | ||
|
|
||
| def text_to_ids(self, text: str) -> List[int]: | ||
| def text_to_ids(self, text: str, add_special_tokens: bool = True) -> List[int]: |
There was a problem hiding this comment.
why do we need to pass special tokens here? Looks like it's not being used by function
There was a problem hiding this comment.
I added it to all tokenizers since we are using this argument in the parent class
| return self.ids_to_text(tokens) | ||
|
|
||
| def text_to_ids(self, text): | ||
| def text_to_ids(self, text, add_special_tokens=True): |
| return document_index, sample_index, shuffle_index | ||
|
|
||
|
|
||
| def _classify_items( |
There was a problem hiding this comment.
what about moving these extra functions to separate file like sft_utils.py?
|
|
||
|
|
||
| @dataclass | ||
| class SFTDatasetConfig(GPTDatasetConfig): |
There was a problem hiding this comment.
do you think it makes sense to move all the configs to separate sft_configs.py?
| @@ -1,201 +0,0 @@ | |||
| # Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. | |||
There was a problem hiding this comment.
is this okay to remove? or should be put in legacy.
…ator test_inter_document_masking_batch runs with sft=False, so get_batch does not call preprocess_sft_batch. After syncing with main it consumed the SFT-path create_sft_data_iterator, which this PR repurposed to emit *un-padded* data (padding, position_ids and max_seqlen are produced by preprocess_sft_batch). With sft=False those fields were absent (position_ids was None), so the test FAILED and the cp=2 case then hung on a NCCL collective timeout -> SIGABRT, crashing the whole tests/unit_tests/data bucket. Add a dedicated create_inter_document_masking_data_iterator that emits data already padded to seq_length with position_ids / cu_seqlens / max_seqlen (the contract main's helper provided before this PR), and point the test at it. Test-only change; production get_batch/utils are unaffected. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: asolergi-nv <asolergibert@nvidia.com>
Sync PR 5017 (SFTDataset) with upstream main (97 commits, up to a046281). Conflicts resolved (2, both in the `from megatron.core.utils import (...)` block of pretrain_gpt.py and pretrain_hybrid.py): keep the PR's `preprocess_sft_batch` import alongside main's new `get_te_version` / `get_torch_version`. Also reconciled a silent name collision in tests/unit_tests/data/test_get_batch.py: main (bf32f44, "Fix inter-document masking crash ... mbs > 1") added its own `create_inter_document_masking_data_iterator` (mbs>1 variant) with the same name this branch had introduced, so it shadowed the branch's version. Dropped the branch's duplicate and pointed `test_inter_document_masking_batch` at main's canonical iterator with `micro_batch_size=1`. Verified on 2 GPUs: tests/unit_tests/data (test_get_batch, test_sft_dataset, test_cp_utils) = 103 passed, 0 failed. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: asolergi-nv <asolergibert@nvidia.com>
aa921f8 to
97b2c55
Compare
|
/ok to test 97b2c55 |
4 participants
Summary
This PR replaces the JSONL-backed
SFTDatasetinmegatron/training/datasets/sft_dataset.pywith a new pre-tokenized, packed implementation undermegatron/core/datasets/sft_dataset.py. The new dataset:.bin/.idx) like pretraining. It's up to the user to apply the chat template and later on pretokenize the data with thetools/preprocess_data.pyscript.ChatTemplateConfig— Nemotron-3 default). Three flags control what counts toward the loss:train_on_assistant_responses_onlytrain_on_thinking_traces(mask<think>…</think>regions)train_on_tool_calls(mask<tool_call>…</tool_call>regions). Tool responses (environment output inside user turns) are always masked.tokens,labels,loss_mask,cu_seqlens) ready for Flash Attention varlen kernels.A new helper,
preprocess_sft_batch, is added tomegatron/core/utils.pyand called from bothpretrain_gpt.pyandpretrain_hybrid.py. It handles the runtime steps that used to happen inside the old__getitem__: padding each sub-sequence to a CP-divisible length, padding or truncating the packed sample toseq_length, and producingposition_ids,cu_seqlens_padded, andmax_seqlen.A standalone debug script,
inspect_sft_file_prefix.py, is added at the repo root for inspecting a.bin/.idxpair against the Nemotron-3 template. Will remove before merge.User-facing reference docs for the dataset live in
megatron/core/datasets/sft.MD.Changes per file
New: dataset and helpers
megatron/core/datasets/sft_dataset.py(new, ~830 lines)SFTDataset,SFTDatasetConfig,ChatTemplateConfig,Nemotron3ChatTemplateConfig, the MFFD packer (pack_samples+_classify_items), and the chat-template segment parser (extract_segments,_split_tool_calls,find_subsequence).megatron/core/datasets/sft.MD(new)megatron/core/utils.py(+289 lines)pad_thd_sequences_for_cp(round each segment up to a CP-divisible length),pad_or_truncate_thd_tensors(force the packed sample to exactlymax_seq_len), andpreprocess_sft_batch(TP-rank-0 batch preprocessing wrapper).inspect_sft_file_prefix.py(new, 332 lines).bin/.idxpair and prints per-segment role assignments using the hard-coded Nemotron-3 chat template. Not imported by training code.Removed: the old SFTDataset
megatron/training/datasets/sft_dataset.pySFTDataset/SFTLowLevelDatasetare superseded by the pre-tokenized core implementation.Training entrypoints wired to the new dataset
pretrain_gpt.pySFTDataset,SFTDatasetConfig,IGNORE_INDEXfrommegatron.core.datasets.sft_dataset. Callpreprocess_sft_batchinsideget_batchwhenargs.sftis set. ReturnSFTDatasetConfigfromcore_gpt_dataset_config_from_argswhen in SFT mode.pretrain_hybrid.pypretrain_gpt.py. Also refactoredcore_gpt_dataset_config_from_argsto build adata_argsdict so the SFT and non-SFT configs share construction.megatron/elastification/pretrain_hybrid_flex.pySFTDatasetimport to point at the new module.Tokenizer
add_special_tokensplumbingThe dataset tokenizes chat-template delimiter strings (
<|im_start|>system\n,<|im_end|>\n,<think>, …) and must do so without the tokenizerauto-inserting BOS/EOS. A new
add_special_tokens: bool = Truekeyword isadded through the tokenizer stack:
megatron/core/tokenizers/text/libraries/abstract_tokenizer.pyadd_special_tokens: bool = Trueto thetext_to_idsABC signature.megatron/core/tokenizers/text/libraries/bytelevel_tokenizer.pyadd_special_tokens.megatron/core/tokenizers/text/libraries/huggingface_tokenizer.pyadd_special_tokens(caller-side responsibility — HF path is unchanged whenTrue).megatron/core/tokenizers/text/libraries/sentencepiece_tokenizer.pyadd_special_tokens.megatron/core/tokenizers/text/libraries/tiktoken_tokenizer.pyadd_special_tokens.megatron/core/tokenizers/text/libraries/sft_tokenizer.pyadd_special_tokenstoself._tokenizer.encode(...); tighten anassertwith a message.megatron/core/tokenizers/text/libraries/null_tokenizer.pyadd_special_tokensand changetext_to_idsto a character-level encoding[ord(c) % vocab_size for c in text]. This makes the null tokenizer usable in the SFT unit tests (it can now tokenize arbitrary template strings reproducibly).megatron/core/tokenizers/text/text_tokenizer.pyadd_special_tokensthroughtokenize(...).megatron/core/tokenizers/utils/build_tokenizer.py{NullTokenizer, NullMultimodalTokenizer} -> library namemapping into aNULL_TOKENIZERSdict.Tests included
New tests
tests/unit_tests/data/test_sft_dataset.py(399 lines) — end-to-endtest of
SFTDatasetviaBlendedMegatronDatasetBuilder..bin/.idxfiles containing 4 conversation shapes(
simple,with_thinking,with_tool_calls,with_thinking_and_tool_calls),with and without system turns.
TestChatTemplateConfigwhose delimiter strings tokenize cleanlyunder the (now character-level)
NullTokenizer.vocab_size ∈ {131072, 20000}and all five(train_on_assistant_responses_only, train_on_thinking_traces, train_on_tool_calls)combinations allowed by the asserts in
SFTDatasetConfig.__post_init__.cu_seqlens, label-shift correctness, and per-segment loss-mask valuesby re-running
extract_segmentson each document inside a packed sample.tests/unit_tests/data/test_cp_utils.py(232 lines) — unit tests forpad_thd_sequences_for_cpinmegatron/core/utils.py._reference_padimplementation.already-divisible (no-op), 2-D input squeezing, dtype preservation
(
int32/int64forcu_seqlens, mixed dtypes per tensor), multi-tensorcalls, single-tensor calls,
divisibility_factor=1, and randomizedparametric tests checking the post-condition that every padded segment
length is a multiple of the divisibility factor.
Updated tests
tests/unit_tests/data/test_get_batch.py— the mock SFT data iteratornow matches the new dataset's output shape (un-padded, no
position_ids,no
max_seqlen). Assertions acrosstest_sft_batchandtest_hybrid_cp_batchupdated to expectcu_seqlens_paddedto bepopulated when
cp_size > 1(now produced bypreprocess_sft_batch).tests/unit_tests/tokenizers/test_tokenizer.py—test_null_tokenizerand
test_detokenize_skip_special_tokens_unsupported_backendupdated toreflect the
NullTokenizer's new character-level encoding.Contribution process
Pre-checks
Code review
Feel free to message or comment @NVIDIA/mcore-oncall to help accelerate your merge into main. The less complex your PR is, the faster it will be approved and merged!
All PRs start as draft. If you open a non-draft PR, it will be automatically converted to draft.
Step 1: Mark PR as "Ready for Review"
.github/CODEOWNERS.Final Review might get declined if these requirements are not fulfilled.
Step 2: Final Review
For PRs that change
megatron/core, once all expert reviewers have approved, theFinal Reviewlabel is applied automatically and final reviewers are assigned.For PRs outside
megatron/core, this step is skipped.Step 3: Approved
Once all required reviewers have approved, the
Approvedlabel is applied automatically.Merge
Any member of mcore-engineers will be able to merge your PR.