diff --git a/docs/user-guide/data-loading.md b/docs/user-guide/data-loading.md index 1f0d544317c..b60cd685cf2 100644 --- a/docs/user-guide/data-loading.md +++ b/docs/user-guide/data-loading.md @@ -72,6 +72,8 @@ If your later training job does not set `--global-batch-size`, or you are prepar This keeps the prepared cache aligned with the sample counts expected by training. +> **Unsupported configurations:** `tools/prepare_cache.py` does not support `--mock-data`, `--sft`, `--fim-data`, or `--step-batch-size-schedule`. Using any of these will cause the script to exit with an error. + ### Step 3: Optionally pre-build per-dataset metadata When blending many datasets, generate the `--per-dataset-sequences-path` JSON ahead of time to avoid one metadata read per file prefix at startup: diff --git a/megatron/core/datasets/blended_dataset.py b/megatron/core/datasets/blended_dataset.py index 802a9770506..9b642ee1ff3 100644 --- a/megatron/core/datasets/blended_dataset.py +++ b/megatron/core/datasets/blended_dataset.py @@ -150,7 +150,10 @@ def _build_indices(self) -> Tuple[numpy.ndarray, numpy.ndarray]: else: cache_hit = False - if not path_to_cache or (not cache_hit and torch.distributed.get_rank() == 0): + if not path_to_cache or ( + not cache_hit + and (not torch.distributed.is_initialized() or torch.distributed.get_rank() == 0) + ): log_single_rank( logger, logging.INFO, f"Build and save the {type(self).__name__} indices" ) diff --git a/megatron/core/datasets/readme.md b/megatron/core/datasets/readme.md index 452bf24e4a2..58721b7471b 100644 --- a/megatron/core/datasets/readme.md +++ b/megatron/core/datasets/readme.md @@ -192,12 +192,24 @@ To query the `BlendedDataset` for the _k_-th sample we do the following To save time during initialization, each index is built/cached sequentially on one process rank and subsequently loaded in parallel on other process ranks. The cached indices are unique to a hash generated in the `BlendedDataset.__init__` function. +## Offline cache preparation + +For GPT-style training, the dataset caches described above can be prepared ahead of time with `tools/prepare_cache.py` instead of waiting for rank 0 to build them during training startup. + +The script reuses the normal dataset construction path used by `pretrain_gpt.py` and `pretrain_mamba.py`, including `GPTDataset`, `BlendedDataset`, and `BlendedMegatronDatasetBuilder`. It accepts the usual dataset arguments, supports blends and per-split dataset definitions, and requires `--data-cache-path` so the generated cache can later be reused by training. + +This is especially useful for large blends or many file prefixes, where building the document, sample, and shuffle indices can take several minutes and leave all GPUs idle while rank 0 performs CPU-only work. + +If the later training job does not specify `--global-batch-size` (which is needed to determine the dataset size and splits), you should specify `--prepare-cache-world-size` to explicitly set the world size used during cache preparation. + +`tools/prepare_cache.py` does not support `--mock-data`, `--sft`, `--fim-data`, or `--step-batch-size-schedule`. + ## Fast DataLoader initialization -Especially for large-scale runs, DataLoader initialization can take several minutes, since it involves opening and memory-mapping multiple files and can significantly stress the filesystem. To speed up this process, we have developed the following three optimizations, controlled by configuration flags": +Especially for large-scale runs, DataLoader initialization can take several minutes, since it involves opening and memory-mapping multiple files and can significantly stress the filesystem. To speed up this process, we have developed the following three optimizations, controlled by configuration flags: - `--dataloader-fast-cache-load`: This option assumes that the dataset cache already exists in the specified `--data-cache-path`. When enabled, it speeds up the creation process by removing synchronization points and file check assertions. - `--dataloader-defer-npy-index-mmap`: This option also assumes that the dataset cache already exists in the specified `--data-cache-path`. When enabled, it defers the memory mapping of the dataset indexes (.npy files) until their first access. We recommend using this configuration together with `--num-workers` > 0 so that the DataLoader prefetches the next batches of data, thereby hiding the cost of index memory mapping. - - `--per-dataset-sequences-path`: With this configuration, we specify the JSON file generated by the `tools/build_sequences_per_dataset.py` script. This script generates a single file containing the required metadata from all the specified file prefixes. This configuration is especially useful when dealing with hundreds to thousands of file prefixes, since it requires only a single `open` operation instead of one per file prefix. \ No newline at end of file + - `--per-dataset-sequences-path`: With this configuration, we specify the JSON file generated by the `tools/build_sequences_per_dataset.py` script. This script generates a single file containing the required metadata from all the specified file prefixes. This configuration is especially useful when dealing with hundreds to thousands of file prefixes, since it requires only a single `open` operation instead of one per file prefix. diff --git a/pretrain_gpt.py b/pretrain_gpt.py index 02ca5dd72c0..bc6eaeee51f 100644 --- a/pretrain_gpt.py +++ b/pretrain_gpt.py @@ -17,7 +17,7 @@ warnings.filterwarnings("ignore", category=FutureWarning) from functools import partial -from typing import List, Optional, Tuple +from typing import Any, List, Optional, Tuple import torch @@ -281,7 +281,7 @@ def is_dataset_built_on_rank(vp_stage=None, is_packed_sequence=False): ) -def core_gpt_dataset_config_from_args(args): +def core_gpt_dataset_config_from_args(args: Any) -> GPTDatasetConfig: tokenizer = build_tokenizer(args) # Sometimes --data-path is too long, instead we parse it from a file. @@ -318,7 +318,7 @@ def core_gpt_dataset_config_from_args(args): "defer_npy_index_mmap": args.dataloader_defer_npy_index_mmap, "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, + "sequence_parallel_size": args.tensor_model_parallel_size * args.sequence_parallel, "hybrid_context_parallel": args.hybrid_context_parallel, } @@ -404,7 +404,7 @@ def get_embedding_ranks(pp_ranks: List[int]): set_startup_timestamps(program_start=_PROGRAM_START_TIME, main_entry=_MAIN_ENTRY_TIME) # Temporary for transition to core datasets - train_valid_test_datasets_provider.is_distributed = True + setattr(train_valid_test_datasets_provider, "is_distributed", True) # Optionally enable inprocess restart on pretrain pretrain, store = inprocess_restart.maybe_wrap_for_inprocess_restart(pretrain) diff --git a/pretrain_hybrid.py b/pretrain_hybrid.py index f073e8e9ab3..7eba74698aa 100644 --- a/pretrain_hybrid.py +++ b/pretrain_hybrid.py @@ -16,7 +16,7 @@ warnings.filterwarnings("ignore", category=FutureWarning) from functools import partial -from typing import List, Optional, Tuple +from typing import Any, List, Optional, Tuple import torch @@ -274,7 +274,7 @@ def is_dataset_built_on_rank(vp_stage=None, is_packed_sequence=False): return is_first_or_last_pipeline_stage(vp_stage) -def core_gpt_dataset_config_from_args(args): +def core_gpt_dataset_config_from_args(args: Any) -> GPTDatasetConfig: tokenizer = build_tokenizer(args) # Sometimes --data-path is too long, instead we parse it from a file. @@ -293,6 +293,8 @@ def core_gpt_dataset_config_from_args(args): blend=blend, blend_per_split=blend_per_split, split=args.split, + multiple_validation_sets=args.multiple_validation_sets, + full_validation=args.full_validation, num_dataset_builder_threads=args.num_dataset_builder_threads, path_to_cache=args.data_cache_path, mmap_bin_files=args.mmap_bin_files, @@ -352,7 +354,7 @@ def train_valid_test_datasets_provider(train_val_test_num_samples, vp_stage=None set_startup_timestamps(program_start=_PROGRAM_START_TIME, main_entry=_MAIN_ENTRY_TIME) # Temporary for transition to core datasets - train_valid_test_datasets_provider.is_distributed = True + setattr(train_valid_test_datasets_provider, "is_distributed", True) # Optionally enable inprocess restart on pretrain pretrain, store = inprocess_restart.maybe_wrap_for_inprocess_restart(pretrain) diff --git a/tests/unit_tests/data/test_prepare_cache.py b/tests/unit_tests/data/test_prepare_cache.py new file mode 100644 index 00000000000..164001121b2 --- /dev/null +++ b/tests/unit_tests/data/test_prepare_cache.py @@ -0,0 +1,288 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +import os +import random +from argparse import Namespace + +import pytest +import torch + +from megatron.core.datasets.blended_dataset import BlendedDataset +from megatron.core.datasets.blended_megatron_dataset_builder import BlendedMegatronDatasetBuilder +from megatron.core.datasets.gpt_dataset import GPTDataset +from megatron.core.datasets.indexed_dataset import DType, IndexedDatasetBuilder +from megatron.core.datasets.utils import compile_helpers +from megatron.core.tokenizers.utils.build_tokenizer import build_tokenizer +from tests.unit_tests.dist_checkpointing import TempNamedDir +from tests.unit_tests.test_utilities import Utils +from tools.prepare_cache import ( + _normalize_prepare_cache_args, + build_dataset_caches, + core_gpt_dataset_config_from_args, +) + + +def _build_null_tokenizer(vocab_size: int = 2048): + return build_tokenizer( + Namespace( + vocab_size=vocab_size, + tokenizer_type="NullTokenizer", + padded_vocab_size=None, + rank=0, + make_vocab_size_divisible_by=128, + tensor_model_parallel_size=1, + ) + ) + + +def _initialize_test_environment() -> None: + if torch.distributed.is_available(): + Utils.initialize_distributed() + if torch.distributed.get_rank() == 0: + compile_helpers() + torch.distributed.barrier() + else: + compile_helpers() + + +def _create_file_prefixes(tokenizer, dataset_dir, number_of_files: int = 4) -> list[str]: + os.makedirs(dataset_dir, exist_ok=True) + + file_prefixes = [] + for i in range(number_of_files): + file_prefix = os.path.join(dataset_dir, f"file_{i}") + builder = IndexedDatasetBuilder( + file_prefix + ".bin", dtype=DType.optimal_dtype(tokenizer.vocab_size) + ) + + for j in range(32): + tokens = [int((i * 97 + j * 13 + k) % tokenizer.vocab_size) for k in range(64)] + builder.add_document(tokens, [len(tokens)]) + + builder.finalize(file_prefix + ".idx") + file_prefixes.append(file_prefix) + + return file_prefixes + + +def _create_shared_file_prefixes(tokenizer, dataset_dir, number_of_files: int = 4) -> list[str]: + if not torch.distributed.is_initialized() or torch.distributed.get_rank() == 0: + file_prefixes = _create_file_prefixes(tokenizer, dataset_dir, number_of_files) + else: + file_prefixes = [os.path.join(dataset_dir, f"file_{i}") for i in range(number_of_files)] + + if torch.distributed.is_initialized(): + torch.distributed.barrier() + + random.seed(1234) # NOTE(asolergi-nv): re-sync random state across all ranks + + return file_prefixes + + +def _build_prepare_cache_args(file_prefixes, data_cache_path, **overrides): + args = dict( + seed=1234, + seq_length=16, + split="70,20,10", + data_path=file_prefixes, + train_data_path=None, + valid_data_path=None, + test_data_path=None, + per_split_data_args_path=None, + data_args_path=None, + per_dataset_sequences_path=None, + data_cache_path=str(data_cache_path), + mmap_bin_files=True, + reset_position_ids=False, + reset_attention_mask=False, + eod_mask_loss=False, + create_attention_mask_in_dataloader=False, + object_storage_cache_path=None, + mid_level_dataset_surplus=0.005, + allow_ambiguous_pad_tokens=False, + dataloader_fast_cache_load=True, + dataloader_defer_npy_index_mmap=True, + context_parallel_size=1, + data_parallel_size=4, + tensor_model_parallel_size=1, + sequence_parallel=False, + hybrid_context_parallel=False, + multiple_validation_sets=False, + full_validation=False, + num_dataset_builder_threads=1, + tokenizer_type="NullTokenizer", + vocab_size=2048, + padded_vocab_size=None, + make_vocab_size_divisible_by=128, + rank=0, + world_size=4, + train_samples=None, + train_iters=4, + skip_train=False, + eval_iters=2, + eval_interval=2, + start_eval_at_iter=None, + global_batch_size=8, + phase_transition_iterations=None, + iteration=0, + mock_data=False, + sft=False, + fim_data=False, + step_batch_size_schedule=None, + ) + args.update(overrides) + return Namespace(**args) + + +def test_prepare_cache_builds_blended_dataset_cache(tmp_path_dist_ckpt): + _initialize_test_environment() + + tokenizer = _build_null_tokenizer() + + with TempNamedDir( + tmp_path_dist_ckpt / "test_prepare_cache_builds_blended_dataset_cache", sync=True + ) as temp_dir: + file_prefixes = _create_shared_file_prefixes(tokenizer, os.path.join(temp_dir, "dataset")) + args = _build_prepare_cache_args(file_prefixes, temp_dir / "cache") + + summary = build_dataset_caches(args) + + assert args.dataloader_fast_cache_load is False + assert args.dataloader_defer_npy_index_mmap is False + assert summary["train_valid_test_num_samples"] == (32, 48, 16) + assert summary["train_dataset_length"] == 32 + assert summary["valid_dataset_length"] == 48 + assert summary["test_dataset_length"] == 16 + assert list((temp_dir / "cache").glob("*document_index.npy")) + assert list((temp_dir / "cache").glob("*dataset_index.npy")) + + +def test_prepare_cache_world_size_override(): + args = Namespace(rank=11, world_size=1, prepare_cache_world_size=8) + + _normalize_prepare_cache_args(args) + + assert args.rank == 0 + assert args.world_size == 8 + + +def test_prepare_cache_builds_and_hits_per_split_dataset_cache(tmp_path_dist_ckpt): + _initialize_test_environment() + + tokenizer = _build_null_tokenizer() + + with TempNamedDir( + tmp_path_dist_ckpt / "test_prepare_cache_builds_and_hits_per_split_dataset_cache", sync=True + ) as temp_dir: + file_prefixes = _create_shared_file_prefixes(tokenizer, os.path.join(temp_dir, "dataset")) + args = _build_prepare_cache_args( + None, + temp_dir / "cache", + split=None, + data_path=None, + train_data_path=[50, file_prefixes[0], 50, file_prefixes[1]], + valid_data_path=[file_prefixes[2]], + test_data_path=[file_prefixes[3]], + ) + + summary = build_dataset_caches(args) + + assert summary["train_valid_test_num_samples"] == (32, 48, 16) + assert list((temp_dir / "cache").glob("*description.txt")) + + slow_args = _build_prepare_cache_args( + None, + temp_dir / "cache", + split=None, + data_path=None, + train_data_path=[50, file_prefixes[0], 50, file_prefixes[1]], + valid_data_path=[file_prefixes[2]], + test_data_path=[file_prefixes[3]], + dataloader_fast_cache_load=False, + dataloader_defer_npy_index_mmap=False, + ) + slow_config = core_gpt_dataset_config_from_args(slow_args) + train_slow, valid_slow, test_slow = BlendedMegatronDatasetBuilder( + GPTDataset, list(summary["train_valid_test_num_samples"]), lambda: True, slow_config + ).build() + + fast_args = _build_prepare_cache_args( + None, + temp_dir / "cache", + split=None, + data_path=None, + train_data_path=[50, file_prefixes[0], 50, file_prefixes[1]], + valid_data_path=[file_prefixes[2]], + test_data_path=[file_prefixes[3]], + dataloader_fast_cache_load=True, + dataloader_defer_npy_index_mmap=True, + ) + fast_config = core_gpt_dataset_config_from_args(fast_args) + train_fast, valid_fast, test_fast = BlendedMegatronDatasetBuilder( + GPTDataset, list(summary["train_valid_test_num_samples"]), lambda: True, fast_config + ).build() + + assert isinstance(train_fast, BlendedDataset) + assert train_fast.dataset_index is None + assert train_fast.dataset_sample_index is None + assert isinstance(valid_fast, GPTDataset) + assert valid_fast.document_index is None + assert valid_fast.sample_index is None + assert valid_fast.shuffle_index is None + assert isinstance(test_fast, GPTDataset) + assert test_fast.document_index is None + assert test_fast.sample_index is None + assert test_fast.shuffle_index is None + + assert summary["train_dataset_length"] == len(train_slow) == len(train_fast) == 32 + assert summary["valid_dataset_length"] == len(valid_slow) == len(valid_fast) + assert summary["test_dataset_length"] == len(test_slow) == len(test_fast) + assert summary["valid_dataset_length"] >= summary["train_valid_test_num_samples"][1] + assert summary["test_dataset_length"] >= summary["train_valid_test_num_samples"][2] + assert torch.all(train_slow[0]["tokens"] == train_fast[0]["tokens"]) + assert torch.all(valid_slow[0]["tokens"] == valid_fast[0]["tokens"]) + assert torch.all(test_slow[0]["tokens"] == test_fast[0]["tokens"]) + + assert train_fast.dataset_index is not None + assert train_fast.dataset_sample_index is not None + assert valid_fast.document_index is not None + assert valid_fast.sample_index is not None + assert valid_fast.shuffle_index is not None + assert test_fast.document_index is not None + assert test_fast.sample_index is not None + assert test_fast.shuffle_index is not None + + +@pytest.mark.parametrize( + ("flag_name", "flag_value", "message"), + [ + ("mock_data", True, "--mock-data"), + ("sft", True, "--sft"), + ("fim_data", True, "--fim-data"), + ("step_batch_size_schedule", [(0, 8)], "--step-batch-size-schedule"), + ], +) +def test_prepare_cache_rejects_unsupported_modes(tmp_path, flag_name, flag_value, message): + args = _build_prepare_cache_args([], tmp_path / "cache", **{flag_name: flag_value}) + + with pytest.raises(ValueError, match=message): + build_dataset_caches(args) + + +def test_prepare_cache_builds_with_train_samples(tmp_path_dist_ckpt): + _initialize_test_environment() + + tokenizer = _build_null_tokenizer() + + with TempNamedDir( + tmp_path_dist_ckpt / "test_prepare_cache_builds_with_train_samples", sync=True + ) as temp_dir: + file_prefixes = _create_shared_file_prefixes(tokenizer, os.path.join(temp_dir, "dataset")) + args = _build_prepare_cache_args( + file_prefixes, temp_dir / "cache", train_iters=None, train_samples=32 + ) + + summary = build_dataset_caches(args) + + assert args.train_iters == 32 // args.global_batch_size + assert summary["train_valid_test_num_samples"][0] == 32 diff --git a/tools/prepare_cache.py b/tools/prepare_cache.py new file mode 100644 index 00000000000..98d57db498a --- /dev/null +++ b/tools/prepare_cache.py @@ -0,0 +1,214 @@ +# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +"""Prepare GPT dataset caches ahead of training. + +Unsupported configurations: + --mock-data, --sft, --fim-data, --step-batch-size-schedule +""" + +import argparse +import json +from typing import Any, Dict, List, Optional, Tuple + +from megatron.core.datasets.blended_megatron_dataset_builder import BlendedMegatronDatasetBuilder +from megatron.core.datasets.gpt_dataset import GPTDataset, GPTDatasetConfig +from megatron.core.datasets.utils import compile_helpers +from megatron.core.tokenizers.utils.build_tokenizer import build_tokenizer +from megatron.training import get_train_valid_test_num_samples +from megatron.training.arguments import parse_args, validate_args +from megatron.training.global_vars import set_args, unset_global_variables +from megatron.training.training import update_train_iters +from megatron.training.utils import get_blend_and_blend_per_split + +try: + from megatron.post_training.arguments import add_modelopt_args + + has_nvidia_modelopt = True +except ImportError: + has_nvidia_modelopt = False + + +def add_prepare_cache_args(parser: argparse.ArgumentParser) -> argparse.ArgumentParser: + """Add cache-preparation specific arguments.""" + + group = parser.add_argument_group(title="prepare cache") + group.add_argument( + "--prepare-cache-world-size", + type=int, + default=None, + help=( + "Optional override for the effective world size used to derive data-parallel size and " + "dataset sample counts during cache preparation." + ), + ) + return parser + + +def _extra_args_provider(parser: argparse.ArgumentParser) -> argparse.ArgumentParser: + parser = add_prepare_cache_args(parser) + if has_nvidia_modelopt: + parser = add_modelopt_args(parser) + return parser + + +def _normalize_prepare_cache_args(args: Any) -> None: + """Apply cache-preparation specific argument normalization.""" + + args.rank = 0 + + if args.prepare_cache_world_size is not None: + if args.prepare_cache_world_size <= 0: + raise ValueError("--prepare-cache-world-size must be positive") + args.world_size = args.prepare_cache_world_size + + +def _validate_prepare_cache_args(args: Any) -> None: + """Validate options that are intentionally unsupported for offline cache prep.""" + + if args.data_cache_path is None: + raise ValueError("--data-cache-path must be provided for cache preparation") + if args.mock_data: + raise ValueError("--mock-data is not supported by tools/prepare_cache.py") + if getattr(args, "sft", False): + raise ValueError("--sft is not supported by tools/prepare_cache.py") + if getattr(args, "fim_data", False): + raise ValueError("--fim-data is not supported by tools/prepare_cache.py") + if getattr(args, "step_batch_size_schedule", None) is not None: + raise ValueError( + "--step-batch-size-schedule is not supported by tools/prepare_cache.py" + ) + + +def _disable_cache_load_only_flags(args: Any) -> Dict[str, bool]: + """Disable flags that only make sense when consuming an existing cache.""" + + ignored = { + "dataloader_fast_cache_load": bool(args.dataloader_fast_cache_load), + "dataloader_defer_npy_index_mmap": bool(args.dataloader_defer_npy_index_mmap), + } + args.dataloader_fast_cache_load = False + args.dataloader_defer_npy_index_mmap = False + return ignored + + +def _get_dataset_length(dataset: Optional[Any]) -> Optional[Any]: + if dataset is None: + return None + if isinstance(dataset, list): + return [len(ds) if ds is not None else None for ds in dataset] + return len(dataset) + + +def _print_effective_configuration( + args: Any, train_valid_test_num_samples: Any, ignored_flags: Dict[str, bool] +) -> None: + print("> preparing dataset cache with the following effective values:") + print(f" world size: {args.world_size}") + print(f" data parallel size: {args.data_parallel_size}") + print(f" global batch size: {args.global_batch_size}") + print(f" cache path: {args.data_cache_path}") + print(" > datasets target sizes (minimum size):") + print(f" train: {train_valid_test_num_samples[0]}") + print(f" validation: {train_valid_test_num_samples[1]}") + print(f" test: {train_valid_test_num_samples[2]}") + if ignored_flags["dataloader_fast_cache_load"]: + print("> ignoring --dataloader-fast-cache-load during cache preparation") + if ignored_flags["dataloader_defer_npy_index_mmap"]: + print("> ignoring --dataloader-defer-npy-index-mmap during cache preparation") + + +def core_gpt_dataset_config_from_args(args: Any) -> GPTDatasetConfig: + """Build the explicit GPTDatasetConfig used for offline cache preparation.""" + + tokenizer = build_tokenizer(args) + + blend: Optional[Tuple[List[str], Optional[List[float]]]] + blend_per_split: Optional[List[Optional[Tuple[List[str], Optional[List[float]]]]]] + blend, blend_per_split = get_blend_and_blend_per_split(args) + + sequences_per_dataset = None + if args.per_dataset_sequences_path is not None: + with open(args.per_dataset_sequences_path, "r") as f: + sequences_per_dataset = json.load(f) + + return GPTDatasetConfig( + random_seed=args.seed, + sequence_length=args.seq_length, + blend=blend, + blend_per_split=blend_per_split, + split=args.split, + multiple_validation_sets=args.multiple_validation_sets, + full_validation=args.full_validation, + num_dataset_builder_threads=args.num_dataset_builder_threads, + path_to_cache=args.data_cache_path, + mmap_bin_files=args.mmap_bin_files, + tokenizer=tokenizer, + reset_position_ids=args.reset_position_ids, + reset_attention_mask=args.reset_attention_mask, + eod_mask_loss=args.eod_mask_loss, + create_attention_mask=args.create_attention_mask_in_dataloader, + object_storage_cache_path=args.object_storage_cache_path, + mid_level_dataset_surplus=args.mid_level_dataset_surplus, + allow_ambiguous_pad_tokens=args.allow_ambiguous_pad_tokens, + fast_cache_load=args.dataloader_fast_cache_load, + sequences_per_dataset=sequences_per_dataset, + defer_npy_index_mmap=args.dataloader_defer_npy_index_mmap, + 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, + hybrid_context_parallel=args.hybrid_context_parallel, + ) + + +def build_dataset_caches(args: Any) -> Dict[str, Any]: + """Build the dataset caches for the plain GPTDataset path.""" + + _validate_prepare_cache_args(args) + ignored_flags = _disable_cache_load_only_flags(args) + + unset_global_variables() + set_args(args) + + try: + # Derive train_iters from --train-samples when needed (pretrain() does the same). + update_train_iters(args) + train_valid_test_num_samples = get_train_valid_test_num_samples() + _print_effective_configuration(args, train_valid_test_num_samples, ignored_flags) + + compile_helpers() + + config = core_gpt_dataset_config_from_args(args) + train_ds, valid_ds, test_ds = BlendedMegatronDatasetBuilder( + GPTDataset, train_valid_test_num_samples, lambda: True, config + ).build() + + print("> finished preparing dataset cache") + print(f" train dataset length: {_get_dataset_length(train_ds)}") + print(f" validation dataset length: {_get_dataset_length(valid_ds)}") + print(f" test dataset length: {_get_dataset_length(test_ds)}") + + return { + "world_size": args.world_size, + "data_parallel_size": args.data_parallel_size, + "global_batch_size": args.global_batch_size, + "train_valid_test_num_samples": tuple(train_valid_test_num_samples), + "train_dataset_length": _get_dataset_length(train_ds), + "valid_dataset_length": _get_dataset_length(valid_ds), + "test_dataset_length": _get_dataset_length(test_ds), + } + finally: + unset_global_variables() + + +def main() -> Dict[str, Any]: + args = parse_args( + extra_args_provider=_extra_args_provider, + ignore_unknown_args=False, + ) + _normalize_prepare_cache_args(args) + validate_args(args, defaults={"tokenizer_type": "GPT2BPETokenizer"}) + return build_dataset_caches(args) + + +if __name__ == "__main__": + main()