diff --git a/megatron/training/datasets/varlen_dataset.py b/megatron/training/datasets/varlen_dataset.py new file mode 100644 index 00000000000..c2533f795bb --- /dev/null +++ b/megatron/training/datasets/varlen_dataset.py @@ -0,0 +1,548 @@ +# 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** — 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, ...) + * **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``. + * **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 + 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 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, +) +from megatron.training.datasets.utils import load_json_arg + +# 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 _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]], 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: + 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" + 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), " + "pretrain-text (text)." + ) + + +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 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: + 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`` / + ``pretrain-text`` (the raw ``text``-column fallback).""" + 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 + # 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 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 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)])] + + assert not self.config.reset_position_ids + assert not self.config.create_attention_mask and not self.config.reset_attention_mask + + # 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() + + # 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. + 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) + + valid_len = len(tokens_list) - 1 + + # 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). + if self.config.varlen_sbhd_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[valid_len:] = 0.0 # mask the right-padded tail by position + 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[valid_len:] = 0.0 # mask the right-padded tail by position + 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. + + ``--varlen-sbhd-validation`` is intentionally not implemented for mock + data; it is guarded against in argument validation. + """ + + @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 = load_json_arg(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 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 + # 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) + + # 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. + 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[original_seq_len:] = 0.0 # mask the right-padded tail by position + 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/tests/unit_tests/data/test_varlen_dataset.py b/tests/unit_tests/data/test_varlen_dataset.py new file mode 100644 index 00000000000..a7ec459d23e --- /dev/null +++ b/tests/unit_tests/data/test_varlen_dataset.py @@ -0,0 +1,594 @@ +# 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 +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, + _messages_passthrough, + _raw_text_loader, + _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]}]}) + + +# ---------------------------------------------------------------------------- +# 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]}) + + +# ---------------------------------------------------------------------------- +# 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"]) + + +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) +# ---------------------------------------------------------------------------- + + +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) + + +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..." + + +# ---------------------------------------------------------------------------- +# 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, sbhd=False): + return SimpleNamespace( + tokenizer=tokenizer, + sequence_length=seq_length, + reset_position_ids=False, + create_attention_mask=False, + reset_attention_mask=False, + varlen_sbhd_validation=sbhd, + data_parallel_size=dp, + context_parallel_size=cp, + hybrid_context_parallel=False, + 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_sbhd_pads_to_seq_length_and_masks_tail(): + tok = _FakeTokenizer(eod=0, pad=None) + ds = _make_varlen(["abc"], _make_config(tok, seq_length=8, sbhd=True)) + out = ds[0] + # 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() + # 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