From b9a65b2cf8fc5707f0636e2eb3a2f72c9cfe28de Mon Sep 17 00:00:00 2001 From: ilml Date: Mon, 20 Jul 2026 19:28:53 +0000 Subject: [PATCH 1/3] Add sequence_packing_scheduler config field, CLI arg, and validation Split 2/10 from #3386 (sequence packing / THD E2E support). Adds the sequence_packing_scheduler knob to ModelParallelConfig, its TransformerConfig validation (TE>=2.9 pin, variable_seq_lengths, alltoall dispatcher), the explicit CLI args, and validate_args checks. Feature-flagged no-op until the scheduler lands. Original changes by @xiaoyao0115 in #3386. Co-Authored-By: Claude Fable 5 Signed-off-by: ilml --- megatron/core/model_parallel_config.py | 8 +++- .../core/transformer/transformer_config.py | 34 +++++++++++++++++ megatron/training/arguments.py | 37 +++++++++++++++---- .../models/test_hybrid_moe_model.py | 1 + 4 files changed, 72 insertions(+), 8 deletions(-) diff --git a/megatron/core/model_parallel_config.py b/megatron/core/model_parallel_config.py index 157ae1437f5..87c7f768182 100644 --- a/megatron/core/model_parallel_config.py +++ b/megatron/core/model_parallel_config.py @@ -114,7 +114,7 @@ class ModelParallelConfig: can handle without overflowing the memory. Typically, a good starting point is to set this to maximum sequence length / context parallel size. This is used to calculate the number and length of sub-samples assigned to - each rank when using hybrid_context_parallel. + each rank when sequence_packing_scheduler is not None. """ hybrid_context_parallel: bool = False @@ -124,6 +124,12 @@ class ModelParallelConfig: Please set max_seqlen_per_dp_cp_rank when using hybrid_context_parallel. """ + sequence_packing_scheduler: Optional[Literal['dp_balanced']] = None + """ + Scheduler for sequence packing and hybrid context parallel. + dp_balanced: DP-balanced scheduler for sequence packing. + """ + expert_model_parallel_size: int = 1 """Distributes Moe Experts across sub data parallel dimension.""" diff --git a/megatron/core/transformer/transformer_config.py b/megatron/core/transformer/transformer_config.py index e783a017056..9d1fa920289 100644 --- a/megatron/core/transformer/transformer_config.py +++ b/megatron/core/transformer/transformer_config.py @@ -3200,6 +3200,40 @@ def _scope_to_str(s): "Disable MoE capacity/expert padding." ) + if self.sequence_packing_scheduler is not None: + # Check TE version. + if not HAVE_PACKAGING: + raise ImportError( + "packaging is not installed. Please install it with `pip install packaging`." + ) + # TODO: remove this after we fix the convergence issue with TE < 2.9. + if not ( + is_te_min_version("2.9.0") or get_te_version() == PkgVersion("2.9.0.dev0+5b3092a") + ): + raise ValueError( + "SFT sequence packing requires Transformer Engine >= 2.9.0 " + f"but got {get_te_version()} (TE < 2.9.0 may have convergence issues)." + ) + + # Needed for passing variable sequences between pp stages. + self.variable_seq_lengths = True + + # TODO(tailaim): add support for other dispatcher types + assert self.moe_token_dispatcher_type == "alltoall", ( + f"sequence_packing only supports moe_token_dispatcher_type='alltoall', " + f"got '{self.moe_token_dispatcher_type}'" + ) + + supported_schedulers = ['dp_balanced'] + if ( + self.sequence_packing_scheduler is not None + and self.sequence_packing_scheduler not in supported_schedulers + ): + raise ValueError( + f"Unsupported scheduler: {self.sequence_packing_scheduler}. " + f"Available schedulers: {supported_schedulers}" + ) + @dataclass class MLATransformerConfig(TransformerConfig): diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py index 2f619468b86..4c072b4824e 100644 --- a/megatron/training/arguments.py +++ b/megatron/training/arguments.py @@ -1246,13 +1246,6 @@ def validate_args(args, defaults={}): if args.rl_use_sequence_packing: args.consumed_train_bins = 0 - # Support for variable sequence lengths across batches/microbatches. - # set it if the dataloader supports generation of variable sequence lengths - # across batches/microbatches. Due to additional communication overhead - # during pipeline parallelism, it should not be set if sequence length - # is constant during training. - args.variable_seq_lengths = False - # Iteration-based training. # Skip these checks when skip_train is set: LR config is irrelevant. if args.train_iters and not args.skip_train: @@ -1430,6 +1423,12 @@ def validate_args(args, defaults={}): assert args.dataloader_type == 'single', 'Hybrid context parallelism only supported with single dataloader type' assert args.calculate_per_token_loss, 'Hybrid context parallelism must be used with --calculate-per-token-loss' + # Support for variable sequence lengths across batches/microbatches. + # set it if the dataloader supports generation of variable sequence lengths + # across batches/microbatches. Due to additional communication overhead + # during pipeline parallelism, it should not be set if sequence length + # is constant during training. + args.variable_seq_lengths = False # disable async_tensor_model_parallel_allreduce when # model parallel memory optimization is enabled if (args.tensor_model_parallel_size > 1 or args.context_parallel_size > 1) \ @@ -1642,6 +1641,19 @@ def validate_args(args, defaults={}): if args.ckpt_format == "fsdp_dtensor": assert args.use_megatron_fsdp, "--ckpt-format fsdp_dtensor is only tested with Megatron FSDP." + # Packed-sequence buffer-size check. Placed after varlen scheduler + # auto-select so it validates the final resolved scheduler. + if args.sequence_packing_scheduler is not None: + args.variable_seq_lengths = True + assert args.max_seqlen_per_dp_cp_rank is not None, ( + "--max-seqlen-per-dp-cp-rank must be set when using sequence packing" + ) + total_cp_ranks = 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 + \ bool(args.data_path) + \ @@ -2336,6 +2348,9 @@ def _add_network_size_args(parser): "gtp_weight_remat_size", # internal/derived: controlled only via --expert-tensor-parallel-num-weight-shards "expert_gtp_weight_remat_size", + "max_seqlen_per_dp_cp_rank", + "hybrid_context_parallel", + "sequence_packing_scheduler", ] transformer_factory = ArgumentGroupFactory(TransformerConfig, exclude=exclude) transformer_group = transformer_factory.build_group(parser, "transformer configuration") @@ -3172,6 +3187,14 @@ def _add_distributed_args(parser): 'all layers will share the same communication type. Users can also ' 'specify separated types for each layer like ' '--cp-comm-type p2p p2p a2a a2a a2a+p2p a2a+p2p') + group.add_argument('--max-seqlen-per-dp-cp-rank', type=int, default=None, + help='Maximum sequence length per CP rank. This is used to calculate the ' + 'number of sub-samples assigned to each CP rank when using heterogeneous context parallel.') + group.add_argument('--hybrid-context-parallel', action='store_true', default=False, + help='Enables hybrid context parallel. This is used to balance the workload ' + 'of each CP rank when we use packed samples with variable sequence lengths. ' + 'Requires --max-seqlen-per-dp-cp-rank to be set.') + group.add_argument('--sequence-packing-scheduler', type=str, default=None, choices=['dp_balanced']) group.add_argument('--fake-process-group', action='store_true', default=False, help='If set, initialize with fake distributed process group and all distributed communication operations will be skipped. \ This is quite useful for profiling memory usage of distributed training with just one GPU. \ diff --git a/tests/unit_tests/models/test_hybrid_moe_model.py b/tests/unit_tests/models/test_hybrid_moe_model.py index 037dbae3ee4..2332e35a82f 100644 --- a/tests/unit_tests/models/test_hybrid_moe_model.py +++ b/tests/unit_tests/models/test_hybrid_moe_model.py @@ -352,6 +352,7 @@ "moe_single_grouped_weight": False, "moe_single_grouped_bias": False, "moe_hybridep_pad_uneven_dispatch_inputs": False, + "sequence_packing_scheduler": None, } # Fields to ignore entirely (ephemeral, environment-specific, very large). SKIP_FIELDS = set() From 6a133aafe5c3f0eafc6c459665bddfb3b578966c Mon Sep 17 00:00:00 2001 From: ilml Date: Mon, 20 Jul 2026 19:29:27 +0000 Subject: [PATCH 2/3] Add mock SFT dataset and generalize SFT padding divisor Split 6/10 from #3386 (sequence packing / THD E2E support). Adds MockSFTDataset/MockSFTLowLevelDataset with file- and distribution-mode sequence-length configs, the load_json_arg helper, the --sft-mock-dataset-config-json arg with lognormal default, and generalizes the SFT padding divisor beyond cp>1. Original changes by @xiaoyao0115 in #3386. Co-Authored-By: Claude Fable 5 Signed-off-by: ilml --- megatron/training/arguments.py | 35 +++- megatron/training/datasets/sft_dataset.py | 208 +++++++++++++++++++++- megatron/training/datasets/utils.py | 29 +++ pretrain_gpt.py | 8 +- 4 files changed, 268 insertions(+), 12 deletions(-) create mode 100644 megatron/training/datasets/utils.py diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py index 4c072b4824e..330bf021871 100644 --- a/megatron/training/arguments.py +++ b/megatron/training/arguments.py @@ -1429,6 +1429,17 @@ def validate_args(args, defaults={}): # during pipeline parallelism, it should not be set if sequence length # is constant during training. args.variable_seq_lengths = False + if args.mock_data and args.sft and args.sft_mock_dataset_config_json is None: + args.sft_mock_dataset_config_json = json.dumps( + { + "mode": "distribution", + "type": "lognormal", + "min_seq_len": args.seq_length // 2, + "max_seq_len": args.seq_length, + "mean_seq_len": args.seq_length // 4 * 3, + "lognormal_sigma": 1.1, + } + ) # disable async_tensor_model_parallel_allreduce when # model parallel memory optimization is enabled if (args.tensor_model_parallel_size > 1 or args.context_parallel_size > 1) \ @@ -3742,8 +3753,28 @@ def _add_kitchen_quantization_arguments(parser: argparse.ArgumentParser): def _add_sft_args(parser): group = parser.add_argument_group(title='sft') group.add_argument('--sft', action="store_true", help='Megatron SFT training') - group.add_argument('--sft-tokenizer-prompt-format', type=str, default="nemotron-h-aligned", - help='SFT prompt format.') + group.add_argument( + '--sft-tokenizer-prompt-format', + type=str, + default="nemotron-h-aligned", + help='SFT prompt format.', + ) + group.add_argument( + '--sft-mock-dataset-config-json', + type=str, + default=None, + help='This config provides the necessary information for the mock dataset. ' + 'Accepts either an inline JSON literal or a path to a JSON file containing ' + 'the same schema. You can either specify a CSV file that contains sequence lengths, ' + 'where each line stores the length of a sequence, for example: ' + '{"mode":"file","path":"/path/to/file"}. Alternatively, you can specify a distribution ' + '(currently only supporting lognormal distribution) along with the required parameters, ' + 'for example, {"mode":"distribution","type":"lognormal","min_seq_len":1024,' + '"max_seq_len":2048,"mean_seq_len":1536,"lognormal_sigma":1.1}, where sigma controls ' + 'the variability of the lognormal distribution. ' + 'If not specified and --mock-data is set, defaults to a lognormal distribution with ' + 'min_seq_len=seq_length//2, max_seq_len=seq_length, mean_seq_len=seq_length*3//4, lognormal_sigma=1.1.', + ) return parser def _add_logits_distillation_args(parser): diff --git a/megatron/training/datasets/sft_dataset.py b/megatron/training/datasets/sft_dataset.py index 3f93927387d..80c625d665c 100644 --- a/megatron/training/datasets/sft_dataset.py +++ b/megatron/training/datasets/sft_dataset.py @@ -1,15 +1,18 @@ # Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. -import atexit, json +import atexit from collections import Counter -from typing import Any, Dict, Optional +import math +from typing import Any, Dict, List, Optional, Union import numpy as np +import pandas as pd import torch from megatron.core.datasets.gpt_dataset import GPTDatasetConfig 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 @@ -61,6 +64,8 @@ def __init__( config: GPTDatasetConfig, ) -> None: super().__init__(dataset, dataset_path, indices, num_samples, index_split, config) + # Pre-calculate padding divisor to avoid redundant computation in get_item + self.padding_divisor = self._calculate_padding_divisor() @staticmethod def numel_low_level_dataset(low_level_dataset: LowLevelDataset) -> int: @@ -88,6 +93,26 @@ def _split_conversations(self, merged_conversations): split_conversations.append(current) return split_conversations + def _calculate_padding_divisor(self) -> int: + """ + Calculate the divisor used for sequence padding. + tp_pad = tp_size * 2 if tp_size > 1 else 1 + cp_pad = cp_size * 2 if cp_size > 1 else 1 + cp_pad = cp_pad * dp_size if hybrid_cp else cp_pad + divisor = cp_pad * tp_pad + """ + if self.config.hybrid_context_parallel: + # Hybrid CP: consider both CP and DP + cp_pad = self.config.data_parallel_size * self.config.context_parallel_size * 2 + else: + # Standard CP: only consider CP + cp_pad = self.config.context_parallel_size * 2 if self.config.context_parallel_size > 1 else 1 + tp_pad = self.config.sequence_parallel_size if self.config.sequence_parallel_size > 0 else 1 + divisor = cp_pad * tp_pad + # TODO(tailaim): do we need to pad for FP8 execution? + # divisor = ((divisor + 15) // 16) * 16 + return divisor + def __getitem__(self, idx: int) -> Dict[str, Any]: tokenizer = self.config.tokenizer @@ -124,12 +149,11 @@ def extend_with_padding(tokens, targets, positions, pad_len): assert not self.config.reset_position_ids pack_positions.extend(range(len(tokens_list))) - if self.config.context_parallel_size > 1: - pad_granularity = self.config.context_parallel_size * 2 - mod_token_count = len(pack_tokens) % pad_granularity - if mod_token_count != 0: - pad_len = pad_granularity - mod_token_count - extend_with_padding(pack_tokens, pack_targets, pack_positions, pad_len) + pad_granularity = self.padding_divisor + mod_token_count = len(pack_tokens) % pad_granularity + if mod_token_count != 0: + pad_len = pad_granularity - mod_token_count + extend_with_padding(pack_tokens, pack_targets, pack_positions, pad_len) # TODO(duncan): Consider also padding to multiple of number of tokens here. This might # be needed for efficiency (and potentially set via command-line argument). @@ -199,3 +223,171 @@ def extend_with_padding(tokens, targets, positions, pad_len): 'cu_seqlens': padded_cu_seqlens, 'max_seqlen': max_seqlen, } + + +class MockSFTLowLevelDataset: + """The low-level mock dataset for SFT + + Args: + mode (str): Either 'file' or 'distribution'. + **kwargs: Additional arguments depending on mode. + For mode='file': path (str) - path to a CSV file with sequence lengths. + For mode='distribution': type (str), min_seq_len (int), max_seq_len (int), + mean_seq_len (int), and distribution-specific params (e.g. lognormal_sigma). + """ + + seed: int = 0 + """The hard-coded random seed to use to set the NumPy RNG""" + + size: int = 1000000 + """The hard-coded number of sequence to generate""" + + def __init__(self, mode: str, **kwargs) -> None: + np.random.seed(self.seed) + + if mode == "file": + self.sequence_lengths = np.array(pd.read_csv(kwargs["path"])).flatten() + self.size = len(self.sequence_lengths) + elif mode == "distribution": + min_seq_len = kwargs["min_seq_len"] + max_seq_len = kwargs["max_seq_len"] + mean_seq_len = kwargs["mean_seq_len"] + if kwargs["type"] == "lognormal": + lognormal_sigma = kwargs["lognormal_sigma"] + self.sequence_lengths = self.generate_lognormal_samples( + self.size, mean_seq_len, lognormal_sigma, min_seq_len, max_seq_len + ) + else: + raise ValueError(f"Unsupported distribution type {kwargs['type']}") + else: + raise ValueError(f"Unsupported mode '{mode}', must be 'file' or 'distribution'") + + def generate_lognormal_samples(self, size, mean, sigma, min_seq_len, max_seq_len): + mu = np.log(mean) - sigma**2 / 2 + samples = np.random.lognormal(mu, sigma, size) + samples = np.clip(samples, min_seq_len, max_seq_len) + return samples.astype(int) + + def __len__(self) -> int: + return self.size + + def __getitem__(self, idx: int) -> List[np.ndarray]: + # the length of sample is 'length', but only length-1 elements are generated here, + # because an eod token will be appended at the end later in SFTDataset + + length = self.sequence_lengths[idx % self.size] + sample = np.arange(1, length, dtype=np.int64) + return sample + + +class MockSFTDataset(SFTDataset): + """The mock dataset used during SFT""" + + 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 build_low_level_dataset(dataset_path: str, config: GPTDatasetConfig) -> LowLevelDataset: + if config.sft_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.sft_mock_dataset_config_json) + return MockSFTLowLevelDataset(**mock_config) + + def __len__(self) -> int: + return self.num_samples + + def __getitem__(self, idx: int) -> Dict[str, Any]: + + tokenizer = self.config.tokenizer + pack_length = self.config.sequence_length + eod = tokenizer.eod + pad = tokenizer.pad + + tokens = self.dataset[int(self.indices[idx % len(self.indices)])] + + def extend_with_padding(tokens, targets, positions, pad_len): + tokens.extend([pad] * pad_len) + targets.extend([pad] * pad_len) + positions.extend(range(positions[-1] + 1, positions[-1] + 1 + pad_len)) + + # Convert tokens to list and add EOD + tokens_list = tokens.tolist() + if tokens_list[-1] != eod: + tokens_list.append(eod) + targets_list = list(tokens_list) + + pack_tokens = list(tokens_list) + pack_targets = list(targets_list) + pack_positions = list(range(len(tokens_list))) + cu_seqlens = [0] + + # Pad to padding_divisor alignment + if self.padding_divisor > 1: + mod_token_count = len(pack_tokens) % self.padding_divisor + if mod_token_count != 0: + pad_len = self.padding_divisor - mod_token_count + extend_with_padding(pack_tokens, pack_targets, pack_positions, pad_len) + + # Record padded boundary after padding + cu_seqlens.append(len(pack_tokens)) + + # Handle any necessary truncation + if len(pack_tokens) >= pack_length + 1: # +1 here to account for later alignment + max_body = pack_length - 1 + pack_tokens = pack_tokens[:max_body] + pack_targets = pack_targets[:max_body] + pack_tokens.extend([eod, pad]) + pack_targets.extend([eod, pad]) + pack_positions = pack_positions[:pack_length + 1] + cu_seqlens[-1] = len(pack_tokens) - 1 + + # Handle any necessary padding + if len(pack_tokens) < pack_length + 1: # +1 here to account for later alignment + pad_len = pack_length + 1 - len(pack_tokens) + extend_with_padding(pack_tokens, pack_targets, pack_positions, pad_len) + cu_seqlens[-1] = len(pack_tokens) - 1 + + assert len(pack_tokens) == pack_length + 1 + assert len(pack_targets) == pack_length + 1 + assert len(pack_positions) == pack_length + 1 + + # Align and convert to tensors + input_ids = torch.tensor(pack_tokens[:-1], dtype=torch.int64) + labels = torch.tensor(pack_targets[1:], dtype=torch.int64) + position_ids = torch.tensor(pack_positions[:-1], dtype=torch.int64) + + # Loss mask + loss_mask = torch.ones(pack_length, dtype=torch.float32) + loss_mask[labels == pad] = 0.0 + loss_mask[labels == IGNORE_INDEX] = 0.0 + + assert len(cu_seqlens) >= 2 + cu_seqlens = torch.tensor(cu_seqlens, dtype=torch.int32) + # Calculating max_seqlen here because of possible effects of truncation and padding + adjacent_diffs = cu_seqlens[1:] - cu_seqlens[:-1] + max_seqlen = adjacent_diffs.max() # max_seqlen is a 0-D tensor + + return { + 'tokens': input_ids, + 'labels': labels, + 'loss_mask': loss_mask, + 'position_ids': position_ids, + 'cu_seqlens': cu_seqlens, + 'max_seqlen': max_seqlen, + } 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/pretrain_gpt.py b/pretrain_gpt.py index 8b979d16d4a..f6e8f5d904c 100644 --- a/pretrain_gpt.py +++ b/pretrain_gpt.py @@ -74,7 +74,7 @@ def _rank0_only_showwarning(message, category, filename, lineno, file=None, line from megatron.training.argument_utils import gpt_config_from_args, pretrain_cfg_container_from_args 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 SFTDataset +from megatron.training.datasets.sft_dataset import MockSFTDataset, SFTDataset from megatron.training.training import update_seqlen_stats_from_cu_seqlens from megatron.training.utils import get_blend_and_blend_per_split, is_first_or_last_pipeline_stage from model_provider import model_provider @@ -428,6 +428,7 @@ def core_gpt_dataset_config_from_args(args: Any) -> GPTDatasetConfig: "sequence_parallel_size": args.tensor_model_parallel_size * args.sequence_parallel, "hybrid_context_parallel": args.hybrid_context_parallel, "inter_document_masking": args.dataloader_inter_document_masking, + "sft_mock_dataset_config_json": args.sft_mock_dataset_config_json, } # add FIM args to the config @@ -466,7 +467,10 @@ def train_valid_test_datasets_provider(train_val_test_num_samples, vp_stage=None is_packed_sequence = False if args.sft: - dataset_type = SFTDataset + if args.mock_data: + dataset_type = MockSFTDataset + else: + dataset_type = SFTDataset is_packed_sequence = True # SFT always uses packed sequence else: if args.mock_data: From ddb5304bebb8a5b09a8fbdfccbed125c32760283 Mon Sep 17 00:00:00 2001 From: ilml Date: Mon, 20 Jul 2026 19:29:53 +0000 Subject: [PATCH 3/3] Add VarlenLowLevelDataset, VarlenDataset, and MockVarlenDataset Split 9/10 from #3386 (sequence packing / THD E2E support). Adds the variable-length packed (THD) dataset family: HF-hub/parquet/jsonl loading, THD __getitem__ with cu_seqlens, SBHD validation mode, and the mock variant, with unit tests. Also adds hybrid_context_parallel=False to the _make_config test helper (deviation from #3386: fixes a latent AttributeError in _calculate_padding_divisor with SimpleNamespace configs). Original changes by @xiaoyao0115 in #3386. Co-Authored-By: Claude Fable 5 Signed-off-by: ilml --- megatron/training/datasets/varlen_dataset.py | 324 +++++++++++++++++++ tests/unit_tests/data/test_varlen_dataset.py | 302 +++++++++++++++++ 2 files changed, 626 insertions(+) diff --git a/megatron/training/datasets/varlen_dataset.py b/megatron/training/datasets/varlen_dataset.py index 07edbf708c1..c2533f795bb 100644 --- a/megatron/training/datasets/varlen_dataset.py +++ b/megatron/training/datasets/varlen_dataset.py @@ -50,6 +50,21 @@ 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", @@ -222,3 +237,312 @@ def _select_converter( "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 index 46c03e32d09..8c9f619ced9 100644 --- a/tests/unit_tests/data/test_varlen_dataset.py +++ b/tests/unit_tests/data/test_varlen_dataset.py @@ -8,12 +8,22 @@ 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, @@ -289,3 +299,295 @@ def test_select_converter_alpaca_beats_pretrain_text(): 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