From 510a414a346bfffc1ea5312c95a7315955890ea7 Mon Sep 17 00:00:00 2001 From: Daniel Korzekwa Date: Wed, 1 Jul 2026 12:37:21 -0700 Subject: [PATCH 01/23] Implement prepare_data_blend utility. Signed-off-by: Daniel Korzekwa --- examples/dataset/prepare_data_blend.py | 70 ++++++++++++++++++++++++++ 1 file changed, 70 insertions(+) create mode 100644 examples/dataset/prepare_data_blend.py diff --git a/examples/dataset/prepare_data_blend.py b/examples/dataset/prepare_data_blend.py new file mode 100644 index 00000000000..aeb8c1d862e --- /dev/null +++ b/examples/dataset/prepare_data_blend.py @@ -0,0 +1,70 @@ +# 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 token-sized Megatron data blend from a YAML configuration.""" + +import argparse +from pathlib import Path +from typing import Any, cast + +import yaml + +__all__ = ["load_config", "main"] + + +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 + 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``, ``target_tokens``, and + ``sources`` keys. Each source has ``hf_dataset``, ``content_field``, and ``weight``; + it uses ``split`` with an optional ``config``, 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 main() -> None: + """Prepare a data blend from the supplied configuration.""" + parser = _build_parser() + args = parser.parse_args() + load_config(args.config) + print(f"Parsed configuration from {args.config}. Data preparation is not implemented yet.") + + +if __name__ == "__main__": + main() From d22f71cd7acd366fd8f8c066e812890ffd0f881e Mon Sep 17 00:00:00 2001 From: Daniel Korzekwa Date: Wed, 1 Jul 2026 12:49:39 -0700 Subject: [PATCH 02/23] remodve not needed __all__ variable Signed-off-by: Daniel Korzekwa --- examples/dataset/prepare_data_blend.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/examples/dataset/prepare_data_blend.py b/examples/dataset/prepare_data_blend.py index aeb8c1d862e..71bd4957b9d 100644 --- a/examples/dataset/prepare_data_blend.py +++ b/examples/dataset/prepare_data_blend.py @@ -21,8 +21,6 @@ import yaml -__all__ = ["load_config", "main"] - def load_config(path: Path) -> dict[str, Any]: """Load a data-blend YAML configuration as a dictionary. From 52e6035ecd7462ee9a0562df36e5fa36fdf05148 Mon Sep 17 00:00:00 2001 From: Daniel Korzekwa Date: Thu, 2 Jul 2026 08:04:50 -0700 Subject: [PATCH 03/23] implement prepare_data_blend.py Signed-off-by: Daniel Korzekwa --- examples/dataset/prepare_data_blend.py | 100 +++++++++++++++++- .../utils/plugins/megatron_preprocess_data.py | 98 ++++++++++++++--- 2 files changed, 182 insertions(+), 16 deletions(-) diff --git a/examples/dataset/prepare_data_blend.py b/examples/dataset/prepare_data_blend.py index 71bd4957b9d..37c29c4e67a 100644 --- a/examples/dataset/prepare_data_blend.py +++ b/examples/dataset/prepare_data_blend.py @@ -16,10 +16,15 @@ """Prepare a token-sized Megatron data blend from a YAML configuration.""" import argparse +import os +import shutil from pathlib import Path from typing import Any, cast import yaml +from huggingface_hub import hf_hub_download + +from modelopt.torch.utils.plugins.megatron_preprocess_data import megatron_preprocess_data def load_config(path: Path) -> dict[str, Any]: @@ -56,12 +61,103 @@ def _build_parser() -> argparse.ArgumentParser: 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, +) -> 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 + + for index, source in enumerate(sources): + weight = float(source["weight"]) + if index == len(sources) - 1: + source_tokens = total_tokens - allocated_tokens + else: + source_tokens = round(total_tokens * weight / 100) + 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 = [ + 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_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_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) + total_tokens = int(config["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() - load_config(args.config) - print(f"Parsed configuration from {args.config}. Data preparation is not implemented yet.") + blend = prepare_data_blend(args.config) + print(f"Prepared {len(blend)} data paths. See data_blend.txt and config.yaml in the output.") if __name__ == "__main__": diff --git a/modelopt/torch/utils/plugins/megatron_preprocess_data.py b/modelopt/torch/utils/plugins/megatron_preprocess_data.py index 52dae7c140b..04e59b0b6ae 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,16 +289,28 @@ 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 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 @@ -291,7 +323,8 @@ 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 = {} @@ -312,7 +345,8 @@ def process_json_file( return 0, 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 +354,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 +381,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 +394,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 +416,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 +424,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" @@ -400,10 +439,14 @@ def process_hf_split( print(f"\t[SKIP] Output files for {dataset_name} {config}/{split} already exist") return 0, 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 +455,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 +507,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, @@ -488,6 +534,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 split. 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 +563,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 +589,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 +606,7 @@ def megatron_preprocess_data( config, split, hf_max_samples_per_split, + remaining_tokens, hf_streaming, ) final_enc_len += enc_len @@ -566,7 +624,12 @@ 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) @@ -633,6 +696,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 +744,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, From 8d40667bec39acdeb1ca83239f3706c57185e147 Mon Sep 17 00:00:00 2001 From: Daniel Korzekwa Date: Thu, 2 Jul 2026 09:55:07 -0700 Subject: [PATCH 04/23] created a tool for efficient data blend preparation Signed-off-by: Daniel Korzekwa --- examples/dataset/MEGATRON_DATA_PREP.md | 4 + examples/dataset/prepare_data_blend.py | 21 ++- .../NVIDIA-Nemotron-Nano-9B-v2/README.md | 9 +- examples/researcher_guide/README.md | 147 ++++++++++++++++++ .../dataset/test_prepare_data_blend.py | 99 ++++++++++++ .../plugins/test_megatron_preprocess_data.py | 55 +++++++ 6 files changed, 326 insertions(+), 9 deletions(-) create mode 100644 examples/researcher_guide/README.md create mode 100644 tests/examples/dataset/test_prepare_data_blend.py diff --git a/examples/dataset/MEGATRON_DATA_PREP.md b/examples/dataset/MEGATRON_DATA_PREP.md index 99366bec7c6..357f16da7c2 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](../researcher_guide/README.md#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)\] | @@ -11,6 +12,9 @@ The distillation and pre-training scripts in Megatron-Bridge or Megatron-LM expe Use the `megatron_preprocess_data` utility to tokenize any JSONL or Hugging Face dataset. The tokenization scripts below print the list of output prefixes (e.g. `tokenized_qwen3/data1_text`) that you can use for the `data_paths` argument (with relative weights on different files) in Megatron training scripts. +For iterative research, use the [token-budgeted data blend workflow](../researcher_guide/README.md#prepare-token-budgeted-data-blends) +to prepare smaller weighted datasets before scaling to a full distillation run. + **Important Notes:** - For Pretraining / raw-text data (`text` key) — use `--append_eod` so Megatron can tell where documents end when concatenating them into long sequences. diff --git a/examples/dataset/prepare_data_blend.py b/examples/dataset/prepare_data_blend.py index 37c29c4e67a..b5cc2607aa8 100644 --- a/examples/dataset/prepare_data_blend.py +++ b/examples/dataset/prepare_data_blend.py @@ -13,7 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Prepare a token-sized Megatron data blend from a YAML configuration.""" +"""Prepare a weighted Megatron data blend from a YAML configuration.""" import argparse import os @@ -34,7 +34,7 @@ def load_config(path: Path) -> dict[str, Any]: tokenizer: /models/Qwen3-8B output_dir: /datasets/qwen3-blend - target_tokens: 1000000 + target_tokens: 1000000 # Optional; omit to prepare every source in full. sources: - hf_dataset: nvidia/Nemotron-Pretraining-SFT-v1 config: Nemotron-SFT-General @@ -47,9 +47,10 @@ def load_config(path: Path) -> dict[str, Any]: content_field: messages weight: 40 - returns a dictionary with ``tokenizer``, ``output_dir``, ``target_tokens``, and - ``sources`` keys. Each source has ``hf_dataset``, ``content_field``, and ``weight``; - it uses ``split`` with an optional ``config``, or selects repository ``files``. + 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)) @@ -77,7 +78,7 @@ def _prepare_sources( sources: list[dict[str, Any]], output_dir: Path, tokenizer: str, - total_tokens: int, + 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) @@ -86,7 +87,9 @@ def _prepare_sources( for index, source in enumerate(sources): weight = float(source["weight"]) - if index == len(sources) - 1: + 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 / 100) @@ -113,6 +116,7 @@ def _prepare_sources( "hf_dataset": dataset, "hf_name": source.get("config"), "hf_split": source["split"], + "hf_max_samples_per_split": source.get("max_samples"), "hf_streaming": True, } @@ -143,7 +147,8 @@ def prepare_data_blend(config_path: Path) -> list[tuple[float, str]]: config = load_config(config_path) output_dir = Path(config["output_dir"]) output_dir.mkdir(parents=True, exist_ok=True) - total_tokens = int(config["target_tokens"]) + 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) 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..d53d7a9a6d8 100644 --- a/examples/pruning/minitron/NVIDIA-Nemotron-Nano-9B-v2/README.md +++ b/examples/pruning/minitron/NVIDIA-Nemotron-Nano-9B-v2/README.md @@ -62,7 +62,14 @@ 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. +Prepare this blend with the +[token-budgeted data blend workflow](../../../researcher_guide/README.md#prepare-token-budgeted-data-blends). +The complete blend listed below contains approximately 93B tokens. For an initial experiment, set +`target_tokens: 1000000000` to prepare a 1B-token subset with the same source weights, avoiding the time and +storage needed to preprocess the complete blend. Omit `target_tokens` to prepare every configured source in +full, subject to any per-source `max_samples` setting. See +[examples/dataset/MEGATRON_DATA_PREP.md](../../../dataset/MEGATRON_DATA_PREP.md) for additional dataset +tokenization commands. 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 new file mode 100644 index 00000000000..c719a0f1b3c --- /dev/null +++ b/examples/researcher_guide/README.md @@ -0,0 +1,147 @@ +# ModelOpt for Researchers: Fast Experimentation Workflows + +Model optimization research depends on short feedback loops: test a hypothesis cheaply, compare candidates +reproducibly, and spend full-scale compute only on the most promising experiments. This guide collects practical +ModelOpt workflows for that iterative research process. + +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 + +[LM-Eval Harness](../llm_eval/README.md) supports many accuracy benchmarks, but full runs are often too slow for +every iteration of model pruning, distillation, or quantization. Use progressively larger evaluation subsets to +reject weak candidates quickly and reserve full runs for the most promising models. + +In LM-Eval, `--limit N` evaluates the first `N` samples of each individual task. For task groups such as MMLU and +MMLU-Pro, the limit applies to every subject, not to the group as a whole. + +The following table gives a practical progression for LM-Eval's MMLU-Pro task group, which contains 14 subjects +and 12,032 questions. Example times assume Qwen3-8B, a batch size of 4, and subject-level parallelism on eight +H100 80GB GPUs: + +| Limit per subject | Questions evaluated | Worst-case 95% margin of error | Example time | +|-------------------|--------------------:|--------------------------------:|-------------:| +| `10` | 140 | ±8.3 percentage points | ~3 minutes | +| `50` | 700 | ±3.7 percentage points | ~14 minutes | +| `100` | 1,400 | ±2.6 percentage points | ~28 minutes | +| `200` | 2,800 | ±1.9 percentage points | ~56 minutes | +| None | 12,032 | ±0.9 percentage points | 4 hours | + +The example times scale an approximately four-hour full run by the fraction of questions evaluated. Actual time +depends on the model, hardware, batch size, and parallelism. + +The margins of error are conservative planning estimates. They use 50% accuracy, the normal approximation for a +[binomial proportion confidence interval](https://en.wikipedia.org/wiki/Binomial_proportion_confidence_interval#Normal_approximation_interval). + +These estimates treat benchmark questions as independent random samples from a broader population of possible +questions. Because `--limit` selects the first samples, limited scores may also be affected by dataset ordering +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 + +Full distillation datasets are often unnecessarily large for testing a pruning or distillation hypothesis. Use +[`prepare_data_blend.py`](../dataset/prepare_data_blend.py) to prepare a smaller 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 Nano 9B v2 distillation blend](../pruning/minitron/NVIDIA-Nemotron-Nano-9B-v2/README.md#1-data-preparation) +down to one billion tokens while preserving its source weights: + +```yaml +tokenizer: nvidia/NVIDIA-Nemotron-Nano-9B-v2 +output_dir: /datasets/tokenized_nemotron_v2_1b +# Optional; omit this field to prepare every source in full. +target_tokens: 1000000000 +sources: + - hf_dataset: nvidia/Nemotron-Pretraining-SFT-v1 + config: Nemotron-SFT-Code + split: train + max_samples: 10000000 + content_field: text + weight: 5 + - hf_dataset: nvidia/Nemotron-Pretraining-SFT-v1 + config: Nemotron-SFT-General + split: train + max_samples: 10000000 + content_field: text + weight: 20 + - hf_dataset: nvidia/Nemotron-Pretraining-SFT-v1 + config: Nemotron-SFT-MATH + split: train + max_samples: 10000000 + content_field: text + weight: 5 + - hf_dataset: nvidia/Nemotron-Math-v2 + split: high_part00 + content_field: messages + weight: 15 + - hf_dataset: nvidia/Nemotron-Math-v2 + split: high_part01 + content_field: messages + weight: 15 + - 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: 5000000 + content_field: messages + weight: 10 + - 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 +``` + +Run from the repository root: + +```bash +python examples/dataset/prepare_data_blend.py --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. See the +[Megatron data preparation guide](../dataset/MEGATRON_DATA_PREP.md) for dataset-specific details. + +## Planned topics + +Future additions can cover: + +- Iterative pruning and distillation workflows diff --git a/tests/examples/dataset/test_prepare_data_blend.py b/tests/examples/dataset/test_prepare_data_blend.py new file mode 100644 index 00000000000..c71f1c0b801 --- /dev/null +++ b/tests/examples/dataset/test_prepare_data_blend.py @@ -0,0 +1,99 @@ +# 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 pytest +import yaml + +# examples/dataset is not a package; add it to the path to import and test the script in-process. +sys.path.insert(0, str(Path(__file__).parents[3] / "examples/dataset")) + +import prepare_data_blend + + +def _setup_test( + tiny_qwen3_path: str, tmp_path: Path, monkeypatch, target_tokens: int | None +) -> tuple[Path, Path, Mock]: + output_dir = tmp_path / "tokenized" + config = { + "tokenizer": tiny_qwen3_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, + }, + ], + } + if target_tokens is not None: + config["target_tokens"] = target_tokens + config_path = tmp_path / "config.yaml" + config_path.write_text(yaml.safe_dump(config), 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(prepare_data_blend, "hf_hub_download", download) + monkeypatch.setattr(sys, "argv", ["prepare_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_data_blend_with_split_and_files_sources( + tiny_qwen3_path: str, tmp_path: Path, monkeypatch, target_tokens: int | None +): + output_dir, config_path, download = _setup_test( + tiny_qwen3_path, tmp_path, monkeypatch, target_tokens + ) + + # Run in-process so the mocked NVIDIA download is visible; run_example_command uses a subprocess. + prepare_data_blend.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() + assert ("_tokens" in prefix) is (target_tokens is not None) + assert (output_dir / "config.yaml").read_bytes() == config_path.read_bytes() 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..9cb6a82ada3 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 @@ -53,6 +53,61 @@ 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, + "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_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 + + @pytest.mark.parametrize( ("hf_dataset", "hf_split", "json_keys"), [ From aeaba79edf39fb2fd4805af09fc8ed47cbc8f3a9 Mon Sep 17 00:00:00 2001 From: Daniel Korzekwa Date: Thu, 2 Jul 2026 10:10:56 -0700 Subject: [PATCH 05/23] add change log and improve a unit test Signed-off-by: Daniel Korzekwa --- CHANGELOG.rst | 1 + tests/examples/dataset/test_prepare_data_blend.py | 2 +- 2 files changed, 2 insertions(+), 1 deletion(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index ed03fec6ef7..2a8b5d4c830 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -17,6 +17,7 @@ Changelog **New Features** +- Add the `ModelOpt for Researchers: Fast Experimentation Workflows `_ guide, covering efficient model evaluation with smaller benchmark subsets and efficient token-budgeted data-blend preparation. - 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. - Add **streaming** speculative-decoding training (EAGLE3 / DFlash): the draft trains on base-model hidden states produced on the fly by a co-located ``vllm serve`` (no disk dump), moved trainer-side over NIXL RDMA, scaling to multi-node (dedicated serve replicas + DDP trainers). New launcher examples for NVFP4 Kimi-K2.5 / K2.6 on GB200/aarch64 under ``tools/launcher/examples/moonshotai/``. diff --git a/tests/examples/dataset/test_prepare_data_blend.py b/tests/examples/dataset/test_prepare_data_blend.py index c71f1c0b801..f4b69ee33a3 100644 --- a/tests/examples/dataset/test_prepare_data_blend.py +++ b/tests/examples/dataset/test_prepare_data_blend.py @@ -95,5 +95,5 @@ def test_prepare_data_blend_with_split_and_files_sources( for _, prefix in blend: assert Path(prefix + ".bin").exists() assert Path(prefix + ".idx").exists() - assert ("_tokens" in prefix) is (target_tokens is not None) + assert ("_tokens" in blend[0][1]) is (target_tokens is not None) assert (output_dir / "config.yaml").read_bytes() == config_path.read_bytes() From aaccd769c3aa737e6df3c5ac8be5e7db6d444513 Mon Sep 17 00:00:00 2001 From: Daniel Korzekwa Date: Thu, 2 Jul 2026 10:40:18 -0700 Subject: [PATCH 06/23] improve docs Signed-off-by: Daniel Korzekwa --- .../README.md | 9 +++++++- .../NVIDIA-Nemotron-Nano-9B-v2/README.md | 13 +++++------ examples/researcher_guide/README.md | 22 ++++++++++++------- 3 files changed, 27 insertions(+), 17 deletions(-) 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..a19511af1e7 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 @@ -69,7 +69,14 @@ 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. +Prepare this blend with the +[token-budgeted data blend workflow](../../../researcher_guide/README.md#prepare-token-budgeted-data-blends). +The complete blend listed below contains approximately 142B tokens. For an initial experiment, set +`target_tokens: 1000000000` to prepare a 1B-token subset with the same source weights, avoiding the time and +storage needed to preprocess the complete blend. Omit `target_tokens` to prepare every configured source in +full, subject to any per-source `max_samples` setting. See +[examples/dataset/MEGATRON_DATA_PREP.md](../../../dataset/MEGATRON_DATA_PREP.md) for additional dataset +tokenization commands. 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 d53d7a9a6d8..5fab5aebbee 100644 --- a/examples/pruning/minitron/NVIDIA-Nemotron-Nano-9B-v2/README.md +++ b/examples/pruning/minitron/NVIDIA-Nemotron-Nano-9B-v2/README.md @@ -62,14 +62,11 @@ Distillation uses the **30% Pretraining (Code 5, General 20, MATH 5) + 70% Post- ### 1. Data Preparation -Prepare this blend with the -[token-budgeted data blend workflow](../../../researcher_guide/README.md#prepare-token-budgeted-data-blends). -The complete blend listed below contains approximately 93B tokens. For an initial experiment, set -`target_tokens: 1000000000` to prepare a 1B-token subset with the same source weights, avoiding the time and -storage needed to preprocess the complete blend. Omit `target_tokens` to prepare every configured source in -full, subject to any per-source `max_samples` setting. See -[examples/dataset/MEGATRON_DATA_PREP.md](../../../dataset/MEGATRON_DATA_PREP.md) for additional dataset -tokenization commands. +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](../../../researcher_guide/README.md#prepare-token-budgeted-data-blends), +but create a custom YAML configuration using this tutorial's tokenizer, sources, and weights below. The +researcher guide's 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 c719a0f1b3c..0c043aeb9b4 100644 --- a/examples/researcher_guide/README.md +++ b/examples/researcher_guide/README.md @@ -56,12 +56,12 @@ 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 Nano 9B v2 distillation blend](../pruning/minitron/NVIDIA-Nemotron-Nano-9B-v2/README.md#1-data-preparation) +[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: ```yaml -tokenizer: nvidia/NVIDIA-Nemotron-Nano-9B-v2 -output_dir: /datasets/tokenized_nemotron_v2_1b +tokenizer: nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16 +output_dir: /path/to/nemotron_3_nano_30b_distillation_blend_1b # Optional; omit this field to prepare every source in full. target_tokens: 1000000000 sources: @@ -86,11 +86,12 @@ sources: - hf_dataset: nvidia/Nemotron-Math-v2 split: high_part00 content_field: messages - weight: 15 - - hf_dataset: nvidia/Nemotron-Math-v2 - split: high_part01 + weight: 10 + - hf_dataset: nvidia/Nemotron-SFT-Math-v3 + files: + - data/train.jsonl content_field: messages - weight: 15 + weight: 17 - hf_dataset: nvidia/Nemotron-SFT-Competitive-Programming-v2 files: - data/competitive_programming_python_00.jsonl @@ -106,7 +107,7 @@ sources: split: stem max_samples: 5000000 content_field: messages - weight: 10 + weight: 8 - hf_dataset: nvidia/Nemotron-Science-v1 files: - data/MCQ.jsonl @@ -127,6 +128,11 @@ sources: - 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 ``` Run from the repository root: From 360c999dd51db733c143d07be6632c2bf35fae3f Mon Sep 17 00:00:00 2001 From: "coderabbitai[bot]" <136622811+coderabbitai[bot]@users.noreply.github.com> Date: Mon, 6 Jul 2026 18:24:55 +0530 Subject: [PATCH 07/23] =?UTF-8?q?=F0=9F=93=9D=20CodeRabbit=20Chat:=20Updat?= =?UTF-8?q?e=20dataset=20blend=20preparation=20example=20(#1929)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Code changes was requested by @chochowski. * https://github.com/NVIDIA/Model-Optimizer/pull/1888#discussion_r3528105185 The following files were modified: * `examples/dataset/prepare_data_blend.py` Co-authored-by: coderabbitai[bot] <136622811+coderabbitai[bot]@users.noreply.github.com> --- examples/dataset/prepare_data_blend.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/examples/dataset/prepare_data_blend.py b/examples/dataset/prepare_data_blend.py index b5cc2607aa8..7ffebd5fe26 100644 --- a/examples/dataset/prepare_data_blend.py +++ b/examples/dataset/prepare_data_blend.py @@ -84,6 +84,8 @@ def _prepare_sources( 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"]) @@ -92,7 +94,7 @@ def _prepare_sources( elif index == len(sources) - 1: source_tokens = total_tokens - allocated_tokens else: - source_tokens = round(total_tokens * weight / 100) + source_tokens = round(total_tokens * weight / weight_sum) allocated_tokens += source_tokens dataset = source["hf_dataset"] From 78794fbf62def2ddca313a7ef35ef03158dcae2c Mon Sep 17 00:00:00 2001 From: Daniel Korzekwa Date: Mon, 6 Jul 2026 09:11:40 -0700 Subject: [PATCH 08/23] Fix token accounting when resuming cached data blend preparation Signed-off-by: Daniel Korzekwa --- .../utils/plugins/megatron_preprocess_data.py | 14 +++++++++++-- .../plugins/test_megatron_preprocess_data.py | 21 +++++++++++++++++++ 2 files changed, 33 insertions(+), 2 deletions(-) diff --git a/modelopt/torch/utils/plugins/megatron_preprocess_data.py b/modelopt/torch/utils/plugins/megatron_preprocess_data.py index 04e59b0b6ae..17389e5672f 100644 --- a/modelopt/torch/utils/plugins/megatron_preprocess_data.py +++ b/modelopt/torch/utils/plugins/megatron_preprocess_data.py @@ -305,6 +305,16 @@ def _encode_docs(self, encoder: "_Encoder", lines, may_stop_early: bool = False) 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, @@ -342,7 +352,7 @@ 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 = 0, 0 @@ -437,7 +447,7 @@ 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 # Workers encode asynchronously; iterating encoded_docs waits for results in input order. pool, encoded_docs = self._encode_docs( 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 9cb6a82ada3..a7d580fd98d 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 @@ -108,6 +108,27 @@ def test_megatron_preprocess_data_hf_split_stops_at_max_tokens(tmp_path): assert limited_size < full_size +def test_megatron_preprocess_data_hf_splits_resume_uses_cached_token_count(tmp_path): + args = { + "hf_dataset": "Salesforce/wikitext", + "hf_name": "wikitext-2-raw-v1", + "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"), [ From 187af2d8cb5f5c091daf0c31a93113ff4318825c Mon Sep 17 00:00:00 2001 From: Daniel Korzekwa Date: Mon, 6 Jul 2026 10:02:09 -0700 Subject: [PATCH 09/23] include _tokens{max_tokens} in the processed_json bin/idx file names Signed-off-by: Daniel Korzekwa --- modelopt/torch/utils/plugins/megatron_preprocess_data.py | 8 +++++--- .../torch/utils/plugins/test_megatron_preprocess_data.py | 9 ++++----- 2 files changed, 9 insertions(+), 8 deletions(-) diff --git a/modelopt/torch/utils/plugins/megatron_preprocess_data.py b/modelopt/torch/utils/plugins/megatron_preprocess_data.py index 17389e5672f..132a5c17bb7 100644 --- a/modelopt/torch/utils/plugins/megatron_preprocess_data.py +++ b/modelopt/torch/utils/plugins/megatron_preprocess_data.py @@ -325,7 +325,8 @@ def process_json_file( 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": @@ -341,8 +342,9 @@ def process_json_file( 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( 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 a7d580fd98d..b1c583f9f1f 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 @@ -62,23 +62,22 @@ def test_megatron_preprocess_data_jsonl_stops_at_max_tokens(tmp_path): 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, - output_dir=tmp_path / "limited", max_tokens=100, )[0] - full_prefix = megatron_preprocess_data( - **common_args, - output_dir=tmp_path / "full", - )[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 From e65ee79c1e31f5e3ea9951fbde12aa88d22ada01 Mon Sep 17 00:00:00 2001 From: Daniel Korzekwa Date: Mon, 6 Jul 2026 10:19:17 -0700 Subject: [PATCH 10/23] improve docs Signed-off-by: Daniel Korzekwa --- modelopt/torch/utils/plugins/megatron_preprocess_data.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/modelopt/torch/utils/plugins/megatron_preprocess_data.py b/modelopt/torch/utils/plugins/megatron_preprocess_data.py index 132a5c17bb7..3d7e334e6ee 100644 --- a/modelopt/torch/utils/plugins/megatron_preprocess_data.py +++ b/modelopt/torch/utils/plugins/megatron_preprocess_data.py @@ -535,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. @@ -547,7 +552,7 @@ def megatron_preprocess_data( 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 split. The final document may make the result slightly larger. + 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"]. From 73b4e7b90f9eda23d33d08993dcbe92904c201d8 Mon Sep 17 00:00:00 2001 From: Daniel Korzekwa Date: Mon, 6 Jul 2026 10:44:54 -0700 Subject: [PATCH 11/23] Improve docs Signed-off-by: Daniel Korzekwa --- examples/researcher_guide/README.md | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/examples/researcher_guide/README.md b/examples/researcher_guide/README.md index 0c043aeb9b4..964259fa29d 100644 --- a/examples/researcher_guide/README.md +++ b/examples/researcher_guide/README.md @@ -59,6 +59,12 @@ prepare a weighted subset, or omit it to prepare every source in full. This exam [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 tokenizer: nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16 output_dir: /path/to/nemotron_3_nano_30b_distillation_blend_1b From 9ac05403bf21052ccc3e36dbe6cec289655ba802 Mon Sep 17 00:00:00 2001 From: Daniel Korzekwa Date: Wed, 8 Jul 2026 02:33:47 -0700 Subject: [PATCH 12/23] refactoring and update docs for prepare_megatron_data_blend utility Signed-off-by: Daniel Korzekwa --- examples/dataset/MEGATRON_DATA_PREP.md | 109 +++++++++++++++++- .../README.md | 2 +- .../NVIDIA-Nemotron-Nano-9B-v2/README.md | 4 +- examples/researcher_guide/README.md | 106 +---------------- modelopt/torch/utils/plugins/__init__.py | 3 + .../plugins/prepare_megatron_data_blend.py | 10 +- tests/gpu_megatron/conftest.py | 14 +++ .../test_prepare_megatron_data_blend.py} | 18 +-- 8 files changed, 144 insertions(+), 122 deletions(-) rename examples/dataset/prepare_data_blend.py => modelopt/torch/utils/plugins/prepare_megatron_data_blend.py (96%) rename tests/{examples/dataset/test_prepare_data_blend.py => gpu_megatron/torch/utils/plugins/test_prepare_megatron_data_blend.py} (85%) diff --git a/examples/dataset/MEGATRON_DATA_PREP.md b/examples/dataset/MEGATRON_DATA_PREP.md index 357f16da7c2..363d4352444 100644 --- a/examples/dataset/MEGATRON_DATA_PREP.md +++ b/examples/dataset/MEGATRON_DATA_PREP.md @@ -4,7 +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](../researcher_guide/README.md#prepare-token-budgeted-data-blends)\] | +| 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)\] | @@ -12,9 +12,6 @@ The distillation and pre-training scripts in Megatron-Bridge or Megatron-LM expe Use the `megatron_preprocess_data` utility to tokenize any JSONL or Hugging Face dataset. The tokenization scripts below print the list of output prefixes (e.g. `tokenized_qwen3/data1_text`) that you can use for the `data_paths` argument (with relative weights on different files) in Megatron training scripts. -For iterative research, use the [token-budgeted data blend workflow](../researcher_guide/README.md#prepare-token-budgeted-data-blends) -to prepare smaller weighted datasets before scaling to a full distillation run. - **Important Notes:** - For Pretraining / raw-text data (`text` key) — use `--append_eod` so Megatron can tell where documents end when concatenating them into long sequences. @@ -70,6 +67,110 @@ 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 +tokenizer: nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16 +output_dir: /path/to/nemotron_3_nano_30b_distillation_blend_1b +# Optional; omit this field to prepare every source in full. +target_tokens: 1000000000 +sources: + - hf_dataset: nvidia/Nemotron-Pretraining-SFT-v1 + config: Nemotron-SFT-Code + split: train + max_samples: 10000000 + content_field: text + weight: 5 + - hf_dataset: nvidia/Nemotron-Pretraining-SFT-v1 + config: Nemotron-SFT-General + split: train + max_samples: 10000000 + content_field: text + weight: 20 + - hf_dataset: nvidia/Nemotron-Pretraining-SFT-v1 + config: Nemotron-SFT-MATH + split: train + max_samples: 10000000 + 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: 5000000 + 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 a19511af1e7..312049ff28b 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,7 +70,7 @@ Distillation uses the **30% Pretraining (Code 5, General 20, MATH 5) + 70% Post- ### 1. Data Preparation Prepare this blend with the -[token-budgeted data blend workflow](../../../researcher_guide/README.md#prepare-token-budgeted-data-blends). +[token-budgeted data blend workflow](../../../dataset/MEGATRON_DATA_PREP.md#prepare-token-budgeted-data-blends). The complete blend listed below contains approximately 142B tokens. For an initial experiment, set `target_tokens: 1000000000` to prepare a 1B-token subset with the same source weights, avoiding the time and storage needed to preprocess the complete blend. Omit `target_tokens` to prepare every configured source in 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 5fab5aebbee..d1a75e43098 100644 --- a/examples/pruning/minitron/NVIDIA-Nemotron-Nano-9B-v2/README.md +++ b/examples/pruning/minitron/NVIDIA-Nemotron-Nano-9B-v2/README.md @@ -64,9 +64,9 @@ Distillation uses the **30% Pretraining (Code 5, General 20, MATH 5) + 70% Post- 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](../../../researcher_guide/README.md#prepare-token-budgeted-data-blends), +[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 -researcher guide's example configuration targets Nemotron 3 and should not be reused unchanged. +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 964259fa29d..dd0334bf12a 100644 --- a/examples/researcher_guide/README.md +++ b/examples/researcher_guide/README.md @@ -49,108 +49,10 @@ split samples across model copies; see the [LM-Eval examples](../llm_eval/README ## Prepare token-budgeted data blends -Full distillation datasets are often unnecessarily large for testing a pruning or distillation hypothesis. Use -[`prepare_data_blend.py`](../dataset/prepare_data_blend.py) to prepare a smaller 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 -tokenizer: nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16 -output_dir: /path/to/nemotron_3_nano_30b_distillation_blend_1b -# Optional; omit this field to prepare every source in full. -target_tokens: 1000000000 -sources: - - hf_dataset: nvidia/Nemotron-Pretraining-SFT-v1 - config: Nemotron-SFT-Code - split: train - max_samples: 10000000 - content_field: text - weight: 5 - - hf_dataset: nvidia/Nemotron-Pretraining-SFT-v1 - config: Nemotron-SFT-General - split: train - max_samples: 10000000 - content_field: text - weight: 20 - - hf_dataset: nvidia/Nemotron-Pretraining-SFT-v1 - config: Nemotron-SFT-MATH - split: train - max_samples: 10000000 - 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: 5000000 - 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 -``` - -Run from the repository root: - -```bash -python examples/dataset/prepare_data_blend.py --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. See the -[Megatron data preparation guide](../dataset/MEGATRON_DATA_PREP.md) for dataset-specific details. +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 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/examples/dataset/prepare_data_blend.py b/modelopt/torch/utils/plugins/prepare_megatron_data_blend.py similarity index 96% rename from examples/dataset/prepare_data_blend.py rename to modelopt/torch/utils/plugins/prepare_megatron_data_blend.py index 7ffebd5fe26..8f5720d96e4 100644 --- a/examples/dataset/prepare_data_blend.py +++ b/modelopt/torch/utils/plugins/prepare_megatron_data_blend.py @@ -21,11 +21,13 @@ from pathlib import Path from typing import Any, cast +import huggingface_hub import yaml -from huggingface_hub import hf_hub_download from modelopt.torch.utils.plugins.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. @@ -104,7 +106,7 @@ def _prepare_sources( if "files" in source: raw_dir = output_dir.parent / "raw" / dataset.replace("/", "--") paths = [ - hf_hub_download( + huggingface_hub.hf_hub_download( repo_id=dataset, filename=file, repo_type="dataset", @@ -144,7 +146,7 @@ def _prepare_sources( return blend -def prepare_data_blend(config_path: Path) -> list[tuple[float, str]]: +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"]) @@ -163,7 +165,7 @@ def main() -> None: """Prepare a data blend from the supplied configuration.""" parser = _build_parser() args = parser.parse_args() - blend = prepare_data_blend(args.config) + blend = prepare_megatron_data_blend(args.config) print(f"Prepared {len(blend)} data paths. See data_blend.txt and config.yaml in the output.") diff --git a/tests/gpu_megatron/conftest.py b/tests/gpu_megatron/conftest.py index 405a83db06d..ccaed731cda 100644 --- a/tests/gpu_megatron/conftest.py +++ b/tests/gpu_megatron/conftest.py @@ -17,10 +17,24 @@ import pytest import torch from _test_utils.torch.distributed.utils import DistributedWorkerPool +from _test_utils.torch.transformers_models import create_tiny_qwen3_dir from megatron.core.parallel_state import destroy_model_parallel import modelopt.torch.utils.distributed as dist + +@pytest.fixture(scope="session") +def tiny_qwen3_path(tmp_path_factory): + return str( + create_tiny_qwen3_dir( + tmp_path_factory.mktemp("tiny_qwen3"), + with_tokenizer=True, + hidden_size=512, + intermediate_size=512, + ) + ) + + apex_destroy = None with contextlib.suppress(ImportError): from apex.transformer.parallel_state import destroy_model_parallel as apex_destroy diff --git a/tests/examples/dataset/test_prepare_data_blend.py b/tests/gpu_megatron/torch/utils/plugins/test_prepare_megatron_data_blend.py similarity index 85% rename from tests/examples/dataset/test_prepare_data_blend.py rename to tests/gpu_megatron/torch/utils/plugins/test_prepare_megatron_data_blend.py index f4b69ee33a3..ccc74c7b3a3 100644 --- a/tests/examples/dataset/test_prepare_data_blend.py +++ b/tests/gpu_megatron/torch/utils/plugins/test_prepare_megatron_data_blend.py @@ -18,13 +18,11 @@ from pathlib import Path from unittest.mock import Mock +import huggingface_hub import pytest import yaml -# examples/dataset is not a package; add it to the path to import and test the script in-process. -sys.path.insert(0, str(Path(__file__).parents[3] / "examples/dataset")) - -import prepare_data_blend +from modelopt.torch.utils.plugins.prepare_megatron_data_blend import main def _setup_test( @@ -65,21 +63,23 @@ def _setup_test( "".join(json.dumps(conversation) + "\n" for _ in range(20)), encoding="utf-8" ) download = Mock(return_value=str(jsonl_path)) - monkeypatch.setattr(prepare_data_blend, "hf_hub_download", download) - monkeypatch.setattr(sys, "argv", ["prepare_data_blend.py", "--config", str(config_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_data_blend_with_split_and_files_sources( +def test_prepare_megatron_data_blend_with_split_and_files_sources( tiny_qwen3_path: str, tmp_path: Path, monkeypatch, target_tokens: int | None ): output_dir, config_path, download = _setup_test( tiny_qwen3_path, tmp_path, monkeypatch, target_tokens ) - # Run in-process so the mocked NVIDIA download is visible; run_example_command uses a subprocess. - prepare_data_blend.main() + # 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", From 79148ba095e6fbaf58915876d6bdb854fc3491ac Mon Sep 17 00:00:00 2001 From: Daniel Korzekwa Date: Wed, 8 Jul 2026 03:04:14 -0700 Subject: [PATCH 13/23] improve docs Signed-off-by: Daniel Korzekwa --- examples/dataset/MEGATRON_DATA_PREP.md | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/examples/dataset/MEGATRON_DATA_PREP.md b/examples/dataset/MEGATRON_DATA_PREP.md index 363d4352444..d4028a2da86 100644 --- a/examples/dataset/MEGATRON_DATA_PREP.md +++ b/examples/dataset/MEGATRON_DATA_PREP.md @@ -86,8 +86,9 @@ down to one billion tokens while preserving its source weights: > 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_nano_30b_distillation_blend_1b +output_dir: /path/to/nemotron_3_distillation_blend_1b # Optional; omit this field to prepare every source in full. target_tokens: 1000000000 sources: From e43ac5397a676a10fd5887e6cc4a038634eecb1b Mon Sep 17 00:00:00 2001 From: Daniel Korzekwa Date: Wed, 8 Jul 2026 03:18:49 -0700 Subject: [PATCH 14/23] allow for a more readable max_tokens number formatting (target_tokens: 1_000_000_000 instead of target_tokens: 1000000000) Signed-off-by: Daniel Korzekwa --- examples/dataset/MEGATRON_DATA_PREP.md | 2 +- .../NVIDIA-Nemotron-3-Nano-30B-A3B-BF16/README.md | 2 +- .../plugins/test_prepare_megatron_data_blend.py | 13 +++++++++---- 3 files changed, 11 insertions(+), 6 deletions(-) diff --git a/examples/dataset/MEGATRON_DATA_PREP.md b/examples/dataset/MEGATRON_DATA_PREP.md index d4028a2da86..794f45a2e63 100644 --- a/examples/dataset/MEGATRON_DATA_PREP.md +++ b/examples/dataset/MEGATRON_DATA_PREP.md @@ -90,7 +90,7 @@ down to one billion tokens while preserving its source weights: 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: 1000000000 +target_tokens: 1_000_000_000 sources: - hf_dataset: nvidia/Nemotron-Pretraining-SFT-v1 config: Nemotron-SFT-Code 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 312049ff28b..afc4eaf1e98 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 @@ -72,7 +72,7 @@ Distillation uses the **30% Pretraining (Code 5, General 20, MATH 5) + 70% Post- Prepare this blend with the [token-budgeted data blend workflow](../../../dataset/MEGATRON_DATA_PREP.md#prepare-token-budgeted-data-blends). The complete blend listed below contains approximately 142B tokens. For an initial experiment, set -`target_tokens: 1000000000` to prepare a 1B-token subset with the same source weights, avoiding the time and +`target_tokens: 1_000_000_000` to prepare a 1B-token subset with the same source weights, avoiding the time and storage needed to preprocess the complete blend. Omit `target_tokens` to prepare every configured source in full, subject to any per-source `max_samples` setting. See [examples/dataset/MEGATRON_DATA_PREP.md](../../../dataset/MEGATRON_DATA_PREP.md) for additional dataset 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 index ccc74c7b3a3..2df4019b1d0 100644 --- 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 @@ -48,10 +48,11 @@ def _setup_test( }, ], } - if target_tokens is not None: - config["target_tokens"] = target_tokens config_path = tmp_path / "config.yaml" - config_path.write_text(yaml.safe_dump(config), encoding="utf-8") + config_yaml = yaml.safe_dump(config) + if target_tokens is not None: + config_yaml += "target_tokens: 1_000\n" + config_path.write_text(config_yaml, encoding="utf-8") jsonl_path = tmp_path / "competitive_programming_python_00.jsonl" conversation = { "messages": [ @@ -95,5 +96,9 @@ def test_prepare_megatron_data_blend_with_split_and_files_sources( for _, prefix in blend: assert Path(prefix + ".bin").exists() assert Path(prefix + ".idx").exists() - assert ("_tokens" in blend[0][1]) is (target_tokens is not None) + token_suffixes = ["_tokens600", "_tokens400"] if target_tokens is not None else ["", ""] + assert [Path(prefix).name for _, prefix in blend] == [ + f"nanotron--minipile_100_samples_None_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() From 08b5ba8edac344b4bf2b2152d35d4705b27b66d0 Mon Sep 17 00:00:00 2001 From: Daniel Korzekwa Date: Wed, 8 Jul 2026 03:27:34 -0700 Subject: [PATCH 15/23] Immprove logging Signed-off-by: Daniel Korzekwa --- modelopt/torch/utils/plugins/megatron_preprocess_data.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/modelopt/torch/utils/plugins/megatron_preprocess_data.py b/modelopt/torch/utils/plugins/megatron_preprocess_data.py index 3d7e334e6ee..dcbf9a6db1d 100644 --- a/modelopt/torch/utils/plugins/megatron_preprocess_data.py +++ b/modelopt/torch/utils/plugins/megatron_preprocess_data.py @@ -650,6 +650,9 @@ def megatron_preprocess_data( 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)}" From 556c64f9cfb47a11e7aeeb6fc33d9203cf0019f7 Mon Sep 17 00:00:00 2001 From: Daniel Korzekwa Date: Mon, 13 Jul 2026 00:34:29 -0700 Subject: [PATCH 16/23] improve docs Signed-off-by: Daniel Korzekwa --- examples/dataset/MEGATRON_DATA_PREP.md | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/examples/dataset/MEGATRON_DATA_PREP.md b/examples/dataset/MEGATRON_DATA_PREP.md index 794f45a2e63..91b89375907 100644 --- a/examples/dataset/MEGATRON_DATA_PREP.md +++ b/examples/dataset/MEGATRON_DATA_PREP.md @@ -95,19 +95,19 @@ sources: - hf_dataset: nvidia/Nemotron-Pretraining-SFT-v1 config: Nemotron-SFT-Code split: train - max_samples: 10000000 + 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: 10000000 + 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: 10000000 + max_samples: 10_000_000 content_field: text weight: 5 - hf_dataset: nvidia/Nemotron-Math-v2 @@ -132,7 +132,7 @@ sources: - hf_dataset: nvidia/Nemotron-Post-Training-Dataset-v1 config: default split: stem - max_samples: 5000000 + max_samples: 5_000_000 content_field: messages weight: 8 - hf_dataset: nvidia/Nemotron-Science-v1 From 6c201eb6de36c2b5bfa5881be04c47e93f0f2890 Mon Sep 17 00:00:00 2001 From: Daniel Korzekwa Date: Mon, 13 Jul 2026 00:41:06 -0700 Subject: [PATCH 17/23] Improve docs Signed-off-by: Daniel Korzekwa --- .../NVIDIA-Nemotron-3-Nano-30B-A3B-BF16/README.md | 13 +++++-------- 1 file changed, 5 insertions(+), 8 deletions(-) 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 afc4eaf1e98..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 @@ -69,14 +69,11 @@ Distillation uses the **30% Pretraining (Code 5, General 20, MATH 5) + 70% Post- ### 1. Data Preparation -Prepare this blend with the -[token-budgeted data blend workflow](../../../dataset/MEGATRON_DATA_PREP.md#prepare-token-budgeted-data-blends). -The complete blend listed below contains approximately 142B tokens. For an initial experiment, set -`target_tokens: 1_000_000_000` to prepare a 1B-token subset with the same source weights, avoiding the time and -storage needed to preprocess the complete blend. Omit `target_tokens` to prepare every configured source in -full, subject to any per-source `max_samples` setting. See -[examples/dataset/MEGATRON_DATA_PREP.md](../../../dataset/MEGATRON_DATA_PREP.md) for additional dataset -tokenization commands. +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`. From 8af74ea81b03bc92e80a179b57008f4008ebedb7 Mon Sep 17 00:00:00 2001 From: Daniel Korzekwa Date: Mon, 13 Jul 2026 00:53:17 -0700 Subject: [PATCH 18/23] improve imports Signed-off-by: Daniel Korzekwa --- modelopt/torch/utils/plugins/prepare_megatron_data_blend.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/modelopt/torch/utils/plugins/prepare_megatron_data_blend.py b/modelopt/torch/utils/plugins/prepare_megatron_data_blend.py index 8f5720d96e4..fa5875a5e2c 100644 --- a/modelopt/torch/utils/plugins/prepare_megatron_data_blend.py +++ b/modelopt/torch/utils/plugins/prepare_megatron_data_blend.py @@ -24,7 +24,7 @@ import huggingface_hub import yaml -from modelopt.torch.utils.plugins.megatron_preprocess_data import megatron_preprocess_data +from .megatron_preprocess_data import megatron_preprocess_data __all__ = ["prepare_megatron_data_blend"] From 97df80034bd1eef863d9ad7a53f9936a766fd318 Mon Sep 17 00:00:00 2001 From: Daniel Korzekwa Date: Mon, 13 Jul 2026 01:39:25 -0700 Subject: [PATCH 19/23] Improve unit test,use the same hf_dataset as in other testsI Signed-off-by: Daniel Korzekwa --- .../torch/utils/plugins/test_megatron_preprocess_data.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) 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 b1c583f9f1f..8f242933ce8 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 @@ -107,10 +107,10 @@ def test_megatron_preprocess_data_hf_split_stops_at_max_tokens(tmp_path): assert limited_size < full_size -def test_megatron_preprocess_data_hf_splits_resume_uses_cached_token_count(tmp_path): +def test_megatron_preprocess_data_hf_split_resume_uses_cached_token_count(tmp_path): args = { - "hf_dataset": "Salesforce/wikitext", - "hf_name": "wikitext-2-raw-v1", + "hf_dataset": "nanotron/minipile_100_samples", + "hf_split": "train", "hf_max_samples_per_split": 1, "hf_streaming": True, "max_tokens": 100, From a7f91949419576abe7dc4c96ddce9bc3ef5fc3d3 Mon Sep 17 00:00:00 2001 From: Daniel Korzekwa Date: Mon, 13 Jul 2026 04:05:28 -0700 Subject: [PATCH 20/23] fix a bug: do not use hard-coded target_tokens:1000 Signed-off-by: Daniel Korzekwa --- .../torch/utils/plugins/test_prepare_megatron_data_blend.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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 index 2df4019b1d0..fc5cd5616cd 100644 --- 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 @@ -51,7 +51,7 @@ def _setup_test( config_path = tmp_path / "config.yaml" config_yaml = yaml.safe_dump(config) if target_tokens is not None: - config_yaml += "target_tokens: 1_000\n" + 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 = { From 73e34f54e82c67f53692fc90d4e937f55c6dcb9e Mon Sep 17 00:00:00 2001 From: Daniel Korzekwa Date: Mon, 13 Jul 2026 04:14:55 -0700 Subject: [PATCH 21/23] create a tiny tokenizer instead of tiny qwen model for test_prepare_megatron_data_blend test Signed-off-by: Daniel Korzekwa --- tests/gpu_megatron/conftest.py | 15 +++++---------- .../plugins/test_prepare_megatron_data_blend.py | 8 ++++---- 2 files changed, 9 insertions(+), 14 deletions(-) diff --git a/tests/gpu_megatron/conftest.py b/tests/gpu_megatron/conftest.py index 0cc06a62027..b8176adedd0 100644 --- a/tests/gpu_megatron/conftest.py +++ b/tests/gpu_megatron/conftest.py @@ -17,22 +17,17 @@ import pytest import torch from _test_utils.torch.distributed.utils import DistributedWorkerPool -from _test_utils.torch.transformers_models import create_tiny_qwen3_dir +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_qwen3_path(tmp_path_factory): - return str( - create_tiny_qwen3_dir( - tmp_path_factory.mktemp("tiny_qwen3"), - with_tokenizer=True, - hidden_size=512, - intermediate_size=512, - ) - ) +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 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 index fc5cd5616cd..9de371ec702 100644 --- 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 @@ -26,11 +26,11 @@ def _setup_test( - tiny_qwen3_path: str, tmp_path: Path, monkeypatch, target_tokens: int | None + tiny_tokenizer_path: str, tmp_path: Path, monkeypatch, target_tokens: int | None ) -> tuple[Path, Path, Mock]: output_dir = tmp_path / "tokenized" config = { - "tokenizer": tiny_qwen3_path, + "tokenizer": tiny_tokenizer_path, "output_dir": str(output_dir), "sources": [ { @@ -73,10 +73,10 @@ def _setup_test( @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_qwen3_path: str, tmp_path: Path, monkeypatch, target_tokens: int | None + tiny_tokenizer_path: str, tmp_path: Path, monkeypatch, target_tokens: int | None ): output_dir, config_path, download = _setup_test( - tiny_qwen3_path, tmp_path, monkeypatch, target_tokens + tiny_tokenizer_path, tmp_path, monkeypatch, target_tokens ) # Run in-process so the CLI entry point uses the mocked NVIDIA download. From aee018115b2ef345d80a55478667ef50850b88e0 Mon Sep 17 00:00:00 2001 From: Daniel Korzekwa Date: Mon, 13 Jul 2026 06:00:00 -0700 Subject: [PATCH 22/23] fix change log Signed-off-by: Daniel Korzekwa --- CHANGELOG.rst | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 6b83dc5c51b..df3cdefdb5e 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -18,7 +18,8 @@ Changelog - Dropped VILA / NVILA vision-language model support in ``examples/hf_ptq``. VILA's modeling code requires ``transformers<=4.50.0``, which conflicts with ModelOpt's minimum supported ``transformers`` version. The VILA-specific bootstrap (repo clone, ``requirements-vila.txt``) and loading paths in ``example_utils.py`` have been removed. **New Features** -- Add the `ModelOpt for Researchers: Fast Experimentation Workflows `_ guide, covering efficient model evaluation with smaller benchmark subsets and efficient token-budgeted data-blend preparation. + +- 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. From c6b39711c8498ae39af7171b74086500ff45ce21 Mon Sep 17 00:00:00 2001 From: Daniel Korzekwa Date: Mon, 13 Jul 2026 10:51:41 -0700 Subject: [PATCH 23/23] fix broken tests Signed-off-by: Daniel Korzekwa --- .../torch/utils/plugins/test_megatron_preprocess_data.py | 9 +++++---- .../utils/plugins/test_prepare_megatron_data_blend.py | 3 ++- 2 files changed, 7 insertions(+), 5 deletions(-) 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 8f242933ce8..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" 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 index 9de371ec702..f531be27ad4 100644 --- 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 @@ -97,8 +97,9 @@ def test_prepare_megatron_data_blend_with_split_and_files_sources( 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_None_train_text_max100{token_suffixes[0]}", + 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()