Skip to content

Add mock SFT dataset and generalize SFT padding divisor - #5906

Closed
ilml wants to merge 2 commits into
NVIDIA:pull-request/5905from
ilml:split/3386-06-mock-sft-dataset
Closed

Add mock SFT dataset and generalize SFT padding divisor#5906
ilml wants to merge 2 commits into
NVIDIA:pull-request/5905from
ilml:split/3386-06-mock-sft-dataset

Conversation

@ilml

@ilml ilml commented Jul 20, 2026

Copy link
Copy Markdown
Contributor
  • I, the PR author, have personally reviewed every line of this PR.

What does this PR do?

Adds MockSFTDataset/MockSFTLowLevelDataset with file- and distribution-mode sequence-length configs, the load_json_arg helper, --sft-mock-dataset-config-json (with a lognormal default when --mock-data --sft), and generalizes the SFT padding divisor beyond the previous cp>1-only path.

Part 06/10 of splitting #3386 (Add E2E support for THD format; dev-branch PR #2924). Original changes by @xiaoyao0115 in #3386 — split into functionally self-contained PRs to ease review. Hard dependencies (must merge first): #5902, #5903.

Split series (#3386)

Part PR Hard deps (merge first)
01 #5901
02 #5902
03 #5903
04 #6589 #5903
05 #5905 #6589
06 #5906 #5902, #5903
07 #5907 #5902, #6589, #5905
08 #5908
09 #5909 #5906, #5908
10 #5910 #5907, #5909

Branches are stacked linearly (each on the previous) so every PR shows a clean own-diff once its base is retargeted to the copy-pr-bot pull-request/<parent> ref; until then the Files-changed view of a stacked PR includes its ancestors — its own change is the last commit.
Stacked on #5905.

Review notes

  • Module-level import pandas as pd in sft_dataset.py makes pretrain_gpt.py hard-require pandas (currently only a ci extra) — consider lazy-guarding like varlen_dataset.py does.
  • New math/Union imports in sft_dataset.py are unused.
  • The padding-divisor generalization changes existing SFT padding when sequence_parallel_size>0 or hybrid CP is on.
  • MockSFTDataset has no dedicated unit test until the varlen PRs exercise it via MockVarlenDataset.

Issue tracking

Linked issue: Related to #3386

Contribution process

Pre-checks

  • 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

🤖 Generated with Claude Code

@copy-pr-bot

copy-pr-bot Bot commented Jul 20, 2026

Copy link
Copy Markdown

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

ilml and others added 2 commits August 10, 2026 22:09
Split 2/10 from NVIDIA#3386 (sequence packing / THD E2E support). Adds the
sequence_packing_scheduler knob to ModelParallelConfig, its
TransformerConfig validation (TE>=2.9 pin, variable_seq_lengths,
alltoall dispatcher), the explicit CLI args, and validate_args checks.
Feature-flagged no-op until the scheduler lands.

Original changes by @xiaoyao0115 in NVIDIA#3386.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Signed-off-by: ilml <tolong@nvidia.com>
Split 6/10 from NVIDIA#3386 (sequence packing / THD E2E support). Adds
MockSFTDataset/MockSFTLowLevelDataset with file- and distribution-mode
sequence-length configs, the load_json_arg helper, the
--sft-mock-dataset-config-json arg with lognormal default, and
generalizes the SFT padding divisor beyond cp>1.

Original changes by @xiaoyao0115 in NVIDIA#3386.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Signed-off-by: ilml <tolong@nvidia.com>
@ilml

ilml commented Aug 17, 2026

Copy link
Copy Markdown
Contributor Author

/ok to test 7169c46

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants