From ed5b737d142cbe664ccdfba1f8c9de1e8312eb54 Mon Sep 17 00:00:00 2001
From: Daniel Korzekwa <dkorzekwa@nvidia.com>
Date: Fri, 3 Jul 2026 09:57:21 -0700
Subject: [PATCH 01/25] Save hf checkpoint at every valitation iteration during
 distillation.

Signed-off-by: Daniel Korzekwa <dkorzekwa@nvidia.com>
---
 examples/megatron_bridge/README.md            | 14 +++
 examples/megatron_bridge/distill.py           | 93 ++++++++++++++++++-
 .../examples/megatron_bridge/test_distill.py  |  6 +-
 3 files changed, 108 insertions(+), 5 deletions(-)

diff --git a/examples/megatron_bridge/README.md b/examples/megatron_bridge/README.md
index 60dfed7fc7d..e777991929a 100644
--- a/examples/megatron_bridge/README.md
+++ b/examples/megatron_bridge/README.md
@@ -225,6 +225,20 @@ torchrun --nnodes 1 --nproc_per_node 8 distill.py \
 
 `--student_hf_model` should match the base architecture of the student (used as a template for export). For non-Puzzletron (i.e. standard) models, it should be same as `--student_hf_path`.
 
+To export the live student after every validation stage, use `--hf_validation_export_path`:
+
+```bash
+torchrun --nnodes 1 --nproc_per_node 8 distill.py \
+    ... \
+    --eval_interval 100 \
+    --hf_validation_export_path /path/to/save/validation_checkpoints \
+    --student_hf_model Qwen/Qwen3-4B
+```
+
+The exports are saved as `iter_0000100`, `iter_0000200`, and so on. They contain only
+HuggingFace model artifacts and can be evaluated independently while distillation retains its
+normal resumable Megatron checkpoints.
+
 **Separate conversion** -- convert any saved iteration using the Megatron-Bridge conversion script:
 
 ```bash
diff --git a/examples/megatron_bridge/distill.py b/examples/megatron_bridge/distill.py
index 283504f1ec2..f03dde783ec 100644
--- a/examples/megatron_bridge/distill.py
+++ b/examples/megatron_bridge/distill.py
@@ -24,6 +24,7 @@
 import contextlib
 import os
 from dataclasses import fields
+from pathlib import Path
 
 import torch
 from megatron.bridge import AutoBridge
@@ -34,6 +35,7 @@
 from megatron.bridge.recipes.utils.optimizer_utils import (
     distributed_fused_adam_with_cosine_annealing,
 )
+from megatron.bridge.training.callbacks import Callback
 from megatron.bridge.training.config import (
     CheckpointConfig,
     ConfigContainer,
@@ -45,10 +47,13 @@
     TrainingConfig,
 )
 from megatron.bridge.training.distill import distill
+from megatron.bridge.training.gpt_step import forward_step_modelopt
 from megatron.bridge.training.post_training.checkpointing import has_modelopt_state
 from megatron.bridge.training.post_training.distillation import ModelOptDistillConfig
+from megatron.bridge.training.pretrain import pretrain
 from megatron.core.datasets.utils import get_blend_from_list
 from megatron.core.distributed import DistributedDataParallelConfig
+from megatron.core.utils import unwrap_model
 from transformers import AutoConfig
 
 import modelopt.torch.distill as mtd
@@ -159,6 +164,58 @@ def _distill_provide_with_megatron_student(
 DistillationProvider.provide = _distill_provide_with_megatron_student
 
 
+class _HFValidationExportCallback(Callback):
+    """Export the live student to Hugging Face format after each validation stage."""
+
+    def __init__(
+        self,
+        export_dir: str,
+        student_hf_model: str,
+        student_hf_path: str,
+        trust_remote_code: bool,
+    ) -> None:
+        self.export_dir = Path(export_dir)
+        self.student_hf_path = student_hf_path
+        self.trust_remote_code = trust_remote_code
+        self._last_exported_iteration: int | None = None
+        self.bridge = AutoBridge.from_hf_pretrained(
+            student_hf_model, trust_remote_code=trust_remote_code
+        )
+
+    def on_eval_end(self, context) -> None:
+        """Export the student at the iteration that was just validated."""
+        iteration = context.state.train_state.step
+        # The final iteration can be validated both on its regular interval and after training.
+        # Avoid exporting and overwriting the same Hugging Face checkpoint twice.
+        if iteration == self._last_exported_iteration:
+            return
+        output_path = self.export_dir / f"iter_{iteration:07d}"
+        print_rank_0(f"Exporting validation checkpoint {iteration} to {output_path}")
+
+        # DistillationModel is the student with teacher and KD-loss modules attached. Hide the
+        # auxiliary modules temporarily so the Hugging Face export contains only student weights.
+        with contextlib.ExitStack() as stack:
+            for model_chunk in unwrap_model(context.model):
+                if isinstance(model_chunk, mtd.DistillationModel):
+                    stack.enter_context(model_chunk.hide_teacher_model())
+                    stack.enter_context(model_chunk.hide_loss_modules())
+            self.bridge.save_hf_pretrained(
+                context.model,
+                output_path,
+                show_progress=True,
+                strict=True,
+            )
+
+        if dist.rank() == 0:
+            # Preserve the student architecture from student_hf_path, including heterogeneous
+            # layer changes; AutoConfig supports both local paths and Hugging Face model IDs.
+            AutoConfig.from_pretrained(
+                self.student_hf_path, trust_remote_code=self.trust_remote_code
+            ).save_pretrained(output_path)
+        torch.distributed.barrier()
+        self._last_exported_iteration = iteration
+
+
 def get_args():
     """Parse command-line arguments."""
     parser = argparse.ArgumentParser(description="Distillation for Megatron-Bridge.")
@@ -284,6 +341,15 @@ def get_args():
             "If provided, exports last iteration checkpoint to HF format after distillation."
         ),
     )
+    parser.add_argument(
+        "--hf_validation_export_path",
+        type=str,
+        default=None,
+        help=(
+            "Path where HuggingFace checkpoints are exported after each validation stage. "
+            "Each checkpoint is saved in an iter_<iteration> subdirectory."
+        ),
+    )
     parser.add_argument(
         "--student_hf_model",
         type=str,
@@ -298,8 +364,10 @@ def get_args():
     if not args.use_mock_data and not args.data_paths:
         raise ValueError("Must provide either --data_paths or set --use_mock_data.")
 
-    if args.hf_export_path and not args.student_hf_model:
-        raise ValueError("Must provide --student_hf_model if --hf_export_path is provided.")
+    if (args.hf_export_path or args.hf_validation_export_path) and not args.student_hf_model:
+        raise ValueError(
+            "Must provide --student_hf_model when HuggingFace checkpoint export is enabled."
+        )
 
     print_args(args)
 
@@ -444,12 +512,28 @@ def _build_model_provider(hf_path, load_weights=True):
     )
 
     print_rank_0("\nStarting distillation...")
-    distill(config)
+    if args.hf_validation_export_path:
+        assert isinstance(config.model, DistillationProvider), (
+            "Distillation requires a DistillationProvider"
+        )
+        callback = _HFValidationExportCallback(
+            export_dir=args.hf_validation_export_path,
+            student_hf_model=args.student_hf_model,
+            student_hf_path=args.student_hf_path,
+            trust_remote_code=args.trust_remote_code,
+        )
+        # TODO: Use distill(..., callbacks=[callback]) once Megatron-Bridge supports callbacks.
+        pretrain(config, forward_step_modelopt, callbacks=[callback])
+    else:
+        distill(config)
     print_rank_0(
         f"\nDistillation done! Saved checkpoint to {checkpoint_dir}"
         " in megatron distributed checkpoint format.\n"
     )
 
+    # TODO: Extend _HFValidationExportCallback to export --hf_export_path from the in-memory
+    # student at the end of training. This avoids reloading the newly saved Megatron checkpoint
+    # and duplicating large disk I/O.
     if args.hf_export_path:
         print_rank_0(f"Exporting final distilled ckpt to HF format to {args.hf_export_path}")
         # Save rank before destroying process group (dist.rank() won't work after destruction)
@@ -470,7 +554,8 @@ def _build_model_provider(hf_path, load_weights=True):
                 show_progress=True,
                 strict=True,
             )
-            # Copy config.json from student_hf_path (handles both local paths and HF model IDs)
+            # Preserve the student architecture from student_hf_path, including heterogeneous
+            # layer changes; AutoConfig supports both local paths and Hugging Face model IDs.
             AutoConfig.from_pretrained(
                 args.student_hf_path, trust_remote_code=args.trust_remote_code
             ).save_pretrained(args.hf_export_path)
diff --git a/tests/examples/megatron_bridge/test_distill.py b/tests/examples/megatron_bridge/test_distill.py
index 8b5e206567d..402a3ddee5a 100644
--- a/tests/examples/megatron_bridge/test_distill.py
+++ b/tests/examples/megatron_bridge/test_distill.py
@@ -28,6 +28,7 @@ def test_distill_and_convert(tmp_path: Path, num_gpus):
     train_iters = 2
     distill_output_dir = tmp_path / "distill_output"
     distilled_hf_path = tmp_path / "distilled_hf"
+    validation_exports = tmp_path / "validation_exports"
     distill_cmd_parts = extend_cmd_parts(
         ["torchrun", f"--nproc_per_node={num_gpus}", "distill.py", "--use_mock_data"],
         student_hf_path=teacher_hf_path,
@@ -40,16 +41,19 @@ def test_distill_and_convert(tmp_path: Path, num_gpus):
         gbs=4,
         train_iters=train_iters,
         lr_warmup_iters=1,
-        eval_interval=train_iters,
+        eval_interval=1,
         eval_iters=1,
         log_interval=1,
         hf_export_path=distilled_hf_path,
+        hf_validation_export_path=validation_exports,
         student_hf_model=teacher_hf_path,
     )
     run_example_command(distill_cmd_parts, example_path="megatron_bridge")
 
     assert (distill_output_dir / f"checkpoints/iter_{train_iters:07d}").exists()
     assert (distilled_hf_path / "config.json").exists()
+    assert (validation_exports / "iter_0000001/config.json").exists()
+    assert (validation_exports / "iter_0000002/config.json").exists()
 
 
 def test_distill_puzzletron_anymodel(tmp_path: Path, num_gpus):

From b2f20445bc2b2ae122957bf686b19d170681cd89 Mon Sep 17 00:00:00 2001
From: Daniel Korzekwa <dkorzekwa@nvidia.com>
Date: Wed, 1 Jul 2026 12:37:21 -0700
Subject: [PATCH 02/25] Implement prepare_data_blend utility.

Signed-off-by: Daniel Korzekwa <dkorzekwa@nvidia.com>
---
 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 ae8253dab9a9a4792f63d6c0d760809d2dac060f Mon Sep 17 00:00:00 2001
From: Daniel Korzekwa <dkorzekwa@nvidia.com>
Date: Wed, 1 Jul 2026 12:49:39 -0700
Subject: [PATCH 03/25] remodve not needed  __all__ variable

Signed-off-by: Daniel Korzekwa <dkorzekwa@nvidia.com>
---
 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 02fff734dfd9fb44fbddd2025a8a07697f31e080 Mon Sep 17 00:00:00 2001
From: Daniel Korzekwa <dkorzekwa@nvidia.com>
Date: Thu, 2 Jul 2026 08:04:50 -0700
Subject: [PATCH 04/25] implement prepare_data_blend.py

Signed-off-by: Daniel Korzekwa <dkorzekwa@nvidia.com>
---
 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 2083df78e08fa372882f1185343481a9036a988b Mon Sep 17 00:00:00 2001
From: Daniel Korzekwa <dkorzekwa@nvidia.com>
Date: Thu, 2 Jul 2026 09:55:07 -0700
Subject: [PATCH 05/25] created a tool for efficient data blend preparation

Signed-off-by: Daniel Korzekwa <dkorzekwa@nvidia.com>
---
 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 ee33cd6db172326db539b90e2d8206d635d4f998 Mon Sep 17 00:00:00 2001
From: Daniel Korzekwa <dkorzekwa@nvidia.com>
Date: Thu, 2 Jul 2026 10:40:18 -0700
Subject: [PATCH 06/25] improve docs

Signed-off-by: Daniel Korzekwa <dkorzekwa@nvidia.com>
---
 .../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 b27265f65cad805bc03f5049d9bd35075e670fb4 Mon Sep 17 00:00:00 2001
From: Daniel Korzekwa <dkorzekwa@nvidia.com>
Date: Fri, 3 Jul 2026 14:35:34 -0700
Subject: [PATCH 07/25] Improve docs

Signed-off-by: Daniel Korzekwa <dkorzekwa@nvidia.com>
---
 examples/researcher_guide/README.md | 51 +++++++++++++++++++++++++++++
 1 file changed, 51 insertions(+)

diff --git a/examples/researcher_guide/README.md b/examples/researcher_guide/README.md
index 0c043aeb9b4..a52b1b4938f 100644
--- a/examples/researcher_guide/README.md
+++ b/examples/researcher_guide/README.md
@@ -7,6 +7,8 @@ ModelOpt workflows for that iterative research process.
 Current workflows include:
 
 - [Efficient model evaluation](#efficient-evaluation-with-lm-eval-harness) with smaller benchmark subsets.
+- [Downstream evaluation over time during distillation](#track-downstream-quality-over-time-during-distillation)
+  with validation checkpoints.
 - [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
@@ -47,6 +49,55 @@ 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.
 
+## Track downstream quality over time during distillation
+
+Validation KD and CE losses show whether the student is fitting the teacher and validation data, but they do not
+necessarily predict downstream accuracy. Export the live student at each validation interval and evaluate the
+resulting checkpoints to see when downstream quality improves, plateaus, or regresses.
+
+Use `--hf_validation_export_path` with `distill.py` as described in the
+[Megatron-Bridge distillation guide](../megatron_bridge/README.md#converting-to-hugging-face-format-optional).
+
+The export path contains one loadable Hugging Face checkpoint per validation iteration:
+
+```text
+hf_validation/
+├── iter_0000100/
+├── iter_0000200/
+└── iter_0000300/
+```
+
+Evaluate the teacher, pruned student, and each exported checkpoint. Follow the
+[LM-Eval Harness instructions](../llm_eval/README.md#lm-eval-harness) and use the
+[efficient evaluation workflow](#efficient-evaluation-with-lm-eval-harness) to choose limits.
+
+The following experiment pruned Qwen3-8B to 0.7x and distilled it on either Salesforce/wikitext
+(`wikitext-103-v1`) or the math and stem splits of nvidia/Nemotron-Post-Training-Dataset-v2 for 100 iterations,
+exporting every 20 iterations. MMLU used 25 samples per subject and MMLU-Pro used 50.
+
+| Model | Iteration | Validation KD | Validation CE | MMLU | MMLU-Pro |
+|-------|----------:|--------------:|--------------:|-----:|---------:|
+| Teacher: Qwen3-8B | - | - | - | 74.93% (full) | 58.62% (full) |
+| Pruned 0.7x student before distillation | 0 | - | - | 48.69% (full) | 23.09% (full) |
+| Distilled on Salesforce/wikitext (`wikitext-103-v1`) | 20 | 0.3031 | 2.6570 | 59.72% | 25.00% |
+| Distilled on Salesforce/wikitext (`wikitext-103-v1`) | 40 | 0.2935 | 2.6554 | 60.98% | 28.00% |
+| Distilled on Salesforce/wikitext (`wikitext-103-v1`) | 60 | 0.2696 | 2.6412 | 62.46% | 27.57% |
+| Distilled on Salesforce/wikitext (`wikitext-103-v1`) | 80 | 0.2479 | 2.6262 | 62.74% | 27.14% |
+| Distilled on Salesforce/wikitext (`wikitext-103-v1`) | 100 | 0.2343 | 2.6091 | 63.58% | 29.29% |
+| Distilled on nvidia/Nemotron-Post-Training-Dataset-v2 (math and stem) | 20 | 0.1919 | 1.0931 | 58.74% | 13.14% |
+| Distilled on nvidia/Nemotron-Post-Training-Dataset-v2 (math and stem) | 40 | 0.1731 | 1.0692 | 59.58% | 1.71% |
+| Distilled on nvidia/Nemotron-Post-Training-Dataset-v2 (math and stem) | 60 | 0.1510 | 1.0347 | 58.46% | 14.43% |
+| Distilled on nvidia/Nemotron-Post-Training-Dataset-v2 (math and stem) | 80 | 0.1343 | 1.0466 | 60.35% | 12.29% |
+| Distilled on nvidia/Nemotron-Post-Training-Dataset-v2 (math and stem) | 100 | 0.1342 | 1.0550 | 60.56% | 14.29% |
+
+Interesting observations include:
+
+- On WikiText, both validation losses decrease while MMLU and MMLU-Pro trend upward.
+- On Nemotron, KD loss decreases, but CE loss improves only through iteration 60 and then worsens. This may indicate
+  diminishing learning signal on the Nemotron validation data; distill longer to determine whether CE has plateaued.
+- Investigate the Nemotron iteration-40 MMLU-Pro outlier, which appears despite a stable MMLU score, by inspecting
+  saved samples.
+
 ## Prepare token-budgeted data blends
 
 Full distillation datasets are often unnecessarily large for testing a pruning or distillation hypothesis. Use

From d9aead1ab75cbc491db0b85f53301acfaa302db1 Mon Sep 17 00:00:00 2001
From: Daniel Korzekwa <dkorzekwa@nvidia.com>
Date: Sat, 4 Jul 2026 04:06:48 -0700
Subject: [PATCH 08/25] Add capability to compute student losses for iteration
 0 only.

Signed-off-by: Daniel Korzekwa <dkorzekwa@nvidia.com>
---
 examples/megatron_bridge/README.md            |  3 ++
 examples/megatron_bridge/distill.py           | 21 ++++++++++---
 examples/researcher_guide/README.md           |  3 +-
 .../examples/megatron_bridge/test_distill.py  | 31 +++++++++++++++++++
 4 files changed, 52 insertions(+), 6 deletions(-)

diff --git a/examples/megatron_bridge/README.md b/examples/megatron_bridge/README.md
index e777991929a..5be67fccd20 100644
--- a/examples/megatron_bridge/README.md
+++ b/examples/megatron_bridge/README.md
@@ -159,6 +159,9 @@ Tensorboard logging is enabled by default and logs are saved to `<output_dir>/te
 To use Weights & Biases for logging, set the `WANDB_API_KEY` environment variable and pass the `--wandb_project` argument.
 Optionally, you can also pass `--wandb_entity` and `--wandb_exp_name` arguments to group runs under a project and experiment name.
 
