diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py index 2d86c5593ad..29203303ee9 100644 --- a/megatron/training/arguments.py +++ b/megatron/training/arguments.py @@ -79,6 +79,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) @@ -1489,6 +1490,34 @@ def validate_args(args, defaults={}): if args.ckpt_format == "fsdp_dtensor": assert 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: + 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: + # Otherwise fall back to ``dp_balanced`` (static packing). + if args.sequence_packing_scheduler is None: + args.sequence_packing_scheduler = 'dp_balanced' + # Packed-sequence buffer-size check. Placed after varlen scheduler # auto-select so it validates the final resolved scheduler. if args.sequence_packing_scheduler is not None: @@ -3511,6 +3540,47 @@ def _add_sft_args(parser): ) 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``. ' + '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 --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_logits_distillation_args(parser): group = parser.add_argument_group(title='Logits Distillation') diff --git a/pretrain_gpt.py b/pretrain_gpt.py index e2e2c772382..4f66e9dc97d 100644 --- a/pretrain_gpt.py +++ b/pretrain_gpt.py @@ -63,6 +63,7 @@ from megatron.training.arguments import core_transformer_config_from_args, parse_and_validate_args from megatron.training.datasets.fim_dataset import GPTFIMDataset, GPTFIMDatasetConfig from megatron.training.datasets.sft_dataset import MockSFTDataset, SFTDataset +from megatron.training.datasets.varlen_dataset import MockVarlenDataset, VarlenDataset from megatron.training.training import update_seqlen_stats_from_cu_seqlens from megatron.training.utils import get_blend_and_blend_per_split, is_first_or_last_pipeline_stage from model_provider import model_provider @@ -454,6 +455,8 @@ def core_gpt_dataset_config_from_args(args: Any) -> GPTDatasetConfig: "hybrid_context_parallel": args.hybrid_context_parallel, "inter_document_masking": args.dataloader_inter_document_masking, "sft_mock_dataset_config_json": args.sft_mock_dataset_config_json, + "varlen_mock_dataset_config_json": args.varlen_mock_dataset_config_json, + "varlen_sbhd_validation": args.varlen_sbhd_validation, } # add FIM args to the config @@ -497,6 +500,17 @@ def train_valid_test_datasets_provider(train_val_test_num_samples, vp_stage=None else: dataset_type = SFTDataset is_packed_sequence = True # SFT always uses packed sequence + elif args.use_varlen_dataset: + # Variable-length packed (THD) dataset, independent of --sft. + # Reuses SFTDataset's THD packing internally but is gated + # by its own top-level flag. + if args.mock_data: + dataset_type = MockVarlenDataset + else: + dataset_type = VarlenDataset + # SBHD validation mode runs the non-packed pipeline; THD mode + # is the packed-sequence path. + is_packed_sequence = not args.varlen_sbhd_validation else: if args.mock_data: dataset_type = MockGPTDataset diff --git a/tests/unit_tests/data/test_varlen_dataset.py b/tests/unit_tests/data/test_varlen_dataset.py index a7ec459d23e..f60b872a855 100644 --- a/tests/unit_tests/data/test_varlen_dataset.py +++ b/tests/unit_tests/data/test_varlen_dataset.py @@ -592,3 +592,200 @@ def test_mock_getitem_thd_keys_and_pad_fallback(): n = out["tokens"].numel() assert out["labels"].numel() == n and out["loss_mask"].numel() == n assert int(out["padded_seq_len"].item()) % 4 == 0 + + +# ---------------------------------------------------------------------------- +# THD handoff: _unpack_batch contract for VarlenDataset-style samples +# +# VarlenDataset already emits one unpacked sub-sample carrying ``padded_seq_len``, +# so _unpack_batch must short-circuit (no cu_seqlens slicing) and only normalize +# the collate batch dim. SFTDataset-style pre-packed samples (cu_seqlens, no +# padded_seq_len) still take the slicing path. +# ---------------------------------------------------------------------------- + + +def test_unpack_batch_short_circuits_for_varlen_samples(): + from megatron.core.datasets.data_schedule_utils import _unpack_batch + + # Two VarlenDataset-style samples, each already a single sub-sample with a + # leading batch dim (as added by the default collate_fn) and padded_seq_len. + batch = [ + { + "tokens": torch.arange(4, dtype=torch.int64).view(1, 4), + "labels": torch.arange(4, dtype=torch.int64).view(1, 4), + "loss_mask": torch.ones(1, 4), + "position_ids": torch.arange(4, dtype=torch.int64).view(1, 4), + "padded_seq_len": torch.tensor([4], dtype=torch.int32), + }, + { + "tokens": torch.arange(8, dtype=torch.int64).view(1, 8), + "labels": torch.arange(8, dtype=torch.int64).view(1, 8), + "loss_mask": torch.ones(1, 8), + "position_ids": torch.arange(8, dtype=torch.int64).view(1, 8), + "padded_seq_len": torch.tensor([8], dtype=torch.int32), + "original_seq_len": torch.tensor([8], dtype=torch.int32), + }, + ] + out = _unpack_batch(batch) + # Short-circuit: same number of samples (no slicing into sub-samples). + assert len(out) == 2 + # Leading collate batch dim dropped. + assert out[0]["tokens"].shape == (4,) + assert out[1]["tokens"].shape == (8,) + # Missing original_seq_len synthesized from padded_seq_len. + assert "original_seq_len" in out[0] + assert int(out[0]["original_seq_len"].item()) == 4 + # Existing original_seq_len preserved. + assert int(out[1]["original_seq_len"].item()) == 8 + + +def test_unpack_batch_slices_prepacked_cu_seqlens_samples(): + from megatron.core.datasets.data_schedule_utils import _unpack_batch + + # SFTDataset-style pre-packed sample: two sub-sequences [0:3) and [3:5), + # described by cu_seqlens, NO padded_seq_len -> takes the slicing path. + batch = [ + { + "tokens": torch.arange(5, dtype=torch.int64), + "labels": torch.arange(5, dtype=torch.int64), + "loss_mask": torch.ones(5), + "position_ids": torch.arange(5, dtype=torch.int64), + "cu_seqlens": torch.tensor([0, 3, 5], dtype=torch.int32), + } + ] + out = _unpack_batch(batch) + # One packed sample with two sub-sequences -> two unpacked samples. + assert len(out) == 2 + assert out[0]["tokens"].numel() == 3 + assert out[1]["tokens"].numel() == 2 + assert int(out[0]["padded_seq_len"].item()) == 3 + assert int(out[1]["padded_seq_len"].item()) == 2 + + +# ---------------------------------------------------------------------------- +# DataLoader collate selection (distributed; run under torch.distributed.run). +# +# Validates the build_pretraining_data_loader contract for the varlen paths: +# * --varlen-sbhd-validation emits fixed-length [seq_length] samples that the +# DEFAULT collate stacks into a [mbs, seq_length] batch. +# * The THD path (--use-varlen-dataset without SBHD) uses the identity collate +# (variable-length dicts are returned as a list, not stacked). +# ---------------------------------------------------------------------------- + + +def _build_varlen_for_loader(items, config, num_samples): + from megatron.core.datasets.utils import Split + + ds = VarlenDataset.__new__(VarlenDataset) + ds.config = config + ds.dataset = items + ds.indices = np.arange(len(items)) + ds.num_samples = num_samples + ds.index_split = Split.train + return ds + + +def _loader_args(*, use_varlen, sbhd, scheduler, mbs, gbs=None): + return SimpleNamespace( + dataloader_type='single', + micro_batch_size=mbs, + global_batch_size=mbs if gbs is None else gbs, + full_validation=False, + num_workers=0, + use_varlen_dataset=use_varlen, + varlen_sbhd_validation=sbhd, + sequence_packing_scheduler=scheduler, + ) + + +def test_sbhd_validation_dataloader_uses_default_collate(): + from megatron.core import parallel_state + from megatron.training.datasets.data_samplers import build_pretraining_data_loader + from megatron.training.global_vars import destroy_global_vars, set_args + from tests.unit_tests.test_utilities import Utils + + Utils.initialize_model_parallel(1, 1) + try: + tok = _FakeTokenizer(eod=0, pad=7) + seq_len, mbs = 16, 2 + # One global batch needs micro_batch_size * data_parallel_size samples; + # size the dataset off the runtime DP world size so this passes under + # any --nproc-per-node (the CI default is 8 ranks -> dp=8). + dp = parallel_state.get_data_parallel_world_size() + n = mbs * dp * 4 + cfg = _make_config(tok, seq_length=seq_len, sbhd=True) + ds = _build_varlen_for_loader(["hello world"] * n, cfg, num_samples=n) + set_args(_loader_args(use_varlen=True, sbhd=True, scheduler=None, mbs=mbs)) + loader = build_pretraining_data_loader(ds, consumed_samples=0) + batch = next(iter(loader)) + # Default collate stacks fixed-length SBHD samples into a tensor batch. + assert isinstance(batch, dict) + assert batch["tokens"].shape == (mbs, seq_len) + assert batch["labels"].shape == (mbs, seq_len) + assert batch["loss_mask"].shape == (mbs, seq_len) + finally: + destroy_global_vars() + Utils.destroy_model_parallel() + + +def test_thd_dataloader_uses_identity_collate(): + from megatron.core import parallel_state + from megatron.training.datasets.data_samplers import build_pretraining_data_loader + from megatron.training.global_vars import destroy_global_vars, set_args + from tests.unit_tests.test_utilities import Utils + + Utils.initialize_model_parallel(1, 1) + try: + tok = _FakeTokenizer(eod=0, pad=7) + mbs = 2 + dp = parallel_state.get_data_parallel_world_size() + n = mbs * dp * 4 + cfg = _make_config(tok, seq_length=64, sbhd=False) + # Variable-length samples so identity collate is required. + variable = ["a", "abcdef", "xy", "qwerty"] + items = [variable[i % len(variable)] for i in range(n)] + ds = _build_varlen_for_loader(items, cfg, num_samples=n) + set_args(_loader_args(use_varlen=True, sbhd=False, scheduler="dp_balanced", mbs=mbs)) + loader = build_pretraining_data_loader(ds, consumed_samples=0) + batch = next(iter(loader)) + # Identity collate returns the raw list of per-sample dicts (unstacked). + assert isinstance(batch, list) + assert len(batch) == mbs + assert "padded_seq_len" in batch[0] + finally: + destroy_global_vars() + Utils.destroy_model_parallel() + + +def test_packing_scheduler_dataloader_yields_microbatches(): + from megatron.core import parallel_state + from megatron.training.datasets.data_samplers import build_pretraining_data_loader + from megatron.training.global_vars import destroy_global_vars, set_args + from tests.unit_tests.test_utilities import Utils + + Utils.initialize_model_parallel(1, 1) + try: + tok = _FakeTokenizer(eod=0, pad=7) + mbs = 2 + num_microbatches = 3 + dp = parallel_state.get_data_parallel_world_size() + gbs = mbs * dp * num_microbatches + n = gbs * 2 + cfg = _make_config(tok, seq_length=64, dp=dp, cp=1) + variable = ["a", "abcdef", "xy", "qwerty"] + items = [variable[i % len(variable)] for i in range(n)] + ds = _build_varlen_for_loader(items, cfg, num_samples=n) + set_args( + _loader_args(use_varlen=True, sbhd=False, scheduler="dp_balanced", mbs=mbs, gbs=gbs) + ) + loader = build_pretraining_data_loader(ds, consumed_samples=0) + batch = next(iter(loader)) + # The packing scheduler calls next(data_iterator) num_microbatches times; + # each loader step must therefore be one local microbatch, not all + # local samples from the global batch. + assert isinstance(batch, list) + assert len(batch) == mbs + assert "padded_seq_len" in batch[0] + finally: + destroy_global_vars() + Utils.destroy_model_parallel()