From 111524cda708dedc5b2c11764870441fde71a46a Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Wed, 30 Sep 2026 21:30:35 +0000 Subject: [PATCH 1/6] hf_ptq: load from a local checkpoint and carry its files verbatim A Hub model ID reached the loader as-is, so every later step that needed the source checkpoint as files had its own way back to a directory: _resolve_model_path (whose TRANSFORMERS_CACHE glob picked the lexicographically largest commit hash rather than refs/main), a hardcoded sidecar allowlist, _resolved_local_dir, and hf_hub_download fallbacks. Now main() calls ensure_local_checkpoint once, leaving --pyt_ckpt_path as given and storing hf_model_path (the local copy every step reads) and hf_model_name (the Hub ID, or None). The layerwise resume directory is still named after the Hub ID. The unified exporters copy the source's non-model files themselves, so the example no longer does -- except after the deprecated TensorRT-LLM export, which does not. It also stops regenerating tokenizer and processor files: the source's are carried over verbatim. That fixes exports shipping pad_token set to the EOS token, which get_tokenizer sets for calibration, and removes the default padding side and pad token threaded through six functions only to undo that before saving. The notebooks do the same. Signed-off-by: Shengliang Xu --- CHANGELOG.rst | 1 + examples/hf_ptq/example_utils.py | 271 +----------------- examples/hf_ptq/hf_ptq.py | 129 +++------ .../1_FP4-FP8_PTQ_Min-Max_Calibration.ipynb | 11 +- .../notebooks/2_PTQ_AWQ_Calibration.ipynb | 11 +- .../notebooks/3_PTQ_AutoQuantization.ipynb | 12 +- tests/examples/hf_ptq/test_example_utils.py | 89 ------ tests/examples/hf_ptq/test_hf_ptq_args.py | 14 +- .../hf_ptq/test_hf_ptq_vision_quantization.py | 5 +- 9 files changed, 76 insertions(+), 467 deletions(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 2efd87bdda4..fb48228e263 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -129,6 +129,7 @@ Changelog - Fix offline DFlash training failing to reconstruct the target logits when the captured hidden states are stored in a different dtype than the target's weights. - Fix ``hf_ptq`` overwriting a model's existing ``pad_token`` with ``eos_token`` when the model already has a valid padding token. The exported tokenizer now preserves the source model's padding configuration. - Fix Hugging Face exports dropping off-index safetensors such as GLM-4.7's ``mtp.safetensors``. +- Fix ``examples/hf_ptq`` exports setting ``pad_token`` to the EOS token in the tokenizer files; the source checkpoint's tokenizer files are now exported unchanged. - Fix Megatron-Core checkpoint saving for quantized grouped MoE experts when tensor and expert parallelism are both enabled. - Fix unified HuggingFace export of RADIO-based VLMs retaining post-conversion vision and projector names instead of restoring the hub layout; deployment loaders could skip those weights. diff --git a/examples/hf_ptq/example_utils.py b/examples/hf_ptq/example_utils.py index c613de08ef3..0f6c7f1ae95 100755 --- a/examples/hf_ptq/example_utils.py +++ b/examples/hf_ptq/example_utils.py @@ -15,14 +15,13 @@ import argparse import copy -import glob import hashlib import inspect import json import logging import os import warnings -from collections.abc import Iterable, Iterator +from collections.abc import Iterator from contextlib import contextmanager from dataclasses import dataclass from datetime import timedelta @@ -49,71 +48,14 @@ from modelopt.torch.export.model_utils import is_multimodal_model from modelopt.torch.models import hf_model_type from modelopt.torch.models.hf import checkpoint_has_mtp, prepare_model_for_loading -from modelopt.torch.utils.plugins.hf_checkpoint_utils import ( - copy_non_safetensor_files_from_ckpt, - copy_off_index_safetensors, -) -from modelopt.torch.utils.plugins.model_load_utils import record_unplaced_source_keys - -try: - from huggingface_hub import snapshot_download -except ImportError: - snapshot_download = None - from modelopt.torch.utils import distributed as dist_utils -from modelopt.torch.utils.mlflow import EXPERIMENT_JSON, Tool, resolved_recipe_texts, tracked_run +from modelopt.torch.utils.mlflow import Tool, resolved_recipe_texts, tracked_run +from modelopt.torch.utils.plugins.model_load_utils import record_unplaced_source_keys logger = logging.getLogger(__name__) SPECULATIVE_MODEL_LIST = ["Eagle", "Medusa"] -_HF_SIDECAR_DOWNLOAD_ALLOW_PATTERNS = [ - "*.jinja", - "*.json", - "*.md", - "*.model", - "*.py", - "*.tiktoken", - "*.txt", - "LICENSE*", - "NOTICE*", -] -_HF_PTQ_WEIGHT_FILE_PATTERNS = ( - "*.safetensors", - "*.safetensors.index.json", - "*.bin", - "*.bin.index.json", - "*.ckpt", - "*.gguf", - "*.h5", - "*.msgpack", - "*.npy", - "*.npz", - "*.onnx", - "*.pb", - "*.pickle", - "*.pkl", - "*.pt", - "*.pth", - "*.tar", - "*.tar.bz2", - "*.tar.gz", - "*.tar.xz", - "*.tflite", - "*.tgz", - "*.zip", -) -_HF_PTQ_EXPORT_OWNED_FILES = { - EXPERIMENT_JSON, - "config.json", - "hf_quant_config.json", - "quant_config.json", - "quantization_config.json", - "quantize_config.json", - "recipe.yaml", - "recipe.yml", -} - @dataclass class DistributedState: @@ -151,7 +93,7 @@ def cleanup_distributed(args): def validate_fsdp2_supported(args, config): """Raise ``NotImplementedError`` for model/CLI combos the FSDP2 path doesn't support yet.""" issues = [] - if "vila" in args.pyt_ckpt_path.lower(): + if "vila" in args.hf_model_path.lower(): issues.append("VILA (custom builder + non-standard layer layout)") if is_nemotron_vl(config) or _is_multimodal_config(config): issues.append("multimodal / VL models (decoder layers not auto-detectable)") @@ -161,7 +103,7 @@ def validate_fsdp2_supported(args, config): issues.append("speculative decoding (--specdec_offline_dataset)") if getattr(args, "low_memory_mode", False): issues.append("--low_memory_mode (redundant with FSDP2)") - if not issues and checkpoint_has_mtp(config.model_type, args.pyt_ckpt_path): + if not issues and checkpoint_has_mtp(config.model_type, args.hf_model_path): issues.append("Nemotron-H MTP (auxiliary modules are not constructed by the FSDP2 loader)") if issues: @@ -453,18 +395,9 @@ def _unpack_compressed_linear_weights(model, ckpt_path=None): if not ckpt_path: return - from huggingface_hub import hf_hub_download - - is_local = os.path.isdir(ckpt_path) - def _resolve_file(filename): - if is_local: - local = os.path.join(ckpt_path, filename) - return local if os.path.exists(local) else None - try: - return hf_hub_download(repo_id=ckpt_path, filename=filename) - except Exception: - return None + local = os.path.join(ckpt_path, filename) + return local if os.path.exists(local) else None # Load non-expert weights and metadata from safetensors checkpoint_weights = {} @@ -600,25 +533,6 @@ def _fmt_max_memory(max_memory: dict) -> str: return "\n".join(parts) -def _resolved_local_dir(ckpt_path: str) -> str: - """Return the local directory ``ckpt_path`` names, resolving a hub id to its snapshot. - - The export re-reads the source checkpoint by path to carry over the weights the loader could - not place. Recording the hub id instead would leave it reading ``org/model``, which is not a - directory -- so every carried weight would be dropped with a warning. ``from_pretrained`` has - already populated the cache by the time this runs, so the lookup is local and offline. - """ - if Path(ckpt_path).is_dir(): - return str(ckpt_path) - if snapshot_download is None: - return str(ckpt_path) - try: - return snapshot_download(ckpt_path, local_files_only=True) - except Exception: - # No snapshot to point at; the export falls back to its own provenance handling. - return str(ckpt_path) - - def _from_pretrained_recording(auto_class, ckpt_path, *, model_type=None, **kwargs): """``from_pretrained`` that records what the loader could not place. @@ -639,7 +553,7 @@ def _from_pretrained_recording(auto_class, ckpt_path, *, model_type=None, **kwar ckpt_path, output_loading_info=True, **kwargs ) unexpected = loading_info.get("unexpected_keys") or [] - record_unplaced_source_keys(model, _resolved_local_dir(ckpt_path), unexpected) + record_unplaced_source_keys(model, str(ckpt_path), unexpected) if unexpected: print( f"✓ {len(unexpected)} checkpoint key(s) the model has no parameter for " @@ -941,175 +855,13 @@ def generate_excludes_prompt(model) -> bool: return bool(getattr(config, "is_encoder_decoder", False)) and not is_diffusion_gemma(config) -def _resolve_model_path(model_name_or_path: str, trust_remote_code: bool = False) -> str: - """Resolve a model name or path to a local directory path. - - If the input is already a local directory, returns it as-is. - If the input is a HuggingFace model ID, attempts to resolve it to the local cache path. - - Args: - model_name_or_path: Either a local directory path or HuggingFace model ID - trust_remote_code: Whether to trust remote code when loading the model - - Returns: - Local directory path to the model files - """ - # If it's already a local directory, return as-is - if os.path.isdir(model_name_or_path): - return model_name_or_path - - # Try to resolve HuggingFace model ID to local cache path - try: - # First try to load the config to trigger caching - config = AutoConfig.from_pretrained(model_name_or_path, trust_remote_code=trust_remote_code) - - # The config object should have the local path information - # Try different ways to get the cached path - if hasattr(config, "_name_or_path") and os.path.isdir(config._name_or_path): - return config._name_or_path - - # Alternative: use snapshot_download if available - if snapshot_download is not None: - try: - local_path = snapshot_download( - repo_id=model_name_or_path, - allow_patterns=_HF_SIDECAR_DOWNLOAD_ALLOW_PATTERNS, - ) - return local_path - except Exception as e: - print( - f"Warning: Could not download checkpoint sidecars using snapshot_download: {e}" - ) - - # Fallback: try to find in HuggingFace cache - from transformers.utils import TRANSFORMERS_CACHE - - # Look for the model in the cache directory - cache_pattern = os.path.join(TRANSFORMERS_CACHE, "models--*") - cache_dirs = glob.glob(cache_pattern) - - # Convert model name to cache directory format - model_cache_name = model_name_or_path.replace("/", "--") - for cache_dir in cache_dirs: - if model_cache_name in cache_dir: - # Look for the snapshots directory - snapshots_dir = os.path.join(cache_dir, "snapshots") - if os.path.exists(snapshots_dir): - # Get the latest snapshot - snapshot_dirs = [ - d - for d in os.listdir(snapshots_dir) - if os.path.isdir(os.path.join(snapshots_dir, d)) - ] - if snapshot_dirs: - latest_snapshot = max(snapshot_dirs) # Use lexicographically latest - snapshot_path = os.path.join(snapshots_dir, latest_snapshot) - return snapshot_path - - except Exception as e: - print(f"Warning: Could not resolve model path for {model_name_or_path}: {e}") - - # If all else fails, return the original path - # This will cause the copy function to skip with a warning - return model_name_or_path - - -def copy_custom_model_files( - source_path: str, - export_path: str, - trust_remote_code: bool = False, - exclude_files: Iterable[str] | None = None, - copy_off_index_weights: bool = True, -): - """Copy source checkpoint sidecar files to an HF PTQ export. - - The HF PTQ script writes ModelOpt-owned metadata and quantized weights first, then - copies source checkpoint sidecars so tokenizer/processor files, remote-code modules, - README assets, parser plugins, and similar deployment files are preserved for both - native and ``trust_remote_code`` loads. Weight and weight-index files are skipped - to avoid copying the unquantized source weights. Export-owned metadata (``config.json``, - ``hf_quant_config.json``) and stale source quantization metadata are also skipped. - Source tokenizer and processor files intentionally still win because Transformers may - not regenerate all metadata in the source format. The exported ``tokenizer_config.json`` - wins when it has a separate chat template. Callers that write a generation config can - exclude it; the TensorRT-LLM export retains the source generation config. - - Args: - source_path: Path to the original model directory or HuggingFace model ID - export_path: Path to the exported model directory - trust_remote_code: Passed to HuggingFace model-ID resolution; does not control copying. - exclude_files: Additional source file names to skip. - copy_off_index_weights: Copy safetensors the loader never opens (GLM-4.7's - ``mtp.safetensors``). Only the unified-HF export gives them meaning -- it seeds their - tensor names into ``quantization_config.ignore`` -- so a TensorRT-LLM export, whose - checkpoint is ``rank.safetensors`` plus its own ``config.json``, should pass False - rather than carry gigabytes nothing there reads. - """ - # Resolve the source path (handles both local paths and HF model IDs) - resolved_source_path = _resolve_model_path(source_path, trust_remote_code) - - source_dir = Path(resolved_source_path) - export_dir = Path(export_path) - - if not source_dir.exists(): - if resolved_source_path != source_path: - print( - f"Warning: Could not find local cache for HuggingFace model '{source_path}' " - f"(resolved to '{resolved_source_path}')" - ) - else: - print(f"Warning: Source directory '{source_path}' does not exist") - return - - if not export_dir.exists(): - print(f"Warning: Export directory {export_path} does not exist") - return - - exclude_files = _HF_PTQ_EXPORT_OWNED_FILES | set(exclude_files or ()) - if (export_dir / "chat_template.jinja").is_file(): - exclude_files.add("tokenizer_config.json") - - copied_files = copy_non_safetensor_files_from_ckpt( - source_dir, - export_dir, - exclude_files=exclude_files, - exclude_patterns=_HF_PTQ_WEIGHT_FILE_PATTERNS, - ) - - # Safetensors the loader never opens are sidecars too: untouched by quantization and absent - # from the export, so copy them rather than leave them behind. Skipped by the call above, - # which excludes every *.safetensors to avoid re-emitting the unquantized source weights. - copied_weights = ( - copy_off_index_safetensors(source_dir, export_dir) if copy_off_index_weights else [] - ) - copied_files = [*copied_files, *copied_weights] - if copied_files: - for file_name in copied_files: - print(f"Copied checkpoint sidecar file: {file_name}") - print(f"Successfully copied {len(copied_files)} checkpoint sidecar files to {export_path}") - else: - print("No checkpoint sidecar files found to copy") - - def save_source_config(args, export_path) -> None: """Copy the source model's config to the export path, for VLMs the exporters skip.""" print(f"Saving original model config to {export_path}") config_kwargs = {"trust_remote_code": args.trust_remote_code} if args.attn_implementation is not None: config_kwargs["attn_implementation"] = args.attn_implementation - AutoConfig.from_pretrained(args.pyt_ckpt_path, **config_kwargs).save_pretrained(export_path) - - -def save_processor_config(args, export_path) -> None: - """Copy the processor config, without which a VLM checkpoint cannot preprocess images.""" - try: - print(f"Saving processor config to {export_path}") - AutoProcessor.from_pretrained( - args.pyt_ckpt_path, trust_remote_code=args.trust_remote_code - ).save_pretrained(export_path) - except Exception as e: - print(f"Warning: Could not save processor config: {e}") - print("This is normal for some VLM architectures that don't use AutoProcessor") + AutoConfig.from_pretrained(args.hf_model_path, **config_kwargs).save_pretrained(export_path) def _prepare_quant_cfg( @@ -1132,7 +884,10 @@ def _prepare_quant_cfg( ) if needs_checkpoint_path_update(quant_cfg): - quant_cfg, resolved_dir = resolve_checkpoint_dir(quant_cfg, args.pyt_ckpt_path) + # Named after the Hub ID when there is one: a cache path would name it by commit hash. + quant_cfg, resolved_dir = resolve_checkpoint_dir( + quant_cfg, args.hf_model_name or args.hf_model_path + ) print(f"Auto-resolved layerwise checkpoint_dir: {resolved_dir}") if args.cast_mxfp4_to_nvfp4: diff --git a/examples/hf_ptq/hf_ptq.py b/examples/hf_ptq/hf_ptq.py index 8b221a175d8..e7ad8997db7 100755 --- a/examples/hf_ptq/hf_ptq.py +++ b/examples/hf_ptq/hf_ptq.py @@ -35,10 +35,8 @@ from example_utils import ( HF_PTQ, _prepare_quant_cfg, - _resolve_model_path, build_quant_cfg, cleanup_distributed, - copy_custom_model_files, create_vlm_calibration_loop, generate_excludes_prompt, get_model, @@ -49,7 +47,6 @@ mlflow_run, recipe_layerwise_blocks, run_nemotron_vl_preview, - save_processor_config, save_source_config, setup_distributed_args, validate_fsdp2_supported, @@ -104,6 +101,10 @@ ) from modelopt.torch.utils.memory_monitor import launch_memory_monitor from modelopt.torch.utils.mlflow import add_mlflow_args, resolve_mlflow_args +from modelopt.torch.utils.plugins.hf_checkpoint_utils import ( + copy_non_model_files, + ensure_local_checkpoint, +) from modelopt.torch.utils.plugins.model_load_utils import parallel_load_and_prepare_fsdp2 from modelopt.torch.utils.speech_dataset_utils import get_speech_dataset_dataloader from modelopt.torch.utils.vlm_dataset_utils import get_vlm_dataset_dataloader @@ -307,11 +308,11 @@ def load_model(args: argparse.Namespace): raise NotImplementedError(_FSDP2_KV_AUTOQUANT_ERROR) if args.use_fsdp2: hf_config = AutoConfig.from_pretrained( - args.pyt_ckpt_path, trust_remote_code=args.trust_remote_code + args.hf_model_path, trust_remote_code=args.trust_remote_code ) validate_fsdp2_supported(args, hf_config) full_model = parallel_load_and_prepare_fsdp2( - args.pyt_ckpt_path, + args.hf_model_path, args.dist_state.device, args.dist_state.rank, args.dist_state.world_size, @@ -322,7 +323,7 @@ def load_model(args: argparse.Namespace): ) elif args.specdec_offline_dataset is not None or not args.low_memory_mode: full_model = get_model( - args.pyt_ckpt_path, + args.hf_model_path, args.dist_state.device, gpu_mem_percentage=args.gpu_max_mem_percentage, trust_remote_code=args.trust_remote_code, @@ -348,16 +349,16 @@ def load_model(args: argparse.Namespace): quant_cfg, gpu_mem_percentage=args.gpu_max_mem_percentage, quant_gemm=False ): hf_config = AutoConfig.from_pretrained( - args.pyt_ckpt_path, trust_remote_code=args.trust_remote_code + args.hf_model_path, trust_remote_code=args.trust_remote_code ) model_kwargs = {"trust_remote_code": args.trust_remote_code} if args.attn_implementation is not None: model_kwargs["attn_implementation"] = args.attn_implementation with prepare_model_for_loading( - hf_config.model_type, args.pyt_ckpt_path, args.trust_remote_code + hf_config.model_type, args.hf_model_path, args.trust_remote_code ): full_model = AutoModelForCausalLM.from_pretrained( - args.pyt_ckpt_path, + args.hf_model_path, **model_kwargs, ) calibration_only = True @@ -378,8 +379,6 @@ def load_model(args: argparse.Namespace): processor = None tokenizer = None language_model = full_model - default_padding_side = None - default_pad_token = None is_nemotron_vl_model = is_nemotron_vl(full_model) @@ -396,14 +395,14 @@ def load_model(args: argparse.Namespace): if model_type == "whisper": processor = get_processor( - args.pyt_ckpt_path, + args.hf_model_path, model_type, trust_remote_code=args.trust_remote_code, ) elif args.calib_with_images: # For VLM image calibration, we need an AutoProcessor to build multimodal inputs. processor = AutoProcessor.from_pretrained( - args.pyt_ckpt_path, + args.hf_model_path, trust_remote_code=args.trust_remote_code, padding_side="left", ) @@ -411,15 +410,13 @@ def load_model(args: argparse.Namespace): if hasattr(processor, "tokenizer") and processor.tokenizer is not None: tokenizer = processor.tokenizer else: - tokenizer = get_tokenizer(args.pyt_ckpt_path, trust_remote_code=args.trust_remote_code) + tokenizer = get_tokenizer(args.hf_model_path, trust_remote_code=args.trust_remote_code) - default_pad_token = tokenizer.pad_token # Some Nemotron tokenizers may not define pad_token by default; but we use padding=True during calibration. if tokenizer.pad_token is None: tokenizer.pad_token = tokenizer.eos_token - assert tokenizer.pad_token is not None, f"Pad token for {args.pyt_ckpt_path} cannot be set!" + assert tokenizer.pad_token is not None, f"Pad token for {args.hf_model_path} cannot be set!" - default_padding_side = tokenizer.padding_side tokenizer.padding_side = "left" # Plain PTQ quantizes only the language model. Recipes keep the complete VLM so their @@ -452,10 +449,8 @@ def load_model(args: argparse.Namespace): if extracted_lm is not None: language_model = extracted_lm - tokenizer = get_tokenizer(args.pyt_ckpt_path, trust_remote_code=args.trust_remote_code) + tokenizer = get_tokenizer(args.hf_model_path, trust_remote_code=args.trust_remote_code) - default_padding_side = tokenizer.padding_side - default_pad_token = tokenizer.pad_token # Left padding usually provides better calibration result. tokenizer.padding_side = "left" @@ -466,8 +461,6 @@ def load_model(args: argparse.Namespace): calibration_only, processor, tokenizer, - default_padding_side, - default_pad_token, device, ) @@ -640,9 +633,6 @@ def export_quantized( full_model: torch.nn.Module, language_model: torch.nn.Module, model_type: str | None, - tokenizer: PreTrainedTokenizerBase | None, - default_padding_side: str | None, - default_pad_token: str | None, ): # Not inference_mode: the FSDP2 path gathers full params in this context and # inference tensors break the subsequent state_dict() -> param.detach(). @@ -662,7 +652,6 @@ def export_quantized( # the source config would replace; it never writes a processor config. if not args.layerwise_export: save_source_config(args, export_path) - save_processor_config(args, export_path) start_time = time.time() is_tensorrt_llm_export = ( @@ -689,17 +678,10 @@ def export_quantized( inference_tensor_parallel=args.inference_tensor_parallel, inference_pipeline_parallel=args.inference_pipeline_parallel, ) - - # Copy custom model files (Python files and JSON configs) for TensorRT-LLM export - # TRT-LLM checkpoints are rank.safetensors plus their own config; nothing - # there reads an off-index sidecar, and the exclude_modules seeding that gives - # one meaning happens only inside export_hf_checkpoint. - copy_custom_model_files( - args.pyt_ckpt_path, - export_path, - args.trust_remote_code, - copy_off_index_weights=False, - ) + # The unified exporters copy the source's non-model files themselves; this deprecated + # one does not. + if args.dist_state.is_main: + copy_non_model_files(args.hf_model_path, export_path) else: # Check arguments for unified_hf export format and set to default if unsupported arguments are provided assert args.sparsity_fmt == "dense", ( @@ -736,28 +718,6 @@ def export_quantized( "vLLM deployment support is in progress." ) - # Restore default padding and export the tokenizer as well. - if tokenizer is not None: - tokenizer.padding_side = default_padding_side - if default_pad_token is not None: - tokenizer.pad_token = default_pad_token - if args.dist_state.is_main: - tokenizer.save_pretrained(export_path) - - # Copy custom model files (Python files and JSON configs) if trust_remote_code is used. - # This must run AFTER tokenizer.save_pretrained() so original tokenizer files - # from the source checkpoint take precedence over regenerated ones (which may - # differ in format due to newer transformers versions). - if args.dist_state.is_main: - exclude_files = None if is_tensorrt_llm_export else {"generation_config.json"} - copy_custom_model_files( - args.pyt_ckpt_path, - export_path, - args.trust_remote_code, - exclude_files=exclude_files, - copy_off_index_weights=not is_tensorrt_llm_export, - ) - args.checkpoint_exported = True end_time = time.time() print_rank_0( @@ -809,7 +769,7 @@ def pre_quantize( full_model, tokenizer, preview_input_ids, - args.pyt_ckpt_path, + args.hf_model_path, "before quantization", allow_fallback=False, trust_remote_code=args.trust_remote_code, @@ -836,8 +796,6 @@ def post_quantize( generated_ids_before_ptq: Any, is_nemotron_vl_model: bool, first_text_speech_dataset: str | None, - default_padding_side: str | None, - default_pad_token: str | None, calib_dataloader: DataLoader, ): """ @@ -851,15 +809,7 @@ def post_quantize( # Early exit for offline speculative decoding: skip generation comparison and export directly. # The model's get_dummy_inputs() provides the right input format for the export forward pass. if args.specdec_offline_dataset is not None: - export_quantized( - args, - full_model, - language_model, - model_type, - tokenizer, - default_padding_side, - default_pad_token, - ) + export_quantized(args, full_model, language_model, model_type) return assert preview_input_ids is not None, "pre_quantize returns preview inputs for this path" @@ -895,7 +845,7 @@ def post_quantize( full_model, tokenizer, preview_input_ids, - args.pyt_ckpt_path, + args.hf_model_path, "after quantization", allow_fallback=False, trust_remote_code=args.trust_remote_code, @@ -951,15 +901,7 @@ def output_decode(generated_ids, input_shape): f"example outputs after ptq: {output_decode(generated_ids_after_ptq, preview_input_ids.shape[1])}" ) - export_quantized( - args, - full_model, - language_model, - model_type, - tokenizer, - default_padding_side, - default_pad_token, - ) + export_quantized(args, full_model, language_model, model_type) def quantize_main( @@ -970,8 +912,6 @@ def quantize_main( calibration_only: bool, processor: ProcessorMixin | None, tokenizer: PreTrainedTokenizerBase | None, - default_padding_side, - default_pad_token, device: torch.device, ): # Load the recipe up front so we can detect layerwise calibration before batch-size probing. @@ -1159,12 +1099,9 @@ def quantize_main( # to NVFP4StaticQuantizer with a data-derived ``_global_amax``); we just # override that scalar with the closed-form value before export. if args.cast_mxfp4_to_nvfp4: - # The cast reads the source MXFP4 ``*_scales``/``*_blocks`` tensors from a local - # checkpoint directory. ``--pyt_ckpt_path`` may be a HF Hub ID (e.g. - # ``openai/gpt-oss-20b``); resolve it to the local snapshot dir that load_model's - # ``from_pretrained`` already populated so the cast works with the documented command. - source_ckpt_dir = _resolve_model_path(args.pyt_ckpt_path, args.trust_remote_code) - apply_cast_mxfp4_to_nvfp4(language_model, source_ckpt_dir) + # The cast reads the source MXFP4 ``*_scales``/``*_blocks`` tensors from the local + # checkpoint directory main() downloaded. + apply_cast_mxfp4_to_nvfp4(language_model, args.hf_model_path) post_quantize( args, @@ -1178,8 +1115,6 @@ def quantize_main( generated_ids_before_ptq, is_nemotron_vl_model, first_text_speech_dataset, - default_padding_side, - default_pad_token, calib_dataloader, ) @@ -1191,7 +1126,9 @@ def parse_args() -> argparse.Namespace: "--model", help=( "Model name or path to the PyTorch checkpoint to be quantized. " - "Can be a local path or a Huggingface model name." + "Can be a local path or a Huggingface model name; a model name is downloaded in " + "full. To skip files, download it yourself (e.g. `hf download --exclude " + "'original/*'`) and pass the local path." ), required=True, ) @@ -1536,6 +1473,10 @@ def main(args: argparse.Namespace): # cleanup_distributed would leave the other ranks blocked on the first collective # until the NCCL timeout. with mlflow_run(args): + # --pyt_ckpt_path stays as given. Every later step reads the local copy of the whole + # checkpoint in hf_model_path; hf_model_name is the Hub ID, or None for a local path. + args.hf_model_name, args.hf_model_path = ensure_local_checkpoint(args.pyt_ckpt_path) + # launch a memory monitor to read the currently used GPU memory. launch_memory_monitor() @@ -1549,8 +1490,6 @@ def main(args: argparse.Namespace): calibration_only, processor, tokenizer, - default_padding_side, - default_pad_token, device, ) = load_model(args) @@ -1567,8 +1506,6 @@ def main(args: argparse.Namespace): calibration_only, processor, tokenizer, - default_padding_side, - default_pad_token, device, ) finally: diff --git a/examples/hf_ptq/notebooks/1_FP4-FP8_PTQ_Min-Max_Calibration.ipynb b/examples/hf_ptq/notebooks/1_FP4-FP8_PTQ_Min-Max_Calibration.ipynb index 50b53ac38e7..140653ebfe6 100644 --- a/examples/hf_ptq/notebooks/1_FP4-FP8_PTQ_Min-Max_Calibration.ipynb +++ b/examples/hf_ptq/notebooks/1_FP4-FP8_PTQ_Min-Max_Calibration.ipynb @@ -111,8 +111,12 @@ "metadata": {}, "outputs": [], "source": [ - "model = AutoModelForCausalLM.from_pretrained(model_name).cuda()\n", - "tokenizer = AutoTokenizer.from_pretrained(model_name)\n", + "from modelopt.torch.export import ensure_local_checkpoint\n", + "\n", + "# A local copy of the whole checkpoint, so the export can carry over its tokenizer and other files.\n", + "_, model_path = ensure_local_checkpoint(model_name)\n", + "model = AutoModelForCausalLM.from_pretrained(model_path).cuda()\n", + "tokenizer = AutoTokenizer.from_pretrained(model_path)\n", "tokenizer.pad_token = tokenizer.eos_token" ] }, @@ -250,8 +254,7 @@ "from modelopt.torch.export import export_hf_checkpoint\n", "\n", "export_path = \"./quantized_model_min-max/\"\n", - "export_hf_checkpoint(model, export_dir=export_path)\n", - "tokenizer.save_pretrained(export_path)" + "export_hf_checkpoint(model, export_dir=export_path)" ] }, { diff --git a/examples/hf_ptq/notebooks/2_PTQ_AWQ_Calibration.ipynb b/examples/hf_ptq/notebooks/2_PTQ_AWQ_Calibration.ipynb index d02b7ba49b8..2fd88a81301 100644 --- a/examples/hf_ptq/notebooks/2_PTQ_AWQ_Calibration.ipynb +++ b/examples/hf_ptq/notebooks/2_PTQ_AWQ_Calibration.ipynb @@ -113,8 +113,12 @@ "metadata": {}, "outputs": [], "source": [ - "model = AutoModelForCausalLM.from_pretrained(model_name).cuda()\n", - "tokenizer = AutoTokenizer.from_pretrained(model_name)\n", + "from modelopt.torch.export import ensure_local_checkpoint\n", + "\n", + "# A local copy of the whole checkpoint, so the export can carry over its tokenizer and other files.\n", + "_, model_path = ensure_local_checkpoint(model_name)\n", + "model = AutoModelForCausalLM.from_pretrained(model_path).cuda()\n", + "tokenizer = AutoTokenizer.from_pretrained(model_path)\n", "tokenizer.pad_token = tokenizer.eos_token" ] }, @@ -247,8 +251,7 @@ "from modelopt.torch.export import export_hf_checkpoint\n", "\n", "export_path = \"./quantized_model_awq/\"\n", - "export_hf_checkpoint(model, export_dir=export_path)\n", - "tokenizer.save_pretrained(export_path)" + "export_hf_checkpoint(model, export_dir=export_path)" ] }, { diff --git a/examples/hf_ptq/notebooks/3_PTQ_AutoQuantization.ipynb b/examples/hf_ptq/notebooks/3_PTQ_AutoQuantization.ipynb index aba5fdfce77..73efb8b24ca 100644 --- a/examples/hf_ptq/notebooks/3_PTQ_AutoQuantization.ipynb +++ b/examples/hf_ptq/notebooks/3_PTQ_AutoQuantization.ipynb @@ -131,8 +131,12 @@ }, "outputs": [], "source": [ - "model = AutoModelForCausalLM.from_pretrained(MODEL_ID, torch_dtype=DTYPE).to(DEVICE)\n", - "tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)\n", + "from modelopt.torch.export import ensure_local_checkpoint\n", + "\n", + "# A local copy of the whole checkpoint, so the export can carry over its tokenizer and other files.\n", + "_, model_path = ensure_local_checkpoint(MODEL_ID)\n", + "model = AutoModelForCausalLM.from_pretrained(model_path, torch_dtype=DTYPE).to(DEVICE)\n", + "tokenizer = AutoTokenizer.from_pretrained(model_path)\n", "tokenizer.padding_side = \"left\"\n", "tokenizer.pad_token = tokenizer.eos_token\n", "model.eval()" @@ -366,7 +370,7 @@ "metadata": {}, "outputs": [], "source": [ - "from modelopt.torch.export import export_hf_checkpoint\n", + "from modelopt.torch.export import copy_non_model_files, export_hf_checkpoint\n", "from modelopt.torch.export.trtllm import export_tensorrt_llm_checkpoint\n", "\n", "if EXPORT_FMT == \"tensorrt_llm\":\n", @@ -376,6 +380,8 @@ " inference_tensor_parallel=1,\n", " inference_pipeline_parallel=1,\n", " )\n", + " # export_hf_checkpoint copies the source's tokenizer and other non-model files itself.\n", + " copy_non_model_files(model_path, EXPORT_DIR)\n", "else:\n", " export_hf_checkpoint(model, export_dir=EXPORT_DIR)\n", "\n", diff --git a/tests/examples/hf_ptq/test_example_utils.py b/tests/examples/hf_ptq/test_example_utils.py index 35e2f75b9a4..a6e44a0c363 100644 --- a/tests/examples/hf_ptq/test_example_utils.py +++ b/tests/examples/hf_ptq/test_example_utils.py @@ -110,95 +110,6 @@ def prepare(*args, model_class): assert calls == failure_at -def test_copy_custom_model_files_preserves_non_weight_sidecars(tmp_path): - source_dir = tmp_path / "source" - export_dir = tmp_path / "export" - source_dir.mkdir() - export_dir.mkdir() - - source_files = { - "super_v3_reasoning_parser.py": "class Parser: pass\n", - "modeling_custom.py": "class Model: pass\n", - "README.md": "# Source model\n", - "LICENSE": "license text\n", - "chat_template.jinja": "{{ messages }}\n", - "tokenizer_config.json": '{"chat_template": "source"}\n', - "generation_config.json": '{"source": "generation"}\n', - "config.json": '{"source": "config"}\n', - "hf_quant_config.json": '{"source": "quant"}\n', - "quant_config.json": '{"source": "stale quant"}\n', - "quantize_config.json": '{"source": "stale quant"}\n', - "recipe.yaml": "quantize: {}\n", - "model.safetensors.index.json": '{"weight_map": {}}\n', - "model-00001-of-00001.safetensors": "source weights\n", - "model.gguf": "source weights\n", - } - for file_name, contents in source_files.items(): - (source_dir / file_name).write_text(contents) - - (export_dir / "config.json").write_text('{"export": "config"}\n') - (export_dir / "generation_config.json").write_text('{"export": "generation"}\n') - (export_dir / "hf_quant_config.json").write_text('{"export": "quant"}\n') - (export_dir / "chat_template.jinja").write_text("{{ exported_messages }}\n") - (export_dir / "tokenizer_config.json").write_text('{"chat_template": "export"}\n') - - example_utils.copy_custom_model_files(str(source_dir), str(export_dir), trust_remote_code=False) - - for file_name in [ - "super_v3_reasoning_parser.py", - "modeling_custom.py", - "README.md", - "LICENSE", - "chat_template.jinja", - "generation_config.json", - ]: - assert (export_dir / file_name).read_text() == source_files[file_name] - - assert (export_dir / "config.json").read_text() == '{"export": "config"}\n' - assert (export_dir / "hf_quant_config.json").read_text() == '{"export": "quant"}\n' - assert (export_dir / "tokenizer_config.json").read_text() == '{"chat_template": "export"}\n' - assert not (export_dir / "quant_config.json").exists() - assert not (export_dir / "quantize_config.json").exists() - assert not (export_dir / "recipe.yaml").exists() - assert not (export_dir / "model.safetensors.index.json").exists() - assert not (export_dir / "model-00001-of-00001.safetensors").exists() - assert not (export_dir / "model.gguf").exists() - - (export_dir / "generation_config.json").write_text('{"export": "generation"}\n') - example_utils.copy_custom_model_files( - str(source_dir), - str(export_dir), - exclude_files={"generation_config.json"}, - ) - assert (export_dir / "generation_config.json").read_text() == '{"export": "generation"}\n' - - -def test_resolve_model_path_snapshot_download_stays_allowlisted(monkeypatch, tmp_path): - snapshot_dir = tmp_path / "snapshot" - - def fake_snapshot_download(**kwargs): - assert kwargs == { - "repo_id": "org/model", - "allow_patterns": example_utils._HF_SIDECAR_DOWNLOAD_ALLOW_PATTERNS, - } - return str(snapshot_dir) - - def fake_from_pretrained(*args, **kwargs): - assert (args, kwargs) == (("org/model",), {"trust_remote_code": False}) - return SimpleNamespace(_name_or_path="org/model") - - monkeypatch.setattr( - example_utils.AutoConfig, - "from_pretrained", - fake_from_pretrained, - ) - monkeypatch.setattr(example_utils, "snapshot_download", fake_snapshot_download) - - assert example_utils._resolve_model_path("org/model", trust_remote_code=False) == str( - snapshot_dir - ) - - # ---------- get_original_hf_quant_method ------------------------------------- # get_model uses this to detect native MXFP4 checkpoints (e.g. openai/gpt-oss-*) and load # them dequantized to BF16 GptOssExperts (so ModelOpt can quantize/export the experts). diff --git a/tests/examples/hf_ptq/test_hf_ptq_args.py b/tests/examples/hf_ptq/test_hf_ptq_args.py index 51d495bb2cf..5a47727e77f 100644 --- a/tests/examples/hf_ptq/test_hf_ptq_args.py +++ b/tests/examples/hf_ptq/test_hf_ptq_args.py @@ -216,7 +216,8 @@ def test_hf_ptq_runs_fixed_ptq_before_kv_autoquantize(monkeypatch): batch_size=1, auto_quantize_checkpoint=None, kv_auto_quantize_checkpoint=None, - pyt_ckpt_path="dummy", + hf_model_path="dummy", + hf_model_name=None, cast_mxfp4_to_nvfp4=False, layerwise_export=False, specdec_offline_dataset=None, @@ -262,7 +263,8 @@ def test_fixed_ptq_then_kv_rejects_explicit_kv_before_calibration(monkeypatch): args = SimpleNamespace( auto_quantize_checkpoint=None, kv_auto_quantize_checkpoint=None, - pyt_ckpt_path="dummy", + hf_model_path="dummy", + hf_model_name=None, cast_mxfp4_to_nvfp4=False, layerwise_export=False, ) @@ -776,12 +778,6 @@ def test_only_the_main_rank_clears_an_inherited_pointer(monkeypatch, example_uti assert inherited.exists() -def test_experiment_json_is_export_owned(example_utils): - """copy_custom_model_files copies source sidecars including dotfiles, so without this - the source checkpoint's pointer would follow it into every derived checkpoint.""" - assert example_utils.EXPERIMENT_JSON in example_utils._HF_PTQ_EXPORT_OWNED_FILES - - def test_untracked_runs_write_no_experiment_json(monkeypatch, example_utils, tmp_path): _, args = _parse_hf_ptq_args( monkeypatch, "--pyt_ckpt_path", "/models/Qwen3-0.6B", "--export_path", str(tmp_path) @@ -884,8 +880,6 @@ def test_post_quantize_export_survives_a_failed_sanity_generate(monkeypatch): generated_ids_before_ptq=torch.zeros(1, 4, dtype=torch.long), is_nemotron_vl_model=False, first_text_speech_dataset=None, - default_padding_side="right", - default_pad_token=None, calib_dataloader=None, ) diff --git a/tests/examples/hf_ptq/test_hf_ptq_vision_quantization.py b/tests/examples/hf_ptq/test_hf_ptq_vision_quantization.py index 4941a11fe69..9a885687d23 100644 --- a/tests/examples/hf_ptq/test_hf_ptq_vision_quantization.py +++ b/tests/examples/hf_ptq/test_hf_ptq_vision_quantization.py @@ -49,7 +49,8 @@ def test_image_calibration_model_target_follows_recipe( use_fsdp2=False, specdec_offline_dataset=None, low_memory_mode=False, - pyt_ckpt_path="dummy", + hf_model_path="dummy", + hf_model_name=None, dist_state=SimpleNamespace(device=torch.device("cpu")), gpu_max_mem_percentage=0.8, trust_remote_code=False, @@ -184,7 +185,5 @@ def test_vision_recipe_requires_image_calibration(hf_ptq, recipe): calibration_only=False, processor=None, tokenizer=None, - default_padding_side=None, - default_pad_token=None, device=torch.device("cpu"), ) From d9ef23a51e3ad94967d508eb8e9b661c8387815e Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Thu, 1 Oct 2026 00:59:30 +0000 Subject: [PATCH 2/6] hf_ptq: store the local copy as local_checkpoint_path, the ID as hub_model_id hf_model_path read as a Hub path when it is the local directory, and hf_model_name did not say it is the Hub identifier; name both for what they are. Signed-off-by: Shengliang Xu --- examples/hf_ptq/example_utils.py | 10 +++-- examples/hf_ptq/hf_ptq.py | 43 +++++++++++-------- .../1_FP4-FP8_PTQ_Min-Max_Calibration.ipynb | 6 +-- .../notebooks/2_PTQ_AWQ_Calibration.ipynb | 6 +-- .../notebooks/3_PTQ_AutoQuantization.ipynb | 8 ++-- tests/examples/hf_ptq/test_hf_ptq_args.py | 8 ++-- .../hf_ptq/test_hf_ptq_vision_quantization.py | 4 +- 7 files changed, 48 insertions(+), 37 deletions(-) diff --git a/examples/hf_ptq/example_utils.py b/examples/hf_ptq/example_utils.py index 0f6c7f1ae95..52241079d39 100755 --- a/examples/hf_ptq/example_utils.py +++ b/examples/hf_ptq/example_utils.py @@ -93,7 +93,7 @@ def cleanup_distributed(args): def validate_fsdp2_supported(args, config): """Raise ``NotImplementedError`` for model/CLI combos the FSDP2 path doesn't support yet.""" issues = [] - if "vila" in args.hf_model_path.lower(): + if "vila" in args.local_checkpoint_path.lower(): issues.append("VILA (custom builder + non-standard layer layout)") if is_nemotron_vl(config) or _is_multimodal_config(config): issues.append("multimodal / VL models (decoder layers not auto-detectable)") @@ -103,7 +103,7 @@ def validate_fsdp2_supported(args, config): issues.append("speculative decoding (--specdec_offline_dataset)") if getattr(args, "low_memory_mode", False): issues.append("--low_memory_mode (redundant with FSDP2)") - if not issues and checkpoint_has_mtp(config.model_type, args.hf_model_path): + if not issues and checkpoint_has_mtp(config.model_type, args.local_checkpoint_path): issues.append("Nemotron-H MTP (auxiliary modules are not constructed by the FSDP2 loader)") if issues: @@ -861,7 +861,9 @@ def save_source_config(args, export_path) -> None: config_kwargs = {"trust_remote_code": args.trust_remote_code} if args.attn_implementation is not None: config_kwargs["attn_implementation"] = args.attn_implementation - AutoConfig.from_pretrained(args.hf_model_path, **config_kwargs).save_pretrained(export_path) + AutoConfig.from_pretrained(args.local_checkpoint_path, **config_kwargs).save_pretrained( + export_path + ) def _prepare_quant_cfg( @@ -886,7 +888,7 @@ def _prepare_quant_cfg( if needs_checkpoint_path_update(quant_cfg): # Named after the Hub ID when there is one: a cache path would name it by commit hash. quant_cfg, resolved_dir = resolve_checkpoint_dir( - quant_cfg, args.hf_model_name or args.hf_model_path + quant_cfg, args.hub_model_id or args.local_checkpoint_path ) print(f"Auto-resolved layerwise checkpoint_dir: {resolved_dir}") diff --git a/examples/hf_ptq/hf_ptq.py b/examples/hf_ptq/hf_ptq.py index e7ad8997db7..db29d5fbf61 100755 --- a/examples/hf_ptq/hf_ptq.py +++ b/examples/hf_ptq/hf_ptq.py @@ -308,11 +308,11 @@ def load_model(args: argparse.Namespace): raise NotImplementedError(_FSDP2_KV_AUTOQUANT_ERROR) if args.use_fsdp2: hf_config = AutoConfig.from_pretrained( - args.hf_model_path, trust_remote_code=args.trust_remote_code + args.local_checkpoint_path, trust_remote_code=args.trust_remote_code ) validate_fsdp2_supported(args, hf_config) full_model = parallel_load_and_prepare_fsdp2( - args.hf_model_path, + args.local_checkpoint_path, args.dist_state.device, args.dist_state.rank, args.dist_state.world_size, @@ -323,7 +323,7 @@ def load_model(args: argparse.Namespace): ) elif args.specdec_offline_dataset is not None or not args.low_memory_mode: full_model = get_model( - args.hf_model_path, + args.local_checkpoint_path, args.dist_state.device, gpu_mem_percentage=args.gpu_max_mem_percentage, trust_remote_code=args.trust_remote_code, @@ -349,16 +349,16 @@ def load_model(args: argparse.Namespace): quant_cfg, gpu_mem_percentage=args.gpu_max_mem_percentage, quant_gemm=False ): hf_config = AutoConfig.from_pretrained( - args.hf_model_path, trust_remote_code=args.trust_remote_code + args.local_checkpoint_path, trust_remote_code=args.trust_remote_code ) model_kwargs = {"trust_remote_code": args.trust_remote_code} if args.attn_implementation is not None: model_kwargs["attn_implementation"] = args.attn_implementation with prepare_model_for_loading( - hf_config.model_type, args.hf_model_path, args.trust_remote_code + hf_config.model_type, args.local_checkpoint_path, args.trust_remote_code ): full_model = AutoModelForCausalLM.from_pretrained( - args.hf_model_path, + args.local_checkpoint_path, **model_kwargs, ) calibration_only = True @@ -395,14 +395,14 @@ def load_model(args: argparse.Namespace): if model_type == "whisper": processor = get_processor( - args.hf_model_path, + args.local_checkpoint_path, model_type, trust_remote_code=args.trust_remote_code, ) elif args.calib_with_images: # For VLM image calibration, we need an AutoProcessor to build multimodal inputs. processor = AutoProcessor.from_pretrained( - args.hf_model_path, + args.local_checkpoint_path, trust_remote_code=args.trust_remote_code, padding_side="left", ) @@ -410,12 +410,16 @@ def load_model(args: argparse.Namespace): if hasattr(processor, "tokenizer") and processor.tokenizer is not None: tokenizer = processor.tokenizer else: - tokenizer = get_tokenizer(args.hf_model_path, trust_remote_code=args.trust_remote_code) + tokenizer = get_tokenizer( + args.local_checkpoint_path, trust_remote_code=args.trust_remote_code + ) # Some Nemotron tokenizers may not define pad_token by default; but we use padding=True during calibration. if tokenizer.pad_token is None: tokenizer.pad_token = tokenizer.eos_token - assert tokenizer.pad_token is not None, f"Pad token for {args.hf_model_path} cannot be set!" + assert tokenizer.pad_token is not None, ( + f"Pad token for {args.local_checkpoint_path} cannot be set!" + ) tokenizer.padding_side = "left" @@ -449,7 +453,9 @@ def load_model(args: argparse.Namespace): if extracted_lm is not None: language_model = extracted_lm - tokenizer = get_tokenizer(args.hf_model_path, trust_remote_code=args.trust_remote_code) + tokenizer = get_tokenizer( + args.local_checkpoint_path, trust_remote_code=args.trust_remote_code + ) # Left padding usually provides better calibration result. tokenizer.padding_side = "left" @@ -681,7 +687,7 @@ def export_quantized( # The unified exporters copy the source's non-model files themselves; this deprecated # one does not. if args.dist_state.is_main: - copy_non_model_files(args.hf_model_path, export_path) + copy_non_model_files(args.local_checkpoint_path, export_path) else: # Check arguments for unified_hf export format and set to default if unsupported arguments are provided assert args.sparsity_fmt == "dense", ( @@ -769,7 +775,7 @@ def pre_quantize( full_model, tokenizer, preview_input_ids, - args.hf_model_path, + args.local_checkpoint_path, "before quantization", allow_fallback=False, trust_remote_code=args.trust_remote_code, @@ -845,7 +851,7 @@ def post_quantize( full_model, tokenizer, preview_input_ids, - args.hf_model_path, + args.local_checkpoint_path, "after quantization", allow_fallback=False, trust_remote_code=args.trust_remote_code, @@ -1101,7 +1107,7 @@ def quantize_main( if args.cast_mxfp4_to_nvfp4: # The cast reads the source MXFP4 ``*_scales``/``*_blocks`` tensors from the local # checkpoint directory main() downloaded. - apply_cast_mxfp4_to_nvfp4(language_model, args.hf_model_path) + apply_cast_mxfp4_to_nvfp4(language_model, args.local_checkpoint_path) post_quantize( args, @@ -1474,8 +1480,11 @@ def main(args: argparse.Namespace): # until the NCCL timeout. with mlflow_run(args): # --pyt_ckpt_path stays as given. Every later step reads the local copy of the whole - # checkpoint in hf_model_path; hf_model_name is the Hub ID, or None for a local path. - args.hf_model_name, args.hf_model_path = ensure_local_checkpoint(args.pyt_ckpt_path) + # 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.pyt_ckpt_path + ) # launch a memory monitor to read the currently used GPU memory. launch_memory_monitor() diff --git a/examples/hf_ptq/notebooks/1_FP4-FP8_PTQ_Min-Max_Calibration.ipynb b/examples/hf_ptq/notebooks/1_FP4-FP8_PTQ_Min-Max_Calibration.ipynb index 140653ebfe6..bac613057fd 100644 --- a/examples/hf_ptq/notebooks/1_FP4-FP8_PTQ_Min-Max_Calibration.ipynb +++ b/examples/hf_ptq/notebooks/1_FP4-FP8_PTQ_Min-Max_Calibration.ipynb @@ -114,9 +114,9 @@ "from modelopt.torch.export import ensure_local_checkpoint\n", "\n", "# A local copy of the whole checkpoint, so the export can carry over its tokenizer and other files.\n", - "_, model_path = ensure_local_checkpoint(model_name)\n", - "model = AutoModelForCausalLM.from_pretrained(model_path).cuda()\n", - "tokenizer = AutoTokenizer.from_pretrained(model_path)\n", + "_, local_checkpoint_path = ensure_local_checkpoint(model_name)\n", + "model = AutoModelForCausalLM.from_pretrained(local_checkpoint_path).cuda()\n", + "tokenizer = AutoTokenizer.from_pretrained(local_checkpoint_path)\n", "tokenizer.pad_token = tokenizer.eos_token" ] }, diff --git a/examples/hf_ptq/notebooks/2_PTQ_AWQ_Calibration.ipynb b/examples/hf_ptq/notebooks/2_PTQ_AWQ_Calibration.ipynb index 2fd88a81301..60bdc632e08 100644 --- a/examples/hf_ptq/notebooks/2_PTQ_AWQ_Calibration.ipynb +++ b/examples/hf_ptq/notebooks/2_PTQ_AWQ_Calibration.ipynb @@ -116,9 +116,9 @@ "from modelopt.torch.export import ensure_local_checkpoint\n", "\n", "# A local copy of the whole checkpoint, so the export can carry over its tokenizer and other files.\n", - "_, model_path = ensure_local_checkpoint(model_name)\n", - "model = AutoModelForCausalLM.from_pretrained(model_path).cuda()\n", - "tokenizer = AutoTokenizer.from_pretrained(model_path)\n", + "_, local_checkpoint_path = ensure_local_checkpoint(model_name)\n", + "model = AutoModelForCausalLM.from_pretrained(local_checkpoint_path).cuda()\n", + "tokenizer = AutoTokenizer.from_pretrained(local_checkpoint_path)\n", "tokenizer.pad_token = tokenizer.eos_token" ] }, diff --git a/examples/hf_ptq/notebooks/3_PTQ_AutoQuantization.ipynb b/examples/hf_ptq/notebooks/3_PTQ_AutoQuantization.ipynb index 73efb8b24ca..50fb06deb22 100644 --- a/examples/hf_ptq/notebooks/3_PTQ_AutoQuantization.ipynb +++ b/examples/hf_ptq/notebooks/3_PTQ_AutoQuantization.ipynb @@ -134,9 +134,9 @@ "from modelopt.torch.export import ensure_local_checkpoint\n", "\n", "# A local copy of the whole checkpoint, so the export can carry over its tokenizer and other files.\n", - "_, model_path = ensure_local_checkpoint(MODEL_ID)\n", - "model = AutoModelForCausalLM.from_pretrained(model_path, torch_dtype=DTYPE).to(DEVICE)\n", - "tokenizer = AutoTokenizer.from_pretrained(model_path)\n", + "_, local_checkpoint_path = ensure_local_checkpoint(MODEL_ID)\n", + "model = AutoModelForCausalLM.from_pretrained(local_checkpoint_path, torch_dtype=DTYPE).to(DEVICE)\n", + "tokenizer = AutoTokenizer.from_pretrained(local_checkpoint_path)\n", "tokenizer.padding_side = \"left\"\n", "tokenizer.pad_token = tokenizer.eos_token\n", "model.eval()" @@ -381,7 +381,7 @@ " inference_pipeline_parallel=1,\n", " )\n", " # export_hf_checkpoint copies the source's tokenizer and other non-model files itself.\n", - " copy_non_model_files(model_path, EXPORT_DIR)\n", + " copy_non_model_files(local_checkpoint_path, EXPORT_DIR)\n", "else:\n", " export_hf_checkpoint(model, export_dir=EXPORT_DIR)\n", "\n", diff --git a/tests/examples/hf_ptq/test_hf_ptq_args.py b/tests/examples/hf_ptq/test_hf_ptq_args.py index 5a47727e77f..ea374c5a6f3 100644 --- a/tests/examples/hf_ptq/test_hf_ptq_args.py +++ b/tests/examples/hf_ptq/test_hf_ptq_args.py @@ -216,8 +216,8 @@ def test_hf_ptq_runs_fixed_ptq_before_kv_autoquantize(monkeypatch): batch_size=1, auto_quantize_checkpoint=None, kv_auto_quantize_checkpoint=None, - hf_model_path="dummy", - hf_model_name=None, + local_checkpoint_path="dummy", + hub_model_id=None, cast_mxfp4_to_nvfp4=False, layerwise_export=False, specdec_offline_dataset=None, @@ -263,8 +263,8 @@ def test_fixed_ptq_then_kv_rejects_explicit_kv_before_calibration(monkeypatch): args = SimpleNamespace( auto_quantize_checkpoint=None, kv_auto_quantize_checkpoint=None, - hf_model_path="dummy", - hf_model_name=None, + local_checkpoint_path="dummy", + hub_model_id=None, cast_mxfp4_to_nvfp4=False, layerwise_export=False, ) diff --git a/tests/examples/hf_ptq/test_hf_ptq_vision_quantization.py b/tests/examples/hf_ptq/test_hf_ptq_vision_quantization.py index 9a885687d23..d37d3aba3b1 100644 --- a/tests/examples/hf_ptq/test_hf_ptq_vision_quantization.py +++ b/tests/examples/hf_ptq/test_hf_ptq_vision_quantization.py @@ -49,8 +49,8 @@ def test_image_calibration_model_target_follows_recipe( use_fsdp2=False, specdec_offline_dataset=None, low_memory_mode=False, - hf_model_path="dummy", - hf_model_name=None, + local_checkpoint_path="dummy", + hub_model_id=None, dist_state=SimpleNamespace(device=torch.device("cpu")), gpu_max_mem_percentage=0.8, trust_remote_code=False, From a0245d56ce97f547258a203d4085a1863894f752 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 2 Oct 2026 01:38:50 +0000 Subject: [PATCH 3/6] hf_ptq: drop the removed padding arguments from the post_quantize test call test_hf_ptq_model_type, new on main, still passed default_padding_side and default_pad_token, which this branch removes from post_quantize. Signed-off-by: Shengliang Xu --- tests/examples/hf_ptq/test_hf_ptq_model_type.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/tests/examples/hf_ptq/test_hf_ptq_model_type.py b/tests/examples/hf_ptq/test_hf_ptq_model_type.py index a5d6fa04165..0ab585b19d0 100644 --- a/tests/examples/hf_ptq/test_hf_ptq_model_type.py +++ b/tests/examples/hf_ptq/test_hf_ptq_model_type.py @@ -129,8 +129,6 @@ def test_post_quantize_preview_keeps_whole_answer( False, None, None, - None, - None, ) assert "example outputs after ptq: ['7 8']" in capsys.readouterr().out From 29ad72d1651422d448cb9d7169fffa509c9f5d8e Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Tue, 6 Oct 2026 21:00:31 +0000 Subject: [PATCH 4/6] Drop the hf_ptq pad_token changelog entry that #2597 already covers Signed-off-by: Shengliang Xu --- CHANGELOG.rst | 1 - 1 file changed, 1 deletion(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index fb48228e263..2efd87bdda4 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -129,7 +129,6 @@ Changelog - Fix offline DFlash training failing to reconstruct the target logits when the captured hidden states are stored in a different dtype than the target's weights. - Fix ``hf_ptq`` overwriting a model's existing ``pad_token`` with ``eos_token`` when the model already has a valid padding token. The exported tokenizer now preserves the source model's padding configuration. - Fix Hugging Face exports dropping off-index safetensors such as GLM-4.7's ``mtp.safetensors``. -- Fix ``examples/hf_ptq`` exports setting ``pad_token`` to the EOS token in the tokenizer files; the source checkpoint's tokenizer files are now exported unchanged. - Fix Megatron-Core checkpoint saving for quantized grouped MoE experts when tensor and expert parallelism are both enabled. - Fix unified HuggingFace export of RADIO-based VLMs retaining post-conversion vision and projector names instead of restoring the hub layout; deployment loaders could skip those weights. From c61b877193597ec6e6a156856226e4decd083559 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Wed, 7 Oct 2026 06:03:14 +0000 Subject: [PATCH 5/6] hf_ptq: give the FSDP2 MTP test the local checkpoint path main() sets Signed-off-by: Shengliang Xu --- tests/examples/hf_ptq/test_hf_ptq_args.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/examples/hf_ptq/test_hf_ptq_args.py b/tests/examples/hf_ptq/test_hf_ptq_args.py index ea374c5a6f3..32c088c90c4 100644 --- a/tests/examples/hf_ptq/test_hf_ptq_args.py +++ b/tests/examples/hf_ptq/test_hf_ptq_args.py @@ -377,6 +377,7 @@ def test_fsdp2_rejects_mtp_before_distributed_loading(monkeypatch, tmp_path, has monkeypatch.setenv("RANK", "0") hf_ptq, args = _parse_hf_ptq_args(monkeypatch, "--pyt_ckpt_path", str(tmp_path), "--use_fsdp2") args.dist_state = SimpleNamespace(device="cpu", rank=0, world_size=1) + args.hub_model_id, args.local_checkpoint_path = None, str(tmp_path) # as main() sets them keys = ["mtp.layers.0.eh_proj.weight", "mtp.layers.1.final_layernorm.weight"] if has_mtp else [] (tmp_path / "model.safetensors.index.json").write_text( json.dumps({"weight_map": dict.fromkeys(keys, "model-00001.safetensors")}) From 265631c1b4760e6b8392a056489806ff9cb16a9a Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Wed, 7 Oct 2026 08:12:22 +0000 Subject: [PATCH 6/6] hf_ptq: test that main() loads from the local copy and the TensorRT-LLM export copies from it Signed-off-by: Shengliang Xu --- tests/examples/hf_ptq/test_hf_ptq_args.py | 61 +++++++++++++++++++++++ 1 file changed, 61 insertions(+) diff --git a/tests/examples/hf_ptq/test_hf_ptq_args.py b/tests/examples/hf_ptq/test_hf_ptq_args.py index 32c088c90c4..0854725ff90 100644 --- a/tests/examples/hf_ptq/test_hf_ptq_args.py +++ b/tests/examples/hf_ptq/test_hf_ptq_args.py @@ -885,3 +885,64 @@ def test_post_quantize_export_survives_a_failed_sanity_generate(monkeypatch): ) assert len(export_calls) == 1 + + +@pytest.mark.usefixtures("clean_env") +def test_main_makes_a_hub_checkpoint_local_before_loading(monkeypatch, tmp_path): + """--pyt_ckpt_path keeps the Hub ID; the model loads from the local copy resolved first.""" + hf_ptq, args = _parse_hf_ptq_args(monkeypatch, "--pyt_ckpt_path", "org/model") + calls = [] + + def ensure_local_checkpoint(model_name_or_path): + calls.append(("ensure_local_checkpoint", model_name_or_path)) + return "org/model", str(tmp_path) + + class _LoadReachedError(Exception): + pass + + def load_model(args): + calls.append( + ("load_model", args.pyt_ckpt_path, args.hub_model_id, args.local_checkpoint_path) + ) + raise _LoadReachedError + + monkeypatch.setattr(torch.cuda, "is_available", lambda: True) + monkeypatch.setattr(torch.compiler, "set_stance", lambda *args, **kwargs: None) + monkeypatch.setattr(hf_ptq, "launch_memory_monitor", lambda: None) + monkeypatch.setattr(hf_ptq, "ensure_local_checkpoint", ensure_local_checkpoint) + monkeypatch.setattr(hf_ptq, "load_model", load_model) + + with pytest.raises(_LoadReachedError): + hf_ptq.main(args) + + assert calls == [ + ("ensure_local_checkpoint", "org/model"), + ("load_model", "org/model", "org/model", str(tmp_path)), + ] + + +def test_tensorrt_llm_export_copies_non_model_files_from_the_local_copy(monkeypatch, tmp_path): + """The deprecated TensorRT-LLM exporter does not copy them itself; hf_ptq does, from disk.""" + hf_ptq, args = _parse_hf_ptq_args( + monkeypatch, + "--pyt_ckpt_path", + "org/model", + "--export_path", + str(tmp_path / "export"), + "--sparsity_fmt", + "sparsegpt", + ) + args.hub_model_id, args.local_checkpoint_path = "org/model", str(tmp_path / "local") + args.dist_state = SimpleNamespace(is_main=True) + copies = [] + monkeypatch.setattr(hf_ptq, "has_spec_opt", lambda model: False) + monkeypatch.setattr(hf_ptq, "is_multimodal_model", lambda model: False) + monkeypatch.setattr(hf_ptq, "is_trtllm_enc_dec_export", lambda model_type: False) + monkeypatch.setattr(hf_ptq, "remove_hook_from_module", lambda *args, **kwargs: None) + monkeypatch.setattr(hf_ptq, "export_tensorrt_llm_checkpoint", lambda *args, **kwargs: None) + monkeypatch.setattr(hf_ptq, "copy_non_model_files", lambda *args: copies.append(args)) + + model = torch.nn.Linear(1, 1) + hf_ptq.export_quantized(args, model, model, "llama") + + assert copies == [(str(tmp_path / "local"), str(tmp_path / "export"))]