+To measure the initial student's CE and distillation losses, add `--validate_only` to the command.
+This skips training and evaluates the student at iteration 0.
+
 To see all available arguments:
 
 ```bash
diff --git a/examples/megatron_bridge/distill.py b/examples/megatron_bridge/distill.py
index f03dde783ec..40fd2e3490c 100644
--- a/examples/megatron_bridge/distill.py
+++ b/examples/megatron_bridge/distill.py
@@ -324,6 +324,11 @@ def get_args():
     parser.add_argument(
         "--eval_iters", type=int, default=32, help="Number of batches per validation stage"
     )
+    parser.add_argument(
+        "--validate_only",
+        action="store_true",
+        help="Run validation on the initial student without training",
+    )
     # Logging arguments
     parser.add_argument("--log_interval", type=int, default=10, help="Write to log every <N> steps")
     parser.add_argument(
@@ -364,6 +369,8 @@ def get_args():
     if not args.use_mock_data and not args.data_paths:
         raise ValueError("Must provide either --data_paths or set --use_mock_data.")
 
+    if args.validate_only and args.hf_export_path:
+        raise ValueError("--hf_export_path cannot be used with --validate_only.")
     if (args.hf_export_path or args.hf_validation_export_path) and not args.student_hf_model:
         raise ValueError(
             "Must provide --student_hf_model when HuggingFace checkpoint export is enabled."
@@ -468,6 +475,7 @@ def _build_model_provider(hf_path, load_weights=True):
             train_iters=args.train_iters,
             eval_interval=args.eval_interval,
             eval_iters=args.eval_iters,
+            skip_train=args.validate_only,
             global_batch_size=args.gbs,
             micro_batch_size=args.mbs,
             manual_gc=True,
@@ -501,7 +509,7 @@ def _build_model_provider(hf_path, load_weights=True):
         checkpoint=CheckpointConfig(
             save_interval=args.eval_interval,
             save=checkpoint_dir,
-            load=checkpoint_dir,  # Resume from this directory (if exists)
+            load=None if args.validate_only else checkpoint_dir,
             most_recent_k=5,  # Keeps 5 most recent checkpoints (not metric-based)
             ckpt_format="torch_dist",
             async_save=True,
@@ -526,10 +534,13 @@ def _build_model_provider(hf_path, load_weights=True):
         pretrain(config, forward_step_modelopt, callbacks=[callback])
     else:
         distill(config)
-    print_rank_0(
-        f"\nDistillation done! Saved checkpoint to {checkpoint_dir}"
-        " in megatron distributed checkpoint format.\n"
-    )
+    if args.validate_only:
+        print_rank_0("\nValidation-only run done!\n")
+    else:
+        print_rank_0(
+            f"\nDistillation done! Saved checkpoint to {checkpoint_dir}"
+            " in megatron distributed checkpoint format.\n"
+        )
 
     # TODO: Extend _HFValidationExportCallback to export --hf_export_path from the in-memory
     # student at the end of training. This avoids reloading the newly saved Megatron checkpoint
diff --git a/examples/researcher_guide/README.md b/examples/researcher_guide/README.md
index a52b1b4938f..9f4bcf354d9 100644
--- a/examples/researcher_guide/README.md
+++ b/examples/researcher_guide/README.md
@@ -78,12 +78,13 @@ exporting every 20 iterations. MMLU used 25 samples per subject and MMLU-Pro use
 | Model | Iteration | Validation KD | Validation CE | MMLU | MMLU-Pro |
 |-------|----------:|--------------:|--------------:|-----:|---------:|
 | Teacher: Qwen3-8B | - | - | - | 74.93% (full) | 58.62% (full) |
-| Pruned 0.7x student before distillation | 0 | - | - | 48.69% (full) | 23.09% (full) |
+| Pruned 0.7x student (Salesforce/wikitext validation) | 0 | 0.8261 | 3.3458 | 48.69% (full) | 23.09% (full) |
 | Distilled on Salesforce/wikitext (`wikitext-103-v1`) | 20 | 0.3031 | 2.6570 | 59.72% | 25.00% |
 | Distilled on Salesforce/wikitext (`wikitext-103-v1`) | 40 | 0.2935 | 2.6554 | 60.98% | 28.00% |
 | Distilled on Salesforce/wikitext (`wikitext-103-v1`) | 60 | 0.2696 | 2.6412 | 62.46% | 27.57% |
 | Distilled on Salesforce/wikitext (`wikitext-103-v1`) | 80 | 0.2479 | 2.6262 | 62.74% | 27.14% |
 | Distilled on Salesforce/wikitext (`wikitext-103-v1`) | 100 | 0.2343 | 2.6091 | 63.58% | 29.29% |
+| Pruned 0.7x student (Nemotron validation) | 0 | 0.5187 | 1.4739 | 48.69% (full) | 23.09% (full) |
 | Distilled on nvidia/Nemotron-Post-Training-Dataset-v2 (math and stem) | 20 | 0.1919 | 1.0931 | 58.74% | 13.14% |
 | Distilled on nvidia/Nemotron-Post-Training-Dataset-v2 (math and stem) | 40 | 0.1731 | 1.0692 | 59.58% | 1.71% |
 | Distilled on nvidia/Nemotron-Post-Training-Dataset-v2 (math and stem) | 60 | 0.1510 | 1.0347 | 58.46% | 14.43% |
diff --git a/tests/examples/megatron_bridge/test_distill.py b/tests/examples/megatron_bridge/test_distill.py
index 402a3ddee5a..41104669d3f 100644
--- a/tests/examples/megatron_bridge/test_distill.py
+++ b/tests/examples/megatron_bridge/test_distill.py
@@ -56,6 +56,37 @@ def test_distill_and_convert(tmp_path: Path, num_gpus):
     assert (validation_exports / "iter_0000002/config.json").exists()
 
 
+def test_distill_validate_only(tmp_path: Path, num_gpus):
+    teacher_hf_path = create_tiny_qwen3_dir(tmp_path, with_tokenizer=True)
+    output_dir = tmp_path / "validation_output"
+    validation_exports = tmp_path / "validation_exports"
+    cmd_parts = extend_cmd_parts(
+        [
+            "torchrun",
+            f"--nproc_per_node={num_gpus}",
+            "distill.py",
+            "--use_mock_data",
+            "--validate_only",
+        ],
+        student_hf_path=teacher_hf_path,
+        teacher_hf_path=teacher_hf_path,
+        output_dir=output_dir,
+        tp_size=num_gpus,
+        pp_size=1,
+        seq_length=16,
+        mbs=1,
+        gbs=4,
+        train_iters=1,
+        eval_iters=1,
+        hf_validation_export_path=validation_exports,
+        student_hf_model=teacher_hf_path,
+    )
+    run_example_command(cmd_parts, example_path="megatron_bridge")
+
+    assert (validation_exports / "iter_0000000/config.json").exists()
+    assert not (output_dir / "checkpoints/iter_0000001").exists()
+
+
 def test_distill_puzzletron_anymodel(tmp_path: Path, num_gpus):
     """Integration test for distill.py with Puzzletron AnyModel (heterogeneous) checkpoints.
 

From fe0e72cd0cc91cc4234fe093b01de73d02b3557e Mon Sep 17 00:00:00 2001
From: Daniel Korzekwa <dkorzekwa@nvidia.com>
Date: Sat, 4 Jul 2026 11:21:32 -0700
Subject: [PATCH 09/25] update docs

Signed-off-by: Daniel Korzekwa <dkorzekwa@nvidia.com>
---
 examples/researcher_guide/README.md | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/examples/researcher_guide/README.md b/examples/researcher_guide/README.md
index 9f4bcf354d9..4f5faf1f98f 100644
--- a/examples/researcher_guide/README.md
+++ b/examples/researcher_guide/README.md
@@ -91,6 +91,8 @@ exporting every 20 iterations. MMLU used 25 samples per subject and MMLU-Pro use
 | Distilled on nvidia/Nemotron-Post-Training-Dataset-v2 (math and stem) | 80 | 0.1343 | 1.0466 | 60.35% | 12.29% |
 | Distilled on nvidia/Nemotron-Post-Training-Dataset-v2 (math and stem) | 100 | 0.1342 | 1.0550 | 60.56% | 14.29% |
 
+On the same validation samples, the teacher CE was 2.6834 on WikiText and 1.1566 on Nemotron.
+
 Interesting observations include:
 
 - On WikiText, both validation losses decrease while MMLU and MMLU-Pro trend upward.

From 760a2e3564679f07812c63728f7512ae37104040 Mon Sep 17 00:00:00 2001
From: Daniel Korzekwa <dkorzekwa@nvidia.com>
Date: Mon, 6 Jul 2026 01:35:25 -0700
Subject: [PATCH 10/25] Improve Track downstream quality over time during
 distillation tutorial (data up to 100m tokens)

Signed-off-by: Daniel Korzekwa <dkorzekwa@nvidia.com>
---
 examples/megatron_bridge/README.md            |  9 +-
 examples/megatron_bridge/distill.py           | 30 ++++++-
 examples/researcher_guide/README.md           | 82 +++++++++++++------
 .../examples/megatron_bridge/test_distill.py  |  7 +-
 4 files changed, 96 insertions(+), 32 deletions(-)

diff --git a/examples/megatron_bridge/README.md b/examples/megatron_bridge/README.md
index 5be67fccd20..5e7b7cfa7b4 100644
--- a/examples/megatron_bridge/README.md
+++ b/examples/megatron_bridge/README.md
@@ -228,17 +228,20 @@ torchrun --nnodes 1 --nproc_per_node 8 distill.py \
 
 `--student_hf_model` should match the base architecture of the student (used as a template for export). For non-Puzzletron (i.e. standard) models, it should be same as `--student_hf_path`.
 
-To export the live student after every validation stage, use `--hf_validation_export_path`:
+To export the live student periodically during validation, use `--hf_validation_export_path`:
 
 ```bash
 torchrun --nnodes 1 --nproc_per_node 8 distill.py \
     ... \
