megatron.core.datasets SFTDataset and get_batch* functions refactor

[asolergi-nv]

What does this PR do ?

In this PR I introduce a new SFTDataset for text that will live in megatron.core.datasets. In addition, I also include a refactor of various get_batch methods and clear guidance about the expected inputs for the three different regimes (pretraining, SFT, and DynamicCP).

• SFTDataset (@NVIDIA/datasets @duncanriach @arendu):
  • Based on file prefixes, just like in pretraining. We will generate them in the same way using the same tools/preprocess_data.py script, but it is important to note that the user will be responsible for adding the chat template to the raw JSONL file data. We can debate whether it is necessary to provide a basic python script that applies the chat template using the nemotron tokenizer or not.
  • Supports training with MULTIPLE file prefixes using the same code & logic as we do for pretraining with the BlendedMegatronDatasetBuilder
  • We have three indexes (document, sample, and shuffle index) that are memory mapped as in pretraining
    • So we no longer have to load a enormous jsonl file into memory
  • We support training only on assistant answers, as well as toggles to choose whether we want to train on tool calls and reasoning traces.
    • To extract segments of each type (system, user, assistant, tool call, tool answer, reasoning), we need to specify the chat template tokens that delimit each segment. In this PR I include the ones used in the Nemotron3 chat template (Nemotron3ChatTemplateConfig), which I manually verified using the inspect_sft_file_prefix.py script (which I will later remove).
  • We pack sequences using the Modified First-Fit Decreasing algorithm.
  • Moved all the padding / truncation / Constraints to the preprocess_sft_batch method (Next section)
  • Created test_sft_dataset.py that checks functionality.
  • Added SFTDataset documentation
• Created preprocess_sft_batch method.
  • The main challenge with SFT data is handling the different metadata tensors used in attention computation and the various shape constraints that exist.
  • Fundamentally, we do the following:
    • Pad individual sub-sequences for CP divisibility via pad_thd_sequences_for_cp from TE (when cp_size > 1), producing cu_seqlens_padded
    • Pad or truncate the packed sample to exactly max_seq_len tokens
    • Create a loss mask that zeros out padding tokens and prompt tokens
    • Compute per-segment position IDs from cumulative sequence lengths
    • Compute max_seqlen (the length of the longest sub-sequence)
  • This method could be implemented as a DataLoader collator, but I decided to place it here explicitly to make it clearer that this preprocessing step is essential.
• Updated mtp_on_this_rank to use explicit arguments (@NVIDIA/gpt & @NVIDIA/hybrid-mamba ).
• Refactored get_batch, get_batch_on_this_cp_rank, and get_batch_on_this_tp_rank functions to use explicit arguments (@NVIDIA/multi-modal @NVIDIA/core-nemo @NVIDIA/datasets @NVIDIA/core-adlr @xiaoyao0115 @parthmannan).
  • I refactored these methods to use explicit arguments so they can be reused across other projects and added tests to ensure that any future changes either add inputs properly or do not break model inputs.
  • Also refactored the pretrain_mamba.py and pretrain_gpt.py scripts so they support the same dataset-related features.
  • Added some tests to the get_batch methods to ensure future changes do not break this function.
• Deleted megatron/training/datasets/sft_dataset.py (@NVIDIA/datasets)
  • Note that IGNORE_INDEX is already defined in megatron/core/tokenizers/text/libraries/sft_tokenizer.py.
• Added add_special_tokens option to MegatronTokenizerText (@NVIDIA/tokenizers).
  • This is required to obtain the chat template tokens in order to mask out chat template, reasoning, and tool-calling tokens. It is ONLY used for SFT with SFTDataset and SFTTokenizer; the rest of the tokenizers ignore this argument (keeping exactly the same behavior as before this change).
  • NullTokenizer now takes a string and converts it to ASCII values. Previously, it accepted a space-separated list of values, and the tokenized sample was those integers cast directly. In short, it previously only supported string values such as "32 456 123 542"; now it supports ANY string value.

TODO:

• Remove inspect_sft_file_prefix.py
• Re-implement changes from Fix quantize.py script and support packed sequences in pretrain_gpt.py #3564

Contribution process

flowchart LR
    A[Pre-checks] --> B[PR Tests]
    subgraph Code Review/Approval
        C1[Expert Review] --> C2[Final Review]
    end
    B --> C1
    C2 --> D[Merge]

Loading

Pre-checks

• I want this PR in a versioned release and have added the appropriate Milestone (e.g., Core 0.8)
• I have added relevant unit tests
• I have added relevant functional tests
• I have added proper typing to my code Typing guidelines
• I have added relevant documentation
• I have run the autoformatter.sh on my PR

Code review

The following process is enforced via the CODEOWNERS file for changes into megatron/core. For changes outside of megatron/core, it is up to the PR author whether or not to tag the Final Reviewer team.

For MRs into `main` branch

Feel free to message or comment the @mcore-oncall to help accelerate your merge into main. The less complex your PR is, the faster it will be approved and merged!

(Step 1): Add PR label Expert Review

(Step 2): Collect the expert reviewers reviews

1. Attach the Expert Review label when your PR is ready for review.
2. GitHub auto-assigns expert reviewers based on your changes. They will get notified and pick up your PR soon.

⚠️ Only proceed to the next step once all reviewers have approved, merge-conflict are resolved and the CI is passing.
Final Review might get declined if these requirements are not fulfilled.

(Step 3): Final Review

1. Add Final Review label
2. GitHub auto-assigns final reviewers based on your changes. They will get notified and pick up your PR soon.

(Optional Step 4): Cherry-pick into release branch

If this PR also needs to be merged into core_r* release branches, after this PR has been merged, select Cherry-pick to open a new PR into the release branch.

For MRs into `dev` branch

The proposed review process for `dev` branch is under active discussion.

MRs are mergable after one approval by either eharper@nvidia.com or zijiey@nvidia.com.

Merging your PR

Any member of core-adlr and core-nemo will be able to merge your PR.

[asolergi-nv]

/ok to test d6c9b29

[asolergi-nv]

/ok to test f07ae4b

[asolergi-nv]

/ok to test b6e48b1

[asolergi-nv]

/ok to test c093244

[asolergi-nv]

/ok to test 807ee2e

[asolergi-nv]

We are dividing this PR into 2 smaller PRs.
First #4103