Skip to content

Add VarlenDataset family and wire it into GPT pretraining - #6680

Closed
ilml wants to merge 6 commits into
NVIDIA:mainfrom
ilml:split/3386-dataset-varlen
Closed

Add VarlenDataset family and wire it into GPT pretraining#6680
ilml wants to merge 6 commits into
NVIDIA:mainfrom
ilml:split/3386-dataset-varlen

Conversation

@ilml

@ilml ilml commented Aug 19, 2026

Copy link
Copy Markdown
Contributor

What does this PR do?

Adds the varlen dataset family (VarlenLowLevelDataset, VarlenDataset, MockVarlenDataset) and wires it into the pretraining CLI, dataloader, and dataset provider. Final PR of the #3386 split series tracked in the Dynamic Context Parallelism project. Original changes by @xiaoyao0115 in #3386. Replaces closed #6629/#6630.

Depends on #6679 — do not merge before it. This branch carries #6679's commits underneath so imports resolve and CI can run (VarlenDataset uses MockSFTDataset/load_json_arg from the mock-SFT layer, and the CLI wiring reads --sequence-packing-scheduler). Review only the last two commits:

Commit Layer Content
HEAD~1 09 VarlenLowLevelDataset, VarlenDataset, MockVarlenDataset (megatron/training/datasets/varlen_dataset.py) + unit tests
HEAD 10 --use-varlen-dataset CLI, dataloader/provider wiring in pretrain_gpt.py, arg validation

Once #6679 merges, this branch gets rebased and the Files-changed diff collapses to exactly those two commits (~+907/−7).

Contribution process

  • Draft PR per contributing guidelines
  • Commits signed off (DCO)

ilml and others added 6 commits August 19, 2026 20:24
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>
…ry point

Split 5/10 from NVIDIA#3386 (sequence packing / THD E2E support). Adds
BasePackingScheduler, DpBalancedScheduler, PackingSchedulerEnum, the
wrap_data_iterator entry point, docs, and the wrap_dataloader unit
tests.

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>
…aining

Split 7/10 from NVIDIA#3386 (sequence packing / THD E2E support). Wires
wrap_data_iterator into train_step/evaluate with dynamic
num_microbatches, adds global seqlen-stats plumbing for FLOPs
accounting, gates the HybridCP sampler, and restructures
pretrain_gpt.forward_step for packed (THD) batches with padding mask.

Original changes by @xiaoyao0115 in NVIDIA#3386.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Signed-off-by: ilml <tolong@nvidia.com>
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>
…vider

Split 10/10 from NVIDIA#3386 (sequence packing / THD E2E support). Adds
--use-varlen-dataset / --varlen-sbhd-validation /
--varlen-mock-dataset-config-json, their validation (auto-selecting the
dp_balanced scheduler), the dataset-provider wiring in pretrain_gpt,
and the remaining dataloader/unpack unit tests.

Original changes by @xiaoyao0115 in NVIDIA#3386.

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

copy-pr-bot Bot commented Aug 19, 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

ilml commented Aug 19, 2026

Copy link
Copy Markdown
Contributor Author

/ok to test cd0a019

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

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant