Save hf checkpoint at every valitation iteration during distillation. (#1897)

### What does this PR do?

Save hf checkpoint at every valitation iteration during distillation.

Addition functionality:
- added `--validate_only` in `examples/megatron_bridge/distill.py` to
enable computing validation losses for iter 0
- added `--reset_optimizer` in `examples/megatron_bridge/distill.py` to
enable not using presaved optimizer, e.g., when changing the number of
train iters.

### Usage

- examples/megatron_bridge/distill.py
- examples/megatron_bridge/README.md (line 228)


### Testing

- tests/examples/megatron_bridge/test_distill.py

### Before your PR is "*Ready for review*"

- Is this change backward compatible?: ✅ 
- Did you write any new necessary tests?: ✅ 

<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit

* **New Features**
  * Added `--validate_only` for student-only validation at iteration 0.
* Added `--hf_validation_export_path` with
`--hf_validation_export_interval` to export validated student
HuggingFace artifacts during distillation.
* Added `prepare_data_blend.py` for YAML-driven token-budgeted data
blends.
* Added `--max_tokens` to stop Megatron preprocessing after a token
budget.
* **Bug Fixes**
* Validation exports avoid duplicate checkpoints, preserve the student
architecture/config, and export only student artifacts.
* **Documentation**
* Expanded researcher and tutorial guides for iterative workflows and
token-budgeted blends.
* **Tests**
* Updated distillation/blend/max_tokens test coverage, including
validate-only and interval-based exports.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Signed-off-by: Daniel Korzekwa <dkorzekwa@nvidia.com>

CHANGELOG.rst

@@ -26,6 +26,7 @@ Experimental
 
 **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).

examples/megatron_bridge/README.md

@@ -158,6 +158,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
@@ -218,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:

examples/megatron_bridge/distill.py

@@ -167,6 +167,21 @@ 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,
+        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(
@@ -201,6 +216,10 @@ 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.")
+    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)
 
@@ -331,7 +350,11 @@ def main(args: argparse.Namespace):
             manual_gc=True,
             manual_gc_interval=100,
         ),
-        validation=ValidationConfig(eval_iters=args.eval_iters, eval_interval=args.eval_interval),
+        validation=ValidationConfig(
+            eval_iters=args.eval_iters,
+            eval_interval=args.eval_interval,
+            skip_train=args.validate_only,
+        ),
         optimizer=optimizer_config,
         scheduler=scheduler_config,
         ddp=DistributedDataParallelConfig(
@@ -359,7 +382,7 @@ def main(args: argparse.Namespace):
             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,
@@ -370,6 +393,10 @@ def main(args: argparse.Namespace):
 
     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"

examples/megatron_bridge/export_distilled_megatron_to_hf.py

@@ -40,10 +40,22 @@ Example VLM (checkpoint reshards on load, so TP/PP/EP need not match training):
         --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 @@ from modelopt.torch.utils.plugins.mbridge import (
     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__":

examples/researcher_guide/README.md

@@ -7,6 +7,10 @@ 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 +51,106 @@ 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. 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.
+
+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
+[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 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.
+
+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) |
+| Pruned 0.7x student | 48.69% (full) | 23.09% (full) |
+
+### WikiText
+
+- Dataset: Salesforce/wikitext (`wikitext-103-v1`)
+- Teacher CE: 2.6834
+
+| 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
+
+- Dataset: nvidia/Nemotron-Post-Training-Dataset-v2 (math and stem)
+- Teacher CE: 1.1566
+
+| 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% |
+
+### 50/50 WikiText and Nemotron v2 blend
+
+- 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 | 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% |
+| 78.6M | 0.0973 | 1.8503 | 67.72% | 41.57% |
+| 100.0M | 0.0916 | 1.7680 | 68.28% | 41.71% |
+
+### Nemotron 3
+
+- 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 | 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% |
+
+Interesting observations include:
+
+- 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. 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
 
 Preparing complete distillation datasets can consume unnecessary time and storage during early experiments.

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)

tests/examples/megatron_bridge/test_distill.py

@@ -57,6 +57,42 @@ 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,
+    )
+    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()
+
+
 # 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

tests/examples/megatron_bridge/test_export_distilled_megatron_to_hf.py

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