diff --git a/CHANGELOG.rst b/CHANGELOG.rst index f20a33f1c17..9ee5a31b096 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -21,6 +21,7 @@ Changelog - Add ``--sft_in_batch_packing``, ``--moe_aux_loss_coeff`` and ``--dsa_kernel_backend`` to ``examples/megatron_bridge/distill.py`` (``--dsa_kernel_backend`` also to ``quantize.py``): pack ``--sft_hf_dataset`` conversations per micro-batch instead of padding each to ``--seq_length``, override the MoE load-balancing loss weight (e.g. 0 for padded SFT data), and select fused DeepSeek sparse-attention kernels for long sequences. - Add ``--vllm_fakequant_export`` to ``examples/megatron_bridge/export_quantized_megatron_to_hf.py`` to evaluate formats vLLM has no kernels for (e.g. GGML IQ) as BF16 fake-quant checkpoints, and allow ``--ep_size`` > 1 for checkpoints built with ``--no_moe_grouped_gemm``. - GGML GPTQ checkpoints from ``examples/megatron_bridge`` keep their calibrated codes when the Megatron checkpoint is reloaded for QAD or export, and layerwise calibration now runs on Megatron vision-language models. +- The Megatron unified HF export (``export_mcore_gpt_to_hf``) now copies the source checkpoint's non-model files (tokenizer, processor, remote code, chat templates, ...) into the export unchanged, from local disk only. The ``examples/megatron_bridge`` scripts make a Hugging Face Hub model local first with ``modelopt.torch.export.ensure_local_checkpoint``. *Misc* @@ -28,8 +29,12 @@ Changelog **Backward Breaking Changes** +- ``export_mcore_gpt_to_hf`` given a Hugging Face Hub model ID no longer writes the tokenizer, processor, generation config or remote-code files. Pass a local directory instead, e.g. from ``modelopt.torch.export.ensure_local_checkpoint``. + **Deprecations** +- ``modelopt.torch.export.copy_non_safetensor_files_from_ckpt`` is deprecated; use ``copy_non_model_files``. + **Bug Fixes** - Fix ONNX quantization placing a per-tensor Q/DQ pair on a rank-0 (scalar) constant input, which TensorRT's ONNX parser rejects (``convertAxis`` asserts the quantization axis lies in ``[0, nbDims]`` but ``nbDims == 0``). Rank-0 constant operands are now kept unquantized in the int8, fp8, and autotune export paths; the existing single-element Mul/Div/Pad stripping stays gated under ``target_dla``. diff --git a/examples/megatron_bridge/distill.py b/examples/megatron_bridge/distill.py index 333ccbbcb3d..8fbb3ab2fb1 100644 --- a/examples/megatron_bridge/distill.py +++ b/examples/megatron_bridge/distill.py @@ -74,6 +74,7 @@ import modelopt.torch.distill as mtd import modelopt.torch.utils.distributed as dist +from modelopt.torch.export import ensure_local_checkpoint from modelopt.torch.opt.conversion import ModeloptStateManager from modelopt.torch.utils import print_args, print_rank_0, warn_rank_0 from modelopt.torch.utils.mlflow import Tool, masked_args @@ -497,6 +498,12 @@ def _tokenizer_prepends_bos(args) -> bool: def main(args: argparse.Namespace, owns_the_run: bool = True): + # --student_hf_path stays as given. The student is built and exported from the local copy of + # the whole checkpoint in student_local_checkpoint_path, since the HF export reads the source's + # files from local disk only; student_hub_model_id is the Hub ID, or None for a local path. + args.student_hub_model_id, args.student_local_checkpoint_path = ensure_local_checkpoint( + args.student_hf_path + ) student_has_modelopt_state = args.student_megatron_path is not None and has_modelopt_state( args.student_megatron_path ) @@ -505,7 +512,7 @@ def main(args: argparse.Namespace, owns_the_run: bool = True): # reads either layout, so it keeps the faster grouped GEMM. moe_grouped_gemm = ( use_moe_grouped_gemm( - args.student_hf_path, + args.student_local_checkpoint_path, trust_remote_code=args.trust_remote_code, force_sequential=args.no_moe_grouped_gemm, ) @@ -553,7 +560,7 @@ def _build_model_provider(hf_path, load_weights=True, moe_grouped_gemm=True): # the built student inside the patched provide() below). # Only the student's layout is pinned -- it must match --student_megatron_path (see quantize.py). student_provider = _build_model_provider( - args.student_hf_path, + args.student_local_checkpoint_path, load_weights=args.student_megatron_path is None, moe_grouped_gemm=moe_grouped_gemm, ) @@ -591,7 +598,9 @@ def _build_model_provider(hf_path, load_weights=True, moe_grouped_gemm=True): # HF VLM configs expose ``vision_config``; Megatron-Bridge nests the text model under # ``language_model`` (used as ``distill_submodule`` below). - is_vlm = is_vlm_config(args.student_hf_path, trust_remote_code=args.trust_remote_code) + is_vlm = is_vlm_config( + args.student_local_checkpoint_path, trust_remote_code=args.trust_remote_code + ) if is_vlm: warn_rank_0( @@ -808,7 +817,7 @@ def _restore_student_hook(model_chunks): save_vlm_to_hf( full_student, args.hf_export_path, - args.student_hf_path, + args.student_local_checkpoint_path, trust_remote_code=args.trust_remote_code, ) print_rank_0(f"Saved distilled VLM to {args.hf_export_path} in HF format") @@ -825,7 +834,7 @@ def _restore_student_hook(model_chunks): export_llm_to_hf( megatron_path=f"{checkpoint_dir}/iter_{args.train_iters:07d}", hf_export_path=args.hf_export_path, - student_hf_path=args.student_hf_path, + student_hf_path=args.student_local_checkpoint_path, template_hf=args.student_hf_model, trust_remote_code=args.trust_remote_code, ) diff --git a/examples/megatron_bridge/export_distilled_megatron_to_hf.py b/examples/megatron_bridge/export_distilled_megatron_to_hf.py index 83c2b595389..26c481c6de5 100644 --- a/examples/megatron_bridge/export_distilled_megatron_to_hf.py +++ b/examples/megatron_bridge/export_distilled_megatron_to_hf.py @@ -72,7 +72,7 @@ from transformers import AutoConfig import modelopt.torch.utils.distributed as dist -from modelopt.torch.export import copy_hf_ckpt_remote_code +from modelopt.torch.export import copy_non_model_files, ensure_local_checkpoint from modelopt.torch.utils import print_args, print_rank_0 from modelopt.torch.utils.mlflow import Tool, masked_args from modelopt.torch.utils.plugins.mbridge import ( @@ -168,6 +168,7 @@ def export_llm_to_hf( AutoConfig.from_pretrained( student_hf_path, trust_remote_code=trust_remote_code ).save_pretrained(hf_export_path) + copy_non_model_files(student_hf_path, hf_export_path) def save_vlm_to_hf( @@ -194,7 +195,9 @@ def save_vlm_to_hf( ) export_bridge.hf_pretrained.save_artifacts(hf_export_path) export_bridge.save_hf_weights([full_model], hf_export_path) - copy_hf_ckpt_remote_code(student_hf_path, hf_export_path) + # Every rank runs this function (save_hf_weights is collective); one writer is enough. + if dist.is_master(): + copy_non_model_files(student_hf_path, hf_export_path) def get_args() -> argparse.Namespace: @@ -268,6 +271,12 @@ def get_args() -> argparse.Namespace: def main(args: argparse.Namespace): + # --student_hf_path stays as given. Later steps read the local copy of the whole checkpoint in + # student_local_checkpoint_path, since exporters read the source's files from local disk only; + # student_hub_model_id is the Hub ID, or None for a local path. + args.student_hub_model_id, args.student_local_checkpoint_path = ensure_local_checkpoint( + args.student_hf_path + ) checkpoint_export_paths: list[tuple[Path, Path]] = _get_checkpoint_export_paths(args) # This path drops quantization, so a QAD checkpoint would export silently unquantized. # ``has_modelopt_state`` ignores ``kd_loss``, so plain distillation still passes. @@ -277,14 +286,16 @@ def main(args: argparse.Namespace): f"{quantized[0]} is quantized; this script exports full precision only and would drop " "the quantizers. Use export_quantized_megatron_to_hf.py instead." ) - is_vlm = is_vlm_config(args.student_hf_path, trust_remote_code=args.trust_remote_code) + is_vlm = is_vlm_config( + args.student_local_checkpoint_path, trust_remote_code=args.trust_remote_code + ) 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("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, + hf_model_name_or_path=args.student_local_checkpoint_path, trust_remote_code=args.trust_remote_code, # Mirrors distill.py's unquantized branch moe_grouped_gemm=not args.no_moe_grouped_gemm, @@ -309,7 +320,7 @@ def main(args: argparse.Namespace): save_vlm_to_hf( full_model, str(hf_export_path), - args.student_hf_path, + args.student_local_checkpoint_path, trust_remote_code=args.trust_remote_code, ) record_exported_checkpoint(args, hf_export_path, dist.is_master()) @@ -327,7 +338,7 @@ def main(args: argparse.Namespace): export_llm_to_hf( megatron_path=str(megatron_path), hf_export_path=str(hf_export_path), - student_hf_path=args.student_hf_path, + student_hf_path=args.student_local_checkpoint_path, template_hf=args.student_hf_model, trust_remote_code=args.trust_remote_code, ) diff --git a/examples/megatron_bridge/export_quantized_megatron_to_hf.py b/examples/megatron_bridge/export_quantized_megatron_to_hf.py index a035dda679d..08a58102ce7 100644 --- a/examples/megatron_bridge/export_quantized_megatron_to_hf.py +++ b/examples/megatron_bridge/export_quantized_megatron_to_hf.py @@ -52,7 +52,7 @@ import modelopt.torch.quantization as mtq import modelopt.torch.utils.distributed as dist -from modelopt.torch.export import export_mcore_gpt_to_hf +from modelopt.torch.export import ensure_local_checkpoint, export_mcore_gpt_to_hf from modelopt.torch.export.plugins.vllm_fakequant_megatron import export_mcore_gpt_to_hf_vllm_fq from modelopt.torch.utils import print_args, print_rank_0 from modelopt.torch.utils.mlflow import Tool, masked_args @@ -158,16 +158,22 @@ def get_args() -> argparse.Namespace: def main(args: argparse.Namespace): + # --hf_model_name_or_path stays as given. Later steps read the local copy of the whole + # checkpoint in local_checkpoint_path, since exporters read the source's files from local disk + # only; hub_model_id is the Hub ID, or None for a local path. + args.hub_model_id, args.local_checkpoint_path = ensure_local_checkpoint( + args.hf_model_name_or_path + ) trust_remote_code = is_safe_repo( - trust_remote_code=args.trust_remote_code, hf_path=args.hf_model_name_or_path + trust_remote_code=args.trust_remote_code, hf_path=args.local_checkpoint_path ) # Build the model structure from HF _bridge, _provider, model, _unwrapped_model, _tokenizer = load_mbridge_model_from_hf( - hf_model_name_or_path=args.hf_model_name_or_path, + hf_model_name_or_path=args.local_checkpoint_path, trust_remote_code=trust_remote_code, moe_grouped_gemm=use_moe_grouped_gemm( - args.hf_model_name_or_path, + args.local_checkpoint_path, trust_remote_code=trust_remote_code, force_sequential=args.no_moe_grouped_gemm, ), @@ -214,7 +220,7 @@ def main(args: argparse.Namespace): ) export_fn( unwrapped_model, - args.hf_model_name_or_path, + args.local_checkpoint_path, export_extra_modules=export_extra_modules, dtype=torch.bfloat16, export_dir=args.export_unified_hf_path, diff --git a/examples/megatron_bridge/prune_minitron.py b/examples/megatron_bridge/prune_minitron.py index 0abaf810f3e..28e86c8baa1 100644 --- a/examples/megatron_bridge/prune_minitron.py +++ b/examples/megatron_bridge/prune_minitron.py @@ -58,7 +58,7 @@ import modelopt.torch.opt as mto import modelopt.torch.prune as mtp import modelopt.torch.utils.distributed as dist -from modelopt.torch.export import copy_hf_ckpt_remote_code +from modelopt.torch.export import copy_non_model_files, ensure_local_checkpoint from modelopt.torch.nas.plugins.megatron_model_stats import parse_main_layer_chars from modelopt.torch.utils import ( get_supported_datasets, @@ -438,8 +438,15 @@ def main(args: argparse.Namespace): warn_rank_0(f"\nPruned model already exists at {args.output_hf_path}. Exiting...") return + # --hf_model_name_or_path stays as given. Later steps read the local copy of the whole + # checkpoint in local_checkpoint_path, since exporters read the source's files from local disk + # only; hub_model_id is the Hub ID, or None for a local path. + args.hub_model_id, args.local_checkpoint_path = ensure_local_checkpoint( + args.hf_model_name_or_path + ) + bridge, provider, model, unwrapped_model, tokenizer = load_mbridge_model_from_hf( - hf_model_name_or_path=args.hf_model_name_or_path, + hf_model_name_or_path=args.local_checkpoint_path, trust_remote_code=args.trust_remote_code, provider_overrides={ "tensor_model_parallel_size": 1, # Tensor parallelism is not supported @@ -510,7 +517,7 @@ def main(args: argparse.Namespace): # the full VLM forward over image-text pairs. if use_image_calib: processor = AutoProcessor.from_pretrained( - args.hf_model_name_or_path, trust_remote_code=args.trust_remote_code + args.local_checkpoint_path, trust_remote_code=args.trust_remote_code ) forward_loop = get_megatron_vlm_calibration_forward_loop( unwrapped_model, # full VLM (vision encoder + projector + language model) @@ -643,7 +650,8 @@ def score_func(m): bridge.save_megatron_model( model, args.output_megatron_path, - hf_tokenizer_path=args.hf_model_name_or_path, + # Recorded in the checkpoint's config: the Hub ID stays valid on other hosts. + hf_tokenizer_path=args.hub_model_id or args.local_checkpoint_path, hf_tokenizer_kwargs=tokenizer_kwargs, ) args.checkpoint_exported = True @@ -752,7 +760,7 @@ def score_func(m): # from_hf_config can't infer it since AutoConfig consumes the kwarg. pruned_bridge.trust_remote_code = args.trust_remote_code pruned_bridge.save_hf_pretrained( - model, args.output_hf_path, source_path=args.hf_model_name_or_path + model, args.output_hf_path, source_path=args.local_checkpoint_path ) else: dummy_model_cls = AutoModelForImageTextToText if is_vlm else AutoModelForCausalLM @@ -764,7 +772,9 @@ def score_func(m): ) pruned_bridge.save_hf_weights(model, args.output_hf_path) - copy_hf_ckpt_remote_code(args.hf_model_name_or_path, args.output_hf_path) + # Every rank gets here (save_hf_weights is collective); one writer is enough. + if dist.is_master(): + copy_non_model_files(args.local_checkpoint_path, args.output_hf_path) args.checkpoint_exported = True print_rank_0(f"Saved pruned model to {args.output_hf_path} in HF checkpoint format") diff --git a/examples/megatron_bridge/quantize.py b/examples/megatron_bridge/quantize.py index 44bcee9bf00..efdab627b36 100644 --- a/examples/megatron_bridge/quantize.py +++ b/examples/megatron_bridge/quantize.py @@ -76,6 +76,7 @@ QUANT_CFG_CHOICES, RecipeSupersededAction, ) +from modelopt.torch.export import ensure_local_checkpoint from modelopt.torch.utils import print_args, print_rank_0, warn_rank_0 from modelopt.torch.utils.dataset_utils import get_supported_datasets from modelopt.torch.utils.mlflow import Tool, masked_args, resolved_recipe_texts @@ -331,18 +332,23 @@ def get_quant_config(args: argparse.Namespace) -> dict: def main(args: argparse.Namespace): + # --hf_model_name_or_path stays as given. Later steps read the local copy of the whole + # checkpoint in local_checkpoint_path; hub_model_id is the Hub ID, or None for a local path. + args.hub_model_id, args.local_checkpoint_path = ensure_local_checkpoint( + args.hf_model_name_or_path + ) trust_remote_code = is_safe_repo( - trust_remote_code=args.trust_remote_code, hf_path=args.hf_model_name_or_path + trust_remote_code=args.trust_remote_code, hf_path=args.local_checkpoint_path ) moe_grouped_gemm = use_moe_grouped_gemm( - args.hf_model_name_or_path, + args.local_checkpoint_path, trust_remote_code=trust_remote_code, force_sequential=args.no_moe_grouped_gemm, ) bridge, _provider, model, unwrapped_model, tokenizer = load_mbridge_model_from_hf( - hf_model_name_or_path=args.hf_model_name_or_path, + hf_model_name_or_path=args.local_checkpoint_path, trust_remote_code=trust_remote_code, moe_grouped_gemm=moe_grouped_gemm, provider_overrides={ @@ -447,7 +453,7 @@ def forward_loop(_model=None): # VLMs: drive the full VLM forward on image-text pairs so the language model's quantizers # see vision-conditioned activations (we still quantize the LM only). processor = AutoProcessor.from_pretrained( - args.hf_model_name_or_path, trust_remote_code=trust_remote_code + args.local_checkpoint_path, trust_remote_code=trust_remote_code ) forward_loop = get_megatron_vlm_calibration_forward_loop( unwrapped_model, # full VLM (vision encoder + projector + language model) @@ -491,7 +497,8 @@ def forward_loop(_model=None): bridge.save_megatron_model( model, args.export_megatron_path, - hf_tokenizer_path=args.hf_model_name_or_path, + # Recorded in the checkpoint's config: the Hub ID stays valid on other hosts. + hf_tokenizer_path=args.hub_model_id or args.local_checkpoint_path, hf_tokenizer_kwargs={"trust_remote_code": trust_remote_code}, ) args.checkpoint_exported = True diff --git a/modelopt/torch/export/plugins/vllm_fakequant_megatron.py b/modelopt/torch/export/plugins/vllm_fakequant_megatron.py index 1355b833189..de873f6d673 100644 --- a/modelopt/torch/export/plugins/vllm_fakequant_megatron.py +++ b/modelopt/torch/export/plugins/vllm_fakequant_megatron.py @@ -509,6 +509,10 @@ def export_mcore_gpt_to_hf_vllm_fq( pretrained model hosted inside a model repo on huggingface.co; or a *directory* containing model weights saved using [`~PreTrainedModel.save_pretrained`], e.g., `./my_model_directory/`. + The source's non-model files (tokenizer, processor, generation config, remote code) + are copied into the export from a local directory only; for a model ID they are + skipped with a warning. Get a local copy first with + `modelopt.torch.export.ensure_local_checkpoint` to include them. export_extra_modules: If True, export extra modules like medusa_heads or eagle_module. Otherwise, only export the base model. dtype: The weights data type to export the unquantized layers. diff --git a/modelopt/torch/export/unified_export_megatron.py b/modelopt/torch/export/unified_export_megatron.py index 6f3f128d1fc..ee2dfcb24c2 100644 --- a/modelopt/torch/export/unified_export_megatron.py +++ b/modelopt/torch/export/unified_export_megatron.py @@ -39,8 +39,7 @@ from modelopt.torch.utils import import_plugin, warn_rank_0 from modelopt.torch.utils.distributed import gather_state_dicts from modelopt.torch.utils.plugins.hf_checkpoint_utils import ( - copy_hf_ckpt_remote_code, - copy_non_safetensor_files_from_ckpt, + copy_non_model_files, load_multimodal_components, ) @@ -83,7 +82,6 @@ with import_plugin("transformers", verbose=False): import transformers - from transformers import AutoProcessor has_mcore = False @@ -129,6 +127,10 @@ class GPTModelExporter: pretrained model hosted inside a model repo on huggingface.co; or a *directory* containing model weights saved using [`~PreTrainedModel.save_pretrained`], e.g., `./my_model_directory/`. + The source's non-model files (tokenizer, processor, generation config, remote code) + are copied into the export from a local directory only; for a model ID they are + skipped with a warning. Get a local copy first with + `modelopt.torch.export.ensure_local_checkpoint` to include them. export_extra_modules: If True, export extra modules like medusa_heads or eagle_module. Otherwise, only export the base model. dtype: The weights data type to export the unquantized layers. @@ -391,47 +393,20 @@ def save_pretrained( if is_last_stage_main_rank: if is_writer_rank: + self._hf_config.save_pretrained(save_directory) + # Everything else the source ships -- tokenizer, processor, generation config, + # remote code -- is carried over verbatim rather than regenerated, from local disk + # only: a Hub ID's cache holds just what loading needed. if self._hf_pretrained_model_name is not None: if os.path.isdir(self._hf_pretrained_model_name): - copy_non_safetensor_files_from_ckpt( - self._hf_pretrained_model_name, save_directory - ) + copy_non_model_files(self._hf_pretrained_model_name, save_directory) else: - copy_hf_ckpt_remote_code(self._hf_pretrained_model_name, save_directory) - self._hf_config.save_pretrained(save_directory) - try: - generation_config = transformers.GenerationConfig.from_pretrained( - self._hf_pretrained_model_name, - trust_remote_code=self.trust_remote_code, - ) - # Pass it through unvalidated: save_pretrained rejects some shipped configs - # (e.g. GLM-5.3-Flash sets top_p without do_sample) on newer transformers. - generation_config.to_json_file( - os.path.join(save_directory, "generation_config.json") - ) - except OSError: - pass - # Hub-ID / None source: fetch tokenizer files via AutoTokenizer. - if self._hf_pretrained_model_name is None or not os.path.isdir( - self._hf_pretrained_model_name - ): - try: - tokenizer = transformers.AutoTokenizer.from_pretrained( - self._hf_pretrained_model_name, - trust_remote_code=self.trust_remote_code, + warn_rank_0( + f"{self._hf_pretrained_model_name!r} is not a local directory, so its " + "non-model files (tokenizer, processor, remote code, ...) were not " + "copied into the export. Get a local copy first with " + "modelopt.torch.export.ensure_local_checkpoint and pass that directory." ) - tokenizer.save_pretrained(save_directory) - except (OSError, TypeError, ValueError, ImportError): - pass - try: - # Load and save preprocessor config from the original model - processor = AutoProcessor.from_pretrained( - self._hf_pretrained_model_name, trust_remote_code=self.trust_remote_code - ) - if hasattr(processor, "image_processor"): - processor.image_processor.save_pretrained(save_directory) - except (OSError, ValueError, ImportError): - pass # The live MTP export runs EP collectives, so every last-stage main rank joins it; the # collective-free copy from the source checkpoint only runs on the writer. @@ -2337,6 +2312,10 @@ def export_mcore_gpt_to_hf( pretrained model hosted inside a model repo on huggingface.co; or a *directory* containing model weights saved using [`~PreTrainedModel.save_pretrained`], e.g., `./my_model_directory/`. + The source's non-model files (tokenizer, processor, generation config, remote code) + are copied into the export from a local directory only; for a model ID they are + skipped with a warning. Get a local copy first with + `modelopt.torch.export.ensure_local_checkpoint` to include them. export_extra_modules: If True, export extra modules like medusa_heads or eagle_module. Otherwise, only export the base model. dtype: The weights data type to export the unquantized layers. diff --git a/modelopt/torch/utils/plugins/hf_checkpoint_utils.py b/modelopt/torch/utils/plugins/hf_checkpoint_utils.py index dc8138d4286..625d5141ea2 100644 --- a/modelopt/torch/utils/plugins/hf_checkpoint_utils.py +++ b/modelopt/torch/utils/plugins/hf_checkpoint_utils.py @@ -736,7 +736,16 @@ def copy_non_safetensor_files_from_ckpt( Returns: File names copied into ``dst``. + + .. deprecated:: + Use :func:`copy_non_model_files`, which also skips non-safetensors weights, keeps files the + export already wrote, and copies subdirectories. """ + warnings.warn( + "copy_non_safetensor_files_from_ckpt is deprecated; use copy_non_model_files.", + DeprecationWarning, + stacklevel=2, + ) if not os.path.isdir(src): raise ValueError(f"Invalid source path: {src}. It should be a directory.") exclude_files = set(exclude_files or ()) diff --git a/tests/examples/megatron_bridge/conftest.py b/tests/examples/megatron_bridge/conftest.py index 1093fb7b281..d8da852f11a 100644 --- a/tests/examples/megatron_bridge/conftest.py +++ b/tests/examples/megatron_bridge/conftest.py @@ -15,9 +15,13 @@ """Run example steps without a ``torchrun`` launch; see ``_test_utils.examples.megatron_example_runner``.""" import os +import shutil import pytest +import transformers +import yaml from _test_utils.examples.megatron_example_runner import ( + _load_example_module, reset_megatron_global_state, run_example_step, ) @@ -55,3 +59,67 @@ def _isolate_megatron_global_state(): reset_megatron_global_state() os.environ.clear() os.environ.update(env_before) + + +class FakeHubCheckpoint: + """Serves a local checkpoint to example scripts as if it were Hub model ``hub_model_id``. + + ``serve(checkpoint, *scripts, nproc=...)`` patches ``ensure_local_checkpoint`` on each script's + module -- the one the in-process runner caches and runs -- to resolve ``hub_model_id`` (or + ``checkpoint`` itself, for a script that must read its argument before resolving it) to a copy + of ``checkpoint`` holding the file ``marker``, and returns ``hub_model_id`` to pass on the + command line: a step that read it without resolving it fails, and ``marker`` reaching an + export shows the export read the resolved copy. Like the Hugging Face cache would for a real Hub + ID, ``AutoTokenizer.from_pretrained(hub_model_id)`` loads the copy's tokenizer: Megatron-Bridge + loads the tokenizer a checkpoint records. Only a single-rank step runs in this process + and sees the patch, so with ``nproc > 1`` nothing is patched, ``checkpoint`` is returned + unchanged and ``active`` is False: check the marker and the Hub ID only when it is True. + """ + + hub_model_id = "org/tiny-model" + marker = "from_the_resolved_copy.md" + + def __init__(self, monkeypatch, tmp_path): + self._monkeypatch = monkeypatch + self._resolved = tmp_path / "resolved_hub_checkpoint" + self.active = False + + def serve(self, checkpoint, *scripts, nproc: int) -> str: + self.active = nproc == 1 + if not self.active: + return str(checkpoint) + shutil.copytree(checkpoint, self._resolved) + (self._resolved / self.marker).write_text("from the resolved copy\n") + + def ensure_local_checkpoint(model_name_or_path, group=None): + assert str(model_name_or_path) in (self.hub_model_id, str(checkpoint)) + return self.hub_model_id, str(self._resolved) + + for script in scripts: + module = _load_example_module(script, "megatron_bridge") + self._monkeypatch.setattr(module, "ensure_local_checkpoint", ensure_local_checkpoint) + + load_tokenizer = transformers.AutoTokenizer.from_pretrained + + def from_pretrained(*args, **kwargs): + # Megatron-LM passes the name as pretrained_model_name_or_path=. + if args and str(args[0]) == self.hub_model_id: + args = (str(self._resolved), *args[1:]) + if str(kwargs.get("pretrained_model_name_or_path")) == self.hub_model_id: + kwargs["pretrained_model_name_or_path"] = str(self._resolved) + return load_tokenizer(*args, **kwargs) + + self._monkeypatch.setattr(transformers.AutoTokenizer, "from_pretrained", from_pretrained) + return self.hub_model_id + + @staticmethod + def saved_tokenizer_model(megatron_path) -> str: + """The tokenizer a Megatron-Bridge checkpoint records in its ``run_config.yaml``.""" + (run_config,) = sorted(megatron_path.rglob("run_config.yaml")) + return yaml.safe_load(run_config.read_text())["tokenizer"]["tokenizer_model"] + + +@pytest.fixture +def fake_hub_checkpoint(monkeypatch, tmp_path): + """See :class:`FakeHubCheckpoint`.""" + return FakeHubCheckpoint(monkeypatch, tmp_path) diff --git a/tests/examples/megatron_bridge/test_distill.py b/tests/examples/megatron_bridge/test_distill.py index 3bdcc460c27..6d146957426 100644 --- a/tests/examples/megatron_bridge/test_distill.py +++ b/tests/examples/megatron_bridge/test_distill.py @@ -40,9 +40,17 @@ @pytest.mark.timeout(360) # distill + inline HF export + export of every iteration -def test_distill_llm_hf_export(tmp_path, num_gpus): - """Distill with an inline HF export, then re-export every saved iteration standalone.""" +def test_distill_llm_hf_export(tmp_path, num_gpus, fake_hub_checkpoint): + """Distill with an inline HF export, then re-export every saved iteration standalone. + + Both exports read the student from the copy ``ensure_local_checkpoint`` resolved (see + ``fake_hub_checkpoint``). distill.py keeps the local path on its command line, since get_args() + loads the tokenizer from it; the standalone export is passed the Hub ID. + """ teacher_hf_path = create_tiny_qwen3_dir(tmp_path, with_tokenizer=True) + student_export_arg = fake_hub_checkpoint.serve( + teacher_hf_path, "distill.py", "export_distilled_megatron_to_hf.py", nproc=num_gpus + ) train_iters = 2 distill_output_dir = tmp_path / "distill_output" distilled_hf_path = tmp_path / "distilled_hf" @@ -70,11 +78,13 @@ def test_distill_llm_hf_export(tmp_path, num_gpus): checkpoint_root = distill_output_dir / "checkpoints" assert (checkpoint_root / f"iter_{train_iters:07d}").exists() assert (distilled_hf_path / "config.json").exists() + if fake_hub_checkpoint.active: + assert (distilled_hf_path / fake_hub_checkpoint.marker).exists() # export_distilled_megatron_to_hf.py converts any saved iteration, not just the final one. all_export_cmd_parts = extend_cmd_parts( ["torchrun", "--nproc_per_node=1", "export_distilled_megatron_to_hf.py"], - student_hf_path=teacher_hf_path, + student_hf_path=student_export_arg, megatron_path=checkpoint_root, hf_export_path=all_exports, export_iterations="all", @@ -83,6 +93,8 @@ def test_distill_llm_hf_export(tmp_path, num_gpus): for it in range(1, train_iters + 1): assert (all_exports / f"iter_{it:07d}/config.json").exists() + if fake_hub_checkpoint.active: + assert (all_exports / f"iter_{it:07d}" / fake_hub_checkpoint.marker).exists() def test_distill_llm_sft(tmp_path, num_gpus): @@ -241,7 +253,7 @@ def test_distill_validate_only(tmp_path, num_gpus): # NOTE: Qwen3.5-VL-MoE covered by test_qad.py @pytest.mark.timeout(420) # sometimes times out in CI env with default 300s timeout -def test_distill_vlm(tmp_path, num_gpus): +def test_distill_vlm(tmp_path, num_gpus, fake_hub_checkpoint): # Self-distillation of a tiny VLM: only the language model is distilled; the vision tower and the # vision->language projector must be left byte-for-byte untouched. tp_size = num_gpus @@ -293,10 +305,14 @@ def _is_language_model(name: str) -> bool: megatron_ckpt = distill_output_dir / f"checkpoints/iter_{train_iters:07d}" assert megatron_ckpt.exists() - # Separately convert the distilled Megatron checkpoint to HF with the standalone export script + # Separately convert the distilled Megatron checkpoint to HF with the standalone export script, + # passing the student as a Hub ID (see ``fake_hub_checkpoint``). + student_export_arg = fake_hub_checkpoint.serve( + vlm_hf_path, "export_distilled_megatron_to_hf.py", nproc=num_gpus + ) export_cmd_parts = extend_cmd_parts( ["torchrun", f"--nproc_per_node={num_gpus}", "export_distilled_megatron_to_hf.py"], - student_hf_path=vlm_hf_path, + student_hf_path=student_export_arg, megatron_path=megatron_ckpt, hf_export_path=distilled_hf_path, tp_size=tp_size, @@ -305,6 +321,8 @@ def _is_language_model(name: str) -> bool: run_example_command(export_cmd_parts, example_path="megatron_bridge") assert (distilled_hf_path / "config.json").exists() + if fake_hub_checkpoint.active: + assert (distilled_hf_path / fake_hub_checkpoint.marker).exists() # from_pretrained (default strict load) verifies the saved weights match the (unchanged) config. distilled_model = AutoModelForImageTextToText.from_pretrained(distilled_hf_path) diff --git a/tests/examples/megatron_bridge/test_prune_minitron.py b/tests/examples/megatron_bridge/test_prune_minitron.py index 71230f54dcd..2cdee235d06 100644 --- a/tests/examples/megatron_bridge/test_prune_minitron.py +++ b/tests/examples/megatron_bridge/test_prune_minitron.py @@ -70,15 +70,20 @@ ), ], ) -def test_prune_minitron(tmp_path, num_gpus, create_teacher, expected_pruned_config): +def test_prune_minitron( + tmp_path, num_gpus, create_teacher, expected_pruned_config, fake_hub_checkpoint +): + """The teacher is passed as a Hub ID (see ``fake_hub_checkpoint``): it is pruned from its local + copy, and the HF output carries that copy's non-model files.""" teacher_hf_path, teacher_model = create_teacher(tmp_path, num_gpus) + teacher_arg = fake_hub_checkpoint.serve(teacher_hf_path, "prune_minitron.py", nproc=num_gpus) teacher_params = sum(p.numel() for p in teacher_model.parameters()) prune_target_params = int(teacher_params * 0.8) pruned_path = tmp_path / "pruned" prune_command_parts = extend_cmd_parts( ["torchrun", f"--nproc_per_node={num_gpus}", "prune_minitron.py"], - hf_model_name_or_path=teacher_hf_path, + hf_model_name_or_path=teacher_arg, output_hf_path=pruned_path, pp_size=num_gpus, calib_dataset_name="cnn_dailymail", @@ -94,6 +99,8 @@ def test_prune_minitron(tmp_path, num_gpus, create_teacher, expected_pruned_conf run_example_command(prune_command_parts, example_path="megatron_bridge") assert (pruned_path / "config.json").exists() + if fake_hub_checkpoint.active: + assert (pruned_path / fake_hub_checkpoint.marker).exists() pruned_model = AutoModelForCausalLM.from_pretrained(pruned_path) assert sum(p.numel() for p in pruned_model.parameters()) <= prune_target_params for field, expected in expected_pruned_config.items(): diff --git a/tests/examples/megatron_bridge/test_quantize_export.py b/tests/examples/megatron_bridge/test_quantize_export.py index 9d46482f328..a60fb417acd 100644 --- a/tests/examples/megatron_bridge/test_quantize_export.py +++ b/tests/examples/megatron_bridge/test_quantize_export.py @@ -70,16 +70,24 @@ def test_quantize_and_export( recipe, quantize_parallelism, export_parallelism, + fake_hub_checkpoint, ): - """Quantize a tiny model via a YAML recipe and export it to a unified HF checkpoint.""" + """Quantize a tiny model via a YAML recipe and export it to a unified HF checkpoint. + + The model is passed as a Hub ID (see ``fake_hub_checkpoint``): both steps load and export from + its local copy, and the Megatron checkpoint records the Hub ID as its tokenizer. + """ hf_model_path = create_model(tmp_path, with_tokenizer=True, **model_kwargs) + model_arg = fake_hub_checkpoint.serve( + hf_model_path, "quantize.py", "export_quantized_megatron_to_hf.py", nproc=num_gpus + ) megatron_path = tmp_path / "quantized_megatron" hf_export_path = tmp_path / "quantized_hf" # Step 1: quantize and save a Megatron checkpoint quantize_cmd = extend_cmd_parts( ["torchrun", f"--nproc_per_node={num_gpus}", "quantize.py", "--skip_generate"], - hf_model_name_or_path=hf_model_path, + hf_model_name_or_path=model_arg, recipe=recipe, **{quantize_parallelism: num_gpus}, calib_dataset_name="cnn_dailymail", @@ -91,11 +99,16 @@ def test_quantize_and_export( run_example_command(quantize_cmd, example_path="megatron_bridge", setup_free_port=True) assert (megatron_path / "latest_checkpointed_iteration.txt").exists() assert_has_modelopt_state(megatron_path) + if fake_hub_checkpoint.active: + assert ( + fake_hub_checkpoint.saved_tokenizer_model(megatron_path) + == fake_hub_checkpoint.hub_model_id + ) # Step 2: export to HF export_cmd = extend_cmd_parts( ["torchrun", f"--nproc_per_node={num_gpus}", "export_quantized_megatron_to_hf.py"], - hf_model_name_or_path=hf_model_path, + hf_model_name_or_path=model_arg, megatron_path=megatron_path, export_unified_hf_path=hf_export_path, **{export_parallelism: num_gpus}, @@ -104,6 +117,8 @@ def test_quantize_and_export( assert (hf_export_path / "config.json").exists() assert (hf_export_path / "hf_quant_config.json").exists() assert_exported_checkpoint_matches(hf_export_path, hf_model_path) + if fake_hub_checkpoint.active: + assert (hf_export_path / fake_hub_checkpoint.marker).exists() # The exported unified checkpoint should be loadable and runnable by vLLM. The deployment check below # is disabled because it takes too long in CI (likely because of first run) diff --git a/tests/unit/torch/utils/test_hf_checkpoint_utils.py b/tests/unit/torch/utils/test_hf_checkpoint_utils.py index 093e8e1cf66..23f23b0e83c 100644 --- a/tests/unit/torch/utils/test_hf_checkpoint_utils.py +++ b/tests/unit/torch/utils/test_hf_checkpoint_utils.py @@ -56,18 +56,20 @@ def test_copy_non_safetensor_files_from_ckpt_supports_additional_exclusions(tmp_ (src_dir / "reasoning_parser.py").write_text("parser") default_dst = tmp_path / "default" - copy_non_safetensor_files_from_ckpt(src_dir, default_dst) + with pytest.deprecated_call(): + copy_non_safetensor_files_from_ckpt(src_dir, default_dst) assert not (default_dst / "model.safetensors").exists() assert not (default_dst / "model.safetensors.index.json").exists() assert (default_dst / "pytorch_model.bin").exists() assert (default_dst / "stats.npy").exists() filtered_dst = tmp_path / "filtered" - copy_non_safetensor_files_from_ckpt( - src_dir, - filtered_dst, - exclude_patterns=("*.bin", "*.npy"), - ) + with pytest.deprecated_call(): + copy_non_safetensor_files_from_ckpt( + src_dir, + filtered_dst, + exclude_patterns=("*.bin", "*.npy"), + ) assert (filtered_dst / "reasoning_parser.py").exists() assert not (filtered_dst / "pytorch_model.bin").exists() assert not (filtered_dst / "stats.npy").exists() @@ -87,7 +89,7 @@ def copy2(source, *args, **kwargs): return original_copy2(source, *args, **kwargs) monkeypatch.setattr(hf_checkpoint_utils.shutil, "copy2", copy2) - with pytest.warns(UserWarning, match="bad.py"): + with pytest.deprecated_call(), pytest.warns(UserWarning, match="bad.py"): copied_files = copy_non_safetensor_files_from_ckpt(src_dir, tmp_path / "dst") assert copied_files == ["good.py"]