Save hf checkpoint at every valitation iteration during distillation.

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 _restore_student_hook(model_chunks):
             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 _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,
@@ -370,6 +393,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"

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__":