From 4a9a9dc864e0cde4f28797f1883c33ac1a43c87a Mon Sep 17 00:00:00 2001 From: tailaim Date: Sat, 16 May 2026 15:08:47 -0700 Subject: [PATCH 1/7] add varlendataset for thd e2e and benchmark Signed-off-by: tailaim --- megatron/core/datasets/data_schedule_utils.py | 23 + megatron/core/datasets/gpt_dataset.py | 12 + megatron/training/arguments.py | 75 +++ megatron/training/datasets/data_samplers.py | 10 +- megatron/training/datasets/varlen_dataset.py | 510 ++++++++++++++++++ pretrain_gpt.py | 18 +- tests/unit_tests/data/test_varlen_dataset.py | 374 +++++++++++++ 7 files changed, 1020 insertions(+), 2 deletions(-) create mode 100644 megatron/training/datasets/varlen_dataset.py create mode 100644 tests/unit_tests/data/test_varlen_dataset.py diff --git a/megatron/core/datasets/data_schedule_utils.py b/megatron/core/datasets/data_schedule_utils.py index c59b1742c0a..6fb583f71d5 100644 --- a/megatron/core/datasets/data_schedule_utils.py +++ b/megatron/core/datasets/data_schedule_utils.py @@ -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 = [] diff --git a/megatron/core/datasets/gpt_dataset.py b/megatron/core/datasets/gpt_dataset.py index 418a02719df..1f5b30b604b 100644 --- a/megatron/core/datasets/gpt_dataset.py +++ b/megatron/core/datasets/gpt_dataset.py @@ -77,6 +77,18 @@ class GPTDatasetConfig(BlendedMegatronDatasetConfig): """The size of the context parallel group. Needed for padding in packed sequences.""" sft_mock_dataset_config_json: Optional[str] = None + + varlen_mock_dataset_config_json: Optional[str] = None + """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_bshd_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 + BSHD reference run that mirrors the THD path's tokenization but skips all + packing — useful for THD numerical-correctness validation.""" """This config provides the necessary information for the mock dataset.""" def __post_init__(self) -> None: diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py index c7ffe840d8b..f5079142327 100644 --- a/megatron/training/arguments.py +++ b/megatron/training/arguments.py @@ -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) @@ -1692,6 +1693,38 @@ 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_bshd_validation: + # BSHD reference mode: each sample is right-padded to + # sequence_length and shipped through the default non-packed + # pipeline. No scheduler / dynamic-cp involved. + assert not args.dynamic_context_parallel, ( + "--varlen-bshd-validation is incompatible with " + "--dynamic-context-parallel (BSHD mode is not packed)." + ) + assert args.sequence_packing_scheduler is None, ( + "--varlen-bshd-validation does not use a sequence packing " + "scheduler; drop --sequence-packing-scheduler." + ) + 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' + # Data blend checks assert ( args.mock_data @@ -4847,6 +4880,48 @@ 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`` 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-bshd-validation', + action="store_true", + help='Reference BSHD 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 JSON for --use-varlen-dataset --mock-data. ' + '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 diff --git a/megatron/training/datasets/data_samplers.py b/megatron/training/datasets/data_samplers.py index 8b14b975aba..e5ab0be3ab6 100644 --- a/megatron/training/datasets/data_samplers.py +++ b/megatron/training/datasets/data_samplers.py @@ -106,7 +106,15 @@ def close_nvidia_fds(): 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): + # All packing-scheduler paths (dynamic_cp + dp_balanced + future schedulers) + # consume variable-length per-sample tensors that the default stack-based + # collate_fn cannot batch. Use an identity collate so the scheduler sees + # a list of dicts and can pack them itself. + if ( + 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 = {} diff --git a/megatron/training/datasets/varlen_dataset.py b/megatron/training/datasets/varlen_dataset.py new file mode 100644 index 00000000000..f4d8b65cfbb --- /dev/null +++ b/megatron/training/datasets/varlen_dataset.py @@ -0,0 +1,510 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +"""Variable-length packed (THD) dataset for SFT-style instruction data. + +This dataset is the entry point for the ``--use-varlen-dataset`` flag. It is +independent of the ``--sft`` flag (no implicit coupling) but shares the same +THD packing / cu_seqlens / dynamic-CP padding logic by extending the existing +:class:`SFTDataset` family. The variable-length aspect is what matters here: +samples have wildly different lengths and are packed into THD format for +training throughput. + +Compared to :class:`SFTDataset`, this dataset adds: + + * **Multi-source loading** — accepts HuggingFace Hub repo ids + (``owner/repo``), local ``.parquet`` files, and local ``.jsonl/.json`` + files; the latter are read via pandas to sidestep pyarrow's per-chunk + JSON schema inference which fails when sample fields vary across rows. + + * **Auto schema detection** — three common instruction-tuning layouts are + auto-detected by column name and normalized to the messages list format + expected by the parent ``SFTDataset.__getitem__``: + + * **openai-messages** — column ``messages`` (Llama post-training, + HuggingFaceH4/no_robots, ...) + * **sharegpt** — column ``conversations`` (OpenOrca, Vicuna, ...) + * **alpaca / dolly** — at least one of + ``instruction|prompt|query|question`` + one of + ``output|response|completion|answer``, plus optional context field + ``input|context``. + + * **Mock variant** — :class:`MockVarlenDataset` mirrors + :class:`MockSFTDataset` end-to-end (synthetic lognormal sequence-length + distribution / fixed-length file / verification mode from an + ``IndexedDataset``), configured via + ``--varlen-mock-dataset-config-json``. + +Limitations (raise a clear ``ValueError`` instead of silently mishandling): + + * Sample content/value must be a plain string — multi-modal content lists + (image+text parts) are not supported. + * Tree-structured (OpenAssistant oasst1) and preference (chosen/rejected) + datasets are out of scope. + * For HF Hub repos, only ``split="train"`` is loaded. +""" + +import json +import os +from typing import Any, Callable, Dict, Iterable, List, Optional, Tuple + +import numpy as np +import torch + +from megatron.core.datasets.gpt_dataset import GPTDatasetConfig +from megatron.core.datasets.megatron_dataset import LowLevelDataset +from megatron.core.datasets.utils import Split +from megatron.training.datasets.sft_dataset import ( + IGNORE_INDEX, + MockSFTDataset, + MockSFTLowLevelDataset, + SFTDataset, + SFTLowLevelDataset, +) + +# Field-name synonyms (probed in order; first non-empty wins). +_INSTRUCTION_FIELDS: Tuple[str, ...] = ( + "instruction", "prompt", "query", "question", +) +_OUTPUT_FIELDS: Tuple[str, ...] = ( + "output", "response", "completion", "answer", +) +# Supplementary user-turn context: Stanford Alpaca's "input", Dolly's "context". +_EXTRA_INPUT_FIELDS: Tuple[str, ...] = ("input", "context") + +# ShareGPT "from" value -> chat-template "role". Unknown values fall back to +# "user" so downstream tokenization does not crash on unfamiliar speakers. +_SHAREGPT_ROLE_MAP: Dict[str, str] = { + "human": "user", + "user": "user", + "gpt": "assistant", + "assistant": "assistant", + "model": "assistant", + "chatgpt": "assistant", + "bing": "assistant", + "bard": "assistant", + "system": "system", + "tool": "tool", + "function": "tool", + "observation": "tool", +} + + +def _looks_like_hf_id(path: str) -> bool: + """Heuristic: does ``path`` look like an ``owner/repo`` HF dataset id? + + True iff ``path`` contains ``/``, is not an absolute/relative file path, + and does not exist on the local filesystem. + """ + if not path: + return False + if os.path.exists(path): + return False + if path.startswith(("/", "./", "../")): + return False + return "/" in path + + +def _first_present( + sample: Dict[str, Any], fields: Iterable[str] +) -> Optional[str]: + """Return the first non-empty string value among the given fields, or None.""" + for f in fields: + v = sample.get(f) + if v in (None, ""): + continue + if not isinstance(v, str): + raise ValueError( + f"VarlenDataset: field '{f}' must be a string, " + f"got {type(v).__name__}." + ) + return v + return None + + +def _ensure_str_content(content: Any, where: str) -> str: + """Validate that a turn's content is a plain string (reject multi-modal lists).""" + if content is None: + return "" + if not isinstance(content, str): + raise ValueError( + f"VarlenDataset: {where} content must be a string, " + f"got {type(content).__name__}. Multi-modal datasets (e.g. " + "content as a list of image/text parts) are not supported." + ) + return content + + +def _alpaca_to_messages(sample: Dict[str, Any]) -> List[Dict[str, str]]: + """Convert an Alpaca/Dolly-style sample to a 3-turn messages list.""" + instruction = _first_present(sample, _INSTRUCTION_FIELDS) or "" + extra_input = _first_present(sample, _EXTRA_INPUT_FIELDS) or "" + output = _first_present(sample, _OUTPUT_FIELDS) or "" + user_content = ( + f"{instruction}\n\n{extra_input}" if extra_input else instruction + ) + return [ + {"role": "system", "content": ""}, + {"role": "user", "content": user_content}, + {"role": "assistant", "content": output}, + ] + + +def _sharegpt_to_messages(sample: Dict[str, Any]) -> List[Dict[str, str]]: + """Convert a ShareGPT ``conversations`` sample to a messages list. + + Prepends an empty ``system`` turn unless the conversation already starts + with one, so ``SFTDataset._split_conversations`` treats the sample as a + single conversation. + """ + conv = sample.get("conversations") or [] + out: List[Dict[str, str]] = [] + first_speaker = (conv[0].get("from") or "").lower() if conv else "" + if first_speaker != "system": + out.append({"role": "system", "content": ""}) + for turn in conv: + speaker = (turn.get("from") or "").lower() + role = _SHAREGPT_ROLE_MAP.get(speaker, "user") + content = _ensure_str_content(turn.get("value"), f"sharegpt turn role={role}") + out.append({"role": role, "content": content}) + return out + + +def _messages_passthrough(sample: Dict[str, Any]) -> List[Dict[str, str]]: + """Pass through an OpenAI ``messages`` sample, ensuring a leading system turn. + + Strips any keys other than ``role``/``content`` (e.g. ``name``, + ``tool_calls``) since they are not part of the chat-template input + expected by SFTTokenizer. + """ + raw = list(sample.get("messages") or []) + if raw and raw[0].get("role") != "system": + raw = [{"role": "system", "content": ""}] + raw + out: List[Dict[str, str]] = [] + for m in raw: + role = m.get("role") or "user" + content = _ensure_str_content(m.get("content"), f"messages turn role={role}") + out.append({"role": role, "content": content}) + return out + + +def _select_converter( + column_names: List[str], +) -> Tuple[Callable[[Dict[str, Any]], List[Dict[str, str]]], str]: + """Pick a sample->messages converter based on dataset column names. + + Priority: openai-messages > sharegpt > alpaca/dolly. + """ + cols = set(column_names) + if "messages" in cols: + return _messages_passthrough, "openai-messages" + if "conversations" in cols: + return _sharegpt_to_messages, "sharegpt" + has_instr = any(f in cols for f in _INSTRUCTION_FIELDS) + has_out = any(f in cols for f in _OUTPUT_FIELDS) + if has_instr and has_out: + return _alpaca_to_messages, "alpaca" + raise ValueError( + "VarlenDataset cannot infer schema from columns " + f"{sorted(cols)}. Supported schemas: " + f"alpaca/dolly ({'|'.join(_INSTRUCTION_FIELDS)} + " + f"{'|'.join(_OUTPUT_FIELDS)} [+ optional {'|'.join(_EXTRA_INPUT_FIELDS)}]), " + "sharegpt (conversations), openai-messages (messages)." + ) + + +class VarlenLowLevelDataset(SFTLowLevelDataset): + """Low-level loader: HF Hub repo / local parquet / local jsonl, normalized. + + Dataset path interpretation: + + * HF Hub repo id (e.g. ``Yukang/LongAlpaca-12k``) — contains ``/`` and + does not exist on the local filesystem; loaded via + ``datasets.load_dataset(path, split="train")``. + * Local ``.parquet`` — loaded via + ``datasets.load_dataset("parquet", data_files=path, split="all")``; + parquet's footer schema makes chunked loading safe. + * Otherwise local jsonl/json — loaded via pandas + ``read_json(lines=True)`` and wrapped in ``Dataset.from_pandas``. + We avoid ``datasets.load_dataset("json", ...)`` for local files + because its pyarrow-based JSON reader infers schema per parallel + chunk and fails with ``CastError`` when the union of fields varies + between rows (e.g. LongAlpaca-12k). + + A sample->messages converter is selected once at construction time based + on column names and applied per-sample at access time. + """ + + def __init__(self, dataset_path: str) -> None: + try: + from datasets import Dataset, load_dataset + except ImportError as exc: + raise ImportError( + "VarlenDataset requires the `datasets` library " + "(pip install datasets)." + ) from exc + + if _looks_like_hf_id(dataset_path): + self.dataset = load_dataset(dataset_path, split="train") + elif dataset_path.endswith(".parquet"): + self.dataset = load_dataset( + "parquet", data_files=dataset_path, split="all" + ) + else: + try: + import pandas as pd + except ImportError as exc: + raise ImportError( + "VarlenDataset requires `pandas` to load local jsonl " + "files (pip install pandas)." + ) from exc + df = pd.read_json(dataset_path, lines=True) + self.dataset = Dataset.from_pandas(df, preserve_index=False) + + self._converter, self._schema_name = _select_converter( + list(self.dataset.column_names) + ) + + @property + def schema_name(self) -> str: + """Detected schema name: ``alpaca`` / ``sharegpt`` / ``openai-messages``.""" + return self._schema_name + + def __len__(self) -> int: + return len(self.dataset) + + def __getitem__(self, idx: int) -> List[Dict[str, str]]: + return self._converter(self.dataset[idx]) + + +class VarlenDataset(SFTDataset): + """Variable-length single-sample SFT dataset for the packed-sequence path. + + Each ``__getitem__`` returns **one tokenized conversation** in unpacked + form: ``tokens``/``labels``/``loss_mask``/``position_ids`` whose length + equals the sample's actual token count (padded to ``pad_granularity``, + NOT to ``sequence_length``), plus ``original_seq_len``/``padded_seq_len`` + tensors that the upstream packing scheduler consumes directly via + :func:`get_batch_and_global_seqlens`. + + This is the schema described in :class:`BasePackingScheduler.get_required_sample_keys`. + It deliberately skips the multi-conversation pre-packing that + :class:`SFTDataset.__getitem__` does, letting the upstream scheduler + pack variable-length samples across the DP×CP grid with no per-sample + padding waste. + + Truncation: samples longer than ``config.sequence_length`` are truncated + on the right; an EOD token is appended if the truncation removed it. + """ + + def __init__( + self, + dataset: LowLevelDataset, + dataset_path: Optional[str], + indices: np.ndarray, + num_samples: Optional[int], + index_split: Split, + config: GPTDatasetConfig, + ) -> None: + super().__init__(dataset, dataset_path, indices, num_samples, index_split, config) + + @staticmethod + def numel_low_level_dataset(low_level_dataset: LowLevelDataset) -> int: + return len(low_level_dataset) + + @staticmethod + def build_low_level_dataset( + dataset_path: str, config: GPTDatasetConfig + ) -> LowLevelDataset: + return VarlenLowLevelDataset(dataset_path) + + def __getitem__(self, idx: int) -> Dict[str, torch.Tensor]: + tokenizer = self.config.tokenizer + max_len = self.config.sequence_length + eod = tokenizer.eod + pad = tokenizer.pad + + # 1. Pull a single conversation (the low-level dataset emits exactly + # one messages list per index — see VarlenLowLevelDataset). + messages = self.dataset[int(self.indices[idx % len(self.indices)])] + + # 2. Tokenize the conversation once; no multi-conv packing here. + tokens, targets = tokenizer.tokenize_conversation( + messages, return_target=True, add_generation_prompt=False + ) + assert not self.config.reset_position_ids + assert not self.config.create_attention_mask and not self.config.reset_attention_mask + + tokens_list = tokens.tolist() + targets_list = targets.tolist() + + # 3. Right-truncate to ``sequence_length + 1`` (we drop the last token + # after the input/label shift below). Keep an EOD at the end so a + # truncated assistant turn still has a valid stop token. + if len(tokens_list) > max_len + 1: + tokens_list = tokens_list[: max_len + 1] + targets_list = targets_list[: max_len + 1] + if tokens_list[-1] != eod: + tokens_list[-1] = eod + targets_list[-1] = eod + + # 4. Ensure EOD is the last token (unconditional for short samples). + if tokens_list[-1] != eod: + tokens_list.append(eod) + targets_list.append(eod) + + # 5a. BSHD validation mode: right-pad to sequence_length + 1, drop + # packing metadata, return shape [sequence_length]. Useful as a + # numerical reference for THD path verification (no scheduler, + # no dynamic-cp). + if self.config.varlen_bshd_validation: + pad_len = max_len + 1 - len(tokens_list) + if pad_len > 0: + tokens_list.extend([pad] * pad_len) + targets_list.extend([pad] * pad_len) + assert len(tokens_list) == max_len + 1 + input_ids = torch.tensor(tokens_list[:-1], dtype=torch.int64) + labels = torch.tensor(targets_list[1:], dtype=torch.int64) + loss_mask = torch.ones(max_len, dtype=torch.float32) + loss_mask[labels == pad] = 0.0 + loss_mask[labels == IGNORE_INDEX] = 0.0 + return { + 'tokens': input_ids, + 'labels': labels, + 'loss_mask': loss_mask, + 'position_ids': torch.arange(max_len, dtype=torch.int64), + } + + original_seq_len = len(tokens_list) - 1 # length after the shift below + + # 5b. THD path: pad to pad_granularity (dp_size * cp_size * 2 * sp), + # the minimum alignment required by CP slicing. We deliberately + # do NOT pad to sequence_length — the upstream packing scheduler + # will combine variable-length samples up to + # max_seqlen_per_dp_cp_rank. + pad_granularity = self._calculate_padding_divisor() + mod = original_seq_len % pad_granularity + if mod != 0: + pad_len = pad_granularity - mod + tokens_list.extend([pad] * pad_len) + targets_list.extend([pad] * pad_len) + padded_seq_len = len(tokens_list) - 1 + + # 6. Apply the next-token shift. + input_ids = torch.tensor(tokens_list[:-1], dtype=torch.int64) + labels = torch.tensor(targets_list[1:], dtype=torch.int64) + position_ids = torch.arange(padded_seq_len, dtype=torch.int64) + loss_mask = torch.ones(padded_seq_len, dtype=torch.float32) + loss_mask[labels == pad] = 0.0 + loss_mask[labels == IGNORE_INDEX] = 0.0 + + return { + 'tokens': input_ids, + 'labels': labels, + 'loss_mask': loss_mask, + 'position_ids': position_ids, + # The packing scheduler consumes these directly; cu_seqlens / + # max_seqlen are produced downstream in _pack_sequences. + 'original_seq_len': torch.tensor([original_seq_len], dtype=torch.int32), + 'padded_seq_len': torch.tensor([padded_seq_len], dtype=torch.int32), + } + + +class MockVarlenDataset(MockSFTDataset): + """Mock variable-length dataset for benchmarking the varlen path. + + Uses :class:`MockSFTLowLevelDataset` for sequence-length sampling (lognormal + distribution / per-line CSV / IndexedDataset verification mode — same JSON + schema as ``--sft-mock-dataset-config-json``, just consumed via + ``--varlen-mock-dataset-config-json``). + + Output shape mirrors :class:`VarlenDataset.__getitem__` (not the inherited + :meth:`MockSFTDataset.__getitem__`) so the mock and real-data paths + exercise exactly the same downstream pipeline: + + * THD mode: emits **one unpacked sample** padded to ``pad_granularity`` + with ``original_seq_len`` / ``padded_seq_len`` tensors. The upstream + scheduler packs across the DP×CP grid. + * BSHD validation mode (``--varlen-bshd-validation``): right-pads to + ``sequence_length`` with no packing metadata, for THD numerical + verification against a non-packed reference run. + """ + + @staticmethod + def build_low_level_dataset( + dataset_path: str, config: GPTDatasetConfig + ) -> LowLevelDataset: + if config.varlen_mock_dataset_config_json is None: + mock_config = { + "mode": "distribution", + "type": "lognormal", + "min_seq_len": config.sequence_length // 2, + "max_seq_len": config.sequence_length, + "mean_seq_len": config.sequence_length // 4 * 3, + "lognormal_sigma": 1.1, + } + else: + mock_config = json.loads(config.varlen_mock_dataset_config_json) + return MockSFTLowLevelDataset(**mock_config) + + def __getitem__(self, idx: int) -> Dict[str, torch.Tensor]: + tokenizer = self.config.tokenizer + max_len = self.config.sequence_length + eod = tokenizer.eod + pad = tokenizer.pad + + # MockSFTLowLevelDataset returns ``length - 1`` token ids; append EOD + # to make the conversation end on a stop token, mirroring the real + # VarlenDataset path. + raw = self.dataset[int(self.indices[idx % len(self.indices)])] + tokens_list = raw.tolist() + tokens_list.append(eod) + # Mock data uses ``tokens == targets`` (no role masking). + targets_list = list(tokens_list) + + # BSHD validation mode: pad to sequence_length + 1, no packing meta. + if self.config.varlen_bshd_validation: + if len(tokens_list) > max_len + 1: + tokens_list = tokens_list[: max_len - 1] + [eod] + targets_list = targets_list[: max_len - 1] + [eod] + pad_len = max_len + 1 - len(tokens_list) + if pad_len > 0: + tokens_list.extend([pad] * pad_len) + targets_list.extend([pad] * pad_len) + assert len(tokens_list) == max_len + 1 + input_ids = torch.tensor(tokens_list[:-1], dtype=torch.int64) + labels = torch.tensor(targets_list[1:], dtype=torch.int64) + loss_mask = torch.ones(max_len, dtype=torch.float32) + loss_mask[labels == pad] = 0.0 + return { + 'tokens': input_ids, + 'labels': labels, + 'loss_mask': loss_mask, + 'position_ids': torch.arange(max_len, dtype=torch.int64), + } + + # THD mode: unpacked single sample, pad to pad_granularity only. + if len(tokens_list) > max_len + 1: + tokens_list = tokens_list[: max_len - 1] + [eod] + targets_list = targets_list[: max_len - 1] + [eod] + original_seq_len = len(tokens_list) - 1 + + pad_granularity = self._calculate_padding_divisor() + mod = original_seq_len % pad_granularity + if mod != 0: + pad_len = pad_granularity - mod + tokens_list.extend([pad] * pad_len) + targets_list.extend([pad] * pad_len) + padded_seq_len = len(tokens_list) - 1 + + input_ids = torch.tensor(tokens_list[:-1], dtype=torch.int64) + labels = torch.tensor(targets_list[1:], dtype=torch.int64) + loss_mask = torch.ones(padded_seq_len, dtype=torch.float32) + loss_mask[labels == pad] = 0.0 + return { + 'tokens': input_ids, + 'labels': labels, + 'loss_mask': loss_mask, + 'position_ids': torch.arange(padded_seq_len, dtype=torch.int64), + 'original_seq_len': torch.tensor([original_seq_len], dtype=torch.int32), + 'padded_seq_len': torch.tensor([padded_seq_len], dtype=torch.int32), + } diff --git a/pretrain_gpt.py b/pretrain_gpt.py index d337ecde2cd..dc5c83329ad 100644 --- a/pretrain_gpt.py +++ b/pretrain_gpt.py @@ -51,6 +51,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.utils import ( get_batch_on_this_cp_rank, get_batch_on_this_tp_rank, @@ -133,7 +134,9 @@ def get_batch(data_iterator, vp_stage: Optional[int] = None): ) # TODO: this is pretty hacky, find a better way - is_packed_sequence = get_args().sft # SFT always uses packed sequence + is_packed_sequence = args.sft or ( + args.use_varlen_dataset and not args.varlen_bshd_validation + ) if ( not is_first_or_last_pipeline_stage(vp_stage) and not is_packed_sequence @@ -349,6 +352,8 @@ def core_gpt_dataset_config_from_args(args: Any) -> GPTDatasetConfig: "context_parallel_size": args.context_parallel_size, "data_parallel_size": args.data_parallel_size, "sequence_parallel_size": args.tensor_model_parallel_size * args.sequence_parallel, + "varlen_mock_dataset_config_json": args.varlen_mock_dataset_config_json, + "varlen_bshd_validation": args.varlen_bshd_validation, } # add FIM args to the config @@ -392,6 +397,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/dynamic-cp packing internally but is gated + # by its own top-level flag. + if args.mock_data: + dataset_type = MockVarlenDataset + else: + dataset_type = VarlenDataset + # BSHD validation mode runs the SBHD non-packed pipeline; THD mode + # is the packed-sequence path. + is_packed_sequence = not args.varlen_bshd_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 new file mode 100644 index 00000000000..6e9df8b0be0 --- /dev/null +++ b/tests/unit_tests/data/test_varlen_dataset.py @@ -0,0 +1,374 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +"""Unit tests for :mod:`megatron.training.datasets.varlen_dataset`. + +These tests cover the schema-detection and message-normalization helpers and +the :class:`VarlenLowLevelDataset` loader. The end-to-end SFTDataset packing +behavior is exercised by the existing SFT test suite; here we focus on the +varlen-specific contracts (auto-detect schema, normalize to messages, +ValueError on unsupported shapes). +""" + +import json +from pathlib import Path + +import pytest + +# Import via the public module path so this test gets discovered through the +# regular pytest entry point. The functions under test are pure Python and do +# not require torch.distributed. +from megatron.training.datasets.varlen_dataset import ( + VarlenLowLevelDataset, + _alpaca_to_messages, + _looks_like_hf_id, + _messages_passthrough, + _select_converter, + _sharegpt_to_messages, +) + +# ---------------------------------------------------------------------------- +# _looks_like_hf_id heuristic +# ---------------------------------------------------------------------------- + + +@pytest.mark.parametrize( + "path,expected", + [ + ("Yukang/LongAlpaca-12k", True), + ("HuggingFaceH4/no_robots", True), + ("databricks/databricks-dolly-15k", True), + ("/tmp/foo.jsonl", False), + ("./local.jsonl", False), + ("../up.jsonl", False), + ("singlename", False), + ("", False), + (None, False), + ], +) +def test_looks_like_hf_id(path, expected): + assert _looks_like_hf_id(path) is expected + + +# ---------------------------------------------------------------------------- +# Schema converters +# ---------------------------------------------------------------------------- + + +def test_alpaca_canonical_with_input(): + out = _alpaca_to_messages( + {"instruction": "Summarize.", "input": "Long passage", "output": "It says X."} + ) + assert [m["role"] for m in out] == ["system", "user", "assistant"] + assert out[1]["content"] == "Summarize.\n\nLong passage" + assert out[2]["content"] == "It says X." + + +def test_alpaca_without_input(): + out = _alpaca_to_messages({"instruction": "Hi.", "output": "Hello."}) + assert out[0] == {"role": "system", "content": ""} + assert out[1]["content"] == "Hi." + assert out[2]["content"] == "Hello." + + +@pytest.mark.parametrize( + "instr_key,out_key", + [ + ("prompt", "response"), + ("query", "answer"), + ("question", "completion"), + ("instruction", "answer"), + ], +) +def test_alpaca_field_synonyms(instr_key, out_key): + out = _alpaca_to_messages({instr_key: "Q?", out_key: "A."}) + assert out[1]["content"] == "Q?" + assert out[2]["content"] == "A." + + +def test_dolly_instruction_context_response(): + """Dolly-15k: instruction + context + response, all via synonyms.""" + out = _alpaca_to_messages( + { + "instruction": "Who wrote 1984?", + "context": "1984 was written in 1948.", + "response": "George Orwell.", + } + ) + assert out[1]["content"] == "Who wrote 1984?\n\n1984 was written in 1948." + assert out[2]["content"] == "George Orwell." + + +def test_sharegpt_human_gpt(): + out = _sharegpt_to_messages( + { + "conversations": [ + {"from": "human", "value": "hi"}, + {"from": "gpt", "value": "hello"}, + ] + } + ) + assert [m["role"] for m in out] == ["system", "user", "assistant"] + assert out[1]["content"] == "hi" + + +def test_sharegpt_preserves_existing_system_turn(): + out = _sharegpt_to_messages( + { + "conversations": [ + {"from": "system", "value": "be terse"}, + {"from": "human", "value": "hi"}, + {"from": "gpt", "value": "hello"}, + ] + } + ) + assert [m["role"] for m in out] == ["system", "user", "assistant"] + assert out[0]["content"] == "be terse" + + +@pytest.mark.parametrize( + "speaker,expected_role", + [ + ("human", "user"), + ("user", "user"), + ("gpt", "assistant"), + ("assistant", "assistant"), + ("model", "assistant"), + ("chatgpt", "assistant"), + ("tool", "tool"), + ("function", "tool"), + ("alien", "user"), # unknown speakers fall back to user + ], +) +def test_sharegpt_role_map(speaker, expected_role): + out = _sharegpt_to_messages({"conversations": [{"from": speaker, "value": "x"}]}) + # First entry is the prepended system turn; second is the actual content. + assert out[1]["role"] == expected_role + + +def test_messages_passthrough_prepends_system_when_missing(): + out = _messages_passthrough( + { + "messages": [ + {"role": "user", "content": "hi"}, + {"role": "assistant", "content": "hello"}, + ] + } + ) + assert [m["role"] for m in out] == ["system", "user", "assistant"] + + +def test_messages_passthrough_keeps_existing_system(): + out = _messages_passthrough( + { + "messages": [ + {"role": "system", "content": "be terse"}, + {"role": "user", "content": "hi"}, + {"role": "assistant", "content": "ok"}, + ] + } + ) + assert [m["role"] for m in out] == ["system", "user", "assistant"] + assert out[0]["content"] == "be terse" + + +def test_messages_passthrough_strips_extra_keys(): + """OpenAI-style messages may carry ``name`` / ``tool_calls`` etc.; + chat-template input only wants ``role`` and ``content``.""" + out = _messages_passthrough( + { + "messages": [ + {"role": "user", "content": "hi", "name": "alice"}, + { + "role": "assistant", + "content": "hi alice", + "tool_calls": [{"function": "foo"}], + }, + ] + } + ) + for m in out: + assert set(m.keys()) == {"role", "content"} + + +# ---------------------------------------------------------------------------- +# Shape validation: reject multi-modal / non-string content +# ---------------------------------------------------------------------------- + + +def test_messages_rejects_list_content(): + with pytest.raises(ValueError, match="must be a string"): + _messages_passthrough( + { + "messages": [ + {"role": "user", "content": [{"type": "image", "url": "x.png"}]}, + ] + } + ) + + +def test_alpaca_rejects_non_string_field(): + with pytest.raises(ValueError, match="must be a string"): + _alpaca_to_messages({"instruction": ["a", "b"], "output": "x"}) + + +def test_sharegpt_rejects_list_value(): + with pytest.raises(ValueError, match="must be a string"): + _sharegpt_to_messages( + {"conversations": [{"from": "human", "value": [1, 2, 3]}]} + ) + + +# ---------------------------------------------------------------------------- +# Schema selector priority +# ---------------------------------------------------------------------------- + + +def test_select_converter_alpaca(): + fn, name = _select_converter(["instruction", "output", "file"]) + assert name == "alpaca" + assert fn is _alpaca_to_messages + + +def test_select_converter_alpaca_via_synonyms(): + fn, name = _select_converter(["prompt", "response"]) + assert name == "alpaca" + + +def test_select_converter_dolly_columns(): + fn, name = _select_converter(["instruction", "context", "response", "category"]) + assert name == "alpaca" + + +def test_select_converter_sharegpt(): + fn, name = _select_converter(["conversations", "id"]) + assert name == "sharegpt" + assert fn is _sharegpt_to_messages + + +def test_select_converter_messages(): + fn, name = _select_converter(["messages"]) + assert name == "openai-messages" + assert fn is _messages_passthrough + + +def test_select_converter_priority_messages_over_alpaca(): + # When both ``messages`` and alpaca-style columns are present, the more + # explicit ``messages`` schema wins. + fn, name = _select_converter(["messages", "instruction", "output"]) + assert name == "openai-messages" + + +def test_select_converter_unrecognized_columns(): + with pytest.raises(ValueError, match="cannot infer schema"): + _select_converter(["foo", "bar"]) + + +def test_select_converter_alpaca_missing_output(): + """Having an instruction column but no output column is not a match.""" + with pytest.raises(ValueError, match="cannot infer schema"): + _select_converter(["instruction", "category"]) + + +# ---------------------------------------------------------------------------- +# VarlenLowLevelDataset on local jsonl (no HF Hub network needed) +# ---------------------------------------------------------------------------- + + +def _write_jsonl(tmp_path: Path, rows): + p = tmp_path / "data.jsonl" + with p.open("w") as f: + for row in rows: + f.write(json.dumps(row) + "\n") + return str(p) + + +def test_low_level_loads_jsonl_alpaca(tmp_path): + pytest.importorskip("datasets") + pytest.importorskip("pandas") + path = _write_jsonl( + tmp_path, + [ + {"instruction": "i1", "output": "o1"}, + {"instruction": "i2", "output": "o2", "file": "extra"}, + ], + ) + ll = VarlenLowLevelDataset(path) + assert len(ll) == 2 + assert ll.schema_name == "alpaca" + sample = ll[0] + assert [m["role"] for m in sample] == ["system", "user", "assistant"] + assert sample[1]["content"] == "i1" + assert sample[2]["content"] == "o1" + + +def test_low_level_loads_jsonl_sharegpt(tmp_path): + pytest.importorskip("datasets") + pytest.importorskip("pandas") + path = _write_jsonl( + tmp_path, + [ + { + "conversations": [ + {"from": "human", "value": "q1"}, + {"from": "gpt", "value": "a1"}, + ] + }, + { + "conversations": [ + {"from": "human", "value": "q2"}, + {"from": "gpt", "value": "a2"}, + ] + }, + ], + ) + ll = VarlenLowLevelDataset(path) + assert len(ll) == 2 + assert ll.schema_name == "sharegpt" + sample = ll[1] + # system prepended + 2 turns from the conversation + assert [m["role"] for m in sample] == ["system", "user", "assistant"] + assert sample[1]["content"] == "q2" + + +def test_low_level_loads_jsonl_messages(tmp_path): + pytest.importorskip("datasets") + pytest.importorskip("pandas") + path = _write_jsonl( + tmp_path, + [ + { + "messages": [ + {"role": "user", "content": "hi"}, + {"role": "assistant", "content": "hello"}, + ] + }, + ], + ) + ll = VarlenLowLevelDataset(path) + assert ll.schema_name == "openai-messages" + sample = ll[0] + assert [m["role"] for m in sample] == ["system", "user", "assistant"] + + +def test_low_level_jsonl_heterogeneous_columns(tmp_path): + """Real datasets often mix rows that have / lack an optional field. Our + pandas-based loader must accept the union schema without ``CastError``.""" + pytest.importorskip("datasets") + pytest.importorskip("pandas") + rows = [{"instruction": "a", "output": "x"}] * 100 + [ + {"instruction": "b", "output": "y", "file": "extra"} + ] * 100 + path = _write_jsonl(tmp_path, rows) + ll = VarlenLowLevelDataset(path) + assert len(ll) == 200 + # Both halves should normalize to the same messages structure. + assert [m["role"] for m in ll[0]] == ["system", "user", "assistant"] + assert [m["role"] for m in ll[150]] == ["system", "user", "assistant"] + + +def test_low_level_rejects_unknown_schema(tmp_path): + pytest.importorskip("datasets") + pytest.importorskip("pandas") + path = _write_jsonl(tmp_path, [{"foo": "bar"}]) + with pytest.raises(ValueError, match="cannot infer schema"): + VarlenLowLevelDataset(path) From 23327854a3ab7cc4f7324f6bd01f9e229fe2eea8 Mon Sep 17 00:00:00 2001 From: tailaim Date: Thu, 28 May 2026 01:00:12 -0700 Subject: [PATCH 2/7] add support for dataset like https://huggingface.co/datasets/allenai/dolma3_longmino_mix-100B-1125 Signed-off-by: tailaim --- megatron/training/datasets/varlen_dataset.py | 69 ++++++++--- tests/unit_tests/data/test_varlen_dataset.py | 113 ++++++++++++------- 2 files changed, 129 insertions(+), 53 deletions(-) diff --git a/megatron/training/datasets/varlen_dataset.py b/megatron/training/datasets/varlen_dataset.py index f4d8b65cfbb..cb755115909 100644 --- a/megatron/training/datasets/varlen_dataset.py +++ b/megatron/training/datasets/varlen_dataset.py @@ -187,12 +187,32 @@ def _messages_passthrough(sample: Dict[str, Any]) -> List[Dict[str, str]]: return out +def _raw_text_loader(sample: Dict[str, Any]) -> str: + """Return the ``text`` column unchanged for pretrain-style packed runs. + + Unlike the SFT schemas this returns a plain string (no messages list). + :class:`VarlenDataset.__getitem__` dispatches on the return type to pick + a tokenization path that skips chat templating and prompt masking. + """ + text = sample.get("text") or "" + if not isinstance(text, str): + raise ValueError( + f"VarlenDataset (pretrain-text schema): 'text' must be a string, " + f"got {type(text).__name__}." + ) + return text + + def _select_converter( column_names: List[str], -) -> Tuple[Callable[[Dict[str, Any]], List[Dict[str, str]]], str]: - """Pick a sample->messages converter based on dataset column names. - - Priority: openai-messages > sharegpt > alpaca/dolly. +) -> Tuple[Callable[[Dict[str, Any]], Any], str]: + """Pick a sample converter based on dataset column names. + + Priority (most explicit first): openai-messages > sharegpt > alpaca/dolly + > pretrain-text. ``pretrain-text`` is the fallback for datasets that + only carry a single ``text`` column (e.g. Dolma / OLMo midtraining + corpora) — long-context pretraining packed through the same THD path + as SFT. """ cols = set(column_names) if "messages" in cols: @@ -203,12 +223,15 @@ def _select_converter( has_out = any(f in cols for f in _OUTPUT_FIELDS) if has_instr and has_out: return _alpaca_to_messages, "alpaca" + if "text" in cols: + return _raw_text_loader, "pretrain-text" raise ValueError( "VarlenDataset cannot infer schema from columns " f"{sorted(cols)}. Supported schemas: " f"alpaca/dolly ({'|'.join(_INSTRUCTION_FIELDS)} + " f"{'|'.join(_OUTPUT_FIELDS)} [+ optional {'|'.join(_EXTRA_INPUT_FIELDS)}]), " - "sharegpt (conversations), openai-messages (messages)." + "sharegpt (conversations), openai-messages (messages), " + "pretrain-text (text)." ) @@ -320,22 +343,38 @@ def build_low_level_dataset( def __getitem__(self, idx: int) -> Dict[str, torch.Tensor]: tokenizer = self.config.tokenizer max_len = self.config.sequence_length + # HuggingFaceTokenizer returns None for ``pad`` when the underlying + # tokenizer has no explicit pad token (common for raw pretraining + # tokenizers like Qwen3). Fall back to eod for padding — irrelevant + # for loss because loss_mask zeros pad positions out. eod = tokenizer.eod - pad = tokenizer.pad + pad = tokenizer.pad if tokenizer.pad is not None else eod + assert eod is not None, ( + "VarlenDataset requires the tokenizer to expose an EOD/EOS token id." + ) - # 1. Pull a single conversation (the low-level dataset emits exactly - # one messages list per index — see VarlenLowLevelDataset). - messages = self.dataset[int(self.indices[idx % len(self.indices)])] + # 1. Pull a single item from the low-level dataset. For SFT schemas + # (alpaca / sharegpt / openai-messages) this is a messages list; + # for the pretrain-text schema it is a raw string. + item = self.dataset[int(self.indices[idx % len(self.indices)])] - # 2. Tokenize the conversation once; no multi-conv packing here. - tokens, targets = tokenizer.tokenize_conversation( - messages, return_target=True, add_generation_prompt=False - ) assert not self.config.reset_position_ids assert not self.config.create_attention_mask and not self.config.reset_attention_mask - tokens_list = tokens.tolist() - targets_list = targets.tolist() + # 2. Tokenize. SFT schemas go through tokenize_conversation (chat + # template + role-aware target masking); pretrain-text bypasses + # chat templating and uses the plain ``tokenize`` interface, + # treating every token as a target (no prompt masking). + if isinstance(item, str): + ids = list(tokenizer.tokenize(item)) + tokens_list = ids + targets_list = list(ids) + else: + tokens, targets = tokenizer.tokenize_conversation( + item, return_target=True, add_generation_prompt=False + ) + tokens_list = tokens.tolist() + targets_list = targets.tolist() # 3. Right-truncate to ``sequence_length + 1`` (we drop the last token # after the input/label shift below). Keep an EOD at the end so a diff --git a/tests/unit_tests/data/test_varlen_dataset.py b/tests/unit_tests/data/test_varlen_dataset.py index 6e9df8b0be0..bd8cd80dacc 100644 --- a/tests/unit_tests/data/test_varlen_dataset.py +++ b/tests/unit_tests/data/test_varlen_dataset.py @@ -22,6 +22,7 @@ _alpaca_to_messages, _looks_like_hf_id, _messages_passthrough, + _raw_text_loader, _select_converter, _sharegpt_to_messages, ) @@ -100,12 +101,7 @@ def test_dolly_instruction_context_response(): def test_sharegpt_human_gpt(): out = _sharegpt_to_messages( - { - "conversations": [ - {"from": "human", "value": "hi"}, - {"from": "gpt", "value": "hello"}, - ] - } + {"conversations": [{"from": "human", "value": "hi"}, {"from": "gpt", "value": "hello"}]} ) assert [m["role"] for m in out] == ["system", "user", "assistant"] assert out[1]["content"] == "hi" @@ -147,12 +143,7 @@ def test_sharegpt_role_map(speaker, expected_role): def test_messages_passthrough_prepends_system_when_missing(): out = _messages_passthrough( - { - "messages": [ - {"role": "user", "content": "hi"}, - {"role": "assistant", "content": "hello"}, - ] - } + {"messages": [{"role": "user", "content": "hi"}, {"role": "assistant", "content": "hello"}]} ) assert [m["role"] for m in out] == ["system", "user", "assistant"] @@ -178,11 +169,7 @@ def test_messages_passthrough_strips_extra_keys(): { "messages": [ {"role": "user", "content": "hi", "name": "alice"}, - { - "role": "assistant", - "content": "hi alice", - "tool_calls": [{"function": "foo"}], - }, + {"role": "assistant", "content": "hi alice", "tool_calls": [{"function": "foo"}]}, ] } ) @@ -198,11 +185,7 @@ def test_messages_passthrough_strips_extra_keys(): def test_messages_rejects_list_content(): with pytest.raises(ValueError, match="must be a string"): _messages_passthrough( - { - "messages": [ - {"role": "user", "content": [{"type": "image", "url": "x.png"}]}, - ] - } + {"messages": [{"role": "user", "content": [{"type": "image", "url": "x.png"}]}]} ) @@ -213,9 +196,29 @@ def test_alpaca_rejects_non_string_field(): def test_sharegpt_rejects_list_value(): with pytest.raises(ValueError, match="must be a string"): - _sharegpt_to_messages( - {"conversations": [{"from": "human", "value": [1, 2, 3]}]} - ) + _sharegpt_to_messages({"conversations": [{"from": "human", "value": [1, 2, 3]}]}) + + +# ---------------------------------------------------------------------------- +# Pretrain-text schema +# ---------------------------------------------------------------------------- + + +def test_raw_text_loader_returns_string(): + """``text``-column samples are returned as plain strings (not messages).""" + out = _raw_text_loader({"text": "Once upon a time...", "id": "doc-1"}) + assert isinstance(out, str) + assert out == "Once upon a time..." + + +def test_raw_text_loader_handles_empty(): + assert _raw_text_loader({"text": None}) == "" + assert _raw_text_loader({}) == "" + + +def test_raw_text_rejects_non_string(): + with pytest.raises(ValueError, match="must be a string"): + _raw_text_loader({"text": [1, 2, 3]}) # ---------------------------------------------------------------------------- @@ -269,6 +272,30 @@ def test_select_converter_alpaca_missing_output(): _select_converter(["instruction", "category"]) +def test_select_converter_pretrain_text(): + fn, name = _select_converter(["text", "id"]) + assert name == "pretrain-text" + assert fn is _raw_text_loader + + +def test_select_converter_pretrain_text_with_metadata(): + """Real corpora (e.g. Dolma) have ``text`` + ``url`` + ``metadata``.""" + fn, name = _select_converter(["text", "url", "metadata", "id"]) + assert name == "pretrain-text" + + +def test_select_converter_alpaca_beats_pretrain_text(): + """When both ``instruction``/``output`` and ``text`` are present (rare), + the alpaca schema is more specific and should win.""" + fn, name = _select_converter(["text", "instruction", "output"]) + assert name == "alpaca" + + +def test_select_converter_messages_beats_pretrain_text(): + fn, name = _select_converter(["text", "messages"]) + assert name == "openai-messages" + + # ---------------------------------------------------------------------------- # VarlenLowLevelDataset on local jsonl (no HF Hub network needed) # ---------------------------------------------------------------------------- @@ -307,18 +334,8 @@ def test_low_level_loads_jsonl_sharegpt(tmp_path): path = _write_jsonl( tmp_path, [ - { - "conversations": [ - {"from": "human", "value": "q1"}, - {"from": "gpt", "value": "a1"}, - ] - }, - { - "conversations": [ - {"from": "human", "value": "q2"}, - {"from": "gpt", "value": "a2"}, - ] - }, + {"conversations": [{"from": "human", "value": "q1"}, {"from": "gpt", "value": "a1"}]}, + {"conversations": [{"from": "human", "value": "q2"}, {"from": "gpt", "value": "a2"}]}, ], ) ll = VarlenLowLevelDataset(path) @@ -341,7 +358,7 @@ def test_low_level_loads_jsonl_messages(tmp_path): {"role": "user", "content": "hi"}, {"role": "assistant", "content": "hello"}, ] - }, + } ], ) ll = VarlenLowLevelDataset(path) @@ -372,3 +389,23 @@ def test_low_level_rejects_unknown_schema(tmp_path): path = _write_jsonl(tmp_path, [{"foo": "bar"}]) with pytest.raises(ValueError, match="cannot infer schema"): VarlenLowLevelDataset(path) + + +def test_low_level_loads_jsonl_pretrain_text(tmp_path): + """Pretrain-text corpora (Dolma / OLMo midtraining) typically have + ``text`` + extra fields like ``id`` / ``url`` / ``metadata``.""" + pytest.importorskip("datasets") + pytest.importorskip("pandas") + path = _write_jsonl( + tmp_path, + [ + {"text": "Doc one body...", "id": "1", "url": "https://x/1"}, + {"text": "Doc two body...", "id": "2", "url": "https://x/2"}, + ], + ) + ll = VarlenLowLevelDataset(path) + assert ll.schema_name == "pretrain-text" + assert len(ll) == 2 + # Each item is a raw string, NOT a messages list. + assert ll[0] == "Doc one body..." + assert ll[1] == "Doc two body..." From a85be943422835335f9f48865ad271a1871118b2 Mon Sep 17 00:00:00 2001 From: tailaim Date: Thu, 28 May 2026 08:53:28 -0700 Subject: [PATCH 3/7] minor fixes according to the comments Signed-off-by: tailaim --- megatron/core/datasets/gpt_dataset.py | 6 ++++ megatron/training/arguments.py | 26 +++++++++++------- megatron/training/datasets/data_samplers.py | 11 ++++---- megatron/training/datasets/sft_dataset.py | 6 ++-- megatron/training/datasets/utils.py | 29 ++++++++++++++++++++ megatron/training/datasets/varlen_dataset.py | 4 +-- 6 files changed, 61 insertions(+), 21 deletions(-) create mode 100644 megatron/training/datasets/utils.py diff --git a/megatron/core/datasets/gpt_dataset.py b/megatron/core/datasets/gpt_dataset.py index 1f5b30b604b..5345ee95170 100644 --- a/megatron/core/datasets/gpt_dataset.py +++ b/megatron/core/datasets/gpt_dataset.py @@ -101,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_bshd_validation: + assert not self.dynamic_context_parallel, ( + "--varlen-bshd-validation is incompatible with " + "--dynamic-context-parallel (BSHD mode is not packed)." + ) + self.token_dtype_code = ( None if self.tokenizer.vocab_size is None diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py index f5079142327..53a9c922ca1 100644 --- a/megatron/training/arguments.py +++ b/megatron/training/arguments.py @@ -1702,13 +1702,10 @@ def validate_args(args, defaults={}): "select the packed-sequence dataset family. Pick one." ) if args.varlen_bshd_validation: - # BSHD reference mode: each sample is right-padded to - # sequence_length and shipped through the default non-packed - # pipeline. No scheduler / dynamic-cp involved. - assert not args.dynamic_context_parallel, ( - "--varlen-bshd-validation is incompatible with " - "--dynamic-context-parallel (BSHD mode is not packed)." - ) + # ``--dynamic-context-parallel`` ⊥ ``--varlen-bshd-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-bshd-validation does not use a sequence packing " "scheduler; drop --sequence-packing-scheduler." @@ -4873,7 +4870,15 @@ 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.', ) @@ -4909,8 +4914,9 @@ def _add_varlen_dataset_args(parser): '--varlen-mock-dataset-config-json', type=str, default=None, - help='Mock-dataset config JSON for --use-varlen-dataset --mock-data. ' - 'Same schema as --sft-mock-dataset-config-json: either ' + 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 ' diff --git a/megatron/training/datasets/data_samplers.py b/megatron/training/datasets/data_samplers.py index e5ab0be3ab6..2c9761f704f 100644 --- a/megatron/training/datasets/data_samplers.py +++ b/megatron/training/datasets/data_samplers.py @@ -105,13 +105,12 @@ 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. - # All packing-scheduler paths (dynamic_cp + dp_balanced + future schedulers) - # consume variable-length per-sample tensors that the default stack-based - # collate_fn cannot batch. Use an identity collate so the scheduler sees - # a list of dicts and can pack them itself. + # Identity collate for VarlenDataset and packing-scheduler paths; + # they emit one variable-length dict per sample, not stack-able by + # the default collate. if ( - args.dynamic_context_parallel + args.use_varlen_dataset + or args.dynamic_context_parallel or args.sequence_packing_scheduler is not None or getattr(args, "use_vanilla_collate_fn", False) ): diff --git a/megatron/training/datasets/sft_dataset.py b/megatron/training/datasets/sft_dataset.py index 250a0137568..65db30d44c3 100644 --- a/megatron/training/datasets/sft_dataset.py +++ b/megatron/training/datasets/sft_dataset.py @@ -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 @@ -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 @@ -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: diff --git a/megatron/training/datasets/utils.py b/megatron/training/datasets/utils.py new file mode 100644 index 00000000000..1fe6d7ef83e --- /dev/null +++ b/megatron/training/datasets/utils.py @@ -0,0 +1,29 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +"""Shared utilities for training-side dataset helpers.""" + +import json +import os +from typing import Any, Optional + + +def load_json_arg(spec: Optional[str]) -> Optional[Any]: + """Parse a CLI JSON argument that may be either a JSON literal or a path + to a JSON file. + + The argument is interpreted as a file path when ``spec`` points to an + existing regular file on the local filesystem; otherwise it is parsed as + a JSON literal string. Returns ``None`` when ``spec`` itself is ``None``, + so callers can use it transparently for optional CLI flags. + + Used by the ``--sft-mock-dataset-config-json`` and + ``--varlen-mock-dataset-config-json`` flags, which both accept either an + inline JSON snippet or the path to a file containing the same JSON + document. + """ + if spec is None: + return None + if os.path.isfile(spec): + with open(spec, "r") as f: + return json.load(f) + return json.loads(spec) diff --git a/megatron/training/datasets/varlen_dataset.py b/megatron/training/datasets/varlen_dataset.py index cb755115909..62e56cc1fda 100644 --- a/megatron/training/datasets/varlen_dataset.py +++ b/megatron/training/datasets/varlen_dataset.py @@ -43,7 +43,6 @@ * For HF Hub repos, only ``split="train"`` is loaded. """ -import json import os from typing import Any, Callable, Dict, Iterable, List, Optional, Tuple @@ -60,6 +59,7 @@ SFTDataset, SFTLowLevelDataset, ) +from megatron.training.datasets.utils import load_json_arg # Field-name synonyms (probed in order; first non-empty wins). _INSTRUCTION_FIELDS: Tuple[str, ...] = ( @@ -482,7 +482,7 @@ def build_low_level_dataset( "lognormal_sigma": 1.1, } else: - mock_config = json.loads(config.varlen_mock_dataset_config_json) + mock_config = load_json_arg(config.varlen_mock_dataset_config_json) return MockSFTLowLevelDataset(**mock_config) def __getitem__(self, idx: int) -> Dict[str, torch.Tensor]: From 5e9b3e23e69c81efc3734c626dd5ea2fe03095a4 Mon Sep 17 00:00:00 2001 From: xiaoyao0115 <1804647152@qq.com> Date: Fri, 5 Jun 2026 00:50:44 -0700 Subject: [PATCH 4/7] some fixes according to the comments Signed-off-by: xiaoyao0115 <1804647152@qq.com> --- megatron/core/datasets/gpt_dataset.py | 2 +- megatron/training/arguments.py | 30 +- megatron/training/datasets/data_samplers.py | 6 +- megatron/training/datasets/varlen_dataset.py | 68 ++-- tests/unit_tests/data/test_varlen_dataset.py | 336 +++++++++++++++++++ 5 files changed, 395 insertions(+), 47 deletions(-) diff --git a/megatron/core/datasets/gpt_dataset.py b/megatron/core/datasets/gpt_dataset.py index 5345ee95170..9cc4bb9b3f8 100644 --- a/megatron/core/datasets/gpt_dataset.py +++ b/megatron/core/datasets/gpt_dataset.py @@ -77,6 +77,7 @@ class GPTDatasetConfig(BlendedMegatronDatasetConfig): """The size of the context parallel group. Needed for padding in packed sequences.""" 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 """Mock-dataset config (same JSON schema as ``sft_mock_dataset_config_json``) @@ -89,7 +90,6 @@ class GPTDatasetConfig(BlendedMegatronDatasetConfig): ``padded_seq_len``), bypassing the packed-sequence path. Used to obtain a BSHD reference run that mirrors the THD path's tokenization but skips all packing — useful for THD numerical-correctness validation.""" - """This config provides the necessary information for the mock dataset.""" def __post_init__(self) -> None: """Do asserts and set fields post init""" diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py index 53a9c922ca1..566eb0de91a 100644 --- a/megatron/training/arguments.py +++ b/megatron/training/arguments.py @@ -1555,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 ( @@ -1710,6 +1700,12 @@ def validate_args(args, defaults={}): "--varlen-bshd-validation does not use a sequence packing " "scheduler; drop --sequence-packing-scheduler." ) + # BSHD validation is a real-data numerical-reference path only; + # MockVarlenDataset does not implement it. + assert not args.mock_data, ( + "--varlen-bshd-validation is not supported with --mock-data; " + "BSHD validation requires a real dataset." + ) else: # VarlenDataset emits one unpacked sample per __getitem__; it # relies on an upstream packing scheduler to group variable-length @@ -1722,6 +1718,20 @@ def validate_args(args, defaults={}): if args.sequence_packing_scheduler is None: args.sequence_packing_scheduler = 'dp_balanced' + # 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 diff --git a/megatron/training/datasets/data_samplers.py b/megatron/training/datasets/data_samplers.py index 2c9761f704f..0d6a791c0c5 100644 --- a/megatron/training/datasets/data_samplers.py +++ b/megatron/training/datasets/data_samplers.py @@ -107,9 +107,11 @@ def close_nvidia_fds(): maybe_worker_init_fn = worker_init_fn if args.num_workers > 0 else None # Identity collate for VarlenDataset and packing-scheduler paths; # they emit one variable-length dict per sample, not stack-able by - # the default collate. + # the default collate. --varlen-bshd-validation is excluded: it bypasses + # packing and emits fixed-length [seq_length] samples that the default + # collate stacks normally. if ( - args.use_varlen_dataset + (args.use_varlen_dataset and not args.varlen_bshd_validation) or args.dynamic_context_parallel or args.sequence_packing_scheduler is not None or getattr(args, "use_vanilla_collate_fn", False) diff --git a/megatron/training/datasets/varlen_dataset.py b/megatron/training/datasets/varlen_dataset.py index 62e56cc1fda..28c9e1b133c 100644 --- a/megatron/training/datasets/varlen_dataset.py +++ b/megatron/training/datasets/varlen_dataset.py @@ -16,9 +16,11 @@ files; the latter are read via pandas to sidestep pyarrow's per-chunk JSON schema inference which fails when sample fields vary across rows. - * **Auto schema detection** — three common instruction-tuning layouts are - auto-detected by column name and normalized to the messages list format - expected by the parent ``SFTDataset.__getitem__``: + * **Auto schema detection** — four input layouts are auto-detected by column + name. The three instruction-tuning layouts are normalized to the messages + list format expected by the parent ``SFTDataset.__getitem__``; the + ``pretrain-text`` fallback instead returns a raw string handled separately + in :meth:`VarlenDataset.__getitem__`: * **openai-messages** — column ``messages`` (Llama post-training, HuggingFaceH4/no_robots, ...) @@ -27,6 +29,8 @@ ``instruction|prompt|query|question`` + one of ``output|response|completion|answer``, plus optional context field ``input|context``. + * **pretrain-text** — column ``text``; returns the raw string (no + messages list, no role masking), tokenized as plain pretraining text. * **Mock variant** — :class:`MockVarlenDataset` mirrors :class:`MockSFTDataset` end-to-end (synthetic lognormal sequence-length @@ -253,8 +257,10 @@ class VarlenLowLevelDataset(SFTLowLevelDataset): chunk and fails with ``CastError`` when the union of fields varies between rows (e.g. LongAlpaca-12k). - A sample->messages converter is selected once at construction time based - on column names and applied per-sample at access time. + A per-sample converter is selected once at construction time based on + column names and applied at access time. The instruction-tuning schemas + convert to a messages list; the ``pretrain-text`` fallback returns the raw + string instead. """ def __init__(self, dataset_path: str) -> None: @@ -289,7 +295,8 @@ def __init__(self, dataset_path: str) -> None: @property def schema_name(self) -> str: - """Detected schema name: ``alpaca`` / ``sharegpt`` / ``openai-messages``.""" + """Detected schema name: ``alpaca`` / ``sharegpt`` / ``openai-messages`` / + ``pretrain-text`` (the raw ``text``-column fallback).""" return self._schema_name def __len__(self) -> int: @@ -376,6 +383,15 @@ def __getitem__(self, idx: int) -> Dict[str, torch.Tensor]: tokens_list = tokens.tolist() targets_list = targets.tolist() + # 2b. Guard against an empty tokenization (e.g. a blank ``pretrain-text`` + # row where ``tokenizer.tokenize("")`` returns no ids). Represent it + # as a single end-of-document token so the next-token shift still + # yields a valid 1-token sample instead of raising on + # ``tokens_list[-1]`` below or producing a zero-length sequence. + if len(tokens_list) == 0: + tokens_list = [eod, eod] + targets_list = [eod, eod] + # 3. Right-truncate to ``sequence_length + 1`` (we drop the last token # after the input/label shift below). Keep an EOD at the end so a # truncated assistant turn still has a valid stop token. @@ -391,6 +407,8 @@ def __getitem__(self, idx: int) -> Dict[str, torch.Tensor]: tokens_list.append(eod) targets_list.append(eod) + valid_len = len(tokens_list) - 1 + # 5a. BSHD validation mode: right-pad to sequence_length + 1, drop # packing metadata, return shape [sequence_length]. Useful as a # numerical reference for THD path verification (no scheduler, @@ -404,7 +422,7 @@ def __getitem__(self, idx: int) -> Dict[str, torch.Tensor]: input_ids = torch.tensor(tokens_list[:-1], dtype=torch.int64) labels = torch.tensor(targets_list[1:], dtype=torch.int64) loss_mask = torch.ones(max_len, dtype=torch.float32) - loss_mask[labels == pad] = 0.0 + loss_mask[valid_len:] = 0.0 # mask the right-padded tail by position loss_mask[labels == IGNORE_INDEX] = 0.0 return { 'tokens': input_ids, @@ -433,7 +451,7 @@ def __getitem__(self, idx: int) -> Dict[str, torch.Tensor]: labels = torch.tensor(targets_list[1:], dtype=torch.int64) position_ids = torch.arange(padded_seq_len, dtype=torch.int64) loss_mask = torch.ones(padded_seq_len, dtype=torch.float32) - loss_mask[labels == pad] = 0.0 + loss_mask[valid_len:] = 0.0 # mask the right-padded tail by position loss_mask[labels == IGNORE_INDEX] = 0.0 return { @@ -463,9 +481,9 @@ class MockVarlenDataset(MockSFTDataset): * THD mode: emits **one unpacked sample** padded to ``pad_granularity`` with ``original_seq_len`` / ``padded_seq_len`` tensors. The upstream scheduler packs across the DP×CP grid. - * BSHD validation mode (``--varlen-bshd-validation``): right-pads to - ``sequence_length`` with no packing metadata, for THD numerical - verification against a non-packed reference run. + + ``--varlen-bshd-validation`` is intentionally not implemented for mock + data; it is guarded against in argument validation. """ @staticmethod @@ -489,7 +507,7 @@ def __getitem__(self, idx: int) -> Dict[str, torch.Tensor]: tokenizer = self.config.tokenizer max_len = self.config.sequence_length eod = tokenizer.eod - pad = tokenizer.pad + pad = tokenizer.pad if tokenizer.pad is not None else eod # MockSFTLowLevelDataset returns ``length - 1`` token ids; append EOD # to make the conversation end on a stop token, mirroring the real @@ -500,27 +518,9 @@ def __getitem__(self, idx: int) -> Dict[str, torch.Tensor]: # Mock data uses ``tokens == targets`` (no role masking). targets_list = list(tokens_list) - # BSHD validation mode: pad to sequence_length + 1, no packing meta. - if self.config.varlen_bshd_validation: - if len(tokens_list) > max_len + 1: - tokens_list = tokens_list[: max_len - 1] + [eod] - targets_list = targets_list[: max_len - 1] + [eod] - pad_len = max_len + 1 - len(tokens_list) - if pad_len > 0: - tokens_list.extend([pad] * pad_len) - targets_list.extend([pad] * pad_len) - assert len(tokens_list) == max_len + 1 - input_ids = torch.tensor(tokens_list[:-1], dtype=torch.int64) - labels = torch.tensor(targets_list[1:], dtype=torch.int64) - loss_mask = torch.ones(max_len, dtype=torch.float32) - loss_mask[labels == pad] = 0.0 - return { - 'tokens': input_ids, - 'labels': labels, - 'loss_mask': loss_mask, - 'position_ids': torch.arange(max_len, dtype=torch.int64), - } - + # MockVarlenDataset only implements the THD (packed) path; BSHD + # validation is a real-data numerical-reference mode (guarded against + # --mock-data in validate_args). # THD mode: unpacked single sample, pad to pad_granularity only. if len(tokens_list) > max_len + 1: tokens_list = tokens_list[: max_len - 1] + [eod] @@ -538,7 +538,7 @@ def __getitem__(self, idx: int) -> Dict[str, torch.Tensor]: input_ids = torch.tensor(tokens_list[:-1], dtype=torch.int64) labels = torch.tensor(targets_list[1:], dtype=torch.int64) loss_mask = torch.ones(padded_seq_len, dtype=torch.float32) - loss_mask[labels == pad] = 0.0 + loss_mask[original_seq_len:] = 0.0 # mask the right-padded tail by position return { 'tokens': input_ids, 'labels': labels, diff --git a/tests/unit_tests/data/test_varlen_dataset.py b/tests/unit_tests/data/test_varlen_dataset.py index bd8cd80dacc..277917ab22d 100644 --- a/tests/unit_tests/data/test_varlen_dataset.py +++ b/tests/unit_tests/data/test_varlen_dataset.py @@ -11,13 +11,19 @@ import json from pathlib import Path +from types import SimpleNamespace +import numpy as np import pytest +import torch # Import via the public module path so this test gets discovered through the # regular pytest entry point. The functions under test are pure Python and do # not require torch.distributed. +from megatron.training.datasets.sft_dataset import IGNORE_INDEX from megatron.training.datasets.varlen_dataset import ( + MockVarlenDataset, + VarlenDataset, VarlenLowLevelDataset, _alpaca_to_messages, _looks_like_hf_id, @@ -409,3 +415,333 @@ def test_low_level_loads_jsonl_pretrain_text(tmp_path): # Each item is a raw string, NOT a messages list. assert ll[0] == "Doc one body..." assert ll[1] == "Doc two body..." + + +# ---------------------------------------------------------------------------- +# VarlenDataset / MockVarlenDataset __getitem__ (fake tokenizer, no GPU) +# +# These bypass the heavy SFTDataset.__init__ and inject the minimal attributes +# __getitem__ reads, so the EOD handling / position-based loss masking / +# pad-to-divisor / packing-metadata contracts can be unit tested without a +# real tokenizer or torch.distributed. +# ---------------------------------------------------------------------------- + + +class _FakeTokenizer: + """Minimal tokenizer for exercising VarlenDataset.__getitem__. + + ``tokenize`` maps each character to a non-zero id (so plain text never + collides with ``eod``/``pad``); ``tokenize("")`` returns ``[]`` to exercise + the empty-row guard. ``tokenize_conversation`` masks non-assistant turns + with ``IGNORE_INDEX`` in the targets. + """ + + def __init__(self, eod: int = 0, pad=None): + self._eod = eod + self._pad = pad + + @property + def eod(self): + return self._eod + + @property + def pad(self): + return self._pad + + def tokenize(self, text): + return [ord(c) % 100 + 1 for c in text] # always >= 1, never eod (0) + + def tokenize_conversation(self, messages, return_target=True, add_generation_prompt=False): + tokens, targets = [], [] + for m in messages: + ids = self.tokenize(m["content"]) + tokens.extend(ids) + # Only assistant turns contribute to the loss; prompt is masked. + targets.extend(ids if m["role"] == "assistant" else [IGNORE_INDEX] * len(ids)) + return (torch.tensor(tokens, dtype=torch.int64), torch.tensor(targets, dtype=torch.int64)) + + +def _make_config(tokenizer, seq_length=64, *, cp=1, dp=1, sp=1, dynamic_cp=False, bshd=False): + return SimpleNamespace( + tokenizer=tokenizer, + sequence_length=seq_length, + reset_position_ids=False, + create_attention_mask=False, + reset_attention_mask=False, + varlen_bshd_validation=bshd, + dynamic_context_parallel=dynamic_cp, + data_parallel_size=dp, + context_parallel_size=cp, + sequence_parallel_size=sp, + ) + + +def _make_varlen(items, config): + ds = VarlenDataset.__new__(VarlenDataset) + ds.config = config + ds.dataset = items + ds.indices = np.arange(len(items)) + return ds + + +def _make_mock_varlen(token_arrays, config): + ds = MockVarlenDataset.__new__(MockVarlenDataset) + ds.config = config + ds.dataset = token_arrays # each item exposes .tolist() + ds.indices = np.arange(len(token_arrays)) + return ds + + +def test_getitem_thd_pretrain_text_keys_and_shapes(): + tok = _FakeTokenizer(eod=0, pad=7) + ds = _make_varlen(["hello world"], _make_config(tok, seq_length=64)) + out = ds[0] + assert set(out) == { + "tokens", + "labels", + "loss_mask", + "position_ids", + "original_seq_len", + "padded_seq_len", + } + n = out["tokens"].numel() + assert out["labels"].numel() == n + assert out["loss_mask"].numel() == n + assert out["position_ids"].numel() == n + assert int(out["padded_seq_len"].item()) == n + + +def test_getitem_thd_sft_prompt_is_masked(): + tok = _FakeTokenizer(eod=0, pad=7) + messages = [ + {"role": "system", "content": ""}, + {"role": "user", "content": "question"}, + {"role": "assistant", "content": "answer"}, + ] + ds = _make_varlen([messages], _make_config(tok, seq_length=64)) + out = ds[0] + # Prompt (user) tokens are IGNORE_INDEX in labels and must be masked out; + # assistant tokens must contribute to the loss. + labels = out["labels"] + loss_mask = out["loss_mask"] + assert torch.all(loss_mask[labels == IGNORE_INDEX] == 0.0) + assert loss_mask.sum() > 0 # assistant span still contributes + + +def test_getitem_thd_pad_masked_by_position_keeps_real_eod(): + """Regression: with pad falling back to eod, the real end-of-document EOD + target must stay in the loss (masked by position, not by value).""" + tok = _FakeTokenizer(eod=0, pad=None) # pad falls back to eod + # cp=2 -> pad divisor = cp*2 = 4, so a 3-token doc gets a padding tail. + ds = _make_varlen(["abc"], _make_config(tok, seq_length=64, cp=2)) + out = ds[0] + loss_mask = out["loss_mask"].tolist() + labels = out["labels"].tolist() + # tokens=[a,b,c,eod] padded to 4 -> labels=[b,c,eod,eod(pad)] + assert len(loss_mask) == 4 + # index 2 is the real end-of-document EOD target -> kept (would be wrongly + # dropped by value-based ``labels == pad`` masking). + assert labels[2] == tok.eod and loss_mask[2] == 1.0 + # index 3 is the appended pad -> masked. + assert loss_mask[3] == 0.0 + + +def test_getitem_thd_padded_to_divisor(): + tok = _FakeTokenizer(eod=0, pad=7) + ds = _make_varlen(["abcde"], _make_config(tok, seq_length=64, cp=2)) # divisor 4 + out = ds[0] + assert int(out["padded_seq_len"].item()) % 4 == 0 + + +def test_getitem_thd_empty_text_does_not_crash(): + """A blank pretrain-text row tokenizes to [] -> must not crash and must + yield a valid (non-zero-length) sample.""" + tok = _FakeTokenizer(eod=0, pad=7) + ds = _make_varlen([""], _make_config(tok, seq_length=64)) + out = ds[0] + assert out["tokens"].numel() >= 1 + assert out["labels"].numel() == out["tokens"].numel() + assert out["loss_mask"].numel() == out["tokens"].numel() + + +def test_getitem_bshd_pads_to_seq_length_and_masks_tail(): + tok = _FakeTokenizer(eod=0, pad=None) + ds = _make_varlen(["abc"], _make_config(tok, seq_length=8, bshd=True)) + out = ds[0] + # BSHD emits fixed [seq_length] samples with no packing metadata. + assert set(out) == {"tokens", "labels", "loss_mask", "position_ids"} + assert out["tokens"].numel() == 8 + loss_mask = out["loss_mask"].tolist() + # tokens=[a,b,c,eod]: valid_len=3 -> first 3 kept (incl. real eod), rest masked. + assert loss_mask[0:3] == [1.0, 1.0, 1.0] + assert all(v == 0.0 for v in loss_mask[3:]) + + +def test_mock_getitem_thd_keys_and_pad_fallback(): + tok = _FakeTokenizer(eod=0, pad=None) # exercise the eod fallback (no crash) + ds = _make_mock_varlen([np.array([1, 2, 3, 4], dtype=np.int64)], _make_config(tok, cp=2)) + out = ds[0] + assert set(out) == { + "tokens", + "labels", + "loss_mask", + "position_ids", + "original_seq_len", + "padded_seq_len", + } + 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-bshd-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 BSHD) 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, bshd, scheduler, mbs): + return SimpleNamespace( + dataloader_type='single', + micro_batch_size=mbs, + global_batch_size=mbs, + full_validation=False, + dynamic_context_parallel=False, + num_workers=0, + use_varlen_dataset=use_varlen, + varlen_bshd_validation=bshd, + sequence_packing_scheduler=scheduler, + ) + + +def test_bshd_validation_dataloader_uses_default_collate(): + 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 + cfg = _make_config(tok, seq_length=seq_len, bshd=True) + ds = _build_varlen_for_loader(["hello world"] * 8, cfg, num_samples=8) + set_args(_loader_args(use_varlen=True, bshd=True, scheduler=None, mbs=mbs)) + loader = build_pretraining_data_loader(ds, consumed_samples=0) + batch = next(iter(loader)) + # Default collate stacks fixed-length BSHD 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.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 + cfg = _make_config(tok, seq_length=64, bshd=False) + # Variable-length samples so identity collate is required. + ds = _build_varlen_for_loader(["a", "abcdef", "xy", "qwerty"] * 2, cfg, num_samples=8) + set_args(_loader_args(use_varlen=True, bshd=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() From 5051382144abf84b32a59f1aff6361a2c08abc61 Mon Sep 17 00:00:00 2001 From: xiaoyao0115 <1804647152@qq.com> Date: Sat, 6 Jun 2026 00:59:58 -0700 Subject: [PATCH 5/7] fix ut Signed-off-by: xiaoyao0115 <1804647152@qq.com> --- tests/unit_tests/data/test_varlen_dataset.py | 15 +++++++++++++-- 1 file changed, 13 insertions(+), 2 deletions(-) diff --git a/tests/unit_tests/data/test_varlen_dataset.py b/tests/unit_tests/data/test_varlen_dataset.py index 277917ab22d..8bafb67d639 100644 --- a/tests/unit_tests/data/test_varlen_dataset.py +++ b/tests/unit_tests/data/test_varlen_dataset.py @@ -700,6 +700,7 @@ def _loader_args(*, use_varlen, bshd, scheduler, mbs): def test_bshd_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 @@ -708,8 +709,13 @@ def test_bshd_validation_dataloader_uses_default_collate(): 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, bshd=True) - ds = _build_varlen_for_loader(["hello world"] * 8, cfg, num_samples=8) + ds = _build_varlen_for_loader(["hello world"] * n, cfg, num_samples=n) set_args(_loader_args(use_varlen=True, bshd=True, scheduler=None, mbs=mbs)) loader = build_pretraining_data_loader(ds, consumed_samples=0) batch = next(iter(loader)) @@ -724,6 +730,7 @@ def test_bshd_validation_dataloader_uses_default_collate(): 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 @@ -732,9 +739,13 @@ def test_thd_dataloader_uses_identity_collate(): 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, bshd=False) # Variable-length samples so identity collate is required. - ds = _build_varlen_for_loader(["a", "abcdef", "xy", "qwerty"] * 2, cfg, num_samples=8) + 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, bshd=False, scheduler="dp_balanced", mbs=mbs)) loader = build_pretraining_data_loader(ds, consumed_samples=0) batch = next(iter(loader)) From 14d94d4b2c1f2d13119ae40dd9bbde83a7b8e2b3 Mon Sep 17 00:00:00 2001 From: xiaoyao0115 <1804647152@qq.com> Date: Sun, 7 Jun 2026 20:52:01 -0700 Subject: [PATCH 6/7] minor fixes Signed-off-by: xiaoyao0115 <1804647152@qq.com> --- megatron/training/datasets/data_samplers.py | 98 ++------------------ pretrain_gpt.py | 1 + tests/unit_tests/data/test_varlen_dataset.py | 47 +++++++++- 3 files changed, 54 insertions(+), 92 deletions(-) diff --git a/megatron/training/datasets/data_samplers.py b/megatron/training/datasets/data_samplers.py index 0d6a791c0c5..bfaf47580f1 100644 --- a/megatron/training/datasets/data_samplers.py +++ b/megatron/training/datasets/data_samplers.py @@ -41,12 +41,6 @@ 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), @@ -54,24 +48,15 @@ def build_pretraining_data_loader(dataset, consumed_samples): 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, @@ -199,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. diff --git a/pretrain_gpt.py b/pretrain_gpt.py index dc5c83329ad..bfc70cd2071 100644 --- a/pretrain_gpt.py +++ b/pretrain_gpt.py @@ -352,6 +352,7 @@ def core_gpt_dataset_config_from_args(args: Any) -> GPTDatasetConfig: "context_parallel_size": args.context_parallel_size, "data_parallel_size": args.data_parallel_size, "sequence_parallel_size": args.tensor_model_parallel_size * args.sequence_parallel, + "dynamic_context_parallel": args.dynamic_context_parallel, "varlen_mock_dataset_config_json": args.varlen_mock_dataset_config_json, "varlen_bshd_validation": args.varlen_bshd_validation, } diff --git a/tests/unit_tests/data/test_varlen_dataset.py b/tests/unit_tests/data/test_varlen_dataset.py index 8bafb67d639..01b4dac483c 100644 --- a/tests/unit_tests/data/test_varlen_dataset.py +++ b/tests/unit_tests/data/test_varlen_dataset.py @@ -685,13 +685,13 @@ def _build_varlen_for_loader(items, config, num_samples): return ds -def _loader_args(*, use_varlen, bshd, scheduler, mbs): +def _loader_args(*, use_varlen, bshd, scheduler, mbs, dynamic_cp=False, gbs=None): return SimpleNamespace( dataloader_type='single', micro_batch_size=mbs, - global_batch_size=mbs, + global_batch_size=mbs if gbs is None else gbs, full_validation=False, - dynamic_context_parallel=False, + dynamic_context_parallel=dynamic_cp, num_workers=0, use_varlen_dataset=use_varlen, varlen_bshd_validation=bshd, @@ -756,3 +756,44 @@ def test_thd_dataloader_uses_identity_collate(): finally: destroy_global_vars() Utils.destroy_model_parallel() + + +def test_dcp_dataloader_yields_microbatches_for_scheduler(): + 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, dynamic_cp=True) + 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, + bshd=False, + scheduler="default_dynamic_cp", + mbs=mbs, + dynamic_cp=True, + gbs=gbs, + ) + ) + loader = build_pretraining_data_loader(ds, consumed_samples=0) + batch = next(iter(loader)) + # The DCP 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() From 0e48a3bcbb538b5f356495c218893d3c8a420b1c Mon Sep 17 00:00:00 2001 From: xiaoyao0115 <1804647152@qq.com> Date: Sun, 7 Jun 2026 21:03:01 -0700 Subject: [PATCH 7/7] reanme varlen_bshd_validation to arlen_sbhd_validation Signed-off-by: xiaoyao0115 <1804647152@qq.com> --- megatron/core/datasets/gpt_dataset.py | 10 +++--- megatron/training/arguments.py | 16 +++++----- megatron/training/datasets/data_samplers.py | 4 +-- megatron/training/datasets/varlen_dataset.py | 8 ++--- pretrain_gpt.py | 8 ++--- tests/unit_tests/data/test_varlen_dataset.py | 32 ++++++++++---------- 6 files changed, 39 insertions(+), 39 deletions(-) diff --git a/megatron/core/datasets/gpt_dataset.py b/megatron/core/datasets/gpt_dataset.py index 9cc4bb9b3f8..e34d930adbf 100644 --- a/megatron/core/datasets/gpt_dataset.py +++ b/megatron/core/datasets/gpt_dataset.py @@ -84,11 +84,11 @@ class GPTDatasetConfig(BlendedMegatronDatasetConfig): used by the ``--use-varlen-dataset`` path; kept separate so the varlen path does not implicitly inherit SFT-specific knobs.""" - varlen_bshd_validation: bool = False + 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 - BSHD reference run that mirrors the THD path's tokenization but skips all + 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: @@ -101,10 +101,10 @@ def __post_init__(self) -> None: assert self.reset_attention_mask is not None assert self.eod_mask_loss is not None - if self.varlen_bshd_validation: + if self.varlen_sbhd_validation: assert not self.dynamic_context_parallel, ( - "--varlen-bshd-validation is incompatible with " - "--dynamic-context-parallel (BSHD mode is not packed)." + "--varlen-sbhd-validation is incompatible with " + "--dynamic-context-parallel (SBHD mode is not packed)." ) self.token_dtype_code = ( diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py index 566eb0de91a..ca051838f42 100644 --- a/megatron/training/arguments.py +++ b/megatron/training/arguments.py @@ -1691,20 +1691,20 @@ def validate_args(args, defaults={}): "--use-varlen-dataset and --sft are mutually exclusive; both " "select the packed-sequence dataset family. Pick one." ) - if args.varlen_bshd_validation: - # ``--dynamic-context-parallel`` ⊥ ``--varlen-bshd-validation`` is + 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-bshd-validation does not use a sequence packing " + "--varlen-sbhd-validation does not use a sequence packing " "scheduler; drop --sequence-packing-scheduler." ) - # BSHD validation is a real-data numerical-reference path only; + # SBHD validation is a real-data numerical-reference path only; # MockVarlenDataset does not implement it. assert not args.mock_data, ( - "--varlen-bshd-validation is not supported with --mock-data; " - "BSHD validation requires a real dataset." + "--varlen-sbhd-validation is not supported with --mock-data; " + "SBHD validation requires a real dataset." ) else: # VarlenDataset emits one unpacked sample per __getitem__; it @@ -4911,9 +4911,9 @@ def _add_varlen_dataset_args(parser): 'distribution; see --varlen-mock-dataset-config-json.', ) group.add_argument( - '--varlen-bshd-validation', + '--varlen-sbhd-validation', action="store_true", - help='Reference BSHD mode for THD numerical verification. When set, ' + 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. ' diff --git a/megatron/training/datasets/data_samplers.py b/megatron/training/datasets/data_samplers.py index bfaf47580f1..4c4501d5de9 100644 --- a/megatron/training/datasets/data_samplers.py +++ b/megatron/training/datasets/data_samplers.py @@ -92,11 +92,11 @@ def close_nvidia_fds(): maybe_worker_init_fn = worker_init_fn if args.num_workers > 0 else None # Identity collate for VarlenDataset and packing-scheduler paths; # they emit one variable-length dict per sample, not stack-able by - # the default collate. --varlen-bshd-validation is excluded: it bypasses + # 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_bshd_validation) + (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) diff --git a/megatron/training/datasets/varlen_dataset.py b/megatron/training/datasets/varlen_dataset.py index 28c9e1b133c..616fc7edcb0 100644 --- a/megatron/training/datasets/varlen_dataset.py +++ b/megatron/training/datasets/varlen_dataset.py @@ -409,11 +409,11 @@ def __getitem__(self, idx: int) -> Dict[str, torch.Tensor]: valid_len = len(tokens_list) - 1 - # 5a. BSHD validation mode: right-pad to sequence_length + 1, drop + # 5a. SBHD validation mode: right-pad to sequence_length + 1, drop # packing metadata, return shape [sequence_length]. Useful as a # numerical reference for THD path verification (no scheduler, # no dynamic-cp). - if self.config.varlen_bshd_validation: + if self.config.varlen_sbhd_validation: pad_len = max_len + 1 - len(tokens_list) if pad_len > 0: tokens_list.extend([pad] * pad_len) @@ -482,7 +482,7 @@ class MockVarlenDataset(MockSFTDataset): with ``original_seq_len`` / ``padded_seq_len`` tensors. The upstream scheduler packs across the DP×CP grid. - ``--varlen-bshd-validation`` is intentionally not implemented for mock + ``--varlen-sbhd-validation`` is intentionally not implemented for mock data; it is guarded against in argument validation. """ @@ -518,7 +518,7 @@ def __getitem__(self, idx: int) -> Dict[str, torch.Tensor]: # Mock data uses ``tokens == targets`` (no role masking). targets_list = list(tokens_list) - # MockVarlenDataset only implements the THD (packed) path; BSHD + # MockVarlenDataset only implements the THD (packed) path; SBHD # validation is a real-data numerical-reference mode (guarded against # --mock-data in validate_args). # THD mode: unpacked single sample, pad to pad_granularity only. diff --git a/pretrain_gpt.py b/pretrain_gpt.py index bfc70cd2071..3674d7f88e2 100644 --- a/pretrain_gpt.py +++ b/pretrain_gpt.py @@ -135,7 +135,7 @@ def get_batch(data_iterator, vp_stage: Optional[int] = None): # TODO: this is pretty hacky, find a better way is_packed_sequence = args.sft or ( - args.use_varlen_dataset and not args.varlen_bshd_validation + args.use_varlen_dataset and not args.varlen_sbhd_validation ) if ( not is_first_or_last_pipeline_stage(vp_stage) @@ -354,7 +354,7 @@ def core_gpt_dataset_config_from_args(args: Any) -> GPTDatasetConfig: "sequence_parallel_size": args.tensor_model_parallel_size * args.sequence_parallel, "dynamic_context_parallel": args.dynamic_context_parallel, "varlen_mock_dataset_config_json": args.varlen_mock_dataset_config_json, - "varlen_bshd_validation": args.varlen_bshd_validation, + "varlen_sbhd_validation": args.varlen_sbhd_validation, } # add FIM args to the config @@ -406,9 +406,9 @@ def train_valid_test_datasets_provider(train_val_test_num_samples, vp_stage=None dataset_type = MockVarlenDataset else: dataset_type = VarlenDataset - # BSHD validation mode runs the SBHD non-packed pipeline; THD mode + # SBHD validation mode runs the non-packed pipeline; THD mode # is the packed-sequence path. - is_packed_sequence = not args.varlen_bshd_validation + 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 01b4dac483c..c158a365b46 100644 --- a/tests/unit_tests/data/test_varlen_dataset.py +++ b/tests/unit_tests/data/test_varlen_dataset.py @@ -461,14 +461,14 @@ def tokenize_conversation(self, messages, return_target=True, add_generation_pro return (torch.tensor(tokens, dtype=torch.int64), torch.tensor(targets, dtype=torch.int64)) -def _make_config(tokenizer, seq_length=64, *, cp=1, dp=1, sp=1, dynamic_cp=False, bshd=False): +def _make_config(tokenizer, seq_length=64, *, cp=1, dp=1, sp=1, dynamic_cp=False, sbhd=False): return SimpleNamespace( tokenizer=tokenizer, sequence_length=seq_length, reset_position_ids=False, create_attention_mask=False, reset_attention_mask=False, - varlen_bshd_validation=bshd, + varlen_sbhd_validation=sbhd, dynamic_context_parallel=dynamic_cp, data_parallel_size=dp, context_parallel_size=cp, @@ -564,11 +564,11 @@ def test_getitem_thd_empty_text_does_not_crash(): assert out["loss_mask"].numel() == out["tokens"].numel() -def test_getitem_bshd_pads_to_seq_length_and_masks_tail(): +def test_getitem_sbhd_pads_to_seq_length_and_masks_tail(): tok = _FakeTokenizer(eod=0, pad=None) - ds = _make_varlen(["abc"], _make_config(tok, seq_length=8, bshd=True)) + ds = _make_varlen(["abc"], _make_config(tok, seq_length=8, sbhd=True)) out = ds[0] - # BSHD emits fixed [seq_length] samples with no packing metadata. + # SBHD emits fixed [seq_length] samples with no packing metadata. assert set(out) == {"tokens", "labels", "loss_mask", "position_ids"} assert out["tokens"].numel() == 8 loss_mask = out["loss_mask"].tolist() @@ -666,9 +666,9 @@ def test_unpack_batch_slices_prepacked_cu_seqlens_samples(): # DataLoader collate selection (distributed; run under torch.distributed.run). # # Validates the build_pretraining_data_loader contract for the varlen paths: -# * --varlen-bshd-validation emits fixed-length [seq_length] samples that the +# * --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 BSHD) uses the identity collate +# * The THD path (--use-varlen-dataset without SBHD) uses the identity collate # (variable-length dicts are returned as a list, not stacked). # ---------------------------------------------------------------------------- @@ -685,7 +685,7 @@ def _build_varlen_for_loader(items, config, num_samples): return ds -def _loader_args(*, use_varlen, bshd, scheduler, mbs, dynamic_cp=False, gbs=None): +def _loader_args(*, use_varlen, sbhd, scheduler, mbs, dynamic_cp=False, gbs=None): return SimpleNamespace( dataloader_type='single', micro_batch_size=mbs, @@ -694,12 +694,12 @@ def _loader_args(*, use_varlen, bshd, scheduler, mbs, dynamic_cp=False, gbs=None dynamic_context_parallel=dynamic_cp, num_workers=0, use_varlen_dataset=use_varlen, - varlen_bshd_validation=bshd, + varlen_sbhd_validation=sbhd, sequence_packing_scheduler=scheduler, ) -def test_bshd_validation_dataloader_uses_default_collate(): +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 @@ -714,12 +714,12 @@ def test_bshd_validation_dataloader_uses_default_collate(): # 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, bshd=True) + 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, bshd=True, scheduler=None, mbs=mbs)) + 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 BSHD samples into a tensor batch. + # 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) @@ -741,12 +741,12 @@ def test_thd_dataloader_uses_identity_collate(): mbs = 2 dp = parallel_state.get_data_parallel_world_size() n = mbs * dp * 4 - cfg = _make_config(tok, seq_length=64, bshd=False) + 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, bshd=False, scheduler="dp_balanced", mbs=mbs)) + 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). @@ -779,7 +779,7 @@ def test_dcp_dataloader_yields_microbatches_for_scheduler(): set_args( _loader_args( use_varlen=True, - bshd=False, + sbhd=False, scheduler="default_dynamic_cp", mbs=mbs, dynamic_cp=True,