-    --eval_interval 100 \
+    --eval_interval 20 \
     --hf_validation_export_path /path/to/save/validation_checkpoints \
+    --hf_validation_export_interval 200 \
     --student_hf_model Qwen/Qwen3-4B
 ```
 
-The exports are saved as `iter_0000100`, `iter_0000200`, and so on. They contain only
+The example validates every 20 steps and exports at iterations 200, 400, and so on. The final
+iteration is always exported. `--hf_validation_export_path` and
+`--hf_validation_export_interval` must be specified together. These directories contain only
 HuggingFace model artifacts and can be evaluated independently while distillation retains its
 normal resumable Megatron checkpoints.
 
diff --git a/examples/megatron_bridge/distill.py b/examples/megatron_bridge/distill.py
index 40fd2e3490c..4c3c34e0451 100644
--- a/examples/megatron_bridge/distill.py
+++ b/examples/megatron_bridge/distill.py
@@ -165,7 +165,7 @@ def _distill_provide_with_megatron_student(
 
 
 class _HFValidationExportCallback(Callback):
-    """Export the live student to Hugging Face format after each validation stage."""
+    """Export the live student to Hugging Face format at selected validation stages."""
 
     def __init__(
         self,
@@ -173,10 +173,12 @@ def __init__(
         student_hf_model: str,
         student_hf_path: str,
         trust_remote_code: bool,
+        export_interval: int,
     ) -> None:
         self.export_dir = Path(export_dir)
         self.student_hf_path = student_hf_path
         self.trust_remote_code = trust_remote_code
+        self.export_interval = export_interval
         self._last_exported_iteration: int | None = None
         self.bridge = AutoBridge.from_hf_pretrained(
             student_hf_model, trust_remote_code=trust_remote_code
@@ -185,6 +187,9 @@ def __init__(
     def on_eval_end(self, context) -> None:
         """Export the student at the iteration that was just validated."""
         iteration = context.state.train_state.step
+        train_iters = context.state.cfg.train.train_iters
+        if iteration % self.export_interval != 0 and iteration != train_iters:
+            return
         # The final iteration can be validated both on its regular interval and after training.
         # Avoid exporting and overwriting the same Hugging Face checkpoint twice.
         if iteration == self._last_exported_iteration:
@@ -290,6 +295,11 @@ def get_args():
     parser.add_argument("--lr", type=float, default=1e-4, help="Peak learning rate")
     parser.add_argument("--min_lr", type=float, default=1e-5, help="Minimum learning rate")
     parser.add_argument("--lr_warmup_iters", type=int, default=50, help="Number of LR warmup steps")
+    parser.add_argument(
+        "--reset_optimizer",
+        action="store_true",
+        help="Do not restore optimizer or scheduler state when resuming from a checkpoint",
+    )
     parser.add_argument(
         "--recompute_granularity",
         type=str,
@@ -355,6 +365,15 @@ def get_args():
             "Each checkpoint is saved in an iter_<iteration> subdirectory."
         ),
     )
+    parser.add_argument(
+        "--hf_validation_export_interval",
+        type=int,
+        default=None,
+        help=(
+            "Export a HuggingFace checkpoint every <N> training steps. "
+            "The final iteration is always exported."
+        ),
+    )
     parser.add_argument(
         "--student_hf_model",
         type=str,
@@ -375,6 +394,13 @@ def get_args():
         raise ValueError(
             "Must provide --student_hf_model when HuggingFace checkpoint export is enabled."
         )
+    if args.hf_validation_export_interval is not None and not args.hf_validation_export_path:
+        raise ValueError("--hf_validation_export_interval requires --hf_validation_export_path.")
+    if args.hf_validation_export_path:
+        if args.hf_validation_export_interval is None:
+            raise ValueError(
+                "--hf_validation_export_path requires --hf_validation_export_interval."
+            )
 
     print_args(args)
 
@@ -510,6 +536,7 @@ def _build_model_provider(hf_path, load_weights=True):
             save_interval=args.eval_interval,
             save=checkpoint_dir,
             load=None if args.validate_only else checkpoint_dir,
+            load_optim=not args.reset_optimizer,
             most_recent_k=5,  # Keeps 5 most recent checkpoints (not metric-based)
             ckpt_format="torch_dist",
             async_save=True,
@@ -529,6 +556,7 @@ def _build_model_provider(hf_path, load_weights=True):
             student_hf_model=args.student_hf_model,
             student_hf_path=args.student_hf_path,
             trust_remote_code=args.trust_remote_code,
+            export_interval=args.hf_validation_export_interval,
         )
         # TODO: Use distill(..., callbacks=[callback]) once Megatron-Bridge supports callbacks.
         pretrain(config, forward_step_modelopt, callbacks=[callback])
diff --git a/examples/researcher_guide/README.md b/examples/researcher_guide/README.md
index 4f5faf1f98f..180028ef963 100644
--- a/examples/researcher_guide/README.md
+++ b/examples/researcher_guide/README.md
@@ -71,35 +71,65 @@ Evaluate the teacher, pruned student, and each exported checkpoint. Follow the
 [LM-Eval Harness instructions](../llm_eval/README.md#lm-eval-harness) and use the
 [efficient evaluation workflow](#efficient-evaluation-with-lm-eval-harness) to choose limits.
 
-The following experiment pruned Qwen3-8B to 0.7x and distilled it on either Salesforce/wikitext
-(`wikitext-103-v1`) or the math and stem splits of nvidia/Nemotron-Post-Training-Dataset-v2 for 100 iterations,
-exporting every 20 iterations. MMLU used 25 samples per subject and MMLU-Pro used 50.
-
-| Model | Iteration | Validation KD | Validation CE | MMLU | MMLU-Pro |
-|-------|----------:|--------------:|--------------:|-----:|---------:|
-| Teacher: Qwen3-8B | - | - | - | 74.93% (full) | 58.62% (full) |
-| Pruned 0.7x student (Salesforce/wikitext validation) | 0 | 0.8261 | 3.3458 | 48.69% (full) | 23.09% (full) |
-| Distilled on Salesforce/wikitext (`wikitext-103-v1`) | 20 | 0.3031 | 2.6570 | 59.72% | 25.00% |
-| Distilled on Salesforce/wikitext (`wikitext-103-v1`) | 40 | 0.2935 | 2.6554 | 60.98% | 28.00% |
-| Distilled on Salesforce/wikitext (`wikitext-103-v1`) | 60 | 0.2696 | 2.6412 | 62.46% | 27.57% |
-| Distilled on Salesforce/wikitext (`wikitext-103-v1`) | 80 | 0.2479 | 2.6262 | 62.74% | 27.14% |
-| Distilled on Salesforce/wikitext (`wikitext-103-v1`) | 100 | 0.2343 | 2.6091 | 63.58% | 29.29% |
-| Pruned 0.7x student (Nemotron validation) | 0 | 0.5187 | 1.4739 | 48.69% (full) | 23.09% (full) |
-| Distilled on nvidia/Nemotron-Post-Training-Dataset-v2 (math and stem) | 20 | 0.1919 | 1.0931 | 58.74% | 13.14% |
-| Distilled on nvidia/Nemotron-Post-Training-Dataset-v2 (math and stem) | 40 | 0.1731 | 1.0692 | 59.58% | 1.71% |
-| Distilled on nvidia/Nemotron-Post-Training-Dataset-v2 (math and stem) | 60 | 0.1510 | 1.0347 | 58.46% | 14.43% |
-| Distilled on nvidia/Nemotron-Post-Training-Dataset-v2 (math and stem) | 80 | 0.1343 | 1.0466 | 60.35% | 12.29% |
-| Distilled on nvidia/Nemotron-Post-Training-Dataset-v2 (math and stem) | 100 | 0.1342 | 1.0550 | 60.56% | 14.29% |
-
-On the same validation samples, the teacher CE was 2.6834 on WikiText and 1.1566 on Nemotron.
+The following experiment pruned Qwen3-8B to 0.7x and distilled the same student for approximately 100 million
+tokens on three datasets. All runs used a global batch size of 8 and sequence length of 4,096. MMLU used 25
+questions per subject (1,425 total), and MMLU-Pro used 50 per subject (700 total). The tables show representative
+checkpoints; token counts are derived from the consumed fixed-length training sequences.
+
+| Baseline model | MMLU | MMLU-Pro |
+|----------------|-----:|---------:|
+| Teacher: Qwen3-8B | 74.93% (full) | 58.62% (full) |
+| Pruned 0.7x student | 48.69% (full) | 23.09% (full) |
+
+### WikiText
+
+Distillation and validation used Salesforce/wikitext (`wikitext-103-v1`).
+
+| Training tokens | Validation KD | Validation CE | MMLU | MMLU-Pro |
+|----------------:|--------------:|--------------:|-----:|---------:|
+| 0 | 0.8261 | 3.3458 | 48.69% (full) | 23.09% (full) |
+| 0.7M | 0.3031 | 2.6570 | 59.72% | 25.00% |
+| 3.3M | 0.2343 | 2.6091 | 63.58% | 29.29% |
+| 39.3M | 0.1495 | 2.5665 | 65.89% | 39.86% |
+| 78.6M | 0.1315 | 2.5699 | 66.46% | 39.57% |
+| 100.0M | 0.1291 | 2.5863 | 67.30% | 40.57% |
+
+### Nemotron v2
+
+Distillation and validation used the math and stem splits of
+nvidia/Nemotron-Post-Training-Dataset-v2.
+
+| Training tokens | Validation KD | Validation CE | MMLU | MMLU-Pro |
+|----------------:|--------------:|--------------:|-----:|---------:|
+| 0 | 0.5187 | 1.4739 | 48.69% (full) | 23.09% (full) |
+| 0.7M | 0.1919 | 1.0931 | 58.74% | 13.14% |
+| 3.3M | 0.1342 | 1.0550 | 60.56% | 14.29% |
+| 39.3M | 0.0675 | 1.0296 | 64.63% | 6.14% |
+| 78.6M | 0.0613 | 1.0773 | 65.61% | 7.71% |
+| 100.0M | 0.0582 | 1.0516 | 65.75% | 11.29% |
+
+### Nemotron 3
+
+Distillation and validation used the Nemotron 3 Nano distillation blend described in the
+[data-blend workflow](#prepare-token-budgeted-data-blends).
+
+| Training tokens | Validation KD | Validation CE | MMLU | MMLU-Pro |
+|----------------:|--------------:|--------------:|-----:|---------:|
+| 0 | Not measured | Not measured | 48.69% (full) | 23.09% (full) |
+| 0.7M | 0.2424 | 1.5910 | 57.05% | 24.86% |
+| 3.3M | 0.1604 | 1.5190 | 62.46% | 36.86% |
+| 39.3M | 0.0978 | 1.4144 | 67.23% | 45.00% |
+| 78.6M | 0.0890 | 1.4112 | 67.93% | 47.14% |
+| 100.0M | 0.0845 | 1.4656 | 67.37% | 47.71% |
+
+On the same validation samples, the teacher CE was 2.6834 on WikiText and 1.1566 on Nemotron v2.
 
 Interesting observations include:
 
-- On WikiText, both validation losses decrease while MMLU and MMLU-Pro trend upward.
-- On Nemotron, KD loss decreases, but CE loss improves only through iteration 60 and then worsens. This may indicate
-  diminishing learning signal on the Nemotron validation data; distill longer to determine whether CE has plateaued.
-- Investigate the Nemotron iteration-40 MMLU-Pro outlier, which appears despite a stable MMLU score, by inspecting
-  saved samples.
+- All three datasets recover MMLU to about 66--67% by 100 million tokens.
+- Nemotron 3 produces the strongest MMLU-Pro trajectory, reaching 47.71%, followed by WikiText at 40.57%.
+- Nemotron v2 KD continues to decrease, but its MMLU-Pro score remains below the pruned baseline. This shows why
+  validation loss alone is insufficient for selecting distillation data.
 
 ## Prepare token-budgeted data blends
 
diff --git a/tests/examples/megatron_bridge/test_distill.py b/tests/examples/megatron_bridge/test_distill.py
index 41104669d3f..f9b90e80edd 100644
--- a/tests/examples/megatron_bridge/test_distill.py
+++ b/tests/examples/megatron_bridge/test_distill.py
@@ -25,7 +25,7 @@
 
 def test_distill_and_convert(tmp_path: Path, num_gpus):
     teacher_hf_path = create_tiny_qwen3_dir(tmp_path, with_tokenizer=True)
-    train_iters = 2
+    train_iters = 3
     distill_output_dir = tmp_path / "distill_output"
     distilled_hf_path = tmp_path / "distilled_hf"
     validation_exports = tmp_path / "validation_exports"
@@ -46,14 +46,16 @@ def test_distill_and_convert(tmp_path: Path, num_gpus):
         log_interval=1,
         hf_export_path=distilled_hf_path,
         hf_validation_export_path=validation_exports,
+        hf_validation_export_interval=2,
         student_hf_model=teacher_hf_path,
     )
     run_example_command(distill_cmd_parts, example_path="megatron_bridge")
 
     assert (distill_output_dir / f"checkpoints/iter_{train_iters:07d}").exists()
     assert (distilled_hf_path / "config.json").exists()
-    assert (validation_exports / "iter_0000001/config.json").exists()
     assert (validation_exports / "iter_0000002/config.json").exists()
+    assert (validation_exports / "iter_0000003/config.json").exists()
+    assert not (validation_exports / "iter_0000001").exists()
 
 
 def test_distill_validate_only(tmp_path: Path, num_gpus):
@@ -79,6 +81,7 @@ def test_distill_validate_only(tmp_path: Path, num_gpus):
         train_iters=1,
         eval_iters=1,
         hf_validation_export_path=validation_exports,
+        hf_validation_export_interval=1,
         student_hf_model=teacher_hf_path,
     )
     run_example_command(cmd_parts, example_path="megatron_bridge")

From ae8eb67ddc1dc6f19f884106068e07b8b4c3b5bb Mon Sep 17 00:00:00 2001
From: Daniel Korzekwa <dkorzekwa@nvidia.com>
Date: Mon, 6 Jul 2026 04:22:58 -0700
Subject: [PATCH 11/25] improve Track downstream quality over time during
 distillation  tutorial

Signed-off-by: Daniel Korzekwa <dkorzekwa@nvidia.com>
---
 examples/researcher_guide/README.md | 26 +++++++++++++++++++++-----
 1 file changed, 21 insertions(+), 5 deletions(-)

diff --git a/examples/researcher_guide/README.md b/examples/researcher_guide/README.md
index 180028ef963..6f4a6f10a99 100644
--- a/examples/researcher_guide/README.md
+++ b/examples/researcher_guide/README.md
@@ -72,7 +72,7 @@ Evaluate the teacher, pruned student, and each exported checkpoint. Follow the
 [efficient evaluation workflow](#efficient-evaluation-with-lm-eval-harness) to choose limits.
 
 The following experiment pruned Qwen3-8B to 0.7x and distilled the same student for approximately 100 million
-tokens on three datasets. All runs used a global batch size of 8 and sequence length of 4,096. MMLU used 25
+tokens using four data recipes. All runs used a global batch size of 8 and sequence length of 4,096. MMLU used 25
 questions per subject (1,425 total), and MMLU-Pro used 50 per subject (700 total). The tables show representative
 checkpoints; token counts are derived from the consumed fixed-length training sequences.
 
@@ -108,6 +108,19 @@ nvidia/Nemotron-Post-Training-Dataset-v2.
 | 78.6M | 0.0613 | 1.0773 | 65.61% | 7.71% |
 | 100.0M | 0.0582 | 1.0516 | 65.75% | 11.29% |
 
+### 50/50 WikiText and Nemotron v2 blend
+
+Distillation and validation used an equal-weight blend of WikiText and the Nemotron v2 math and stem data.
+
+| Training tokens | Validation KD | Validation CE | MMLU | MMLU-Pro |
+|----------------:|--------------:|--------------:|-----:|---------:|
+| 0 | Not measured | Not measured | 48.69% (full) | 23.09% (full) |
+| 0.7M | 0.2479 | 1.8363 | 57.89% | 12.57% |
+| 3.3M | 0.1824 | 2.0265 | 62.46% | 23.14% |
+| 39.3M | 0.1164 | 1.9157 | 67.44% | 33.86% |
+| 78.6M | 0.0973 | 1.8503 | 67.72% | 41.57% |
+| 100.0M | 0.0916 | 1.7680 | 68.28% | 41.71% |
+
 ### Nemotron 3
 
 Distillation and validation used the Nemotron 3 Nano distillation blend described in the
@@ -126,10 +139,13 @@ On the same validation samples, the teacher CE was 2.6834 on WikiText and 1.1566
 
 Interesting observations include:
 
-- All three datasets recover MMLU to about 66--67% by 100 million tokens.
-- Nemotron 3 produces the strongest MMLU-Pro trajectory, reaching 47.71%, followed by WikiText at 40.57%.
-- Nemotron v2 KD continues to decrease, but its MMLU-Pro score remains below the pruned baseline. This shows why
-  validation loss alone is insufficient for selecting distillation data.
+- All four data recipes recover MMLU to about 66% to 68% by 100 million tokens. The 50/50 blend is numerically
+  highest at 68.28%.
+- Nemotron 3 produces the strongest MMLU-Pro trajectory, reaching 47.71%.
+- Although Nemotron v2 performs poorly alone, its 50/50 blend with WikiText slightly outperforms WikiText alone
+  on both final benchmarks.
+- Nemotron v2 KD continues to decrease, but its MMLU-Pro score remains below the pruned baseline. Understanding
+  this discrepancy could help design better distillation corpora.
 
 ## Prepare token-budgeted data blends
 

From a69fb12dedbe97ff5b8e96e93f9bd77146679a63 Mon Sep 17 00:00:00 2001
From: Daniel Korzekwa <dkorzekwa@nvidia.com>
Date: Mon, 6 Jul 2026 04:35:42 -0700
Subject: [PATCH 12/25] improve docs

Signed-off-by: Daniel Korzekwa <dkorzekwa@nvidia.com>
---
 examples/researcher_guide/README.md | 20 ++++++++++----------
 1 file changed, 10 insertions(+), 10 deletions(-)

diff --git a/examples/researcher_guide/README.md b/examples/researcher_guide/README.md
index 6f4a6f10a99..5e751b859fa 100644
--- a/examples/researcher_guide/README.md
+++ b/examples/researcher_guide/README.md
@@ -83,7 +83,8 @@ checkpoints; token counts are derived from the consumed fixed-length training se
 
 ### WikiText
 
-Distillation and validation used Salesforce/wikitext (`wikitext-103-v1`).
+- Dataset: Salesforce/wikitext (`wikitext-103-v1`)
+- Teacher CE: 2.6834
 
 | Training tokens | Validation KD | Validation CE | MMLU | MMLU-Pro |
 |----------------:|--------------:|--------------:|-----:|---------:|
@@ -96,8 +97,8 @@ Distillation and validation used Salesforce/wikitext (`wikitext-103-v1`).
 
 ### Nemotron v2
 
-Distillation and validation used the math and stem splits of
-nvidia/Nemotron-Post-Training-Dataset-v2.
+- Dataset: nvidia/Nemotron-Post-Training-Dataset-v2 (math and stem)
+- Teacher CE: 1.1566
 
 | Training tokens | Validation KD | Validation CE | MMLU | MMLU-Pro |
 |----------------:|--------------:|--------------:|-----:|---------:|
@@ -110,11 +111,12 @@ nvidia/Nemotron-Post-Training-Dataset-v2.
 
 ### 50/50 WikiText and Nemotron v2 blend
 
-Distillation and validation used an equal-weight blend of WikiText and the Nemotron v2 math and stem data.
+- Dataset: 50/50 blend of WikiText and Nemotron v2 math and stem
+- Teacher CE: 1.9025
 
 | Training tokens | Validation KD | Validation CE | MMLU | MMLU-Pro |
 |----------------:|--------------:|--------------:|-----:|---------:|
-| 0 | Not measured | Not measured | 48.69% (full) | 23.09% (full) |
+| 0 | 0.6662 | 2.3780 | 48.69% (full) | 23.09% (full) |
 | 0.7M | 0.2479 | 1.8363 | 57.89% | 12.57% |
 | 3.3M | 0.1824 | 2.0265 | 62.46% | 23.14% |
 | 39.3M | 0.1164 | 1.9157 | 67.44% | 33.86% |
@@ -123,20 +125,18 @@ Distillation and validation used an equal-weight blend of WikiText and the Nemot
 
 ### Nemotron 3
 
-Distillation and validation used the Nemotron 3 Nano distillation blend described in the
-[data-blend workflow](#prepare-token-budgeted-data-blends).
+- Dataset: Nemotron 3 Nano [distillation blend](#prepare-token-budgeted-data-blends)
+- Teacher CE: 1.4702
 
 | Training tokens | Validation KD | Validation CE | MMLU | MMLU-Pro |
 |----------------:|--------------:|--------------:|-----:|---------:|
-| 0 | Not measured | Not measured | 48.69% (full) | 23.09% (full) |
+| 0 | 0.6395 | 1.9113 | 48.69% (full) | 23.09% (full) |
 | 0.7M | 0.2424 | 1.5910 | 57.05% | 24.86% |
 | 3.3M | 0.1604 | 1.5190 | 62.46% | 36.86% |
 | 39.3M | 0.0978 | 1.4144 | 67.23% | 45.00% |
 | 78.6M | 0.0890 | 1.4112 | 67.93% | 47.14% |
 | 100.0M | 0.0845 | 1.4656 | 67.37% | 47.71% |
 
-On the same validation samples, the teacher CE was 2.6834 on WikiText and 1.1566 on Nemotron v2.
-
 Interesting observations include:
 
 - All four data recipes recover MMLU to about 66% to 68% by 100 million tokens. The 50/50 blend is numerically

From 8b4c88cec8d1677f313a7151d034d3e9c966679c Mon Sep 17 00:00:00 2001
From: Daniel Korzekwa <dkorzekwa@nvidia.com>
Date: Mon, 6 Jul 2026 04:41:23 -0700
Subject: [PATCH 13/25] improve docs

Signed-off-by: Daniel Korzekwa <dkorzekwa@nvidia.com>
---
 examples/researcher_guide/README.md | 9 +++++++++
 1 file changed, 9 insertions(+)

diff --git a/examples/researcher_guide/README.md b/examples/researcher_guide/README.md
index 5e751b859fa..cc3eb6decb5 100644
--- a/examples/researcher_guide/README.md
+++ b/examples/researcher_guide/README.md
@@ -76,6 +76,15 @@ tokens using four data recipes. All runs used a global batch size of 8 and seque
 questions per subject (1,425 total), and MMLU-Pro used 50 per subject (700 total). The tables show representative
 checkpoints; token counts are derived from the consumed fixed-length training sequences.
 
+Measured compute per data recipe on eight H100 80GB GPUs:
+
+| Stage | Checkpoints | Time | GPU use |
+|-------|------------:|-----:|---------|
+| Distillation to 100M tokens | - | ~2h10m | 8 GPUs |
+| MMLU trajectory | 21 | ~51m | 8 GPUs |
+| MMLU-Pro trajectory | 13 | ~3h50m | Two checkpoints in parallel, 4 GPUs each |
+| Total | - | ~6h50m | Excludes Slurm queue and worker setup |
+
 | Baseline model | MMLU | MMLU-Pro |
 |----------------|-----:|---------:|
 | Teacher: Qwen3-8B | 74.93% (full) | 58.62% (full) |

From 1c8d328c25e406eb506a74fdde9e592f0eae2a22 Mon Sep 17 00:00:00 2001
From: Daniel Korzekwa <dkorzekwa@nvidia.com>
Date: Mon, 6 Jul 2026 04:52:15 -0700
Subject: [PATCH 14/25] improve docs

Signed-off-by: Daniel Korzekwa <dkorzekwa@nvidia.com>
---
 examples/researcher_guide/README.md | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/examples/researcher_guide/README.md b/examples/researcher_guide/README.md
index cc3eb6decb5..00b009be685 100644
--- a/examples/researcher_guide/README.md
+++ b/examples/researcher_guide/README.md
@@ -55,8 +55,9 @@ Validation KD and CE losses show whether the student is fitting the teacher and
 necessarily predict downstream accuracy. Export the live student at each validation interval and evaluate the
 resulting checkpoints to see when downstream quality improves, plateaus, or regresses.
 
-Use `--hf_validation_export_path` with `distill.py` as described in the
-[Megatron-Bridge distillation guide](../megatron_bridge/README.md#converting-to-hugging-face-format-optional).
+Use `--hf_validation_export_path` to choose the output directory and
+`--hf_validation_export_interval` to choose how often `distill.py` saves a Hugging Face checkpoint, as described in
+the [Megatron-Bridge distillation guide](../megatron_bridge/README.md#converting-to-hugging-face-format-optional).
 
 The export path contains one loadable Hugging Face checkpoint per validation iteration:
 

From 69c2339bdd1c6342a50060454cf499d5deff5c8d Mon Sep 17 00:00:00 2001
From: Daniel Korzekwa <dkorzekwa@nvidia.com>
Date: Wed, 15 Jul 2026 06:34:16 -0700
Subject: [PATCH 15/25] revert not needed  changes

Signed-off-by: Daniel Korzekwa <dkorzekwa@nvidia.com>
---
 examples/dataset/MEGATRON_DATA_PREP.md | 3 ---
 1 file changed, 3 deletions(-)

diff --git a/examples/dataset/MEGATRON_DATA_PREP.md b/examples/dataset/MEGATRON_DATA_PREP.md
index 612ee400d24..91b89375907 100644
--- a/examples/dataset/MEGATRON_DATA_PREP.md
+++ b/examples/dataset/MEGATRON_DATA_PREP.md
@@ -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.

From cb8b96838d982b19a985738f5a9b6633fa0e97a7 Mon Sep 17 00:00:00 2001
From: Daniel Korzekwa <dkorzekwa@nvidia.com>
Date: Wed, 15 Jul 2026 06:36:01 -0700
Subject: [PATCH 16/25] remove not needed test (moved to other location)

Signed-off-by: Daniel Korzekwa <dkorzekwa@nvidia.com>
---
 .../dataset/test_prepare_data_blend.py        | 99 -------------------
 1 file changed, 99 deletions(-)
 delete mode 100644 tests/examples/dataset/test_prepare_data_blend.py

diff --git a/tests/examples/dataset/test_prepare_data_blend.py b/tests/examples/dataset/test_prepare_data_blend.py
deleted file mode 100644
index c71f1c0b801..00000000000
--- a/tests/examples/dataset/test_prepare_data_blend.py
+++ /dev/null
@@ -1,99 +0,0 @@
-# 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()

From 5bfb7b16f4d37ddcce7a9799862118de81b07f30 Mon Sep 17 00:00:00 2001
From: Daniel Korzekwa <dkorzekwa@nvidia.com>
Date: Wed, 15 Jul 2026 06:41:47 -0700
Subject: [PATCH 17/25] Delete moved file

Signed-off-by: Daniel Korzekwa <dkorzekwa@nvidia.com>
---
 examples/dataset/prepare_data_blend.py | 169 -------------------------
 1 file changed, 169 deletions(-)
 delete mode 100644 examples/dataset/prepare_data_blend.py

diff --git a/examples/dataset/prepare_data_blend.py b/examples/dataset/prepare_data_blend.py
deleted file mode 100644
index b5cc2607aa8..00000000000
--- a/examples/dataset/prepare_data_blend.py
+++ /dev/null
@@ -1,169 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-"""Prepare a weighted Megatron data blend from a YAML configuration."""
-
-import argparse
-import os
-import shutil
-from pathlib import Path
-from typing import Any, cast
-
-import 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]:
-    """Load a data-blend YAML configuration as a dictionary.
-
-    For example, this YAML::
-
-        tokenizer: /models/Qwen3-8B
-        output_dir: /datasets/qwen3-blend
-        target_tokens: 1000000  # Optional; omit to prepare every source in full.
-        sources:
-          - hf_dataset: nvidia/Nemotron-Pretraining-SFT-v1
-            config: Nemotron-SFT-General
-            split: train
-            content_field: text
-            weight: 60
-          - hf_dataset: nvidia/Nemotron-SFT-Competitive-Programming-v2
-            files:
-              - data/competitive_programming_python_00.jsonl
-            content_field: messages
-            weight: 40
-
-    returns a dictionary with ``tokenizer``, ``output_dir``, and ``sources`` keys, plus
-    optional ``target_tokens``. Each source has ``hf_dataset``, ``content_field``, and
-    ``weight``; it uses ``split`` with optional ``config`` and ``max_samples``, or
-    selects repository ``files``.
-    """
-    with path.open(encoding="utf-8") as stream:
-        return cast("dict[str, Any]", yaml.safe_load(stream))
-
-
-def _build_parser() -> argparse.ArgumentParser:
-    parser = argparse.ArgumentParser(description=__doc__)
-    parser.add_argument("--config", type=Path, required=True, help="Path to the blend YAML file")
-    return parser
-
-
-def _write_data_blend(path: Path, blend: list[tuple[float, str]]) -> None:
-    """Write the weighted Megatron dataset paths."""
-    content = "\n".join(f"{weight:g} {prefix}" for weight, prefix in blend) + "\n"
-    path.write_text(content, encoding="utf-8")
-
-
-def _copy_config(source: Path, destination: Path) -> None:
-    """Copy the input configuration alongside the generated blend."""
-    if source.resolve() != destination.resolve():
-        shutil.copyfile(source, destination)
-
-
-def _prepare_sources(
-    sources: list[dict[str, Any]],
-    output_dir: Path,
-    tokenizer: str,
-    total_tokens: int | None,
-) -> list[tuple[float, str]]:
-    """Tokenize all sources and return their weighted output paths."""
-    workers = min(32, os.cpu_count() or 1)
-    blend: list[tuple[float, str]] = []  # (weight, shared .bin/.idx path without extension)
-    allocated_tokens = 0
-
-    for index, source in enumerate(sources):
-        weight = float(source["weight"])
-        if total_tokens is None:
-            source_tokens = None
-        elif index == len(sources) - 1:
-            source_tokens = total_tokens - allocated_tokens
-        else:
-            source_tokens = round(total_tokens * weight / 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_max_samples_per_split": source.get("max_samples"),
-                "hf_streaming": True,
-            }
-
-        # Each prefix is the path shared by a tokenized Megatron .bin/.idx file pair.
-        prefixes = megatron_preprocess_data(
-            **input_args,
-            output_dir=source_dir,
-            tokenizer_name_or_path=tokenizer,
-            json_keys=content_field,
-            # Plain text lacks chat-template boundary tokens, so terminate each document with EOS.
-            append_eod=content_field == "text",
-            # Join lines in text documents by replacing each newline with a space.
-            strip_newlines=content_field == "text",
-            reasoning_content="inline" if content_field == "messages" else "strip",
-            # Guard against pathological records by capping each tokenized document at 256K tokens.
-            max_sequence_length=256_000,
-            max_tokens=source_tokens,
-            workers=workers,
-        )
-        prefix_weight = weight / len(prefixes)
-        blend.extend((prefix_weight, prefix) for prefix in prefixes)
-
-    return blend
-
-
-def prepare_data_blend(config_path: Path) -> list[tuple[float, str]]:
-    """Download and tokenize the configured weighted data sources."""
-    config = load_config(config_path)
-    output_dir = Path(config["output_dir"])
-    output_dir.mkdir(parents=True, exist_ok=True)
-    target_tokens = config.get("target_tokens")
-    total_tokens = None if target_tokens is None else int(target_tokens)
-    tokenizer = str(config["tokenizer"])
-
-    blend = _prepare_sources(config["sources"], output_dir, tokenizer, total_tokens)
-    _write_data_blend(output_dir / "data_blend.txt", blend)
-    _copy_config(config_path, output_dir / "config.yaml")
-    return blend
-
-
-def main() -> None:
-    """Prepare a data blend from the supplied configuration."""
-    parser = _build_parser()
-    args = parser.parse_args()
-    blend = prepare_data_blend(args.config)
-    print(f"Prepared {len(blend)} data paths. See data_blend.txt and config.yaml in the output.")
-
-
-if __name__ == "__main__":
-    main()

From bf29d42125f6b3879552c1a710ff6a4f2f7842fc Mon Sep 17 00:00:00 2001
From: Daniel Korzekwa <dkorzekwa@nvidia.com>
Date: Wed, 15 Jul 2026 06:49:58 -0700
Subject: [PATCH 18/25] Add TODO in the code

Signed-off-by: Daniel Korzekwa <dkorzekwa@nvidia.com>
---
 examples/megatron_bridge/distill.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/examples/megatron_bridge/distill.py b/examples/megatron_bridge/distill.py
index 4c3c34e0451..61cef59506a 100644
--- a/examples/megatron_bridge/distill.py
+++ b/examples/megatron_bridge/distill.py
@@ -199,6 +199,8 @@ def on_eval_end(self, context) -> None:
 
         # DistillationModel is the student with teacher and KD-loss modules attached. Hide the
         # auxiliary modules temporarily so the Hugging Face export contains only student weights.
+        # TODO: Replace this with a cleaner student-only export path instead of mutating the
+        # distillation wrapper state during export.
         with contextlib.ExitStack() as stack:
             for model_chunk in unwrap_model(context.model):
                 if isinstance(model_chunk, mtd.DistillationModel):

From 51dc323af70c04da54d4533512ea9381a0817a28 Mon Sep 17 00:00:00 2001
From: Daniel Korzekwa <dkorzekwa@nvidia.com>
Date: Wed, 15 Jul 2026 07:05:29 -0700
Subject: [PATCH 19/25] improve comments

Signed-off-by: Daniel Korzekwa <dkorzekwa@nvidia.com>
---
 examples/megatron_bridge/distill.py | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/examples/megatron_bridge/distill.py b/examples/megatron_bridge/distill.py
index 61cef59506a..a50a11d5e28 100644
--- a/examples/megatron_bridge/distill.py
+++ b/examples/megatron_bridge/distill.py
@@ -560,7 +560,10 @@ def _build_model_provider(hf_path, load_weights=True):
             trust_remote_code=args.trust_remote_code,
             export_interval=args.hf_validation_export_interval,
         )
-        # TODO: Use distill(..., callbacks=[callback]) once Megatron-Bridge supports callbacks.
+        # Megatron-Bridge distill(config) currently only checks DistillationProvider and calls
+        # pretrain(config, forward_step_modelopt). Call pretrain directly here only to attach the
+        # export callback; keep this path behaviorally equivalent to distill(config).
+        # TODO: Use distill(..., callbacks=[callback]) once distill(...) in Megatron-Bridge supports callbacks.
         pretrain(config, forward_step_modelopt, callbacks=[callback])
     else:
         distill(config)

From 0fee306c7c038894c65c62a3d72471c25b6d2fde Mon Sep 17 00:00:00 2001
From: Daniel Korzekwa <dkorzekwa@nvidia.com>
Date: Thu, 16 Jul 2026 23:56:41 -0700
Subject: [PATCH 20/25] Make a temporary fix to `activation_func=None` issue
 when exporting intermediate HF checkpoints during distillation.

Signed-off-by: Daniel Korzekwa <dkorzekwa@nvidia.com>
---
 examples/megatron_bridge/distill.py | 57 ++++++++++++++++++++++++++---
 1 file changed, 52 insertions(+), 5 deletions(-)

diff --git a/examples/megatron_bridge/distill.py b/examples/megatron_bridge/distill.py
index a50a11d5e28..4442617eb32 100644
--- a/examples/megatron_bridge/distill.py
+++ b/examples/megatron_bridge/distill.py
@@ -184,6 +184,38 @@ def __init__(
             student_hf_model, trust_remote_code=trust_remote_code
         )
 
+    @contextlib.contextmanager
+    def _detach_distillation_provider_teacher_for_export(
+        self, modules: list[tuple[str, torch.nn.Module]]
+    ):
+        """Protect the live teacher provider while Megatron-Bridge sanitizes export configs.
+
+        Hugging Face export calls Megatron-Bridge config cleanup, which can replace
+        non-pickleable callables such as ``activation_func`` with ``None``. In distillation,
+        ``model.config`` is a ``DistillationProvider`` and its ``__setattr__`` mirrors writes
+        to the teacher provider. Without detaching the teacher here, export can accidentally
+        set the live teacher's ``activation_func=None`` and the next validation/training
+        forward fails inside the teacher MLP.
+        """
+        detached: list[tuple[DistillationProvider, object]] = []
+        seen_config_ids: set[int] = set()
+        for _, root_module in modules:
+            for module in root_module.modules():
+                config = getattr(module, "config", None)
+                if not isinstance(config, DistillationProvider) or id(config) in seen_config_ids:
+                    continue
+                seen_config_ids.add(id(config))
+                teacher = config.teacher
+                if teacher is not None:
+                    config.teacher = None
+                    detached.append((config, teacher))
+
+        try:
+            yield
+        finally:
+            for config, teacher in detached:
+                config.teacher = teacher
+
     def on_eval_end(self, context) -> None:
         """Export the student at the iteration that was just validated."""
         iteration = context.state.train_state.step
@@ -196,16 +228,31 @@ def on_eval_end(self, context) -> None:
             return
         output_path = self.export_dir / f"iter_{iteration:07d}"
         print_rank_0(f"Exporting validation checkpoint {iteration} to {output_path}")
+        model_chunks = unwrap_model(context.model)
+        distill_chunks = [
+            model_chunk
+            for model_chunk in model_chunks
+            if isinstance(model_chunk, mtd.DistillationModel)
+        ]
+        visible_chunks = [
+            (f"model[{idx}]", model_chunk) for idx, model_chunk in enumerate(model_chunks)
+        ]
 
         # DistillationModel is the student with teacher and KD-loss modules attached. Hide the
         # auxiliary modules temporarily so the Hugging Face export contains only student weights.
+        # Also detach the teacher provider while exporting: Megatron-Bridge export sanitizes
+        # ``model.config`` and can turn non-pickleable callables into ``None``. DistillationProvider
+        # mirrors those config writes to its teacher, which was the root cause of the
+        # ``activation_func=None`` failure in the next teacher forward after validation export.
         # TODO: Replace this with a cleaner student-only export path instead of mutating the
         # distillation wrapper state during export.
-        with contextlib.ExitStack() as stack:
-            for model_chunk in unwrap_model(context.model):
-                if isinstance(model_chunk, mtd.DistillationModel):
-                    stack.enter_context(model_chunk.hide_teacher_model())
-                    stack.enter_context(model_chunk.hide_loss_modules())
+        with (
+            self._detach_distillation_provider_teacher_for_export(visible_chunks),
+            contextlib.ExitStack() as stack,
+        ):
+            for model_chunk in distill_chunks:
+                stack.enter_context(model_chunk.hide_teacher_model())
+                stack.enter_context(model_chunk.hide_loss_modules())
             self.bridge.save_hf_pretrained(
                 context.model,
                 output_path,

From 5435743d4be0939186cd030beea33d288ab7bd2e Mon Sep 17 00:00:00 2001
From: Daniel Korzekwa <dkorzekwa@nvidia.com>
Date: Tue, 21 Jul 2026 05:28:20 -0700
Subject: [PATCH 21/25] Do hf checkpoint exporting offline for multiple
 distilaltion mbridge checkpoints, once distillation is finished. Keep all
 mcore distilaltion on disk.

Signed-off-by: Daniel Korzekwa <dkorzekwa@nvidia.com>
---
 CHANGELOG.rst                                 |   1 +
 examples/megatron_bridge/README.md            |  27 ++++
 examples/megatron_bridge/distill.py           |  14 ++-
 .../export_distilled_megatron_to_hf.py        | 118 ++++++++++++++----
 examples/researcher_guide/README.md           |  19 +--
 .../test_export_distilled_megatron_to_hf.py   |  87 +++++++++++++
 6 files changed, 225 insertions(+), 41 deletions(-)
 create mode 100644 tests/examples/megatron_bridge/test_export_distilled_megatron_to_hf.py

diff --git a/CHANGELOG.rst b/CHANGELOG.rst
index a15a715a5fc..eafc3b3f908 100755
--- a/CHANGELOG.rst
+++ b/CHANGELOG.rst
@@ -19,6 +19,7 @@ Changelog
 
 **New Features**
 
+- Add support for retaining all Megatron-Bridge distillation checkpoints via ``distill.py --checkpoint_keep_last -1`` and exporting all or selected iterations with ``export_distilled_megatron_to_hf.py --export_iterations``.
 - 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).
diff --git a/examples/megatron_bridge/README.md b/examples/megatron_bridge/README.md
index 01088031780..45606d691ee 100644
--- a/examples/megatron_bridge/README.md
+++ b/examples/megatron_bridge/README.md
@@ -221,6 +221,33 @@ torchrun --nproc_per_node 1 export_distilled_megatron_to_hf.py \
     --hf_export_path /path/to/save/distilled_hf_ckpt
 ```
 
