# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. import argparse import contextlib import os import random import time import warnings from datetime import timedelta from pathlib import Path from typing import Any import numpy as np import torch from accelerate.hooks import remove_hook_from_module from autoquant_utils import ( _get_autoquant_memory_probe, _quantize_config_explicitly_enables_kv, _recipe_is_auto_quantize, auto_quantize, ) from cast_mxfp4_to_nvfp4 import apply_to_model as apply_cast_mxfp4_to_nvfp4 from example_utils import ( HF_PTQ, _prepare_quant_cfg, build_quant_cfg, create_vlm_calibration_forward, create_vlm_calibration_loop, exports_tensorrt_llm_checkpoint, generate_excludes_prompt, get_model, get_tokenizer, keep_the_experts_implementation, layerwise_export_requested, mlflow_run, preview_generate_kwargs, recipe_layerwise_blocks, resolve_distribution, run_vlm_preview, save_source_config, ) from torch.utils.data import DataLoader from transformers import ( AutoModelForCausalLM, AutoProcessor, PreTrainedTokenizer, PreTrainedTokenizerBase, PreTrainedTokenizerFast, ProcessorMixin, ) import modelopt.torch.opt as mto import modelopt.torch.quantization as mtq import modelopt.torch.sparsity as mts from modelopt.recipe import ModelOptAutoQuantizeRecipe, ModelOptPTQRecipe, load_recipe from modelopt.recipe.presets import ( KV_CACHE_NONE, KV_QUANT_CFG_CHOICES, QUANT_CFG_CHOICES, RecipeSupersededAction, ) from modelopt.torch.export import ( export_hf_checkpoint, export_hf_vllm_fq_checkpoint, export_speculative_decoding, has_spec_opt, save_expert_token_count_table, ) from modelopt.torch.export.layerwise_export import LayerwiseExporter from modelopt.torch.export.model_utils import get_language_model_from_vl, is_multimodal_model from modelopt.torch.export.trtllm import export_tensorrt_llm_checkpoint from modelopt.torch.models import get_ptq_spec, hf_model_type from modelopt.torch.models.hf import ( calibration_sample_input, load_calibration_processor, make_calibration_dataloader, prepare_model_for_calibration, ) from modelopt.torch.quantization.config import need_calibration from modelopt.torch.quantization.plugins.accelerate import init_quantized_weights from modelopt.torch.quantization.utils import is_quantized from modelopt.torch.quantization.utils.layerwise_calib import get_max_layerwise_batch_size from modelopt.torch.speculative.eagle.utils import ( EagleOfflineDataCollator, OfflineSupervisedDataset, ) from modelopt.torch.utils import print_rank_0 from modelopt.torch.utils.dataset_utils import ( create_forward_loop, get_dataset_dataloader, get_max_batch_size, get_supported_datasets, ) from modelopt.torch.utils.distributed import ( distributed_forward_loop, distributed_run, distribution_plan, local_rank, reinstall_a2a_forwards, shard_dataloader, ) from modelopt.torch.utils.fsdp2_utils import fsdp2_root_unsharded 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.moe import expert_parallel_mesh 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 load_pretrained from modelopt.torch.utils.vlm_dataset_utils import get_vlm_dataset_dataloader RAND_SEED = 1234 # Most tokens per calibration batch when layerwise calibration picks the batch size (64 x 512). LAYERWISE_CALIB_TOKENS = 32768 _PREVIEW_STOPS_THE_RUN = ( "The preview after quantization failed, which stops the run before the export, so no " "checkpoint is written; --skip_generate exports without the preview." ) def _layerwise_calib_batch_size(calib_seq: int, num_samples: int) -> int: """Largest default layerwise batch: LAYERWISE_CALIB_TOKENS tokens, but 1 to all of the samples.""" return max(1, min(LAYERWISE_CALIB_TOKENS // calib_seq, num_samples)) mto.enable_huggingface_checkpointing() def extract_and_prepare_language_model_from_vl( full_model: torch.nn.Module, ) -> torch.nn.Module | None: """Extract language model from VL model and disable quantization for non-language components. Args: full_model: The full VLM model Returns: The language model, or None if not a VLM """ language_model_lineage = get_language_model_from_vl(full_model, strict=True) if language_model_lineage is not None: language_model = language_model_lineage.pop(-1) ancestors = language_model_lineage # Apply disabled quant to all modules that are not part of language_model # This excludes them during HF export disabled_quant_cfg = { "quant_cfg": [{"quantizer_name": "*", "enable": False}], "algorithm": "max", } memo = set(ancestors) | {language_model} for ancestor in ancestors: for _, module in ancestor.named_children(): if module not in memo: mtq.quantize(module, disabled_quant_cfg, forward_loop=None) memo.add(module) return language_model return None def fsdp2_call_context(full_model: torch.nn.Module, model: torch.nn.Module): """Context for calling ``model``, the quantization target, of a possibly FSDP2 ``full_model``. Plain PTQ calibrates the extracted language model directly, which skips the FSDP2 root's pre-forward; the root's own parameters (e.g. the embeddings) are gathered around the call. """ return contextlib.nullcontext() if model is full_model else fsdp2_root_unsharded(full_model) class _DeviceDataLoader: """Wrapper around a DataLoader that moves each batch to a target device.""" def __init__(self, dataloader: DataLoader, device: torch.device): self.dataloader = dataloader self.device = device def __iter__(self): for batch in self.dataloader: yield _move_batch_to_device(batch, self.device) def __len__(self): return len(self.dataloader) def _move_batch_to_device(batch: dict, device: torch.device) -> dict: """Recursively move all tensors in a batch dict to the given device.""" def _to_device(value): if isinstance(value, torch.Tensor): return value.to(device) if isinstance(value, dict): return {k: _to_device(v) for k, v in value.items()} return value return {k: _to_device(v) for k, v in batch.items()} def _shard_calib_dataloader(calib_dataloader, rank: int, world_size: int): """This rank's share of the calibration data, in the loader type it came in.""" if isinstance(calib_dataloader, _DeviceDataLoader): return _DeviceDataLoader( shard_dataloader(calib_dataloader.dataloader, rank, world_size), calib_dataloader.device ) if isinstance(calib_dataloader, DataLoader): return shard_dataloader(calib_dataloader, rank, world_size) raise NotImplementedError( f"A {type(calib_dataloader).__name__} calibration loader cannot be split across this " f"run's {world_size} ranks." ) def make_calib_dataloader( args: argparse.Namespace, language_model: torch.nn.Module, processor: ProcessorMixin | None, tokenizer: PreTrainedTokenizerBase | None, device: torch.device, model_type: str | None, autoquant_gradient_recipe: bool = False, ) -> tuple[DataLoader | _DeviceDataLoader, str | None]: # Every distributed plan is data-parallel over the world: each rank calibrates on its own shard, # and the quantizers' amax is synced across the ranks. world_size = distribution_plan(language_model).world_size calib_dataloader = None preview_input_text = None if args.specdec_offline_dataset is not None: offline_data_path = Path(args.specdec_offline_dataset) dumped_files = sorted(str(p) for p in offline_data_path.glob("*.pt")) if not dumped_files: raise ValueError(f"No .pt files found in {args.specdec_offline_dataset}") if args.calib_size[0] > 0: dumped_files = dumped_files[: args.calib_size[0]] dataset = OfflineSupervisedDataset(dumped_files) collator = EagleOfflineDataCollator(train_len=args.calib_seq) raw_loader = DataLoader( dataset, batch_size=args.batch_size, shuffle=False, collate_fn=collator, ) # Wrap to move batches to the target device; device-transfer logic is kept # out of the data collator to avoid interference with dataloader prefetching. calib_dataloader = _DeviceDataLoader(raw_loader, device) elif args.calib_with_images: # VLM image-text calibration path: assume Nemotron VLM dataset by default. assert processor is not None, ( "Please provide a processor (e.g., AutoProcessor) for image calibration." ) assert len(args.calib_size) == 1, ( "Image calibration currently supports a single dataset. " "Please pass --calib_size with one value (e.g., --calib_size 256)." ) calib_dataloader = get_vlm_dataset_dataloader( dataset_name="nemotron_vlm_dataset_v2", processor=processor, batch_size=args.batch_size, num_samples=args.calib_size[0], device=device, require_image=True, subsets=["sparsetables", "plotqa_cot", "wiki_en"], shuffle_buffer_size=10_000, seed=42, use_media_shards=True, max_shards=1, # Sharded here, so each rank drops its own undecodable images and pads to the others' # batch count: an empty batch on one rank would leave the others in a collective. dp_size=world_size, dp_rank=args.dist_state.rank, ) elif ( provided := make_calibration_dataloader( model_type, language_model, processor, args.dataset, args.calib_size, args.batch_size, device, ) ) is not None: # The model type brings its own calibration data (speech features, not token ids). calib_dataloader, preview_input_text = provided else: assert tokenizer is not None and isinstance( tokenizer, PreTrainedTokenizer | PreTrainedTokenizerFast ), "The PreTrainedTokenizer must be set" # Labels are only needed for gradient-based auto_quantize include_labels = autoquant_gradient_recipe calib_dataloader = get_dataset_dataloader( dataset_name=args.dataset, tokenizer=tokenizer, batch_size=args.batch_size, num_samples=args.calib_size, max_sample_length=args.calib_seq, device=device, include_labels=include_labels, ) if world_size > 1 and calib_dataloader is not None and not args.calib_with_images: calib_dataloader = _shard_calib_dataloader( calib_dataloader, args.dist_state.rank, world_size ) return calib_dataloader, preview_input_text def _resolve_kv_auto_quantize_checkpoint(args: argparse.Namespace) -> str | None: """Resolve a KV-primary checkpoint with a one-release legacy fallback.""" if args.kv_auto_quantize_checkpoint is not None: return args.kv_auto_quantize_checkpoint if args.auto_quantize_checkpoint is not None: warnings.warn( "Using --auto_quantize_checkpoint for a KV-cache search is deprecated; use " "--kv_auto_quantize_checkpoint instead.", FutureWarning, ) return args.auto_quantize_checkpoint return None def _validate_recipe_calibration(args: argparse.Namespace, recipe) -> None: """Require image-text calibration when a PTQ recipe enables visual input quantizers.""" if not isinstance(recipe, ModelOptPTQRecipe) or args.calib_with_images: return visual_input_names = { "*visual.*input_quantizer", "*visual.*merger*.input_quantizer", } visual_input_entry = next( ( entry for entry in reversed(recipe.quantize.quant_cfg) if entry.quantizer_name in visual_input_names ), None, ) if visual_input_entry is not None and visual_input_entry.enable: raise ValueError( "Vision Encoder quantization recipes require --calib_with_images so visual input " "quantizers receive activation calibration data." ) def load_model(args: argparse.Namespace): # If low memory mode is enabled, we compress the model while loading the HF checkpoint. calibration_only = False plan, hf_config = resolve_distribution(args) if plan.is_expert_parallel: full_model = load_pretrained( args.local_checkpoint_path, distribution=plan, hf_config=hf_config, trust_remote_code=args.trust_remote_code, attn_implementation=args.attn_implementation, ) elif plan.is_fsdp_dense: if args.experts_implementation is not None: # Set on the config, as build_meta_causal_lm sets a pre-fetched config's # attn_implementation: the meta model is built from it, with no kwarg of its own. The # setter carries it to the sub-configs, such as the text config a VLM's experts read. hf_config._experts_implementation = args.experts_implementation full_model = load_pretrained( args.local_checkpoint_path, distribution=plan, hf_config=hf_config, trust_remote_code=args.trust_remote_code, cpu_offload=args.cpu_offload, attn_implementation=args.attn_implementation, ) elif args.specdec_offline_dataset is not None or not args.low_memory_mode: full_model = get_model( args.local_checkpoint_path, args.dist_state.device, distribution=plan, gpu_mem_percentage=args.gpu_max_mem_percentage, trust_remote_code=args.trust_remote_code, use_seq_device_map=args.use_seq_device_map, attn_implementation=args.attn_implementation, experts_implementation=args.experts_implementation, offload_folder=args.offload_folder, max_cpu_memory_gb=args.max_cpu_memory_gb, max_gpu_memory_gb=args.max_gpu_memory_gb, ) else: assert args.qformat in QUANT_CFG_CHOICES, ( f"Quantization format is not supported for low memory mode. Supported formats: {list(QUANT_CFG_CHOICES)}" ) quant_cfg = QUANT_CFG_CHOICES[args.qformat] if args.kv_cache_qformat != KV_CACHE_NONE: quant_cfg = mtq.utils.update_quant_cfg_with_kv_cache_quant( quant_cfg, KV_QUANT_CFG_CHOICES[args.kv_cache_qformat]["quant_cfg"], ) # Do not use real quant GEMM so the calibration can be more accurate. with init_quantized_weights( quant_cfg, gpu_mem_percentage=args.gpu_max_mem_percentage, quant_gemm=False ): model_kwargs = {"trust_remote_code": args.trust_remote_code} if args.attn_implementation is not None: model_kwargs["attn_implementation"] = args.attn_implementation if args.experts_implementation is not None: model_kwargs["experts_implementation"] = args.experts_implementation full_model = load_pretrained( args.local_checkpoint_path, AutoModelForCausalLM, model_type=hf_config.model_type, distribution=plan, **model_kwargs, ) calibration_only = True # Include auxiliary modules whose activation quantizers require calibration data. prepare_model_for_calibration(full_model) # The checkpoint's Hugging Face model_type. Always taken from the root model: a VLM's language # model reports its own sub-config's type, which varies across transformers versions. model_type = hf_model_type(full_model) if plan.is_fsdp_dense: device = args.dist_state.device else: device = full_model.device if hasattr(full_model, "model"): device = full_model.model.device processor = None tokenizer = None language_model = full_model vlm_wrapper = get_ptq_spec(model_type).vlm_wrapper # Default to image-text calibration for VLM wrappers. Skip for the AutoQuantize recipe path, whose # text-only path does not support image-text calibration yet (auto_quantize() would raise); # auto-enabling it here would make a VLM wrapper's AutoQuantize fail unconditionally. Nor for # FSDP2, where image-text calibration is opt-in via --calib_with_images. if ( vlm_wrapper and not args.calib_with_images and not args.use_fsdp2 and not _recipe_is_auto_quantize(args.recipe) ): print("VLM wrapper detected. Enabling image-text calibration by default.") args.calib_with_images = True # A model type that brings its own calibration data needs only its processor. processor = load_calibration_processor( model_type, args.local_checkpoint_path, args.trust_remote_code ) if processor is None and args.calib_with_images: # For VLM image calibration, we need an AutoProcessor to build multimodal inputs. processor = AutoProcessor.from_pretrained( args.local_checkpoint_path, trust_remote_code=args.trust_remote_code, padding_side="left", ) if hasattr(processor, "tokenizer") and processor.tokenizer is not None: tokenizer = processor.tokenizer else: tokenizer = get_tokenizer( args.local_checkpoint_path, trust_remote_code=args.trust_remote_code ) # Some 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.local_checkpoint_path} cannot be set!" ) tokenizer.padding_side = "left" # Plain PTQ quantizes only the language model. Recipes keep the complete VLM so their # quantizer rules can target vision and language components in one state. if args.recipe is None: extracted_lm = extract_and_prepare_language_model_from_vl(full_model) if extracted_lm is not None: language_model = extracted_lm elif processor is None: if args.specdec_offline_dataset is not None: language_model = full_model else: if args.dataset is None: args.dataset = ["cnn_nemotron_v2_mix"] warnings.warn( "No dataset specified. Defaulting to the 'cnn_nemotron_v2_mix' combo " "(cnn_dailymail + nemotron-post-training-dataset-v2)." ) # Adjust calib_size to match dataset length by extending or truncating as needed args.calib_size = (args.calib_size + [args.calib_size[-1]] * len(args.dataset))[ : len(args.dataset) ] # Plain PTQ quantizes only the extracted language model. The recipe path keeps the outer # CausalLM so recipes / search can see the Qwen3.5/3.6-MoE VLM lm_head; extracting here # would leave modelopt state on the ancestors and make auto_quantize() fail with # "multiple modelopt states". if args.recipe is None: extracted_lm = extract_and_prepare_language_model_from_vl(full_model) if extracted_lm is not None: language_model = extracted_lm 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" return ( full_model, language_model, model_type, calibration_only, processor, tokenizer, device, ) def sparsity_main( args: argparse.Namespace, full_model: torch.nn.Module, tokenizer: PreTrainedTokenizerBase | None, device: torch.device, ): if args.batch_size == 0: # Sparse algorithm takes more GPU memory so we reduce the batch_size by 4. args.batch_size = max(get_max_batch_size(full_model) // 4, 1) args.batch_size = min(args.batch_size, sum(args.calib_size)) print(f"Use calib batch_size {args.batch_size}") # Different calibration datasets are also available, e.g., "pile" and "wikipedia" # Please also check the docstring for the datasets available assert tokenizer is not None and isinstance( tokenizer, PreTrainedTokenizer | PreTrainedTokenizerFast ), "The PreTrainedTokenizer must be set" calib_dataloader = get_dataset_dataloader( dataset_name=args.dataset, tokenizer=tokenizer, batch_size=args.batch_size, num_samples=args.calib_size, max_sample_length=args.calib_seq, device=device, ) full_model = mts.sparsify( full_model, args.sparsity_fmt, config={"data_loader": calib_dataloader, "collect_func": lambda x: x}, ) mts.export(full_model) def mono_quantize( args: argparse.Namespace, quant_cfg: dict[str, Any], full_model: torch.nn.Module, language_model: torch.nn.Module, model_type: str | None, calibration_only: bool, calib_dataloader: DataLoader, vlm_wrapper: bool, ): """Plain quantization of the selected model target to one quantization configuration.""" model_is_already_quantized = is_quantized(language_model) if "awq" in args.qformat: print( "\n####\nAWQ calibration could take longer than other calibration methods. " "Consider reducing calib_size to reduce calibration time.\n####\n" ) if not model_is_already_quantized or calibration_only: # quantize the model use_calibration = need_calibration(quant_cfg) if not use_calibration: warnings.warn("Dynamic quantization. Calibration skipped.") calibrate_loop = None if use_calibration: # Image calibration batches contain multimodal kwargs (for example pixel_values). # They must be consumed by the complete VLM even when only a nested component is the # quantization target; the full forward still exercises that component's quantizers. if distribution_plan(full_model).is_expert_parallel: # The distributed loop first re-plants the torch_a2a forwards that quantizer # insertion clobbers on transformers >= 5.16. It is handed the full model, whose # expert-parallel plan names the experts; text batches still run through the # language model. calibrate_loop = distributed_forward_loop( full_model, calib_dataloader, device=args.dist_state.device, ep_mesh=expert_parallel_mesh(full_model), forward=create_vlm_calibration_forward(full_model) if args.calib_with_images else lambda batch: language_model(**batch), ) elif args.calib_with_images: calibrate_loop = create_vlm_calibration_loop(full_model, calib_dataloader) else: calibrate_loop = create_forward_loop( dataloader=calib_dataloader, allowed_non_tensor_keys={"base_model_outputs"} if args.specdec_offline_dataset is not None else None, ) if args.use_fsdp2: text_loop = calibrate_loop def calibrate_loop(model): with fsdp2_call_context(full_model, model): text_loop(model) if args.layerwise_export: LayerwiseExporter(full_model, args.export_path) if calibration_only: language_model = mtq.calibrate( language_model, quant_cfg["algorithm"], forward_loop=calibrate_loop ) else: language_model = mtq.quantize(language_model, quant_cfg, forward_loop=calibrate_loop) if distribution_plan(full_model).is_expert_parallel: # A config that runs no calibration loop never re-planted the a2a forwards, which the # preview and the export need. reinstall_a2a_forwards(full_model, expert_parallel_mesh(full_model)) # For VL models, update full_model to use the quantized language model if vlm_wrapper and language_model is not full_model: language_model_lineage = get_language_model_from_vl(full_model) if language_model_lineage is not None: print("Updating full_model with quantized language_model...") language_model_lineage[-2].language_model = language_model else: warnings.warn("Skipping quantization: model is already quantized.") def _run_auto_quantize_recipe( args: argparse.Namespace, recipe: ModelOptAutoQuantizeRecipe, full_model: torch.nn.Module, language_model: torch.nn.Module, model_type: str | None, calibration_only: bool, calib_dataloader: DataLoader, vlm_wrapper: bool, ) -> None: """Run the recipe's fixed PTQ, weight search, and KV search in order.""" primary = recipe.auto_quantize followup_kv = recipe.kv_auto_quantize primary_is_kv = primary.constraints.cost_model == "kv_cache" fixed_quantize_config = recipe.quantize if fixed_quantize_config is not None and (primary_is_kv or followup_kv is not None): if _quantize_config_explicitly_enables_kv(fixed_quantize_config.model_dump()): raise ValueError( "The fixed quantize stage explicitly enables K/V quantizers before KV-cache " "AutoQuantize. Disable them in the fixed stage." ) if primary_is_kv and fixed_quantize_config is not None: quant_cfg = _prepare_quant_cfg(args, fixed_quantize_config.model_dump(), full_model) mono_quantize( args, quant_cfg, full_model, language_model, model_type, calibration_only, calib_dataloader, vlm_wrapper, ) fixed_quantize_config = None auto_quantize( args, full_model, calib_dataloader, aq_config=primary, full_model=full_model, fixed_quantize_config=fixed_quantize_config, allow_uniform_kv=followup_kv is None, checkpoint=( _resolve_kv_auto_quantize_checkpoint(args) if primary_is_kv else args.auto_quantize_checkpoint ), ) if followup_kv is not None: auto_quantize( args, full_model, calib_dataloader, aq_config=followup_kv, full_model=full_model, allow_uniform_kv=False, # The weight search owns --auto_quantize_checkpoint, so a follow-up KV search must # never use the KV-primary legacy fallback and collide with the weight state. checkpoint=args.kv_auto_quantize_checkpoint, ) def export_quantized( args: argparse.Namespace, full_model: torch.nn.Module, language_model: torch.nn.Module, model_type: str | None, ): plan = distribution_plan(full_model) if not (args.dist_state.is_main or plan.is_fsdp_dense or plan.is_expert_parallel): # Replicated over the ranks, the model is whole on each: rank 0 writes it, and the rest # would race it into the same directory. A sharded model needs every rank for its part. return # Not inference_mode: the FSDP2 path gathers full params in this context and # inference tensors break the subsequent state_dict() -> param.detach(). with torch.no_grad(): export_path = args.export_path # Early exit for speculative decoding checkpoints # No tokenizer saving needed for spec ckpts if has_spec_opt(full_model): if plan.is_fsdp_dense or plan.is_expert_parallel: # export_speculative_decoding writes the whole checkpoint from one process and # knows no plan: each rank of a sharded model would write it over the others. raise NotImplementedError( f"Exporting a speculative-decoding model sharded as {plan!r} is not " "supported; quantize it in a single process." ) export_speculative_decoding(full_model, export_dir=export_path) args.checkpoint_exported = True print(f"Quantized speculative decoding checkpoint exported to: {export_path}") return if is_multimodal_model(full_model): # Per-layer export writes its own config.json with quantization_config, which # the source config would replace; it never writes a processor config. Rank 0's # export writes config.json, so no other rank may write the source config after it. if not args.layerwise_export and args.dist_state.is_main: save_source_config(args, export_path) start_time = time.time() if exports_tensorrt_llm_checkpoint(args, model_type): if ( args.inference_tensor_parallel != 1 or args.inference_pipeline_parallel != 1 ) and args.qformat == "nvfp4_svdquant": raise NotImplementedError("Svdquant does not support multiple GPUs yet.") warnings.warn( "Still exporting TensorRT-LLM checkpoints for models not supported by the TensorRT-LLM torch runtime." ) # Move meta tensor back to device before exporting. remove_hook_from_module(language_model, recurse=True) # The deprecated exporter detects its TensorRT-LLM decoder_type from the model class. export_tensorrt_llm_checkpoint( language_model, export_dir=export_path, inference_tensor_parallel=args.inference_tensor_parallel, inference_pipeline_parallel=args.inference_pipeline_parallel, ) # 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.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", ( f"Sparsity format {args.sparsity_fmt} not supported by unified export api." ) if args.inference_tensor_parallel != 1 or args.inference_pipeline_parallel != 1: warnings.warn( "Unified HF export format does not specify inference tensor parallel or pipeline parallel. " "They will be set at deployment time." ) # Load any missing weights from non-standard safetensors (handled in get_model for non-low-memory mode) # Store the MTP layer prefixes on the model for later exclusion from quantization if args.vllm_fakequant_export: # save_pretrained inside the exporter writes model-backed state only; weights # the loader could not place (an MTP head, an auxiliary tower) are carried over # as an extra shard afterward -- see _carry_over_unplaced_weights. export_hf_vllm_fq_checkpoint( full_model, export_dir=export_path, inplace_mem_efficient=True ) else: # Weights the loader could not place (an MTP head, an auxiliary tower) are # carried over by the exporter from the keys recorded at load time; nothing # architecture-specific is needed here. export_hf_checkpoint( full_model, export_dir=export_path, ) if args.qformat == "w4a16_nvfp4": warnings.warn( "TensorRT-LLM and SGLang do not support this format. " "vLLM deployment support is in progress." ) args.checkpoint_exported = True end_time = time.time() print_rank_0( f"Quantized model exported to: {export_path}. Total time used {end_time - start_time}s" ) def pre_quantize( args: argparse.Namespace, full_model: torch.nn.Module, model_type: str | None, tokenizer: PreTrainedTokenizerBase | None, calib_dataloader: DataLoader | None, vlm_wrapper: bool, ): """ Processing before the quantization. Currently we run one round of generation for a sample prompt, to be compared with post-quantize generation. """ # Offline specdec models skip pre-quantize preview (no tokenizer or standard dataloader) if args.specdec_offline_dataset is not None: return None, None, None # Only run single sample for preview assert calib_dataloader is not None, "calib_dataloader is required for pre-quantize preview" batch = next(iter(calib_dataloader)) # A speech batch carries ``input_features`` and no ``input_ids``. input_key = "input_ids" if "input_ids" in batch else "input_features" preview_input_ids = batch[input_key][0:1] # Pass attention_mask to generate(): HF cannot infer it when pad_token == eos_token. preview_attention_mask = batch["attention_mask"][0:1] if "attention_mask" in batch else None if input_key == "input_ids" and preview_attention_mask is not None: preview_input_ids = preview_input_ids[:, preview_attention_mask[0].bool()] preview_attention_mask = torch.ones_like(preview_input_ids) # Generate preview before quantization if args.skip_generate: generated_ids_before_ptq = None elif (reason := get_ptq_spec(model_type).preview_unsupported) is not None: print(f"Skipping the generation preview: {reason}") generated_ids_before_ptq = None elif vlm_wrapper and tokenizer is not None: generated_ids_before_ptq = run_vlm_preview( full_model, tokenizer, preview_input_ids, args.local_checkpoint_path, "before quantization", allow_fallback=False, trust_remote_code=args.trust_remote_code, ) else: with keep_the_experts_implementation(full_model): generated_ids_before_ptq = full_model.generate( preview_input_ids, attention_mask=preview_attention_mask, max_new_tokens=100, **preview_generate_kwargs(full_model), ) return preview_input_ids, preview_attention_mask, generated_ids_before_ptq def post_quantize( args: argparse.Namespace, full_model: torch.nn.Module, language_model: torch.nn.Module, model_type: str | None, tokenizer: PreTrainedTokenizerBase | None, processor: ProcessorMixin | None, preview_input_ids: torch.Tensor | None, preview_attention_mask: torch.Tensor | None, generated_ids_before_ptq: Any, vlm_wrapper: bool, preview_input_text: str | None, calib_dataloader: DataLoader, ): """ Processing after the quantization, then export. For offline speculative decoding models, skip generation comparison and proceed directly to export. For standard models, run one round of generation using the quantized model for a sample prompt and compare it with pre-quantize generation. """ # 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) return assert preview_input_ids is not None, "pre_quantize returns preview inputs for this path" if args.verbose and args.dist_state.is_main: try: mtq.print_quant_summary(full_model, args.export_path) save_expert_token_count_table(full_model, args.export_path) except Exception as e: print(f"Error saving quant summary: {e}") print("Continuing with generation...") # Run some samples torch.cuda.empty_cache() generated_ids_after_ptq = None post_ptq_unsupported = get_ptq_spec(model_type).post_ptq_generation_unsupported if generated_ids_before_ptq is None: pass elif post_ptq_unsupported is None and not vlm_wrapper: # Our fake quantizer may not be fully compatible with torch.compile. # This is a best-effort sanity check: e.g. a `device_map="auto"` load that offloads # part of the model to CPU (seen on unified-memory single-GPU hosts) can make a # quantized layer run on CPU, which some kernels (e.g. NVFP4 dynamic block # quantization) don't support. Don't let that discard the completed calibration. lockstep = preview_generate_kwargs(full_model) try: with keep_the_experts_implementation(full_model): generated_ids_after_ptq = full_model.generate( preview_input_ids, attention_mask=preview_attention_mask, max_new_tokens=100, **lockstep, ) except Exception as e: if lockstep: # The ranks decode in lockstep: a rank that skipped the preview would go on to the # export's collectives while the others still wait in generate's, so the run stops. raise RuntimeError(_PREVIEW_STOPS_THE_RUN) from e warnings.warn(f"Post-quantization generation sanity check failed, skipping it: {e}") elif vlm_wrapper and tokenizer is not None: try: generated_ids_after_ptq = run_vlm_preview( full_model, tokenizer, preview_input_ids, args.local_checkpoint_path, "after quantization", allow_fallback=False, trust_remote_code=args.trust_remote_code, ) except Exception as e: raise RuntimeError(_PREVIEW_STOPS_THE_RUN) from e else: warnings.warn( f"{post_ptq_unsupported or 'The preview needs a tokenizer'}. Skipping generation sample." ) def input_decode(input_ids): if preview_input_text is not None: return preview_input_text elif tokenizer is not None: return tokenizer.batch_decode(input_ids) else: raise ValueError("The processor or tokenizer must be set") def output_decode(generated_ids, input_shape): # Some `.generate()` returns a ModelOutput dataclass (e.g. DiffusionGemma); # unwrap to the token tensor so downstream slicing works uniformly. if hasattr(generated_ids, "sequences"): generated_ids = generated_ids.sequences if generate_excludes_prompt(full_model): if tokenizer is None and processor is not None: return processor.tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0] elif tokenizer is not None: return tokenizer.batch_decode(generated_ids, skip_special_tokens=True) elif tokenizer is not None: return tokenizer.batch_decode(generated_ids[:, input_shape:]) else: raise ValueError("The processor or tokenizer must be set") if generated_ids_after_ptq is not None: print("--------") if vlm_wrapper: # For VLM wrappers, generated_ids are text strings from model.chat() print("VLM wrapper text-only generation results:") print(f"Text response before quantization: {generated_ids_before_ptq}") print("--------") print(f"Text response after quantization: {generated_ids_after_ptq}") print("--------") print("Note: Additional VL tests with images were run separately above") else: # For regular LLMs, generated_ids are token tensors that need decoding print(f"example test input: {input_decode(preview_input_ids)}") print("--------") print( f"example outputs before ptq: {output_decode(generated_ids_before_ptq, preview_input_ids.shape[1])}" ) print("--------") print( 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) def quantize_main( args: argparse.Namespace, full_model: torch.nn.Module, language_model: torch.nn.Module, model_type: str | None, calibration_only: bool, processor: ProcessorMixin | None, tokenizer: PreTrainedTokenizerBase | None, device: torch.device, ): # Load the recipe up front so we can detect layerwise calibration before batch-size probing. recipe = None if args.recipe is not None: print(f"Use recipe {args.recipe} for quantization") recipe = load_recipe(args.recipe) if not isinstance(recipe, ModelOptPTQRecipe | ModelOptAutoQuantizeRecipe): raise TypeError( f"Expected PTQ or AutoQuantize recipe, but got {type(recipe).__name__} " f"from {args.recipe}" ) _validate_recipe_calibration(args, recipe) # AutoQuantize is recipe-driven: everything downstream reads the resolved AutoQuantizeConfig. if isinstance(recipe, ModelOptAutoQuantizeRecipe): aq_config = recipe.auto_quantize else: aq_config = None layerwise_cfgs = recipe_layerwise_blocks(recipe) is_layerwise = any(cfg.get("enable", False) for cfg in layerwise_cfgs) # The value is a placeholder, replaced with --export_path below; presence is the switch. args.layerwise_export = layerwise_export_requested(layerwise_cfgs) if not args.layerwise_export and any( cfg.get("export_dir") is not None for cfg in layerwise_cfgs ): warnings.warn( "layerwise.export_dir is set but layerwise.enable is not, so there is no " "per-layer pass to write the shards: the whole-model export runs instead, which " "holds the full state dict in host memory." ) if args.layerwise_export: if isinstance(recipe, ModelOptAutoQuantizeRecipe): # Only the mono-quantize path retargets export_dir and runs the refusals; # auto_quantize would export to the placeholder and skip the real export. raise NotImplementedError( "layerwise.export_dir is not supported with an AutoQuantize recipe; " "use a PTQ recipe, or drop export_dir and export afterwards." ) if not args.skip_generate: print("Layerwise export: forcing --skip_generate, the model is left in export form.") args.skip_generate = True if args.batch_size == 0: # For VL models with image-text calibration, skip automatic batch size detection # since get_max_batch_size can't handle multimodal inputs if args.use_fsdp2 and aq_config is not None: print_rank_0( "FSDP2 AutoQuant uses batch_size=1 by default; set --batch_size to override." ) args.batch_size = 1 elif args.calib_with_images: print("Image-text calibration enabled. Using default batch_size=1 for calibration.") args.batch_size = 1 # Speculative decoding offline model dost not support get_max_batch_size() because of # the customized dataloader, so we set batch_size to 1 to avoid OOM. elif args.specdec_offline_dataset is not None: print( "Offline speculative decoding calibration enabled. Using default batch_size=1 for calibration." ) args.batch_size = 1 # Layerwise calibration processes one layer at a time; auto batch-size probing runs a # full-model forward which defeats the point and can OOM on very large models. A batch # then only holds one layer's activations, so probe one decoder layer of each kind, up # to a token budget: batch 1 makes an MoE layer run its per-expert loop once per sample. elif is_layerwise: max_batch_size = _layerwise_calib_batch_size(args.calib_seq, sum(args.calib_size)) args.batch_size = get_max_layerwise_batch_size( language_model, args.calib_seq, max_batch_size, sum(args.calib_size), device=device ) print( f"Layerwise calibration enabled. Using batch_size={args.batch_size}, the most " f"that fits a decoder layer's forward, up to {max_batch_size} " f"({LAYERWISE_CALIB_TOKENS} tokens)." ) else: # Calibration/sparsification will actually take much more memory than regular inference # due to intermediate tensors for fake quantization. Setting sample_memory_usage_ratio # to 2 to avoid OOM for AWQ/SmoothQuant fake quantization as it will take more memory than inference. sample_memory_usage_ratio = ( 2 if "awq" in args.qformat or "smoothquant" in args.qformat else 1.1 ) # A model type that takes something other than token ids supplies its own sample; for # the rest the sample input is set up inside get_max_batch_size(). sample_input_single_batch = calibration_sample_input(model_type, language_model) run_auto_quant = aq_config is not None with fsdp2_call_context(full_model, language_model): args.batch_size = get_max_batch_size( language_model, max_sample_length=args.calib_seq, sample_memory_usage_ratio=sample_memory_usage_ratio if not run_auto_quant else 1.0, sample_input_single_batch=sample_input_single_batch, enable_grad=( aq_config is not None and aq_config.auto_quantize_method in ("gradient", "aumann_shapley") ), forward_step=(_get_autoquant_memory_probe(recipe) if run_auto_quant else None), ) args.batch_size = min(args.batch_size, sum(args.calib_size)) print(f"Use calib batch_size {args.batch_size}") calib_dataloader, preview_input_text = make_calib_dataloader( args, language_model, processor, tokenizer, device, model_type, autoquant_gradient_recipe=( aq_config is not None and aq_config.auto_quantize_method == "gradient" ), ) vlm_wrapper = get_ptq_spec(model_type).vlm_wrapper preview_input_ids, preview_attention_mask, generated_ids_before_ptq = pre_quantize( args, full_model, model_type, tokenizer, calib_dataloader, vlm_wrapper ) if aq_config is not None: assert isinstance(recipe, ModelOptAutoQuantizeRecipe) _run_auto_quantize_recipe( args, recipe, full_model, language_model, model_type, calibration_only, calib_dataloader, vlm_wrapper, ) else: # mono quantization if recipe is not None: quant_cfg = recipe.quantize.model_dump() else: assert len(args.qformat.split(",")) == 1, ( "Plain quantization supports only one quantization format." ) assert args.qformat in QUANT_CFG_CHOICES, ( f"Unsupported quantization format: {args.qformat}, choices are: {list(QUANT_CFG_CHOICES)}" ) quant_cfg = QUANT_CFG_CHOICES[args.qformat] quant_cfg = build_quant_cfg( quant_cfg, args.awq_block_size, args.moe_calib_experts_ratio, ) enable_quant_kv_cache = args.kv_cache_qformat != KV_CACHE_NONE print(f"{'Enable' if enable_quant_kv_cache else 'Disable'} KV cache quantization") # Check if any bmm_quantizer is in the quant_cfg. If so, we need to enable the bmm_quantizer. if enable_quant_kv_cache: quant_cfg = mtq.update_quant_cfg_with_kv_cache_quant( quant_cfg, KV_QUANT_CFG_CHOICES[args.kv_cache_qformat]["quant_cfg"], ) quant_cfg = _prepare_quant_cfg(args, quant_cfg, full_model) if quant_cfg: mono_quantize( args, quant_cfg, full_model, language_model, model_type, calibration_only, calib_dataloader, vlm_wrapper, ) else: print(f"qformat: {args.qformat}. No quantization applied, export {device} model") # If asked, run the closed-form MXFP4 -> NVFP4 cast: read the source MXFP4 # *_scales tensors and pin each NVFP4 weight quantizer's scale_2 to 2^m. # Runs after calibration (max_calibrate has already promoted weight quantizers # 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 the local # checkpoint directory main() downloaded. apply_cast_mxfp4_to_nvfp4(language_model, args.local_checkpoint_path) post_quantize( args, full_model, language_model, model_type, tokenizer, processor, preview_input_ids, preview_attention_mask, generated_ids_before_ptq, vlm_wrapper, preview_input_text, calib_dataloader, ) def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser(description=__doc__) parser.add_argument( "--pyt_ckpt_path", "--model", help=( "Model name or path to the PyTorch checkpoint to be quantized. " "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, ) parser.add_argument( "--recipe", help=( "PTQ or AutoQuantize recipe YAML file or name without suffix (e.g. " "general/ptq/nvfp4_default-kv_fp8_cast, general/auto_quantize/nvfp4_fp8_at_4p8bits). " "KV cache behavior depends on the recipe type: PTQ recipes configure it in quant_cfg " "and ignore --kv_cache_qformat; weight AutoQuantize recipes use their kv_cache setting " "or fall back to --kv_cache_qformat; KV-cache AutoQuantize recipes select per-layer K/V " "formats from candidate_formats and ignore --kv_cache_qformat." ), default=None, ) parser.add_argument("--device", default="cuda") parser.add_argument( "--qformat", action=RecipeSupersededAction, help="Quantization format for single-format PTQ. For mixed-precision search, use an " "AutoQuantize recipe via --recipe. (deprecated: use --recipe)", default="fp8", ) parser.add_argument( "--batch_size", help="Batch size for calibration. Default to 0 as we calculate max batch size on-the-fly", type=int, default=0, ) parser.add_argument( "--calib_size", help=( "Number of samples for calibration. If a comma separated list of values is provided, " "each value will be used as the calibration size for the corresponding dataset. " "This argument will be parsed and converted as a list of ints." ), type=str, default="1024", ) parser.add_argument( "--calib_seq", help="Maximum sequence length for calibration.", type=int, default=512, ) parser.add_argument("--export_path", default="exported_model") parser.add_argument( "--dataset", help=( f"name of a dataset, or a comma separated list of datasets. " f"dataset choices are {get_supported_datasets()}" ), type=str, default=None, ) parser.add_argument( "--specdec_offline_dataset", help=( "If set, the model is a speculative decoding model," "which uses offline dataset for calibration. " ), default=None, ) parser.add_argument( "--calib_with_images", action="store_true", help=( "Calibrate with image-text pairs (for VLMs). " "This uses nemotron_vlm_dataset_v2 with default subsets (sparsetables, plotqa_cot, wiki_en)." ), ) parser.add_argument("--inference_tensor_parallel", type=int, default=1) parser.add_argument("--inference_pipeline_parallel", type=int, default=1) parser.add_argument("--awq_block_size", default=0, type=int) parser.add_argument( "--sparsity_fmt", help="Sparsity format.", default="dense", choices=["dense", "sparsegpt"], ) parser.add_argument( "--kv_cache_qformat", action=RecipeSupersededAction, required=False, default="fp8_cast", choices=[KV_CACHE_NONE, *KV_QUANT_CFG_CHOICES], help=( "Specify KV cache quantization format. Default: fp8_cast. " "Formats whose preset pins use_constant_amax on the KV bmm quantizer " "(e.g. fp8_cast, nvfp4_cast) set the amax to FP8 range without data-driven " "calibration; all other formats (fp8, nvfp4, ...) use data-driven calibration. " "With --recipe, the source depends on the recipe type: a PTQ recipe is " "authoritative for KV cache and ignores this flag; an AutoQuantize recipe " "falls back to this flag unless it sets an explicit kv_cache field. (deprecated: use --recipe)" ), ) parser.add_argument( "--export_fmt", required=False, default="hf", choices=["tensorrt_llm", "hf"], help="Deprecated. Please avoid using this argument.", ) parser.add_argument( "--trust_remote_code", help="Set trust_remote_code for Huggingface models and tokenizers", default=False, action="store_true", ) parser.add_argument( "--gpu_max_mem_percentage", help=( "Specify the percentage of available GPU memory to use for loading the model when " "device_map is set to sequential. " "By default, 80%% of the available GPU memory is used." ), type=float, default=0.8, ) parser.add_argument( "--use_seq_device_map", help=( "Use device_map=sequential to load the model onto GPUs. This ensures the model is loaded " "utilizing the percentage of available GPU memory as specified by the value passed with gpu_max_mem flag." "Helpful in cases where device_map=auto loads model unevenly on GPUs causing GPU OOM during quantization." ), default=False, action="store_true", ) parser.add_argument( "--use_fsdp2", action="store_true", help=( "Run calibration under PyTorch FSDP2 (requires torchrun); takes precedence over " "--use_seq_device_map. Includes experimental weight and KV-cache AutoQuantize. " "Supports causal LMs and VLMs, calibrated on text or, with --calib_with_images, on " "image-text samples (required for a VLM whose forward needs images). Not supported: " "VILA / encoder-decoder / pack-quantized / speculative / sparsity / MTP." ), ) parser.add_argument( "--cpu_offload", action="store_true", help="With --use_fsdp2, keep decoder shards on CPU between forwards (frees GPU memory, adds PCIe traffic).", ) parser.add_argument( "--verbose", help="Print verbose output (e.g. quantization summary). Disable by --no-verbose.", default=True, action=argparse.BooleanOptionalAction, ) parser.add_argument( "--skip_generate", help=( "Skip pre/post-quantization preview calls that invoke model.generate(). " "Note: this does not skip calibration or batch-size probing. " "For very large models, pair with --batch_size 1 to avoid max-batch probing." ), default=False, action="store_true", ) parser.add_argument( "--low_memory_mode", help=( "Use low memory mode for quantization." "This is an experimental feature and may not work for all quantization formats." ), default=False, action="store_true", ) parser.add_argument( "--attn_implementation", help=( "Specify the attention implementation to use. " "This arg will be passed to the HF model loading if specified." ), default=None, type=str, ) parser.add_argument( "--experts_implementation", help=( "Fused MoE experts backend to load the model with, passed to the HF model loading as " "experts_implementation. By default transformers picks grouped_mm where the model " "supports it; eager calibrates exactly as earlier releases did. Quantization redirects " "every other backend to one of these two." ), default=None, choices=["eager", "grouped_mm"], ) parser.add_argument( "--auto_quantize_checkpoint", type=str, default=None, help=( "Path to checkpoint file for saving/restoring weight AutoQuantize search state " "(sensitivity scores, costs, etc.). Used with an AutoQuantize --recipe." ), ) parser.add_argument( "--kv_auto_quantize_checkpoint", type=str, default=None, help=( "Path for saving/restoring any KV-cache AutoQuantize search checkpoint. Use a new " "path whenever the preceding weight/activation quantization stage changes. " "KV-primary recipes temporarily accept --auto_quantize_checkpoint as a deprecated " "fallback." ), ) parser.add_argument( "--moe_calib_experts_ratio", type=float, default=None, help=( "Fraction of experts to calibrate during forward pass (ratio in (0.0, 1.0]). " "Only used for MOE models; used to reduce the number of experts calibrated during the forward pass. " "Does not impact non-MOE models." ), ) parser.add_argument( "--vllm_fakequant_export", default=False, action="store_true", help="Export as vLLM fake-quant checkpoint (produces vllm_fq_modelopt_state.pth " "for use with vllm_serve_fakequant.py).", ) parser.add_argument( "--cast_mxfp4_to_nvfp4", action="store_true", default=False, help=( "After calibration, override NVFP4 weight quantizers' global_amax with " "the closed-form value derived from the source MXFP4 *_scales. " "Per-block _amax is computed from the loaded BF16 weights (data-derived). " "Use when --pyt_ckpt_path points at an MXFP4 HF checkpoint (e.g. " "openai/gpt-oss-20b) and the target qformat is NVFP4-family." ), ) parser.add_argument( "--offload_folder", type=str, default=None, help=( "Path to a local folder for disk-offloaded model weights. " "When set, activates disk-offload mode: model weights that exceed the GPU+CPU " "budgets are streamed from disk during calibration and export. " "Pair with --max_cpu_memory_gb to cap CPU RAM usage. " "Incompatible with --low_memory_mode and --use_seq_device_map." ), ) parser.add_argument( "--max_cpu_memory_gb", type=float, default=None, help=( "Maximum CPU RAM budget in GiB for disk-offload model loading. " "Only effective when --offload_folder is set. " "Weights beyond this limit are streamed from disk." ), ) parser.add_argument( "--max_gpu_memory_gb", type=float, default=None, help=( "Maximum GPU memory budget per device in GiB for disk-offload model loading. " "Only effective when --offload_folder is set. " "Defaults to 80%% of available GPU memory when not specified." ), ) add_mlflow_args(parser, HF_PTQ) args = parser.parse_args() # Flipped by export_quantized once a checkpoint is actually on disk. The MLflow pointer # is gated on it rather than on --export_path existing, which proves nothing. args.checkpoint_exported = False resolve_mlflow_args(args, parser, HF_PTQ) if args.moe_calib_experts_ratio is not None and not (0.0 < args.moe_calib_experts_ratio <= 1.0): parser.error("--moe_calib_experts_ratio must be in the range (0.0, 1.0].") if args.calib_seq <= 0: parser.error("--calib_seq must be positive.") if args.specdec_offline_dataset is not None and args.sparsity_fmt != "dense": parser.error("--specdec_offline_dataset is only supported with --sparsity_fmt dense (PTQ).") if args.specdec_offline_dataset is not None and args.low_memory_mode: parser.error("--specdec_offline_dataset is not compatible with --low_memory_mode.") # The low-memory loader pre-instruments quantizers from --qformat/--kv_cache_qformat # via init_quantized_weights(), so it cannot honor a --recipe (which is authoritative # for the quant layout in quantize_main). Reject the combination rather than silently # instrumenting a layout that diverges from the recipe. if args.low_memory_mode and args.recipe is not None: parser.error( "--low_memory_mode does not support --recipe; the low-memory loader initializes " "quantizers from --qformat/--kv_cache_qformat." ) if args.use_fsdp2 and args.use_seq_device_map: warnings.warn("--use_seq_device_map is ignored when --use_fsdp2 is set.") args.use_seq_device_map = False if args.use_fsdp2 and os.environ.get("RANK") is None: parser.error("--use_fsdp2 requires launching with torchrun") if args.cpu_offload and not args.use_fsdp2: parser.error("--cpu_offload requires --use_fsdp2") if args.use_fsdp2 and args.sparsity_fmt != "dense": parser.error(f"--use_fsdp2 does not support --sparsity_fmt {args.sparsity_fmt}.") if args.use_fsdp2 and args.vllm_fakequant_export: parser.error("--use_fsdp2 does not support --vllm_fakequant_export.") if args.use_fsdp2 and args.cast_mxfp4_to_nvfp4: parser.error("--use_fsdp2 does not support --cast_mxfp4_to_nvfp4.") if args.offload_folder is not None and args.low_memory_mode: parser.error("--offload_folder (disk-offload) is not compatible with --low_memory_mode.") if args.offload_folder is not None and args.use_seq_device_map: parser.error( "--offload_folder (disk-offload) is not compatible with --use_seq_device_map; " "device_map=auto is used for disk-offload to let accelerate place layers across " "GPU, CPU, and disk." ) if args.offload_folder is not None and args.device == "cpu": parser.error( "--offload_folder (disk-offload) is not compatible with --device cpu; " "device_map=cpu makes accelerate ignore the memory budgets and offload folder, " "loading the whole model into RAM." ) if args.offload_folder is None and ( args.max_cpu_memory_gb is not None or args.max_gpu_memory_gb is not None ): parser.error( "--max_cpu_memory_gb/--max_gpu_memory_gb only apply to disk-offload loading; " "pass --offload_folder to enable it." ) return args # Derived state and the tracking settings themselves; everything else argparse parsed is a # parameter of the run. Deriving the list means a new flag is tracked without touching this. def main(args: argparse.Namespace): if not torch.cuda.is_available(): raise OSError("GPU is required for inference.") random.seed(RAND_SEED) np.random.seed(RAND_SEED) # The run owns a distributed launch's process group, and ends it even if loading raises. with distributed_run(timeout=timedelta(hours=3), device=args.device) as dist_state: args.dist_state = dist_state _main_inner(args) def _main_inner(args: argparse.Namespace): with mlflow_run(args): # --pyt_ckpt_path stays as given. Every later step reads 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.pyt_ckpt_path) # launch a memory monitor to read the currently used GPU memory, one per node: it reads all # of the node's GPUs. if args.dist_state.world_size == 1 or local_rank() == 0: launch_memory_monitor() # Force eager execution for all model types. torch.compiler.set_stance("force_eager") ( full_model, language_model, model_type, calibration_only, processor, tokenizer, device, ) = load_model(args) if args.sparsity_fmt != "dense": # Sparse sparsity_main(args, full_model, tokenizer, device) else: # Quantize quantize_main( args, full_model, language_model, model_type, calibration_only, processor, tokenizer, device, ) if __name__ == "__main__": args = parse_args() if args.export_fmt != "hf": warnings.warn("Deprecated. --export_fmt forced to hf.") args.dataset = args.dataset.split(",") if isinstance(args.dataset, str) else args.dataset args.calib_size = [int(num_sample) for num_sample in args.calib_size.split(",")] if args.specdec_offline_dataset is not None and len(args.calib_size) != 1: raise ValueError( "--specdec_offline_dataset expects a single --calib value, not a comma-separated list." ) if args.cast_mxfp4_to_nvfp4: qformats = [q.strip() for q in args.qformat.split(",")] if not all("nvfp4" in q for q in qformats): raise ValueError( "--cast_mxfp4_to_nvfp4 requires NVFP4-family --qformat values " f"(got {args.qformat!r}). Use e.g. --qformat nvfp4 or nvfp4_mlp_only." ) main(args)