diff --git a/CHANGELOG.rst b/CHANGELOG.rst index aadd0f0563f..2ea1b2f3db6 100755 --- a/CHANGELOG.rst +++ b/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 `_. - 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 `_) 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 e69af367a5b..45606d691ee 100644 --- a/examples/megatron_bridge/README.md +++ b/examples/megatron_bridge/README.md @@ -158,6 +158,9 @@ Tensorboard logging is enabled by default and logs are saved to `/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_` subdirectory: + +```bash +torchrun --nproc_per_node 1 export_distilled_megatron_to_hf.py \ + --student_hf_path \ + --megatron_path /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 96880d2ae82..b0f6f1af865 100644 --- a/examples/megatron_bridge/distill.py +++ b/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 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 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" 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_." + ), + ) + parser.add_argument( + "--export_iterations", + nargs="+", + default=None, + help=( + "Export checkpoints from the checkpoint root passed to --megatron_path. Use " + "'all' for every iter_ 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 dd0334bf12a..15512dcde5e 100644 --- a/examples/researcher_guide/README.md +++ b/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. 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 60a71d3a812..1a5bf9568bf 100644 --- a/tests/examples/megatron_bridge/test_distill.py +++ b/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 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..3a862d1400c --- /dev/null +++ b/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()