NVIDIA · danielkorzekwa · Jul 13, 2026 · Jul 1, 2026 · Jul 1, 2026 · Jul 2, 2026
@@ -19,6 +19,7 @@ Changelog
 
 **New Features**
 
+- Add the ``prepare_megatron_data_blend`` utility to prepare weighted Megatron data blends from YAML configs, including optional token-budgeted subsets for distillation workflows. See the `Megatron data preparation guide <https://github.com/NVIDIA/Model-Optimizer/tree/main/examples/dataset/MEGATRON_DATA_PREP.md#prepare-token-budgeted-data-blends>`_.
 - 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 <https://arxiv.org/abs/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.

@@ -4,6 +4,7 @@
 | :---: | :---: | :---: |
 | From JSONL files | Tokenize local JSONL files | \[[Link](#from-jsonl-files)\] |
 | From Hugging Face Hub | Stream or download HF datasets and tokenize | \[[Link](#from-hugging-face-hub)\] |
+| Token-budgeted data blends | Prepare weighted subsets for fast experiments | \[[Link](#prepare-token-budgeted-data-blends)\] |
 | `reasoning_content` for Post-Training v3 | Control how chain-of-thought traces are handled | \[[Link](#reasoning_content-for-post-training-v3-datasets)\] |
 | Nemotron Pre/Post-Training Datasets | Ready-to-run commands for all Nemotron datasets | \[[Link](#ready-to-run-tokenization-commands)\] |
 
@@ -66,6 +67,111 @@ For very large datasets (tens of millions of documents), or datasets with comple
 > Re-runs read from cache and are much faster.
 > Streaming re-downloads on every run with no cache, so it is slower for full-dataset processing.
 
+## Prepare token-budgeted data blends
+
+For iterative research, prepare smaller weighted datasets before scaling to a full distillation run.
+Use [`prepare_megatron_data_blend`](../../modelopt/torch/utils/plugins/prepare_megatron_data_blend.py) to
+prepare a weighted blend with a shared token budget. The utility supports Hugging Face configurations and splits
+as well as specific JSONL files stored in a Hugging Face dataset repository.
+
+Define the tokenizer, output directory, and source weights in YAML. Set the optional `target_tokens` field to
+prepare a weighted subset, or omit it to prepare every source in full. This example scales the
+[Nemotron 3 Nano distillation blend](../megatron_bridge/tutorials/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16/README.md#1-data-preparation)
+down to one billion tokens while preserving its source weights:
+
+> [!IMPORTANT]
+> When `target_tokens` is set, JSONL records specified with `files` are consumed from the beginning
+> of each file rather than selected randomly. Pre-shuffle JSONL files to obtain a random subset.
+> Hugging Face dataset splits are shuffled deterministically; streaming datasets use an
+> approximate buffer shuffle.
+
+```yaml
+# Nemotron 3 models share this tokenizer, so the tokenized blend can be reused across the family.
+tokenizer: nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16
+output_dir: /path/to/nemotron_3_distillation_blend_1b
+# Optional; omit this field to prepare every source in full.
+target_tokens: 1_000_000_000
+sources:
+  - hf_dataset: nvidia/Nemotron-Pretraining-SFT-v1
+    config: Nemotron-SFT-Code
+    split: train
+    max_samples: 10_000_000
+    content_field: text
+    weight: 5
+  - hf_dataset: nvidia/Nemotron-Pretraining-SFT-v1
+    config: Nemotron-SFT-General
+    split: train
+    max_samples: 10_000_000
+    content_field: text
+    weight: 20
+  - hf_dataset: nvidia/Nemotron-Pretraining-SFT-v1
+    config: Nemotron-SFT-MATH
+    split: train
+    max_samples: 10_000_000
+    content_field: text
+    weight: 5
+  - hf_dataset: nvidia/Nemotron-Math-v2
+    split: high_part00
+    content_field: messages
+    weight: 10
+  - hf_dataset: nvidia/Nemotron-SFT-Math-v3
+    files:
+      - data/train.jsonl
+    content_field: messages
+    weight: 17
+  - hf_dataset: nvidia/Nemotron-SFT-Competitive-Programming-v2
+    files:
+      - data/competitive_programming_python_00.jsonl
+    content_field: messages
+    weight: 15
+  - hf_dataset: nvidia/Nemotron-SFT-Competitive-Programming-v2
+    files:
+      - data/competitive_programming_cpp_00.jsonl
+    content_field: messages
+    weight: 5
+  - hf_dataset: nvidia/Nemotron-Post-Training-Dataset-v1
+    config: default
+    split: stem
+    max_samples: 5_000_000
+    content_field: messages
+    weight: 8
+  - hf_dataset: nvidia/Nemotron-Science-v1
+    files:
+      - data/MCQ.jsonl
+    content_field: messages
+    weight: 3
+  - hf_dataset: nvidia/Nemotron-Science-v1
+    files:
+      - data/RQA.jsonl
+    content_field: messages
+    weight: 2
+  - hf_dataset: nvidia/Nemotron-SFT-Instruction-Following-Chat-v2
+    files:
+      - data/reasoning_on.jsonl
+    content_field: messages
+    weight: 3
+  - hf_dataset: nvidia/Nemotron-SFT-Instruction-Following-Chat-v2
+    files:
+      - data/reasoning_off.jsonl
+    content_field: messages
+    weight: 2
+  - hf_dataset: nvidia/Nemotron-Agentic-v1
+    files:
+      - data/tool_calling.jsonl
+    content_field: messages
+    weight: 5
+```
+
+With ModelOpt installed, run:
+
+```bash
+python -m modelopt.torch.utils.plugins.prepare_megatron_data_blend --config blend.yaml
+```
+
+The output contains tokenized Megatron `.bin`/`.idx` files, `data_blend.txt` with the weighted paths for training,
+and `config.yaml` recording how the blend was generated. The final token count can slightly exceed the target
+because the final document from each source is kept whole.
+
 ## `reasoning_content` for Post-Training v3 Datasets
 
 v3 datasets include a `reasoning_content` field in assistant messages (chain-of-thought separate from

@@ -70,6 +70,10 @@ Distillation uses the **30% Pretraining (Code 5, General 20, MATH 5) + 70% Post-
 ### 1. Data Preparation
 
 See [examples/dataset/MEGATRON_DATA_PREP.md](../../../dataset/MEGATRON_DATA_PREP.md) for tokenization commands for all datasets used in this blend.
+To prepare a token-limited subset, follow the
+[token-budgeted data blend workflow](../../../dataset/MEGATRON_DATA_PREP.md#prepare-token-budgeted-data-blends),
+but create a custom YAML configuration using this tutorial's tokenizer, sources, and weights below. The
+example configuration targets Nemotron 3 and should not be reused unchanged.
 
 For this experiment: `TOKENIZER=nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16`, `OUTPUT_DIR=tokenized_nemotron_3`.
 

@@ -63,6 +63,10 @@ Distillation uses the **30% Pretraining (Code 5, General 20, MATH 5) + 70% Post-
 ### 1. Data Preparation
 
 See [examples/dataset/MEGATRON_DATA_PREP.md](../../../dataset/MEGATRON_DATA_PREP.md) for tokenization commands for all datasets used in this blend.
+To prepare a token-limited subset, follow the
+[token-budgeted data blend workflow](../../../dataset/MEGATRON_DATA_PREP.md#prepare-token-budgeted-data-blends),
+but create a custom YAML configuration using this tutorial's tokenizer, sources, and weights below. The
+example configuration targets Nemotron 3 and should not be reused unchanged.
 
 For this experiment: `TOKENIZER=nvidia/NVIDIA-Nemotron-Nano-9B-v2`, `OUTPUT_DIR=tokenized_nemotron_v2`.
 

@@ -4,9 +4,14 @@ Model optimization research depends on short feedback loops: test a hypothesis c
 reproducibly, and spend full-scale compute only on the most promising experiments. This guide collects practical
 ModelOpt workflows for that iterative research process.
 
-The guide starts with efficient model evaluation and will grow as additional research workflows are documented.
-It complements the feature-specific [examples](../) by connecting them into experimentation strategies rather
-than replacing their detailed instructions.
+Current workflows include:
+
+- [Efficient model evaluation](#efficient-evaluation-with-lm-eval-harness) with smaller benchmark subsets.
+- [Efficient data blend preparation](#prepare-token-budgeted-data-blends) for distillation experiments.
+
+The guide will grow as additional research workflows are documented. It complements the feature-specific
+[examples](../) by connecting them into experimentation strategies rather than replacing their detailed
+instructions.
 
 ## Efficient evaluation with LM-Eval Harness
 
@@ -42,6 +47,13 @@ and should not be reported as final benchmark results.
 Add `--log_samples` for paired per-question analysis. When multiple GPUs are available, use data parallelism to
 split samples across model copies; see the [LM-Eval examples](../llm_eval/README.md) for commands.
 
+## Prepare token-budgeted data blends
+
+Preparing complete distillation datasets can consume unnecessary time and storage during early experiments.
+ModelOpt can preserve source weights while preparing only a requested token budget. See
+[Prepare token-budgeted data blends](../dataset/MEGATRON_DATA_PREP.md#prepare-token-budgeted-data-blends) for the
+configuration format, commands, and generated outputs.
+
 ## Planned topics
 
 Future additions can cover:

@@ -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"):