diff --git a/CHANGELOG.rst b/CHANGELOG.rst index ad8417a3c0e..df3cdefdb5e 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -19,6 +19,7 @@ Changelog **New Features** +- Add the ``prepare_megatron_data_blend`` utility to prepare weighted Megatron data blends from YAML configs, including optional token-budgeted subsets for distillation workflows. See the `Megatron data preparation guide `_. - Add Learned Scale Quantization (LSQ) and Dual-LSQ support for quantization-aware distillation, including learnable ``amax`` parameters, tied-scale and pre-scale options, focused NVFP4 recipes, and scale-only training. - Add the **D-PACE** loss objective for DFlash speculative-decoding training (`arXiv:2605.18810 `_) and make it the default (``dflash_loss_objective: dpace``). It replaces the static exponential position decay with dynamic, confidence-derived per-position weights that adapt to whichever block positions currently limit acceptance. Smoothing is controlled by ``dflash_dpace_alpha`` (default 0.5); set ``dflash_loss_objective: decay`` to restore the previous static schedule. Training-only and detached from the gradient (no architecture or inference change). - Add the ``day0-release`` agent skill (``.agents/skills/day0-release/``), a deterministic end-to-end driver that chains the PTQ → evaluation → comparison skills (the evaluation stage deploys the checkpoint itself) with an enforced gate after each stage and returns a publish decision (ACCEPT / REGRESSION / ANOMALOUS / INFEASIBLE). Ships three GPU-free, unit-tested gate scripts (``gate_ptq.py``, ``gate_run.py``, ``gate_compare.py``) that validate checkpoint coverage, evaluation-run completeness, and baseline-vs-candidate accuracy threshold. v1 reports and stops on regression; the recipe-search loop is deferred. diff --git a/examples/dataset/MEGATRON_DATA_PREP.md b/examples/dataset/MEGATRON_DATA_PREP.md index 99366bec7c6..91b89375907 100644 --- a/examples/dataset/MEGATRON_DATA_PREP.md +++ b/examples/dataset/MEGATRON_DATA_PREP.md @@ -4,6 +4,7 @@ | :---: | :---: | :---: | | From JSONL files | Tokenize local JSONL files | \[[Link](#from-jsonl-files)\] | | From Hugging Face Hub | Stream or download HF datasets and tokenize | \[[Link](#from-hugging-face-hub)\] | +| Token-budgeted data blends | Prepare weighted subsets for fast experiments | \[[Link](#prepare-token-budgeted-data-blends)\] | | `reasoning_content` for Post-Training v3 | Control how chain-of-thought traces are handled | \[[Link](#reasoning_content-for-post-training-v3-datasets)\] | | Nemotron Pre/Post-Training Datasets | Ready-to-run commands for all Nemotron datasets | \[[Link](#ready-to-run-tokenization-commands)\] | @@ -66,6 +67,111 @@ For very large datasets (tens of millions of documents), or datasets with comple > Re-runs read from cache and are much faster. > Streaming re-downloads on every run with no cache, so it is slower for full-dataset processing. +## Prepare token-budgeted data blends + +For iterative research, prepare smaller weighted datasets before scaling to a full distillation run. +Use [`prepare_megatron_data_blend`](../../modelopt/torch/utils/plugins/prepare_megatron_data_blend.py) to +prepare a weighted blend with a shared token budget. The utility supports Hugging Face configurations and splits +as well as specific JSONL files stored in a Hugging Face dataset repository. + +Define the tokenizer, output directory, and source weights in YAML. Set the optional `target_tokens` field to +prepare a weighted subset, or omit it to prepare every source in full. This example scales the +[Nemotron 3 Nano distillation blend](../megatron_bridge/tutorials/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16/README.md#1-data-preparation) +down to one billion tokens while preserving its source weights: + +> [!IMPORTANT] +> When `target_tokens` is set, JSONL records specified with `files` are consumed from the beginning +> of each file rather than selected randomly. Pre-shuffle JSONL files to obtain a random subset. +> Hugging Face dataset splits are shuffled deterministically; streaming datasets use an +> approximate buffer shuffle. + +```yaml +# Nemotron 3 models share this tokenizer, so the tokenized blend can be reused across the family. +tokenizer: nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16 +output_dir: /path/to/nemotron_3_distillation_blend_1b +# Optional; omit this field to prepare every source in full. +target_tokens: 1_000_000_000 +sources: + - hf_dataset: nvidia/Nemotron-Pretraining-SFT-v1 + config: Nemotron-SFT-Code + split: train + max_samples: 10_000_000 + content_field: text + weight: 5 + - hf_dataset: nvidia/Nemotron-Pretraining-SFT-v1 + config: Nemotron-SFT-General + split: train + max_samples: 10_000_000 + content_field: text + weight: 20 + - hf_dataset: nvidia/Nemotron-Pretraining-SFT-v1 + config: Nemotron-SFT-MATH + split: train + max_samples: 10_000_000 + content_field: text + weight: 5 + - hf_dataset: nvidia/Nemotron-Math-v2 + split: high_part00 + content_field: messages + weight: 10 + - hf_dataset: nvidia/Nemotron-SFT-Math-v3 + files: + - data/train.jsonl + content_field: messages + weight: 17 + - hf_dataset: nvidia/Nemotron-SFT-Competitive-Programming-v2 + files: + - data/competitive_programming_python_00.jsonl + content_field: messages + weight: 15 + - hf_dataset: nvidia/Nemotron-SFT-Competitive-Programming-v2 + files: + - data/competitive_programming_cpp_00.jsonl + content_field: messages + weight: 5 + - hf_dataset: nvidia/Nemotron-Post-Training-Dataset-v1 + config: default + split: stem + max_samples: 5_000_000 + content_field: messages + weight: 8 + - hf_dataset: nvidia/Nemotron-Science-v1 + files: + - data/MCQ.jsonl + content_field: messages + weight: 3 + - hf_dataset: nvidia/Nemotron-Science-v1 + files: + - data/RQA.jsonl + content_field: messages + weight: 2 + - hf_dataset: nvidia/Nemotron-SFT-Instruction-Following-Chat-v2 + files: + - data/reasoning_on.jsonl + content_field: messages + weight: 3 + - hf_dataset: nvidia/Nemotron-SFT-Instruction-Following-Chat-v2 + files: + - data/reasoning_off.jsonl + content_field: messages + weight: 2 + - hf_dataset: nvidia/Nemotron-Agentic-v1 + files: + - data/tool_calling.jsonl + content_field: messages + weight: 5 +``` + +With ModelOpt installed, run: + +```bash +python -m modelopt.torch.utils.plugins.prepare_megatron_data_blend --config blend.yaml +``` + +The output contains tokenized Megatron `.bin`/`.idx` files, `data_blend.txt` with the weighted paths for training, +and `config.yaml` recording how the blend was generated. The final token count can slightly exceed the target +because the final document from each source is kept whole. + ## `reasoning_content` for Post-Training v3 Datasets v3 datasets include a `reasoning_content` field in assistant messages (chain-of-thought separate from diff --git a/examples/megatron_bridge/tutorials/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16/README.md b/examples/megatron_bridge/tutorials/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16/README.md index 22a31ff9db7..cb4e5f9b650 100644 --- a/examples/megatron_bridge/tutorials/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16/README.md +++ b/examples/megatron_bridge/tutorials/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16/README.md @@ -70,6 +70,10 @@ Distillation uses the **30% Pretraining (Code 5, General 20, MATH 5) + 70% Post- ### 1. Data Preparation See [examples/dataset/MEGATRON_DATA_PREP.md](../../../dataset/MEGATRON_DATA_PREP.md) for tokenization commands for all datasets used in this blend. +To prepare a token-limited subset, follow the +[token-budgeted data blend workflow](../../../dataset/MEGATRON_DATA_PREP.md#prepare-token-budgeted-data-blends), +but create a custom YAML configuration using this tutorial's tokenizer, sources, and weights below. The +example configuration targets Nemotron 3 and should not be reused unchanged. For this experiment: `TOKENIZER=nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16`, `OUTPUT_DIR=tokenized_nemotron_3`. diff --git a/examples/pruning/minitron/NVIDIA-Nemotron-Nano-9B-v2/README.md b/examples/pruning/minitron/NVIDIA-Nemotron-Nano-9B-v2/README.md index 55d4706175d..d1a75e43098 100644 --- a/examples/pruning/minitron/NVIDIA-Nemotron-Nano-9B-v2/README.md +++ b/examples/pruning/minitron/NVIDIA-Nemotron-Nano-9B-v2/README.md @@ -63,6 +63,10 @@ Distillation uses the **30% Pretraining (Code 5, General 20, MATH 5) + 70% Post- ### 1. Data Preparation See [examples/dataset/MEGATRON_DATA_PREP.md](../../../dataset/MEGATRON_DATA_PREP.md) for tokenization commands for all datasets used in this blend. +To prepare a token-limited subset, follow the +[token-budgeted data blend workflow](../../../dataset/MEGATRON_DATA_PREP.md#prepare-token-budgeted-data-blends), +but create a custom YAML configuration using this tutorial's tokenizer, sources, and weights below. The +example configuration targets Nemotron 3 and should not be reused unchanged. For this experiment: `TOKENIZER=nvidia/NVIDIA-Nemotron-Nano-9B-v2`, `OUTPUT_DIR=tokenized_nemotron_v2`. diff --git a/examples/researcher_guide/README.md b/examples/researcher_guide/README.md index bb9b32f5677..dd0334bf12a 100644 --- a/examples/researcher_guide/README.md +++ b/examples/researcher_guide/README.md @@ -4,9 +4,14 @@ Model optimization research depends on short feedback loops: test a hypothesis c reproducibly, and spend full-scale compute only on the most promising experiments. This guide collects practical ModelOpt workflows for that iterative research process. -The guide starts with efficient model evaluation and will grow as additional research workflows are documented. -It complements the feature-specific [examples](../) by connecting them into experimentation strategies rather -than replacing their detailed instructions. +Current workflows include: + +- [Efficient model evaluation](#efficient-evaluation-with-lm-eval-harness) with smaller benchmark subsets. +- [Efficient data blend preparation](#prepare-token-budgeted-data-blends) for distillation experiments. + +The guide will grow as additional research workflows are documented. It complements the feature-specific +[examples](../) by connecting them into experimentation strategies rather than replacing their detailed +instructions. ## Efficient evaluation with LM-Eval Harness @@ -42,6 +47,13 @@ and should not be reported as final benchmark results. Add `--log_samples` for paired per-question analysis. When multiple GPUs are available, use data parallelism to split samples across model copies; see the [LM-Eval examples](../llm_eval/README.md) for commands. +## Prepare token-budgeted data blends + +Preparing complete distillation datasets can consume unnecessary time and storage during early experiments. +ModelOpt can preserve source weights while preparing only a requested token budget. See +[Prepare token-budgeted data blends](../dataset/MEGATRON_DATA_PREP.md#prepare-token-budgeted-data-blends) for the +configuration format, commands, and generated outputs. + ## Planned topics Future additions can cover: diff --git a/modelopt/torch/utils/plugins/__init__.py b/modelopt/torch/utils/plugins/__init__.py index f2f02852906..da40fe9e565 100644 --- a/modelopt/torch/utils/plugins/__init__.py +++ b/modelopt/torch/utils/plugins/__init__.py @@ -29,6 +29,9 @@ with import_plugin("megatron_preprocess_data"): from .megatron_preprocess_data import * +with import_plugin("prepare_megatron_data_blend"): + from .prepare_megatron_data_blend import * + # NOTE: Dont pre-import megatron bridge plugin here to avoid circular dependency issues. # We dont register anything so this isnt a problem. # with import_plugin("megatron bridge"): diff --git a/modelopt/torch/utils/plugins/megatron_preprocess_data.py b/modelopt/torch/utils/plugins/megatron_preprocess_data.py index 52dae7c140b..dcbf9a6db1d 100644 --- a/modelopt/torch/utils/plugins/megatron_preprocess_data.py +++ b/modelopt/torch/utils/plugins/megatron_preprocess_data.py @@ -85,6 +85,7 @@ import argparse import gzip +import itertools import json import multiprocessing import time @@ -261,7 +262,26 @@ def _print_processing_stats( flush=True, ) - def _encode_docs(self, encoder: "_Encoder", lines): + @staticmethod + def _encode_in_batches(pool, encoder: "_Encoder", lines, batch_size: int): + """Encode finite batches for a run bounded by ``max_tokens``. + + Once the token target is reached, the caller stops without consuming the remaining input. + Batches run one at a time, while documents within each batch are encoded in parallel. Unlike + ``imap``, ``map`` collects every result from the current batch before returning. Therefore, + no worker is writing a pending result when the caller stops at the token target. + The final batch may encode documents that the caller does not use after reaching its limit. + """ + lines = iter(lines) + while True: + batch = list(itertools.islice(lines, batch_size)) + if not batch: + break + + encoded_batch = pool.map(encoder.encode, batch, chunksize=1) + yield from encoded_batch + + def _encode_docs(self, encoder: "_Encoder", lines, may_stop_early: bool = False): """Tokenize ``lines``, forking worker processes only when ``workers > 1``. ``multiprocessing.Pool`` always ``fork()``s, even for a single worker. Forking a @@ -269,21 +289,44 @@ def _encode_docs(self, encoder: "_Encoder", lines): this is called in-process after GPU work) is unsafe and can segfault the children. The single-worker path avoids the fork entirely by tokenizing inline in this process. + When ``may_stop_early`` is true, wait for finite batches so no worker results are pending + if the caller stops consuming documents after reaching its token limit. + Returns ``(pool, encoded_docs)``; ``pool`` is ``None`` in the inline case. """ if self.workers == 1: encoder.initializer() return None, map(encoder.encode, lines) pool = multiprocessing.Pool(self.workers, initializer=encoder.initializer) - return pool, pool.imap(encoder.encode, lines, 32) + if may_stop_early: + batch_size = self.workers * 4 # Balance throughput against unused final-batch work. + encoded_docs = self._encode_in_batches(pool, encoder, lines, batch_size) + else: + encoded_docs = pool.imap(encoder.encode, lines, 32) + return pool, encoded_docs + + @staticmethod + def _cached_token_count(prefixes: list[str]) -> int: + """Return the number of tokens represented by cached Megatron datasets.""" + token_count = 0 + for prefix in prefixes: + # Avoid memory-mapping the token file because a cached split can be empty. + dataset = indexed_dataset.IndexedDataset(prefix, mmap=False) + token_count += int(dataset.sequence_lengths.sum()) + return token_count def process_json_file( - self, input_file_name: str | Path, output_dir: str | Path, encoder: _Encoder + self, + input_file_name: str | Path, + output_dir: str | Path, + encoder: _Encoder, + max_tokens: int | None = None, ) -> tuple[int, list[str]]: input_path = Path(input_file_name) stem = input_path.stem if input_path.suffix != ".gz" else Path(input_path.stem).stem output_prefix = Path(output_dir) / stem - prefixes = [f"{output_prefix}_{key}" for key in self.json_keys] + token_tag = f"_tokens{max_tokens}" if max_tokens is not None else "" + prefixes = [f"{output_prefix}_{key}{token_tag}" for key in self.json_keys] print(f"\nOpening {input_file_name}") if input_path.suffix == ".gz": @@ -291,15 +334,17 @@ def process_json_file( else: fin = open(input_path, encoding="utf-8") - pool, encoded_docs = self._encode_docs(encoder, fin) + # Workers encode asynchronously; iterating encoded_docs waits for results in input order. + pool, encoded_docs = self._encode_docs(encoder, fin, may_stop_early=max_tokens is not None) output_bin_files = {} output_idx_files = {} builders = {} for key in self.json_keys: - output_bin_files[key] = f"{output_prefix}_{key}.bin" - output_idx_files[key] = f"{output_prefix}_{key}.idx" + prefix = f"{output_prefix}_{key}{token_tag}" + output_bin_files[key] = f"{prefix}.bin" + output_idx_files[key] = f"{prefix}.idx" if Path(output_bin_files[key]).exists() and Path(output_idx_files[key]).exists(): continue builders[key] = indexed_dataset.IndexedDatasetBuilder( @@ -309,10 +354,11 @@ def process_json_file( if not builders: print(f"\t[SKIP] Output files corresponding to {input_file_name} already exist") - return 0, prefixes + return self._cached_token_count(prefixes), prefixes start_time = time.time() - total_doc_len, total_enc_len, final_enc_len = 0, 0, 0 + total_doc_len, total_enc_len = 0, 0 + final_enc_len = 0 # Tokens written to output, including appended EOD tokens. for i, (doc, sentence_lens, (doc_len, enc_len)) in enumerate(encoded_docs, start=1): total_doc_len += doc_len total_enc_len += enc_len @@ -320,6 +366,8 @@ def process_json_file( for key in doc: builders[key].add_document(doc[key], sentence_lens[key]) self._print_processing_stats(i, total_doc_len, total_enc_len, start_time) + if max_tokens is not None and final_enc_len >= max_tokens: + break self._print_processing_stats(i, total_doc_len, total_enc_len, start_time, force_print=True) fin.close() @@ -345,6 +393,7 @@ def process_hf_split( config: str | None, split: str, max_samples: int | None = None, + max_tokens: int | None = None, streaming: bool = False, ) -> tuple[int, list[str]]: """Load a HF dataset split and tokenize directly without writing an intermediate JSONL. @@ -357,11 +406,12 @@ def process_hf_split( """ print(f"\nLoading HF dataset {dataset_name=}, {config=}, {split=}, {streaming=}") ds = load_dataset(path=dataset_name, name=config, split=split, streaming=streaming) - if max_samples is not None: + if max_samples is not None or max_tokens is not None: # Shuffle first so the selected subset is random, not a biased prefix. # Non-streaming: global index shuffle (memory-mapped, efficient) then .select(N). # Streaming: buffer shuffle (approximate) then .take(N). ds = ds.shuffle(seed=42) + if max_samples is not None: if streaming: ds = ds.take(max_samples) else: @@ -378,6 +428,7 @@ def process_hf_split( safe_name = dataset_name.replace("/", "--") sample_tag = f"_max{max_samples}" if max_samples is not None else "" + token_tag = f"_tokens{max_tokens}" if max_tokens is not None else "" output_prefix = Path(output_dir) / f"{safe_name}_{config}_{split}" prefixes = [] @@ -385,7 +436,7 @@ def process_hf_split( output_idx_files = {} builders = {} for key in self.json_keys: - prefix = f"{output_prefix}_{key}{sample_tag}" + prefix = f"{output_prefix}_{key}{sample_tag}{token_tag}" prefixes.append(prefix) output_bin_files[key] = f"{prefix}.bin" output_idx_files[key] = f"{prefix}.idx" @@ -398,12 +449,16 @@ def process_hf_split( if not builders: print(f"\t[SKIP] Output files for {dataset_name} {config}/{split} already exist") - return 0, prefixes + return self._cached_token_count(prefixes), prefixes - pool, encoded_docs = self._encode_docs(encoder, self._iter_hf_as_json(ds)) + # Workers encode asynchronously; iterating encoded_docs waits for results in input order. + pool, encoded_docs = self._encode_docs( + encoder, self._iter_hf_as_json(ds), may_stop_early=max_tokens is not None + ) start_time = time.time() - total_doc_len, total_enc_len, final_enc_len = 0, 0, 0 + total_doc_len, total_enc_len = 0, 0 + final_enc_len = 0 # Tokens written to output, including appended EOD tokens. i = 0 for i, (doc, sentence_lens, (doc_len, enc_len)) in enumerate(encoded_docs, start=1): total_doc_len += doc_len @@ -412,6 +467,8 @@ def process_hf_split( for key in doc: builders[key].add_document(doc[key], sentence_lens[key]) self._print_processing_stats(i, total_doc_len, total_enc_len, start_time) + if max_tokens is not None and final_enc_len >= max_tokens: + break if i: self._print_processing_stats( @@ -462,6 +519,7 @@ def megatron_preprocess_data( hf_split: str | None = None, hf_max_samples_per_split: int | None = None, hf_streaming: bool = False, + max_tokens: int | None = None, # Other arguments output_dir: str | Path, tokenizer_name_or_path: str, @@ -477,6 +535,11 @@ def megatron_preprocess_data( Exactly one of ``input_dir``, ``jsonl_paths``, or ``hf_dataset`` must be provided. + Important: When ``max_tokens`` is set, JSONL records are consumed from the beginning of each + file rather than selected randomly. Pre-shuffle JSONL files to obtain a random subset. Hugging + Face datasets are shuffled deterministically before applying the limit; streaming datasets use + an approximate buffer shuffle. + Args: input_dir: Directory containing JSONL files to tokenize. jsonl_paths: One or more paths to JSONL files. @@ -488,6 +551,8 @@ def megatron_preprocess_data( downloaded — useful for very large pretraining datasets or datasets with complex nested message schemas that cause Arrow type-cast errors in non-streaming mode. Note: streaming does not cache to disk, so re-runs re-download. Defaults to False. + max_tokens: Stop after processing at least this many tokens across the source files or + selected Hugging Face splits. The final document may make the result slightly larger. output_dir: Path to directory to save binary output files. tokenizer_name_or_path: Name or path of the Hugging Face tokenizer to use. json_keys: Key or list of keys to extract from json. Defaults to ["text"]. @@ -515,11 +580,16 @@ def megatron_preprocess_data( raise ValueError( "Exactly one of `input_dir`, `jsonl_paths`, or `hf_dataset` must be provided." ) - if hf_streaming and hf_max_samples_per_split is None and _is_main_or_first_worker(): + if ( + hf_streaming + and hf_max_samples_per_split is None + and max_tokens is None + and _is_main_or_first_worker() + ): warnings.warn( - "--hf_streaming is set but --hf_max_samples_per_split is not. " - "Streaming without a sample cap re-downloads the full dataset on every run with no " - "disk cache, which is slower than the cached non-streaming path.", + "--hf_streaming is set but neither --hf_max_samples_per_split nor --max_tokens is " + "set. Streaming without a sample or token cap re-downloads the full dataset on " + "every run with no disk cache, which is slower than the cached non-streaming path.", stacklevel=2, ) @@ -536,12 +606,16 @@ def megatron_preprocess_data( ) partition = _Partition(vocab_size, json_keys, log_interval, workers) + # Tokens written across all input files or Hugging Face splits. final_enc_len = 0 all_prefixes: list[str] = [] overall_start = time.time() if hf_dataset is not None: for config, split in _enumerate_hf_splits(hf_dataset, hf_name, hf_split): + remaining_tokens = None if max_tokens is None else max_tokens - final_enc_len + if remaining_tokens is not None and remaining_tokens <= 0: + break enc_len, prefixes = partition.process_hf_split( output_dir, encoder, @@ -549,6 +623,7 @@ def megatron_preprocess_data( config, split, hf_max_samples_per_split, + remaining_tokens, hf_streaming, ) final_enc_len += enc_len @@ -566,10 +641,18 @@ def megatron_preprocess_data( file_names = list(jsonl_paths) # type: ignore[arg-type] for name in file_names: - enc_len, prefixes = partition.process_json_file(name, output_dir, encoder) + remaining_tokens = None if max_tokens is None else max_tokens - final_enc_len + if remaining_tokens is not None and remaining_tokens <= 0: + break + enc_len, prefixes = partition.process_json_file( + name, output_dir, encoder, remaining_tokens + ) final_enc_len += enc_len all_prefixes.extend(prefixes) + if max_tokens is not None and final_enc_len >= max_tokens: + print(f"\n>>> Early stopping: {max_tokens=} achieved with {final_enc_len} tokens.") + elapsed = (time.time() - overall_start) / 60 print( f"\n\n>>> Total number of tokens currently processed: {num2hrb(final_enc_len)}" @@ -633,6 +716,12 @@ def main(): parser.add_argument( "--max_sequence_length", type=int, default=None, help="Maximum sequence length" ) + parser.add_argument( + "--max_tokens", + type=int, + default=None, + help="Stop after processing at least this many tokens", + ) parser.add_argument("--workers", type=int, default=8, help="Number of worker processes") parser.add_argument("--log_interval", type=int, default=100000, help="Log interval") parser.add_argument( @@ -675,6 +764,7 @@ def main(): json_keys=args.json_keys, append_eod=args.append_eod, max_sequence_length=args.max_sequence_length, + max_tokens=args.max_tokens, workers=args.workers, log_interval=args.log_interval, reasoning_content=args.reasoning_content, diff --git a/modelopt/torch/utils/plugins/prepare_megatron_data_blend.py b/modelopt/torch/utils/plugins/prepare_megatron_data_blend.py new file mode 100644 index 00000000000..fa5875a5e2c --- /dev/null +++ b/modelopt/torch/utils/plugins/prepare_megatron_data_blend.py @@ -0,0 +1,173 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Prepare a weighted Megatron data blend from a YAML configuration.""" + +import argparse +import os +import shutil +from pathlib import Path +from typing import Any, cast + +import huggingface_hub +import yaml + +from .megatron_preprocess_data import megatron_preprocess_data + +__all__ = ["prepare_megatron_data_blend"] + + +def load_config(path: Path) -> dict[str, Any]: + """Load a data-blend YAML configuration as a dictionary. + + For example, this YAML:: + + tokenizer: /models/Qwen3-8B + output_dir: /datasets/qwen3-blend + target_tokens: 1000000 # Optional; omit to prepare every source in full. + sources: + - hf_dataset: nvidia/Nemotron-Pretraining-SFT-v1 + config: Nemotron-SFT-General + split: train + content_field: text + weight: 60 + - hf_dataset: nvidia/Nemotron-SFT-Competitive-Programming-v2 + files: + - data/competitive_programming_python_00.jsonl + content_field: messages + weight: 40 + + returns a dictionary with ``tokenizer``, ``output_dir``, and ``sources`` keys, plus + optional ``target_tokens``. Each source has ``hf_dataset``, ``content_field``, and + ``weight``; it uses ``split`` with optional ``config`` and ``max_samples``, or + selects repository ``files``. + """ + with path.open(encoding="utf-8") as stream: + return cast("dict[str, Any]", yaml.safe_load(stream)) + + +def _build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--config", type=Path, required=True, help="Path to the blend YAML file") + return parser + + +def _write_data_blend(path: Path, blend: list[tuple[float, str]]) -> None: + """Write the weighted Megatron dataset paths.""" + content = "\n".join(f"{weight:g} {prefix}" for weight, prefix in blend) + "\n" + path.write_text(content, encoding="utf-8") + + +def _copy_config(source: Path, destination: Path) -> None: + """Copy the input configuration alongside the generated blend.""" + if source.resolve() != destination.resolve(): + shutil.copyfile(source, destination) + + +def _prepare_sources( + sources: list[dict[str, Any]], + output_dir: Path, + tokenizer: str, + total_tokens: int | None, +) -> list[tuple[float, str]]: + """Tokenize all sources and return their weighted output paths.""" + workers = min(32, os.cpu_count() or 1) + blend: list[tuple[float, str]] = [] # (weight, shared .bin/.idx path without extension) + allocated_tokens = 0 + # Weights are relative, not required to sum to 100 (matching data_blend.txt semantics). + weight_sum = sum(float(source["weight"]) for source in sources) + + for index, source in enumerate(sources): + weight = float(source["weight"]) + if total_tokens is None: + source_tokens = None + elif index == len(sources) - 1: + source_tokens = total_tokens - allocated_tokens + else: + source_tokens = round(total_tokens * weight / weight_sum) + allocated_tokens += source_tokens + + dataset = source["hf_dataset"] + source_dir = output_dir / f"{index:02d}_{dataset.replace('/', '--')}" + content_field = source["content_field"] + input_args: dict[str, Any] + if "files" in source: + raw_dir = output_dir.parent / "raw" / dataset.replace("/", "--") + paths = [ + huggingface_hub.hf_hub_download( + repo_id=dataset, + filename=file, + repo_type="dataset", + local_dir=raw_dir, + ) + for file in source["files"] + ] + input_args = {"jsonl_paths": paths} + else: + input_args = { + "hf_dataset": dataset, + "hf_name": source.get("config"), + "hf_split": source["split"], + "hf_max_samples_per_split": source.get("max_samples"), + "hf_streaming": True, + } + + # Each prefix is the path shared by a tokenized Megatron .bin/.idx file pair. + prefixes = megatron_preprocess_data( + **input_args, + output_dir=source_dir, + tokenizer_name_or_path=tokenizer, + json_keys=content_field, + # Plain text lacks chat-template boundary tokens, so terminate each document with EOS. + append_eod=content_field == "text", + # Join lines in text documents by replacing each newline with a space. + strip_newlines=content_field == "text", + reasoning_content="inline" if content_field == "messages" else "strip", + # Guard against pathological records by capping each tokenized document at 256K tokens. + max_sequence_length=256_000, + max_tokens=source_tokens, + workers=workers, + ) + prefix_weight = weight / len(prefixes) + blend.extend((prefix_weight, prefix) for prefix in prefixes) + + return blend + + +def prepare_megatron_data_blend(config_path: Path) -> list[tuple[float, str]]: + """Download and tokenize the configured weighted data sources.""" + config = load_config(config_path) + output_dir = Path(config["output_dir"]) + output_dir.mkdir(parents=True, exist_ok=True) + target_tokens = config.get("target_tokens") + total_tokens = None if target_tokens is None else int(target_tokens) + tokenizer = str(config["tokenizer"]) + + blend = _prepare_sources(config["sources"], output_dir, tokenizer, total_tokens) + _write_data_blend(output_dir / "data_blend.txt", blend) + _copy_config(config_path, output_dir / "config.yaml") + return blend + + +def main() -> None: + """Prepare a data blend from the supplied configuration.""" + parser = _build_parser() + args = parser.parse_args() + blend = prepare_megatron_data_blend(args.config) + print(f"Prepared {len(blend)} data paths. See data_blend.txt and config.yaml in the output.") + + +if __name__ == "__main__": + main() diff --git a/tests/gpu_megatron/conftest.py b/tests/gpu_megatron/conftest.py index d84ea99765b..b8176adedd0 100644 --- a/tests/gpu_megatron/conftest.py +++ b/tests/gpu_megatron/conftest.py @@ -17,10 +17,19 @@ import pytest import torch from _test_utils.torch.distributed.utils import DistributedWorkerPool +from _test_utils.torch.transformers_models import get_tiny_tokenizer from megatron.core.parallel_state import destroy_model_parallel import modelopt.torch.utils.distributed as dist + +@pytest.fixture(scope="session") +def tiny_tokenizer_path(tmp_path_factory): + tokenizer_path = tmp_path_factory.mktemp("tiny_tokenizer") + get_tiny_tokenizer().save_pretrained(tokenizer_path) + return str(tokenizer_path) + + apex_destroy = None with contextlib.suppress(ImportError): from apex.transformer.parallel_state import destroy_model_parallel as apex_destroy diff --git a/tests/gpu_megatron/torch/utils/plugins/test_megatron_preprocess_data.py b/tests/gpu_megatron/torch/utils/plugins/test_megatron_preprocess_data.py index c0453484056..ec9b11ff2e3 100644 --- a/tests/gpu_megatron/torch/utils/plugins/test_megatron_preprocess_data.py +++ b/tests/gpu_megatron/torch/utils/plugins/test_megatron_preprocess_data.py @@ -19,14 +19,15 @@ import pytest -from modelopt.torch.utils.dataset_utils import download_hf_dataset_as_jsonl from modelopt.torch.utils.plugins.megatron_preprocess_data import megatron_preprocess_data def test_megatron_preprocess_data_with_jsonl_path(tmp_path): - input_jsonl = download_hf_dataset_as_jsonl("nanotron/minipile_100_samples", tmp_path / "raw") - assert len(input_jsonl) == 1, "Expected 1 JSONL file" - input_jsonl = Path(input_jsonl[0]) + input_jsonl = tmp_path / "minipile.jsonl" + input_jsonl.write_text( + "".join(json.dumps({"text": f"Sample document {index}."}) + "\n" for index in range(4)), + encoding="utf-8", + ) assert input_jsonl.stat().st_size > 0, "Input JSONL file should not be empty" @@ -53,6 +54,81 @@ def test_megatron_preprocess_data_with_jsonl_path(tmp_path): assert Path(prefixes[0] + ".idx").stat().st_size > 0, "Index file should not be empty" +def test_megatron_preprocess_data_jsonl_stops_at_max_tokens(tmp_path): + input_path = tmp_path / "data.jsonl" + input_path.write_text( + "".join(json.dumps({"text": f"Document {index} " * 100}) + "\n" for index in range(10)), + encoding="utf-8", + ) + + common_args = { + "jsonl_paths": input_path, + "output_dir": tmp_path, + "tokenizer_name_or_path": "gpt2", + "json_keys": "text", + "workers": 2, + } + limited_prefix = megatron_preprocess_data( + **common_args, + max_tokens=100, + )[0] + full_prefix = megatron_preprocess_data(**common_args)[0] + + # The .bin file stores token IDs, so its byte size reflects how many tokens were written. + limited_size = Path(limited_prefix + ".bin").stat().st_size + full_size = Path(full_prefix + ".bin").stat().st_size + assert limited_prefix == str(tmp_path / "data_text_tokens100") + assert full_prefix == str(tmp_path / "data_text") + assert limited_size < full_size + + +def test_megatron_preprocess_data_hf_split_stops_at_max_tokens(tmp_path): + common_args = { + "hf_dataset": "nanotron/minipile_100_samples", + "hf_split": "train", + "hf_max_samples_per_split": 100, + "hf_streaming": True, + "tokenizer_name_or_path": "gpt2", + "json_keys": "text", + "workers": 2, + } + limited_prefix = megatron_preprocess_data( + **common_args, + output_dir=tmp_path / "limited", + max_tokens=100, + )[0] + full_prefix = megatron_preprocess_data( + **common_args, + output_dir=tmp_path / "full", + )[0] + + # The .bin file stores token IDs, so its byte size reflects how many tokens were written. + limited_size = Path(limited_prefix + ".bin").stat().st_size + full_size = Path(full_prefix + ".bin").stat().st_size + assert limited_size < full_size + + +def test_megatron_preprocess_data_hf_split_resume_uses_cached_token_count(tmp_path): + args = { + "hf_dataset": "nanotron/minipile_100_samples", + "hf_split": "train", + "hf_max_samples_per_split": 1, + "hf_streaming": True, + "max_tokens": 100, + "output_dir": tmp_path, + "tokenizer_name_or_path": "gpt2", + "json_keys": "text", + "max_sequence_length": 16, + "workers": 2, + } + + prefixes = megatron_preprocess_data(**args) + cached_files = set(tmp_path.iterdir()) + + assert megatron_preprocess_data(**args) == prefixes + assert set(tmp_path.iterdir()) == cached_files + + @pytest.mark.parametrize( ("hf_dataset", "hf_split", "json_keys"), [ diff --git a/tests/gpu_megatron/torch/utils/plugins/test_prepare_megatron_data_blend.py b/tests/gpu_megatron/torch/utils/plugins/test_prepare_megatron_data_blend.py new file mode 100644 index 00000000000..f531be27ad4 --- /dev/null +++ b/tests/gpu_megatron/torch/utils/plugins/test_prepare_megatron_data_blend.py @@ -0,0 +1,105 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import json +import sys +from pathlib import Path +from unittest.mock import Mock + +import huggingface_hub +import pytest +import yaml + +from modelopt.torch.utils.plugins.prepare_megatron_data_blend import main + + +def _setup_test( + tiny_tokenizer_path: str, tmp_path: Path, monkeypatch, target_tokens: int | None +) -> tuple[Path, Path, Mock]: + output_dir = tmp_path / "tokenized" + config = { + "tokenizer": tiny_tokenizer_path, + "output_dir": str(output_dir), + "sources": [ + { + "hf_dataset": "nanotron/minipile_100_samples", + "split": "train", + "max_samples": 100, + "content_field": "text", + "weight": 60, + }, + { + "hf_dataset": "nvidia/Nemotron-SFT-Competitive-Programming-v2", + "files": ["data/competitive_programming_python_00.jsonl"], + "content_field": "messages", + "weight": 40, + }, + ], + } + config_path = tmp_path / "config.yaml" + config_yaml = yaml.safe_dump(config) + if target_tokens is not None: + config_yaml += f"target_tokens: {target_tokens:_}\n" + config_path.write_text(config_yaml, encoding="utf-8") + jsonl_path = tmp_path / "competitive_programming_python_00.jsonl" + conversation = { + "messages": [ + {"role": "user", "content": "Write a Python function that adds two integers."}, + {"role": "assistant", "content": "def add(a, b):\n return a + b"}, + ] + } + jsonl_path.write_text( + "".join(json.dumps(conversation) + "\n" for _ in range(20)), encoding="utf-8" + ) + download = Mock(return_value=str(jsonl_path)) + monkeypatch.setattr(huggingface_hub, "hf_hub_download", download) + monkeypatch.setattr( + sys, "argv", ["prepare_megatron_data_blend.py", "--config", str(config_path)] + ) + return output_dir, config_path, download + + +@pytest.mark.parametrize("target_tokens", [1_000, None], ids=["token-budget", "all-data"]) +def test_prepare_megatron_data_blend_with_split_and_files_sources( + tiny_tokenizer_path: str, tmp_path: Path, monkeypatch, target_tokens: int | None +): + output_dir, config_path, download = _setup_test( + tiny_tokenizer_path, tmp_path, monkeypatch, target_tokens + ) + + # Run in-process so the CLI entry point uses the mocked NVIDIA download. + main() + + download.assert_called_once_with( + repo_id="nvidia/Nemotron-SFT-Competitive-Programming-v2", + filename="data/competitive_programming_python_00.jsonl", + repo_type="dataset", + local_dir=tmp_path / "raw/nvidia--Nemotron-SFT-Competitive-Programming-v2", + ) + blend = [ + line.split(maxsplit=1) + for line in (output_dir / "data_blend.txt").read_text(encoding="utf-8").splitlines() + ] + assert [weight for weight, _ in blend] == ["60", "40"] + for _, prefix in blend: + assert Path(prefix + ".bin").exists() + assert Path(prefix + ".idx").exists() + token_suffixes = ["_tokens600", "_tokens400"] if target_tokens is not None else ["", ""] + # HF split prefixes use {dataset}_{config}_{split}_{field}_max{samples}. + assert [Path(prefix).name for _, prefix in blend] == [ + f"nanotron--minipile_100_samples_default_train_text_max100{token_suffixes[0]}", + f"competitive_programming_python_00_messages{token_suffixes[1]}", + ] + assert (output_dir / "config.yaml").read_bytes() == config_path.read_bytes()