Save hf checkpoint at every valitation iteration during distillation.

examples/dataset/MEGATRON_DATA_PREP.md

@@ -4,13 +4,17 @@
 | :---: | :---: | :---: |
 | 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)\] |
 
 The distillation and pre-training scripts in Megatron-Bridge or Megatron-LM expect data pre-tokenized in Megatron's binary indexed format (`.bin` / `.idx`).
 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.

examples/dataset/prepare_data_blend.py

@@ -0,0 +1,169 @@
+# 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()

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

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)

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`.
 

examples/pruning/minitron/NVIDIA-Nemotron-Nano-9B-v2/README.md

@@ -63,6 +63,10 @@ Distillation uses the **30% Pretraining (Code 5, General 20, MATH 5) + 70% Post-
 ### 1. Data Preparation
 
 See [examples/dataset/MEGATRON_DATA_PREP.md](../../../dataset/MEGATRON_DATA_PREP.md) for tokenization commands for all datasets used in this blend.
+To prepare a token-limited subset, follow the
+[token-budgeted data blend workflow](../../../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`.