+Use `--export_iterations` to export multiple saved checkpoints, for example to evaluate how model
+quality changes during distillation. To export multiple iterations, keep those Megatron checkpoints
+during distillation. The default is to keep the last 5 checkpoints; set `--checkpoint_keep_last -1`
+to keep all saved checkpoints.
+
+Then export all retained checkpoints, with one Hugging Face checkpoint written per
+`iter_<iteration>` subdirectory:
+
+```bash
+torchrun --nproc_per_node 1 export_distilled_megatron_to_hf.py \
+    --student_hf_path <student_hf_model_or_path> \
+    --megatron_path <distill_output_dir>/checkpoints \
+    --hf_export_path /path/to/save/hf_validation_checkpoints \
+    --export_iterations all
+```
+
+The export path contains one loadable Hugging Face checkpoint per exported iteration:
+
+```text
+hf_validation/
+├── iter_0000100/
+├── iter_0000200/
+└── iter_0000300/
+```
+
+To export selected iterations instead, use `--export_iterations 200 400 600`.
+
 ### Quantization Aware Distillation (QAD)
 
 To recover the accuracy lost during [Post-Training Quantization](#post-training-quantization), distill the quantized model (student) from the original, unquantized model (teacher). Pass the quantized **Megatron checkpoint** produced by `quantize.py` via `--student_megatron_path` (the ModelOpt quantizers are restored automatically, so distillation trains the fake-quantized student), while `--student_hf_path` provides the student architecture and `--teacher_hf_path` points to the original unquantized model. We also use a smaller learning rate for QAD:
diff --git a/examples/megatron_bridge/distill.py b/examples/megatron_bridge/distill.py
index dfd404fab30..ab7f2a1cd0e 100644
--- a/examples/megatron_bridge/distill.py
+++ b/examples/megatron_bridge/distill.py
@@ -166,6 +166,16 @@ def get_args():
     parser.add_argument(
         "--eval_iters", type=int, default=32, help="Number of batches per validation stage"
     )
+    parser.add_argument(
+        "--checkpoint_keep_last",
+        type=int,
+        default=5,
+        help=(
+            "Keep only the most recent <N> Megatron checkpoints. Set to -1 to disable "
+            "checkpoint rotation and keep all validation checkpoints, for example for Hugging Face "
+            "export and downstream evaluation."
+        ),
+    )
     # Logging arguments
     parser.add_argument("--log_interval", type=int, default=10, help="Write to log every <N> steps")
     parser.add_argument(
@@ -200,6 +210,8 @@ def get_args():
 
     if args.student_hf_model is None:
         args.student_hf_model = args.student_hf_path
+    if args.checkpoint_keep_last < -1:
+        raise ValueError("--checkpoint_keep_last must be >= -1.")
 
     print_args(args)
 
@@ -361,7 +373,7 @@ def _restore_student_hook(model_chunks):
             save_interval=args.eval_interval,
             save=checkpoint_dir,
             load=checkpoint_dir,  # Resume from this directory (if exists)
-            most_recent_k=5,  # Keeps 5 most recent checkpoints (not metric-based)
+            most_recent_k=args.checkpoint_keep_last,  # Keeps most recent checkpoints (-1 keeps all)
             ckpt_format="torch_dist",
             async_save=True,
             fully_parallel_save=True,
diff --git a/examples/megatron_bridge/export_distilled_megatron_to_hf.py b/examples/megatron_bridge/export_distilled_megatron_to_hf.py
index bb8726463ed..1ba76ea7ae4 100644
--- a/examples/megatron_bridge/export_distilled_megatron_to_hf.py
+++ b/examples/megatron_bridge/export_distilled_megatron_to_hf.py
@@ -40,10 +40,22 @@
         --megatron_path /tmp/distill-out/checkpoints/iter_0000500 \
         --hf_export_path /tmp/distilled-vlm-hf_iter_0000500
 
+Example selected validation iterations:
+
+    torchrun --nproc_per_node 1 export_distilled_megatron_to_hf.py \
+        --student_hf_path Qwen/Qwen3-0.6B \
+        --megatron_path /tmp/distill-out/checkpoints \
+        --hf_export_path /tmp/distilled-hf-validation \
+        --export_iterations all
+
+    Replace ``all`` with explicit iteration numbers, e.g. ``200 400 600``, to export only
+    selected checkpoints.
+
 See `README.md` in this directory for more details.
 """
 
 import argparse
+from pathlib import Path
 
 import torch
 from megatron.bridge import AutoBridge
@@ -57,6 +69,42 @@
     load_modelopt_megatron_checkpoint,
 )
 
+# Megatron-Bridge checkpoint iteration directories use names like ``iter_0000100``.
+_ITER_DIR_PREFIX = "iter_"
+
+
+def _iteration_dir_name(iteration: int) -> str:
+    return f"{_ITER_DIR_PREFIX}{iteration:07d}"
+
+
+def _get_checkpoint_export_paths(args: argparse.Namespace) -> list[tuple[Path, Path]]:
+    """Return ``(Megatron checkpoint path, HF export path)`` pairs for this invocation."""
+    megatron_path = Path(args.megatron_path)
+    hf_export_path = Path(args.hf_export_path)
+
+    if not args.export_iterations:
+        return [(megatron_path, hf_export_path)]
+
+    if len(args.export_iterations) == 1 and args.export_iterations[0].lower() == "all":
+        checkpoint_dirs = [
+            path
+            for path in megatron_path.iterdir()
+            if path.is_dir() and path.name.startswith(_ITER_DIR_PREFIX)
+        ]
+        checkpoint_dirs = sorted(checkpoint_dirs)
+    else:
+        iterations = sorted({int(iteration) for iteration in args.export_iterations})
+        checkpoint_dirs = [
+            megatron_path / _iteration_dir_name(iteration) for iteration in iterations
+        ]
+        for checkpoint_dir in checkpoint_dirs:
+            if not checkpoint_dir.is_dir():
+                raise ValueError(f"Checkpoint not found: {checkpoint_dir}")
+
+    return [
+        (checkpoint_dir, hf_export_path / checkpoint_dir.name) for checkpoint_dir in checkpoint_dirs
+    ]
+
 
 def export_llm_to_hf(
     megatron_path: str,
@@ -132,13 +180,29 @@ def get_args() -> argparse.Namespace:
         "--megatron_path",
         type=str,
         required=True,
-        help="Distilled Megatron checkpoint to convert (an iter_* directory or its parent).",
+        help=(
+            "Distilled Megatron checkpoint to convert, or checkpoint root when using "
+            "--export_iterations."
+        ),
     )
     parser.add_argument(
         "--hf_export_path",
         type=str,
         required=True,
-        help="Directory to write the exported HuggingFace checkpoint to.",
+        help=(
+            "Directory to write the exported HuggingFace checkpoint to. When exporting multiple "
+            "checkpoints, each checkpoint is written under this root as iter_<iteration>."
+        ),
+    )
+    parser.add_argument(
+        "--export_iterations",
+        nargs="+",
+        default=None,
+        help=(
+            "Export checkpoints from the checkpoint root passed to --megatron_path. Use "
+            "'all' for every iter_<iteration> checkpoint, or pass selected iteration numbers, "
+            "for example: --export_iterations all or --export_iterations 100 200 300."
+        ),
     )
     parser.add_argument(
         "--student_hf_model",
@@ -161,6 +225,7 @@ def get_args() -> argparse.Namespace:
 
 
 def main(args: argparse.Namespace):
+    checkpoint_export_paths: list[tuple[Path, Path]] = _get_checkpoint_export_paths(args)
     is_vlm = hasattr(
         AutoConfig.from_pretrained(args.student_hf_path, trust_remote_code=args.trust_remote_code),
         "vision_config",
@@ -169,9 +234,7 @@ def main(args: argparse.Namespace):
     if is_vlm:
         # Build the full VLM (vision tower / projector + original LM from HF), then overwrite the LM
         # with the distilled checkpoint weights, then export the assembled VLM.
-        print_rank_0(
-            f"Reassembling distilled VLM and exporting to HF format at {args.hf_export_path}"
-        )
+        print_rank_0("Reassembling distilled VLM and exporting to HF format")
         _bridge, _provider, _model, full_model, _tokenizer = load_mbridge_model_from_hf(
             hf_model_name_or_path=args.student_hf_path,
             trust_remote_code=args.trust_remote_code,
@@ -187,34 +250,37 @@ def main(args: argparse.Namespace):
             init_model_parallel=True,
             load_weights=True,  # vision tower / projector + original LM; the LM is overwritten below
         )
-        # Load only the distilled language-model weights (skip ModelOpt-state restore -- the kd_loss
-        # mode / teacher are irrelevant for export and would otherwise require a teacher model).
-        load_modelopt_megatron_checkpoint(
-            [full_model.language_model], args.megatron_path, restore_modelopt_state=False
-        )
-        save_vlm_to_hf(
-            full_model,
-            args.hf_export_path,
-            args.student_hf_path,
-            trust_remote_code=args.trust_remote_code,
-        )
-        print_rank_0(f"Saved distilled VLM to {args.hf_export_path} in HF format")
+        for megatron_path, hf_export_path in checkpoint_export_paths:
+            # Load only the distilled language-model weights (skip ModelOpt-state restore -- the kd_loss
+            # mode / teacher are irrelevant for export and would otherwise require a teacher model).
+            load_modelopt_megatron_checkpoint(
+                [full_model.language_model], str(megatron_path), restore_modelopt_state=False
+            )
+            save_vlm_to_hf(
+                full_model,
+                str(hf_export_path),
+                args.student_hf_path,
+                trust_remote_code=args.trust_remote_code,
+            )
+            print_rank_0(f"Saved distilled VLM to {hf_export_path} in HF format")
     else:
-        print_rank_0(f"Exporting distilled checkpoint to HF format at {args.hf_export_path}")
+        print_rank_0("Exporting distilled checkpoint(s) to HF format")
         # Save rank before destroying process group (dist.rank() won't work after destruction).
         is_rank_0 = dist.rank() == 0
         # export_ckpt creates its own temporary process group; destroy this one first so cleanup
         # does not hang on a barrier once rank 0 has left.
         dist.cleanup()
         if is_rank_0:
-            export_llm_to_hf(
-                megatron_path=args.megatron_path,
-                hf_export_path=args.hf_export_path,
-                student_hf_path=args.student_hf_path,
-                template_hf=args.student_hf_model,
-                trust_remote_code=args.trust_remote_code,
-            )
-        print_rank_0(f"Exported HuggingFace checkpoint to {args.hf_export_path}")
+            for megatron_path, hf_export_path in checkpoint_export_paths:
+                print(f"Exporting {megatron_path} to HF format at {hf_export_path}")
+                export_llm_to_hf(
+                    megatron_path=str(megatron_path),
+                    hf_export_path=str(hf_export_path),
+                    student_hf_path=args.student_hf_path,
+                    template_hf=args.student_hf_model,
+                    trust_remote_code=args.trust_remote_code,
+                )
+                print(f"Exported HuggingFace checkpoint to {hf_export_path}")
 
 
 if __name__ == "__main__":
diff --git a/examples/researcher_guide/README.md b/examples/researcher_guide/README.md
index 681f83d2ee4..f44e72ebe58 100644
--- a/examples/researcher_guide/README.md
+++ b/examples/researcher_guide/README.md
@@ -54,21 +54,12 @@ split samples across model copies; see the [LM-Eval examples](../llm_eval/README
 ## Track downstream quality over time during distillation
 
 Validation KD and CE losses show whether the student is fitting the teacher and validation data, but they do not
-necessarily predict downstream accuracy. Export the live student at each validation interval and evaluate the
-resulting checkpoints to see when downstream quality improves, plateaus, or regresses.
+necessarily predict downstream accuracy. Keep the Megatron checkpoints saved at validation intervals, export them
+to Hugging Face format, and evaluate the resulting checkpoints to see when downstream quality improves, plateaus,
+or regresses.
 
-Use `--hf_validation_export_path` to choose the output directory and
-`--hf_validation_export_interval` to choose how often `distill.py` saves a Hugging Face checkpoint, as described in
-the [Megatron-Bridge distillation guide](../megatron_bridge/README.md#converting-to-hugging-face-format-optional).
-
-The export path contains one loadable Hugging Face checkpoint per validation iteration:
-
-```text
-hf_validation/
-├── iter_0000100/
-├── iter_0000200/
-└── iter_0000300/
-```
+See the [Megatron-Bridge distillation guide](../megatron_bridge/README.md#converting-to-hugging-face-format-optional)
+for how to retain and export intermediate distillation checkpoints.
 
 Evaluate the teacher, pruned student, and each exported checkpoint. Follow the
 [LM-Eval Harness instructions](../llm_eval/README.md#lm-eval-harness) and use the
diff --git a/tests/examples/megatron_bridge/test_export_distilled_megatron_to_hf.py b/tests/examples/megatron_bridge/test_export_distilled_megatron_to_hf.py
new file mode 100644
index 00000000000..f415c99bd6a
--- /dev/null
+++ b/tests/examples/megatron_bridge/test_export_distilled_megatron_to_hf.py
@@ -0,0 +1,87 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024 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.
+"""Tests for export_distilled_megatron_to_hf.py."""
+
+from pathlib import Path
+
+from _test_utils.examples.run_command import extend_cmd_parts, run_example_command
+from _test_utils.torch.transformers_models import create_tiny_qwen3_dir
+
+
+def test_export_distilled_megatron_iterations(tmp_path: Path, num_gpus):
+    teacher_hf_path = create_tiny_qwen3_dir(tmp_path, with_tokenizer=True)
+    train_iters = 2
+    distill_output_dir = tmp_path / "distill_output"
+    selected_exports = tmp_path / "selected_exports"
+    all_exports = tmp_path / "all_exports"
+    distill_cmd_parts = extend_cmd_parts(
+        ["torchrun", f"--nproc_per_node={num_gpus}", "distill.py", "--use_mock_data"],
+        student_hf_path=teacher_hf_path,
+        teacher_hf_path=teacher_hf_path,
+        output_dir=distill_output_dir,
+        tp_size=num_gpus,
+        pp_size=1,
+        seq_length=16,
+        mbs=1,
+        gbs=4,
+        train_iters=train_iters,
+        lr_warmup_iters=1,
+        eval_interval=1,
+        eval_iters=1,
+        log_interval=1,
+        checkpoint_keep_last=-1,
+    )
+    run_example_command(distill_cmd_parts, example_path="megatron_bridge")
+
+    checkpoint_root = distill_output_dir / "checkpoints"
+    assert (checkpoint_root / "iter_0000001").exists()
+    assert (checkpoint_root / f"iter_{train_iters:07d}").exists()
+
+    selected_export_cmd_parts = [
+        "torchrun",
+        "--nproc_per_node=1",
+        "export_distilled_megatron_to_hf.py",
+        "--student_hf_path",
+        str(teacher_hf_path),
+        "--megatron_path",
+        str(checkpoint_root),
+        "--hf_export_path",
+        str(selected_exports),
+        "--export_iterations",
+        "1",
+        "2",
+    ]
+    run_example_command(selected_export_cmd_parts, example_path="megatron_bridge")
+
+    assert (selected_exports / "iter_0000001/config.json").exists()
+    assert (selected_exports / "iter_0000002/config.json").exists()
+
+    all_export_cmd_parts = [
+        "torchrun",
+        "--nproc_per_node=1",
+        "export_distilled_megatron_to_hf.py",
+        "--student_hf_path",
+        str(teacher_hf_path),
+        "--megatron_path",
+        str(checkpoint_root),
+        "--hf_export_path",
+        str(all_exports),
+        "--export_iterations",
+        "all",
+    ]
+    run_example_command(all_export_cmd_parts, example_path="megatron_bridge")
+
+    assert (all_exports / "iter_0000001/config.json").exists()
+    assert (all_exports / "iter_0000002/config.json").exists()

From 04624d1aa3185b44ec71cd5c3e42a48e0dffe381 Mon Sep 17 00:00:00 2001
From: Daniel Korzekwa <dkorzekwa@nvidia.com>
Date: Tue, 21 Jul 2026 05:57:04 -0700
Subject: [PATCH 22/25] Update docs.

Signed-off-by: Daniel Korzekwa <dkorzekwa@nvidia.com>
---
 examples/researcher_guide/README.md | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/examples/researcher_guide/README.md b/examples/researcher_guide/README.md
index f44e72ebe58..15512dcde5e 100644
--- a/examples/researcher_guide/README.md
+++ b/examples/researcher_guide/README.md
@@ -147,8 +147,9 @@ Interesting observations include:
 - Nemotron 3 produces the strongest MMLU-Pro trajectory, reaching 47.71%.
 - Although Nemotron v2 performs poorly alone, its 50/50 blend with WikiText slightly outperforms WikiText alone
   on both final benchmarks.
-- Nemotron v2 KD continues to decrease, but its MMLU-Pro score remains below the pruned baseline. Understanding
-  this discrepancy could help design better distillation corpora.
+- Nemotron v2 KD continues to decrease, but its MMLU-Pro score remains below the pruned baseline. The severe
+  MMLU-Pro regression appears to come from overfitting to Nemotron v2-style responses, even though MMLU-Pro prompts
+  ask the model to answer in a specific multiple-choice style.
 
 ## Prepare token-budgeted data blends
 

From a7e7f917cdca154bf777afba42467400ccaacbe3 Mon Sep 17 00:00:00 2001
From: Daniel Korzekwa <dkorzekwa@nvidia.com>
Date: Wed, 22 Jul 2026 01:12:31 -0700
Subject: [PATCH 23/25] improve tests. make them to run faster

Signed-off-by: Daniel Korzekwa <dkorzekwa@nvidia.com>
---
 .../test_export_distilled_megatron_to_hf.py   | 67 +++++++++++++------
 1 file changed, 47 insertions(+), 20 deletions(-)

diff --git a/tests/examples/megatron_bridge/test_export_distilled_megatron_to_hf.py b/tests/examples/megatron_bridge/test_export_distilled_megatron_to_hf.py
index f415c99bd6a..3a862d1400c 100644
--- a/tests/examples/megatron_bridge/test_export_distilled_megatron_to_hf.py
+++ b/tests/examples/megatron_bridge/test_export_distilled_megatron_to_hf.py
@@ -14,17 +14,63 @@
 # limitations under the License.
 """Tests for export_distilled_megatron_to_hf.py."""
 
+import argparse
+import sys
 from pathlib import Path
 
 from _test_utils.examples.run_command import extend_cmd_parts, run_example_command
 from _test_utils.torch.transformers_models import create_tiny_qwen3_dir
 
+# TODO: Move pure checkpoint path-selection logic out of the example script and into src so it
+# can be imported by unit tests without modifying sys.path.
+sys.path.insert(0, str(Path(__file__).resolve().parents[3] / "examples" / "megatron_bridge"))
+
+from export_distilled_megatron_to_hf import _get_checkpoint_export_paths
+
+
+def test_checkpoint_export_paths_for_selected_iterations(tmp_path: Path):
+    checkpoint_root = tmp_path / "checkpoints"
+    checkpoint_root.mkdir()
+    (checkpoint_root / "iter_0000001").mkdir()
+    (checkpoint_root / "iter_0000002").mkdir()
+    hf_export_root = tmp_path / "hf_validation"
+
+    args = argparse.Namespace(
+        megatron_path=str(checkpoint_root),
+        hf_export_path=str(hf_export_root),
+        export_iterations=["2", "1", "2"],
+    )
+
+    assert _get_checkpoint_export_paths(args) == [
+        (checkpoint_root / "iter_0000001", hf_export_root / "iter_0000001"),
+        (checkpoint_root / "iter_0000002", hf_export_root / "iter_0000002"),
+    ]
+
+
+def test_checkpoint_export_paths_for_all_iterations(tmp_path: Path):
+    checkpoint_root = tmp_path / "checkpoints"
+    checkpoint_root.mkdir()
+    (checkpoint_root / "iter_0000002").mkdir()
+    (checkpoint_root / "not_an_iteration").mkdir()
+    (checkpoint_root / "iter_0000001").mkdir()
+    hf_export_root = tmp_path / "hf_validation"
+
+    args = argparse.Namespace(
+        megatron_path=str(checkpoint_root),
+        hf_export_path=str(hf_export_root),
+        export_iterations=["all"],
+    )
+
+    assert _get_checkpoint_export_paths(args) == [
+        (checkpoint_root / "iter_0000001", hf_export_root / "iter_0000001"),
+        (checkpoint_root / "iter_0000002", hf_export_root / "iter_0000002"),
+    ]
+
 
 def test_export_distilled_megatron_iterations(tmp_path: Path, num_gpus):
     teacher_hf_path = create_tiny_qwen3_dir(tmp_path, with_tokenizer=True)
     train_iters = 2
     distill_output_dir = tmp_path / "distill_output"
-    selected_exports = tmp_path / "selected_exports"
     all_exports = tmp_path / "all_exports"
     distill_cmd_parts = extend_cmd_parts(
         ["torchrun", f"--nproc_per_node={num_gpus}", "distill.py", "--use_mock_data"],
@@ -49,25 +95,6 @@ def test_export_distilled_megatron_iterations(tmp_path: Path, num_gpus):
     assert (checkpoint_root / "iter_0000001").exists()
     assert (checkpoint_root / f"iter_{train_iters:07d}").exists()
 
-    selected_export_cmd_parts = [
-        "torchrun",
-        "--nproc_per_node=1",
-        "export_distilled_megatron_to_hf.py",
-        "--student_hf_path",
-        str(teacher_hf_path),
-        "--megatron_path",
-        str(checkpoint_root),
-        "--hf_export_path",
-        str(selected_exports),
-        "--export_iterations",
-        "1",
-        "2",
-    ]
-    run_example_command(selected_export_cmd_parts, example_path="megatron_bridge")
-
-    assert (selected_exports / "iter_0000001/config.json").exists()
-    assert (selected_exports / "iter_0000002/config.json").exists()
-
     all_export_cmd_parts = [
         "torchrun",
         "--nproc_per_node=1",

From 89389e9055d2f903380aae8cb81774c2f53b381a Mon Sep 17 00:00:00 2001
From: Daniel Korzekwa <dkorzekwa@nvidia.com>
Date: Thu, 23 Jul 2026 04:06:48 -0700
Subject: [PATCH 24/25] Add --validate_only to distill.py

Signed-off-by: Daniel Korzekwa <dkorzekwa@nvidia.com>
---
 examples/megatron_bridge/distill.py           | 12 +++++++
 .../examples/megatron_bridge/test_distill.py  | 34 +++++++++++++++++++
 2 files changed, 46 insertions(+)

diff --git a/examples/megatron_bridge/distill.py b/examples/megatron_bridge/distill.py
index ab7f2a1cd0e..8d4dc3b58d6 100644
--- a/examples/megatron_bridge/distill.py
+++ b/examples/megatron_bridge/distill.py
@@ -166,6 +166,11 @@ def get_args():
     parser.add_argument(
         "--eval_iters", type=int, default=32, help="Number of batches per validation stage"
     )
+    parser.add_argument(
+        "--validate_only",
+        action="store_true",
+        help="Skip training and run validation at iteration 0.",
+    )
     parser.add_argument(
         "--checkpoint_keep_last",
         type=int,
@@ -212,6 +217,8 @@ def get_args():
         args.student_hf_model = args.student_hf_path
     if args.checkpoint_keep_last < -1:
         raise ValueError("--checkpoint_keep_last must be >= -1.")
+    if args.validate_only and (args.eval_interval <= 0 or args.eval_iters <= 0):
+        raise ValueError("--validate_only requires --eval_interval > 0 and --eval_iters > 0.")
 
     print_args(args)
 
@@ -339,6 +346,7 @@ def _restore_student_hook(model_chunks):
             train_iters=args.train_iters,
             eval_interval=args.eval_interval,
             eval_iters=args.eval_iters,
+            skip_train=True if args.validate_only else None,
             global_batch_size=args.gbs,
             micro_batch_size=args.mbs,
             manual_gc=True,
@@ -384,6 +392,10 @@ def _restore_student_hook(model_chunks):
 
     print_rank_0("\nStarting distillation...")
     distill(config)
+    if args.validate_only:
+        print_rank_0("\nValidation-only run done! Skipped training and checkpoint export.\n")
+        return
+
     print_rank_0(
         f"\nDistillation done! Saved checkpoint to {checkpoint_dir}"
         " in megatron distributed checkpoint format.\n"
diff --git a/tests/examples/megatron_bridge/test_distill.py b/tests/examples/megatron_bridge/test_distill.py
index 60a71d3a812..7d8ac9e9d42 100644
--- a/tests/examples/megatron_bridge/test_distill.py
+++ b/tests/examples/megatron_bridge/test_distill.py
@@ -57,6 +57,40 @@ def test_distill_llm(tmp_path, num_gpus):
     assert (distilled_hf_path / "config.json").exists()
 
 
+def test_distill_validate_only(tmp_path, num_gpus):
+    teacher_hf_path = create_tiny_qwen3_dir(tmp_path, with_tokenizer=True)
+    train_iters = 2
+    output_dir = tmp_path / "distill_output"
+    hf_export_path = tmp_path / "distilled_hf"
+    distill_cmd_parts = extend_cmd_parts(
+        [
+            "torchrun",
+            f"--nproc_per_node={num_gpus}",
+            "distill.py",
+            "--use_mock_data",
+            "--validate_only",
+        ],
+        student_hf_path=teacher_hf_path,
+        teacher_hf_path=teacher_hf_path,
+        output_dir=output_dir,
+        tp_size=num_gpus,
+        pp_size=1,
+        seq_length=16,
+        mbs=1,
+        gbs=4,
+        train_iters=train_iters,
+        lr_warmup_iters=1,
+        eval_interval=train_iters,
+        eval_iters=1,
+        log_interval=1,
+        hf_export_path=hf_export_path,
+    )
+    run_example_command(distill_cmd_parts, example_path="megatron_bridge")
+
+    assert not (output_dir / f"checkpoints/iter_{train_iters:07d}").exists()
+    assert not (hf_export_path / "config.json").exists()
+
+
 # NOTE: Qwen3.5-VL-MoE covered by test_qad.py
 def test_distill_vlm(tmp_path, num_gpus):
     # Self-distillation of a tiny VLM: only the language model is distilled; the vision tower and the

From 6acf370b475129ca8b37dd91955f6907cb0eca84 Mon Sep 17 00:00:00 2001
From: Daniel Korzekwa <dkorzekwa@nvidia.com>
Date: Thu, 23 Jul 2026 04:46:16 -0700
Subject: [PATCH 25/25] improve tests

Signed-off-by: Daniel Korzekwa <dkorzekwa@nvidia.com>
---
 tests/_test_utils/examples/run_command.py      | 4 ++--
 tests/examples/megatron_bridge/test_distill.py | 4 +++-
 2 files changed, 5 insertions(+), 3 deletions(-)

diff --git a/tests/_test_utils/examples/run_command.py b/tests/_test_utils/examples/run_command.py
index eeb3df69108..44d8f06f603 100644
--- a/tests/_test_utils/examples/run_command.py
+++ b/tests/_test_utils/examples/run_command.py
@@ -77,7 +77,7 @@ def run_example_command(
     env: dict[str, str] | None = None,
     hf_max_retries: int = _HF_MAX_RETRIES,
     hf_retry_delay_s: int = _HF_RETRY_DELAY_S,
-):
+) -> str | None:
     """Run an example command, retrying transient HuggingFace access errors."""
     print(f"[{example_path}] Running command: {cmd_parts}")
     env = env or os.environ.copy()
@@ -88,7 +88,7 @@ def run_example_command(
             env["MASTER_PORT"] = str(get_free_port())  # fresh port per attempt
         returncode, output = _run_capturing(cmd_parts, cwd, env)
         if returncode == 0:
-            return
+            return output
         transient = any(marker in output for marker in _HF_TRANSIENT_MARKERS)
         if not transient or attempt == hf_max_retries:
             raise subprocess.CalledProcessError(returncode, cmd_parts)
diff --git a/tests/examples/megatron_bridge/test_distill.py b/tests/examples/megatron_bridge/test_distill.py
index 7d8ac9e9d42..1a5bf9568bf 100644
--- a/tests/examples/megatron_bridge/test_distill.py
+++ b/tests/examples/megatron_bridge/test_distill.py
@@ -85,8 +85,10 @@ def test_distill_validate_only(tmp_path, num_gpus):
         log_interval=1,
         hf_export_path=hf_export_path,
     )
-    run_example_command(distill_cmd_parts, example_path="megatron_bridge")
+    output = run_example_command(distill_cmd_parts, example_path="megatron_bridge")
 
+    assert "skipping training ..." in output
+    assert "iteration 0 on validation set" in output
     assert not (output_dir / f"checkpoints/iter_{train_iters:07d}").exists()
     assert not (hf_export_path / "config.json").exists()