Skip to content

varlendataset for thd e2e and benchmark #4832

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
xiaoyao0115 merged 7 commits into from
Jun 8, 2026
Merged

varlendataset for thd e2e and benchmark #4832

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
4a9a9dc
add varlendataset for thd e2e and benchmark
xiaoyao0115 May 16, 2026
2332785
add support for dataset like https://huggingface.co/datasets/allenai/…
xiaoyao0115 May 28, 2026
a85be94
minor fixes according to the comments
xiaoyao0115 May 28, 2026
5e9b3e2
some fixes according to the comments
xiaoyao0115 Jun 5, 2026
5051382
fix ut
xiaoyao0115 Jun 6, 2026
14d94d4
minor fixes
xiaoyao0115 Jun 8, 2026
0e48a3b
reanme varlen_bshd_validation to arlen_sbhd_validation
xiaoyao0115 Jun 8, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
23 changes: 23 additions & 0 deletions megatron/core/datasets/data_schedule_utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -43,7 +43,30 @@ def _unpack_batch(batch: List[Dict[str, torch.Tensor]]) -> List[Dict[str, torch.
Since each sub-sample may be routed to different DPxCP ranks,
we unpack the sample here to avoid unnecessarily transferring
the entire packed sample.

Two input shapes are accepted:

* **Pre-packed** (e.g. :class:`SFTDataset`): each sample carries a
``cu_seqlens`` tensor and the tokens of multiple sub-samples
concatenated together. We slice them apart and synthesize
``original_seq_len`` / ``padded_seq_len`` from the cu_seqlens deltas.

* **Already unpacked** (e.g. :class:`VarlenDataset`): each sample is a
single sub-sample that already carries ``padded_seq_len`` (and
usually ``original_seq_len``). We just normalize the leading batch
dimension introduced by the default collate_fn and return as-is.
"""
# Short-circuit for datasets that already emit one sub-sample per index.
if batch and "padded_seq_len" in batch[0]:
for sample in batch:
for key in sample.keys():
if sample[key].ndim == 2 and sample[key].shape[0] == 1:
# Drop the redundant batch dim added by collate_fn.
sample[key] = sample[key].squeeze(0)
if "original_seq_len" not in sample:
sample["original_seq_len"] = sample["padded_seq_len"].clone()
return batch

batch_unpacked = []
dev = batch[0]["tokens"].device
original_seq_lens = []
Expand Down
18 changes: 18 additions & 0 deletions megatron/core/datasets/gpt_dataset.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -79,6 +79,18 @@ class GPTDatasetConfig(BlendedMegatronDatasetConfig):
sft_mock_dataset_config_json: Optional[str] = None
"""This config provides the necessary information for the mock dataset."""

varlen_mock_dataset_config_json: Optional[str] = None

@yuzhongw-nvidia yuzhongw-nvidia May 19, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Can we support either a JSON string or a JSON file path for sft_mock_dataset_config_json and varlen_mock_dataset_config_json? It is a little bit annoying for users to pass a JSON string via the CLI.

@xiaoyao0115 xiaoyao0115 May 28, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

ok, i'll make the change~

"""Mock-dataset config (same JSON schema as ``sft_mock_dataset_config_json``)
used by the ``--use-varlen-dataset`` path; kept separate so the varlen path
does not implicitly inherit SFT-specific knobs."""

varlen_sbhd_validation: bool = False
"""When True, :class:`VarlenDataset.__getitem__` emits SBHD samples padded
to ``sequence_length`` (no ``cu_seqlens`` / ``original_seq_len`` /
``padded_seq_len``), bypassing the packed-sequence path. Used to obtain a
SBHD reference run that mirrors the THD path's tokenization but skips all
packing — useful for THD numerical-correctness validation."""

def __post_init__(self) -> None:
"""Do asserts and set fields post init"""
super().__post_init__()
Expand All @@ -89,6 +101,12 @@ def __post_init__(self) -> None:
assert self.reset_attention_mask is not None
assert self.eod_mask_loss is not None

if self.varlen_sbhd_validation:
assert not self.dynamic_context_parallel, (
"--varlen-sbhd-validation is incompatible with "
"--dynamic-context-parallel (SBHD mode is not packed)."
)

self.token_dtype_code = (
None
if self.tokenizer.vocab_size is None
Expand Down
113 changes: 102 additions & 11 deletions megatron/training/arguments.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -89,6 +89,7 @@ def add_megatron_arguments(parser: argparse.ArgumentParser):
parser = _add_msc_args(parser)
parser = _add_kitchen_quantization_arguments(parser)
parser = _add_sft_args(parser)
parser = _add_varlen_dataset_args(parser)

parser = _add_fault_injector_args(parser)

Expand Down Expand Up @@ -1554,16 +1555,6 @@ def validate_args(args, defaults={}):
f"to {args.data_parallel_size * args.context_parallel_size}."
)

if args.sequence_packing_scheduler is not None:
if args.sequence_packing_scheduler == 'dp_balanced':
total_cp_ranks = args.context_parallel_size
else:
total_cp_ranks = args.data_parallel_size * args.context_parallel_size
assert total_cp_ranks * args.max_seqlen_per_dp_cp_rank >= args.seq_length, (
f'Packed sequence buffer size ({total_cp_ranks * args.max_seqlen_per_dp_cp_rank}) '
f'must be >= single sequence max length ({args.seq_length})'
)

# disable async_tensor_model_parallel_allreduce when
# model parallel memory optimization is enabled
if (
Expand Down Expand Up @@ -1692,6 +1683,55 @@ def validate_args(args, defaults={}):
args.use_megatron_fsdp
), "--ckpt-format fsdp_dtensor is only tested with Megatron FSDP."

# --use-varlen-dataset: independent of --sft. Cannot be combined with --sft
# because they are mutually-exclusive top-level dataset selectors that both
# drive the packed-sequence (THD) path.
if args.use_varlen_dataset:
assert not args.sft, (
"--use-varlen-dataset and --sft are mutually exclusive; both "
"select the packed-sequence dataset family. Pick one."
)
if args.varlen_sbhd_validation:
# ``--dynamic-context-parallel`` ⊥ ``--varlen-sbhd-validation`` is
# checked in ``GPTDatasetConfig.__post_init__``; only the
# scheduler check stays here, since ``sequence_packing_scheduler``
# is a training-framework flag not stored on the dataset config.
assert args.sequence_packing_scheduler is None, (
"--varlen-sbhd-validation does not use a sequence packing "
"scheduler; drop --sequence-packing-scheduler."
)
# SBHD validation is a real-data numerical-reference path only;
# MockVarlenDataset does not implement it.
assert not args.mock_data, (
"--varlen-sbhd-validation is not supported with --mock-data; "
"SBHD validation requires a real dataset."
)
else:
# VarlenDataset emits one unpacked sample per __getitem__; it
# relies on an upstream packing scheduler to group variable-length
# samples into THD batches. Auto-pick a default scheduler when
# the user did not request one explicitly:
# * ``--dynamic-context-parallel`` is already wired to
# ``default_dynamic_cp`` upstream (see the dynamic-cp block
# earlier in ``validate_args``).
# * Otherwise fall back to ``dp_balanced`` (static packing).
if args.sequence_packing_scheduler is None:
args.sequence_packing_scheduler = 'dp_balanced'

@yuzhongw-nvidia yuzhongw-nvidia Jun 1, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The normal --use-varlen-dataset path auto-selects dp_balanced after the generic sequence-packing validation has already run. See https://github.com/NVIDIA/Megatron-LM/pull/4832/changes#r3333048405.


# Packed-sequence buffer-size check. Placed after all scheduler auto-select
# logic (dynamic-cp and --use-varlen-dataset both set the scheduler above)
# so it validates the final resolved scheduler; the varlen path picks its
# default after the earlier generic validation has run.
if args.sequence_packing_scheduler is not None:
if args.sequence_packing_scheduler == 'dp_balanced':
total_cp_ranks = args.context_parallel_size
else:
total_cp_ranks = args.data_parallel_size * args.context_parallel_size
assert total_cp_ranks * args.max_seqlen_per_dp_cp_rank >= args.seq_length, (
f'Packed sequence buffer size ({total_cp_ranks * args.max_seqlen_per_dp_cp_rank}) '
f'must be >= single sequence max length ({args.seq_length})'
)

# Data blend checks
assert (
args.mock_data
Expand Down Expand Up @@ -4840,13 +4880,64 @@ def _add_sft_args(parser):
'--sft-mock-dataset-config-json',
type=str,
default=None,
help='This config provides the necessary information for the mock dataset. You can either specify a CSV file that contains sequence lengths, where each line stores the length of a sequence, for example: {"mode":"file","path":"/path/to/file"}. Alternatively, you can specify a distribution (currently only supporting lognormal distribution) along with the required parameters, for example, {"mode":"distribution","type":"lognormal","min_seq_len":1024,"max_seq_len":2048,"mean_seq_len":1536,"lognormal_sigma":1.1}, where sigma controls the variability of the lognormal distribution. '
help='This config provides the necessary information for the mock dataset. '
'Accepts either an inline JSON literal or a path to a JSON file containing '
'the same schema. You can either specify a CSV file that contains sequence lengths, '
'where each line stores the length of a sequence, for example: '
'{"mode":"file","path":"/path/to/file"}. Alternatively, you can specify a distribution '
'(currently only supporting lognormal distribution) along with the required parameters, '
'for example, {"mode":"distribution","type":"lognormal","min_seq_len":1024,'
'"max_seq_len":2048,"mean_seq_len":1536,"lognormal_sigma":1.1}, where sigma controls '
'the variability of the lognormal distribution. '
'If not specified and --mock-data is set, defaults to a lognormal distribution with '
'min_seq_len=seq_length//2, max_seq_len=seq_length, mean_seq_len=seq_length*3//4, lognormal_sigma=1.1.',
)
return parser


def _add_varlen_dataset_args(parser):
group = parser.add_argument_group(title='varlen dataset')
group.add_argument(
'--use-varlen-dataset',
action="store_true",
help='Train with VarlenDataset, a variable-length packed (THD) dataset '
'that consumes instruction-tuning data from a HuggingFace Hub repo id, '
'a local parquet file, or a local jsonl file. Schema (alpaca / sharegpt '
'/ openai-messages) is auto-detected from the dataset columns. '
'Mutually exclusive with --sft. Auto-picks a sequence packing '
'scheduler when none is given: ``dp_balanced`` by default, '
'``default_dynamic_cp`` when ``--dynamic-context-parallel`` is set. '
'Combine with --mock-data for a synthetic lognormal sequence-length '
'distribution; see --varlen-mock-dataset-config-json.',
)
group.add_argument(
'--varlen-sbhd-validation',
action="store_true",
help='Reference SBHD mode for THD numerical verification. When set, '
'VarlenDataset emits SBHD-style samples right-padded to '
'--seq-length (no cu_seqlens, no packing scheduler), so the run can '
'be compared against the THD path to validate correctness. '
'Incompatible with --dynamic-context-parallel and '
'--sequence-packing-scheduler.',
)
group.add_argument(
'--varlen-mock-dataset-config-json',
type=str,
default=None,
help='Mock-dataset config for --use-varlen-dataset --mock-data. '
'Accepts either an inline JSON literal or a path to a JSON file containing '
'the same schema as --sft-mock-dataset-config-json: either '
'{"mode":"file","path":"/path/to/lengths.csv"}, '
'{"mode":"distribution","type":"lognormal","min_seq_len":1024,'
'"max_seq_len":2048,"mean_seq_len":1536,"lognormal_sigma":1.1}, or '
'{"mode":"verification","data_path":"/prefix/of/IndexedDataset"}. '
'If not specified, defaults to a lognormal distribution with '
'min_seq_len=seq_length//2, max_seq_len=seq_length, '
'mean_seq_len=seq_length*3//4, lognormal_sigma=1.1.',
)
return parser


def _add_fault_injector_args(parser):
from megatron.training.config import FaultInjectorConfig

Expand Down
111 changes: 20 additions & 91 deletions megatron/training/datasets/data_samplers.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -41,37 +41,22 @@ def build_pretraining_data_loader(dataset, consumed_samples):
if is_eval
else args.micro_batch_size
)
global_batch_size = (
getattr(args, 'eval_global_batch_size', args.global_batch_size)
if is_eval
else args.global_batch_size
)

if split == Split.valid and args.full_validation:
batch_sampler = MegatronFullValidationSampler(
total_samples=len(dataset),
data_parallel_rank=mpu.get_data_parallel_rank(),
data_parallel_size=mpu.get_data_parallel_world_size(),
)
elif args.dataloader_type == 'single':
if args.dynamic_context_parallel:
batch_sampler = HybridCPMegatronPretrainingSampler(
total_samples=len(dataset),
consumed_samples=consumed_samples,
micro_batch_size=micro_batch_size,
global_batch_size=global_batch_size,
data_parallel_rank=mpu.get_data_parallel_rank(),
data_parallel_size=mpu.get_data_parallel_world_size(),
)
else:
# Megatron sampler
batch_sampler = MegatronPretrainingSampler(
total_samples=len(dataset),
consumed_samples=consumed_samples,
micro_batch_size=micro_batch_size,
data_parallel_rank=mpu.get_data_parallel_rank(),
data_parallel_size=mpu.get_data_parallel_world_size(),
)
# Packing schedulers consume one microbatch at a time and form
# global/DCP batches themselves.
batch_sampler = MegatronPretrainingSampler(
total_samples=len(dataset),
consumed_samples=consumed_samples,
micro_batch_size=micro_batch_size,
data_parallel_rank=mpu.get_data_parallel_rank(),
data_parallel_size=mpu.get_data_parallel_world_size(),
)
elif args.dataloader_type == 'cyclic':
batch_sampler = MegatronPretrainingRandomSampler(
dataset,
Expand Down Expand Up @@ -105,8 +90,17 @@ def close_nvidia_fds():
DistributedSignalHandler(args.exit_signal).__enter__()

maybe_worker_init_fn = worker_init_fn if args.num_workers > 0 else None
# Torch dataloader.
if args.dynamic_context_parallel or getattr(args, "use_vanilla_collate_fn", False):
# Identity collate for VarlenDataset and packing-scheduler paths;
# they emit one variable-length dict per sample, not stack-able by
# the default collate. --varlen-sbhd-validation is excluded: it bypasses
# packing and emits fixed-length [seq_length] samples that the default
# collate stacks normally.
if (
(args.use_varlen_dataset and not args.varlen_sbhd_validation)
or args.dynamic_context_parallel
or args.sequence_packing_scheduler is not None
or getattr(args, "use_vanilla_collate_fn", False)
):
extra_kwargs = {"collate_fn": lambda x: x}
else:
extra_kwargs = {}
Expand Down Expand Up @@ -190,71 +184,6 @@ def __iter__(self):
start_idx, end_idx = self.get_start_end_idx()
yield batch[start_idx:end_idx]


class HybridCPMegatronPretrainingSampler(MegatronPretrainingSampler):
"""
Data sampler for hybrid context parallel (Hybrid CP) format.
This data sampler pulls in the entire global batch at once across all data parallel ranks.
This helps provide the Hybrid CP Dataloader Wrapper to schedule and load balance sub-samples
of the entire global batch.
"""

def __init__(
self,
total_samples,
consumed_samples,
micro_batch_size,
global_batch_size,
data_parallel_rank,
data_parallel_size,
drop_last=True,
):
super().__init__(
total_samples,
consumed_samples,
micro_batch_size,
data_parallel_rank,
data_parallel_size,
drop_last,
)
self.global_batch_size = global_batch_size
self.data_parallel_size = data_parallel_size
self.num_micro_batches = self.global_batch_size // self.micro_batch_times_data_parallel_size

def __len__(self):
return self.total_samples

def get_start_end_idx_global_batch(self):
start_idx = [
self.data_parallel_rank * self.micro_batch_size
+ i * self.micro_batch_size * self.data_parallel_size
for i in range(self.num_micro_batches)
]
end_idx = [start_idx[i] + self.micro_batch_size for i in range(self.num_micro_batches)]
return start_idx, end_idx

def __iter__(self):
batch = []
# Last batch will be dropped if drop_last is not set False
for idx in range(self.consumed_samples, self.total_samples):
batch.append(idx)
if len(batch) == self.micro_batch_times_data_parallel_size * self.num_micro_batches:
start_idx, end_idx = self.get_start_end_idx_global_batch()
global_batch_idx = []
for i in range(self.num_micro_batches):
global_batch_idx.extend(batch[start_idx[i] : end_idx[i]])
yield global_batch_idx
batch = []

# Check the last partial batch and see drop_last is set
if len(batch) > 0 and not self.drop_last:
start_idx, end_idx = self.get_start_end_idx_global_batch()
global_batch_idx = []
for i in range(self.num_micro_batches):
global_batch_idx.extend(batch[start_idx[i] : end_idx[i]])
yield global_batch_idx


class MegatronFullValidationSampler:
"""Sampler for full validation that handles small datasets gracefully.

Expand Down
6 changes: 3 additions & 3 deletions megatron/training/datasets/sft_dataset.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,8 +1,7 @@
# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.

import atexit, json
import atexit
from collections import Counter
import json
import math
from typing import Any, Dict, Optional, List, Union

Expand All @@ -14,6 +13,7 @@
from megatron.core.datasets.indexed_dataset import IndexedDataset
from megatron.core.datasets.megatron_dataset import LowLevelDataset, MegatronDataset
from megatron.core.datasets.utils import Split
from megatron.training.datasets.utils import load_json_arg

IGNORE_INDEX = -100

Expand Down Expand Up @@ -335,7 +335,7 @@ def build_low_level_dataset(dataset_path: str, config: GPTDatasetConfig) -> LowL
"lognormal_sigma": 1.1,
}
else:
mock_config = json.loads(config.sft_mock_dataset_config_json)
mock_config = load_json_arg(config.sft_mock_dataset_config_json)
return MockSFTLowLevelDataset(**mock_config)

def __len__(self) -> int:
Expand Down
Loading
Loading