Add mock SFT dataset and varlen dataset family by ilml · Pull Request #6734 · NVIDIA/Megatron-LM

ilml · 2026-08-21T03:03:57Z

What does this PR do?

Adds the mock SFT dataset (with a generalized SFT padding divisor) and the varlen dataset family (VarlenLowLevelDataset, VarlenDataset, MockVarlenDataset) with unit tests. The datasets are inert until the follow-up training-integration PR wires them into pretraining.

Part of the #3386 split series tracked in the Dynamic Context Parallelism project. Original changes by @xiaoyao0115 in #3386.

Note: Replaces #6685, which was auto-closed when its base mirror ref was deleted on #6684's merge. All of @jaredcasper's review comments on #6685 are addressed here (and canonically in #5902): sequence-packing field checks moved into ModelParallelConfig.__post_init__, argument formatting matched to the file's style with no reformat of untouched lines.

Series status

Remaining PRs of the #3386 split, tracked in the Dynamic Context Parallelism project. All PRs target main — no stacked bases. Reviews run in parallel; merges go in order:

Order	PR	Content
1	#5902	config field, CLI arg, validation
2	this PR (#6734)	mock SFT + varlen datasets
3	#6742	training-loop + varlen wiring (`pretrain_gpt` + `pretrain_hybrid`)

#6734 carries a temporary copy of #5902's commit, and #6742 carries #6734's commits (review its last three commits) — the duplicates drop out automatically as parents merge and the branches are rebased.

Merged so far: #5901, #5903, #6625, #5908, #6684, #5902.

Contribution process

Draft PR per contributing guidelines
Commits signed off (DCO)

@xiaoyao0115

Split 2/10 from NVIDIA#3386 (sequence packing / THD E2E support). Adds the sequence_packing_scheduler knob to ModelParallelConfig, its validation (field checks in ModelParallelConfig post_init; TE>=2.9 pin and alltoall dispatcher in TransformerConfig), 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>

@xiaoyao0115

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>

@xiaoyao0115

Split 9/10 from NVIDIA#3386 (sequence packing / THD E2E support). Adds the variable-length packed (THD) dataset family: HF-hub/parquet/jsonl loading, THD __getitem__ with cu_seqlens, SBHD validation mode, and the mock variant, with unit tests. Also adds hybrid_context_parallel=False to the _make_config test helper (deviation from NVIDIA#3386: fixes a latent AttributeError in _calculate_padding_divisor with SimpleNamespace configs). Original changes by @xiaoyao0115 in NVIDIA#3386. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: ilml <tolong@nvidia.com>

ilml · 2026-08-21T03:03:59Z

/ok to test d9fdfa6

copy-pr-bot · 2026-08-21T03:04:01Z

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 · 2026-08-21T23:44:49Z

/ok to test 946476b

The module-level import broke the minimal-install import check: pandas is deliberately only in the ci dependency group, not a core dependency. Guard the import inside the mode == "file" branch with a targeted ImportError, matching the existing pattern in varlen_dataset.py. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: ilml <tolong@nvidia.com>

janEbert

It seems like Pandas is just used to load CSV/JSON. In the case of JSON, I think you can load from JSON directly with datasets.load_dataset.
For CSV, I'd also say just use the csv package from stdlib, unless it's too slow compared to using Pandas.

Then we avoid introducing the additional dependency, since Pandas isn't used anywhere else yet (except CI).

ilml · 2026-08-24T16:56:05Z

/ok to test 4ccffde

Per review: pandas was only used to read the mock sequence-length CSV and local jsonl files, and is not a package dependency (it is only in the ci group). - MockSFTLowLevelDataset(mode="file") reads the CSV with the stdlib csv module, skipping non-numeric cells such as a header row. - VarlenLowLevelDataset parses local jsonl with the stdlib json module and builds the table with Dataset.from_list. This keeps the property the pandas path existed for: one pyarrow pass over the whole file, so rows with differing field sets unify into one schema instead of failing datasets.load_dataset("json")'s per-chunk inference with a CastError (e.g. LongAlpaca-12k). - Remove the now-stale pytest.importorskip("pandas") gates so the tests also run in environments without pandas. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: ilml <tolong@nvidia.com>

ilml · 2026-08-26T17:35:15Z

/ok to test 41de47f

The pre-existing allgather-vs-variable_seq_lengths ValueError ran for dense models too, where the (unused) dispatcher default is allgather. Sequence packing derives variable_seq_lengths=True, so every dense packing config failed validation. Gate the check on num_moe_experts, matching the packing-specific dispatcher assert below, and update the MoE test to expect the general check's ValueError, which fires before the packing assert for allgather. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: ilml <tolong@nvidia.com>

ilml · 2026-08-26T19:17:01Z

/ok to test 0b38f86

ilml · 2026-08-27T05:19:00Z

/ok to test 300c9c45b4ebf69d4e1a9dc8fe69449a82e94e7c

# Conflicts: # megatron/training/arguments.py

nemo-automation-bot · 2026-08-27T11:06:22Z

🔄 Merge queue validation started!

You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/33065907111

ilml and others added 3 commits August 21, 2026 03:03

copy-pr-bot Bot temporarily deployed to test August 21, 2026 03:05 Inactive

ilml marked this pull request as ready for review August 21, 2026 03:06

ilml requested review from a team as code owners August 21, 2026 03:06

copy-pr-bot Bot temporarily deployed to public August 21, 2026 03:07 Inactive

github-actions Bot added the complexity: medium label Aug 21, 2026

This was referenced Aug 21, 2026

Integrate sequence-packing scheduler into training loop and wire varlen dataset into GPT pretraining #6742

Open

Add sequence_packing_scheduler config field, CLI arg, and validation #5902

Merged

cuichenx approved these changes Aug 21, 2026

View reviewed changes

zhongbozhu approved these changes Aug 21, 2026

View reviewed changes

copy-pr-bot Bot temporarily deployed to test August 21, 2026 23:46 Inactive

copy-pr-bot Bot temporarily deployed to public August 21, 2026 23:48 Inactive

janEbert reviewed Aug 24, 2026

View reviewed changes

copy-pr-bot Bot deployed to test August 24, 2026 16:57 Active

copy-pr-bot Bot deployed to test August 25, 2026 16:41 Active

copy-pr-bot Bot deployed to public August 25, 2026 16:44 Active

copy-pr-bot Bot deployed to test August 26, 2026 17:37 Active

copy-pr-bot Bot deployed to public August 26, 2026 17:42 Active

Merge remote-tracking branch 'origin/main' into split/3386-dataset-r2

0b38f86

copy-pr-bot Bot deployed to test August 26, 2026 19:18 Active

copy-pr-bot Bot deployed to public August 26, 2026 19:21 Active

jaredcasper approved these changes Aug 26, 2026

View reviewed changes

ilml enabled auto-merge August 27, 2026 05:17

svcnvidia-nemo-ci added the Approved All necessary approvals have been made label Aug 27, 2026

Merge remote-tracking branch 'origin/main' into split/3386-dataset-r2

300c9c4

# Conflicts: # megatron/training/arguments.py

copy-pr-bot Bot deployed to test August 27, 2026 05:20 Active

copy-pr-bot Bot deployed to public August 27, 2026 05:22 Active

ilml added this pull request to the merge queue Aug 27, 2026

Merged via the queue into NVIDIA:main with commit d2931db Aug 27, 2026
92 checks passed

ilml deleted the split/3386-dataset-r2 branch August 27, 2026 12:32

Navigation Menu

Uh oh!

Add mock SFT dataset and varlen dataset family - #6734

Add mock SFT dataset and varlen dataset family#6734
ilml merged 11 commits into
NVIDIA:mainfrom
ilml:split/3386-dataset-r2

ilml commented Aug 21, 2026 •

edited

Loading

Uh oh!

ilml commented Aug 21, 2026

Uh oh!

copy-pr-bot Bot commented Aug 21, 2026

Uh oh!

ilml commented Aug 21, 2026

Uh oh!

janEbert left a comment

Uh oh!

ilml commented Aug 24, 2026

Uh oh!

ilml commented Aug 26, 2026

Uh oh!

ilml commented Aug 26, 2026

Uh oh!

ilml commented Aug 27, 2026

Uh oh!

nemo-automation-bot Bot commented Aug 27, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

6 participants

Uh oh!

Conversation

ilml commented Aug 21, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What does this PR do?

Series status

Contribution process

Uh oh!

ilml commented Aug 21, 2026

Uh oh!

copy-pr-bot Bot commented Aug 21, 2026

Uh oh!

ilml commented Aug 21, 2026

Uh oh!

janEbert left a comment

Choose a reason for hiding this comment

Uh oh!

ilml commented Aug 24, 2026

Uh oh!

ilml commented Aug 26, 2026

Uh oh!

ilml commented Aug 26, 2026

Uh oh!

ilml commented Aug 27, 2026

Uh oh!

nemo-automation-bot Bot commented Aug 27, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

6 participants

ilml commented Aug 21, 2026 •

edited

Loading