diff --git a/CHANGELOG.rst b/CHANGELOG.rst index d4a29905480..16c48f5bb27 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -17,6 +17,8 @@ Changelog - Add ``mtq.temporarily_fold_weights`` for repeated frozen-weight inference and ``mtq.preserve_quantizer_attributes_context`` for restoring temporary quantizer property and type changes. Temporary folding snapshots affected fake-quant weights on a configurable device and restores them with their quantizer state; retained pre-quant scales are inactive, while shared weights, shared quantizers, and ``SequentialQuantizer`` weights are unsupported. - Add the ``nvfp4_act_headroom`` calibration algorithm for NVFP4 **activation** global scales. Instead of setting the global scale from the largest per-block amax seen during calibration (plain ``max``, which leaves no room above it so any larger activation saturates), it anchors the scale to a low percentile of the per-block amax distribution, leaving the rest of the FP8 block-scale range as headroom: ``amax = max(rho * anchor, upper)``, where ``anchor`` and ``upper`` are the per-block amaxes at ``anchor_percentile`` (default 1) and ``upper_percentile`` (default 99.99; set to 100 to never clip calibration data), and ``rho`` (default 16384) is the headroom factor. Applies only to NVFP4 dynamic-block input quantizers; ``SequentialQuantizer`` activation quantizers raise. Weight scales are an orthogonal axis selected by a nested ``weight_scale_algorithm`` (``max`` by default, or ``mse`` / ``local_hessian``), so one recipe can combine a weight calibration with this activation policy in a single pass. Ships ``modelopt_recipes/general/ptq/nvfp4_act_headroom-kv_fp8_cast.yaml``, which mirrors ``nvfp4_default-kv_fp8_cast`` with only the calibration algorithm swapped and exports a standard NVFP4 checkpoint. +- Add PTQ support for Step-3.7 (``stepfun-ai/Step-3.7-Flash``), whose routed experts were previously left unquantized. Quantize with the new ``huggingface/step3p7/ptq/nvfp4_experts_only-kv_fp8_cast`` or ``huggingface/step3p7/ptq/nvfp4_mlp_only-kv_fp8`` recipes rather than the general ones, which select experts by module names Step does not use. + *Megatron Framework (M-LM / M-Bridge)* - Add ``clamp_kv_cache_scales`` to ``export_mcore_gpt_to_hf``. Set it to ``False`` when exporting a QAT Megatron-Core model to preserve its learned FP8 KV-cache scales; the default retains the existing minimum scale of 1.0. @@ -34,6 +36,7 @@ Changelog **Backward Breaking Changes** +- ``get_te_hybrid_stack_spec`` was removed from ``modelopt.torch.nas.plugins.megatron``; it had no use outside tests. Use ``modelopt.torch.utils.plugins.megatron_layer_specs.te_hybrid_stack_spec_sequential_mlp`` for the SequentialMLP layout, or ``megatron.core.models.hybrid.hybrid_layer_specs.hybrid_stack_spec`` for grouped GEMM. - Migrate the FAR3D ONNX PTQ example to the shared evaluator and ModelOpt containers and ``quantize_vovnet.py``. Only the encoder supports INT8 and FP8; decoder calibration, quantization, and related CLI flags are removed, and the decoder remains in its exported mixed FP16/FP32 precision. - Image-text calibration with ``--calib_with_images`` now forwards multimodal batches through the complete VLM for all VLM families, so existing non-Nemotron commands may produce different language-model activation ranges and output scales. Recipe-based VLM PTQ also targets the complete VLM: vision modules stay in high precision by default and are quantized only when a model-specific recipe enables them, so custom recipes must explicitly exclude vision modules when required. - Move the checkpoint-mirror recipe tier from ``huggingface/models///`` to the top-level ``models///``, keyed by each recipe's canonical Hugging Face Hub id — so the Step 3.5 Flash recipe moves to ``models/stepfun-ai/Step-3.5-Flash/ptq/`` and the NVIDIA Nemotron recipes gain the ``NVIDIA-`` prefix (e.g. ``models/nvidia/NVIDIA-Nemotron-3-Super-120B-A12B-BF16/ptq/nvfp4-mse``). Update any saved ``--recipe`` paths for these checkpoint recipes accordingly; the per-``model_type`` recipes under ``huggingface/`` are unchanged. @@ -45,6 +48,7 @@ Changelog - Remove in-trainer quantization via ``QuantizationArguments.quant_cfg`` / ``--quant_cfg`` (deprecated in 0.45); use ``--recipe``. New recipes ``general/ptq/mxfp4_mlp_weight_only`` and ``general/ptq/nvfp4_mlp_weight_only`` replace ``MXFP4_MLP_WEIGHT_ONLY_CFG`` / ``NVFP4_MLP_WEIGHT_ONLY_CFG`` in the ``examples/gpt-oss`` QAT flow. - Remove the ``QuantizationArgumentsWithConfig`` alias in ``modelopt.torch.quantization.plugins.transformers_trainer`` (deprecated in 0.45). Use ``QuantizationArguments``. - Transformer Engine ``TEGroupedLinear`` (fused MoE experts) now uses **per-expert** weight quantization (one ``amax`` per expert) instead of a single shared ``amax``, so ModelOpt checkpoints containing quantized ``TEGroupedLinear`` modules saved before 0.47 are **not compatible** with 0.47. Re-run PTQ to regenerate compatible checkpoints. +- ``mtq.quantize`` now raises when a config asks for weight quantization but none of its weight-quantizer patterns match the model, instead of calibrating and exporting a silently unquantized checkpoint (``"quant_algo": null``). Configs that quantize activations or the KV cache only are unaffected, as are patterns that match and are then disabled by a later entry. If this fires, use the recipe for that architecture under ``modelopt_recipes/huggingface//`` or fix the module patterns. Set ``MODELOPT_SKIP_WEIGHT_QUANT_CHECK=1`` to disable the check process-wide, e.g. for a pipeline-parallel rank whose local stage legitimately has none of the targeted modules. **Deprecations** @@ -53,8 +57,18 @@ Changelog **Bug Fixes** +- Fix shared ONNX export metadata and Diffusers attention policy: every ``NVFP4QuantExporter`` post-process now upgrades the default-domain opset to at least 23, all FP8 custom-op exports re-run ONNX shape/type inference after setting output metadata, and quantized SDPA derives FP8 MHA enablement from the live Q/K/V quantizers instead of honoring a caller-set ``_disable_fp8_mha`` attribute. +- Fix ONNX FP16 conversion failing to preserve public output types when type inference changes a graph output declaration before output casts are inserted. +- Fix ``examples/hf_ptq/hf_ptq.py`` discarding a completed PTQ run (no checkpoint exported) when the optional post-quantization sanity-check ``generate()`` call raised, for example because ``device_map="auto"`` placed part of the model on CPU. That failure is now caught and only skips the sanity check; export proceeds regardless. +- Hybrid (e.g. Nemotron-H) checkpoints saved by the ``examples/megatron_bridge`` scripts now record their layer spec in ``run_config.yaml`` in a form that reloads, so they can be converted to HuggingFace; a checkpoint saved by an earlier release still needs its ``model.hybrid_stack_spec`` block replaced by hand. +- Fail fast on non-finite AutoQuantize output gradients with an actionable error before accumulating sensitivity scores, without changing attention backend settings. +- Fix ONNX INT8 entropy calibration failing or producing invalid quantization parameters for FP16 activations. +- Fix HuggingFace checkpoint export failing with ``activation scaling factor 0.0 not positive`` when a dynamic-block quantizer (such as an NVFP4 input quantizer) ends calibration with an amax of zero because the calibration data never activated that layer or expert. Such a quantizer now exports a positive fallback scale and warns instead of crashing, matching what static quantizers already did; if you see the warning, check whether the layer is expected to be inactive and consider a larger calibration size. +- Speed up ``mtq.quantize`` on FSDP2-sharded fused-MoE models. Promoting static-block weight quantizers gathered each expert's slice of the fused weight across ranks even though only quantizer state is read, adding a collective per expert to calibration. +- Fix ONNX AutoCast failing on models with external initializers larger than 2 GiB. - Avoid querying CUDA/Blackwell capability when ``NVFP4QTensor.quantize`` uses its CPU path or has the optional TensorRT-LLM fast path disabled. - Fix NVFP4 ONNX export to quantize FP4 weights with the published FP8 block scales, matching eager ModelOpt packed weights. Block scales below ``2**-9`` are now clamped to that minimum, and non-finite or negative scales raise an error. +- Fix FP8 ONNX export of BF16 models during real-weight compression. - Fix Megatron-Bridge Quantization Aware Distillation of a vision-language model silently discarding the ModelOpt state, so the distilled checkpoint restored no quantizers and exported as an unquantized model. Re-run QAD to regenerate any affected checkpoint. - Fix Megatron-Core HuggingFace export silently omitting fused (grouped GEMM) MoE experts for architectures without an ``experts.linear_fc1`` rule (e.g. ``Qwen3MoeForCausalLM``), which produced a valid-looking checkpoint containing no expert weights. The exporter now raises instead of writing that checkpoint; the scripts also avoid the situation by selecting ``SequentialMLP`` for those architectures. - Fix GatedDeltaNet (Qwen3.5) quantizer exclusions on Megatron-Core: the recipe patterns name the HuggingFace ``linear_attn`` module, so the ``conv1d`` was calibrated and the alpha / beta gate projections were exported in FP8. ``conv1d`` now has a ``self_attention`` alias in the default disabled-quantizer units, and the alpha / beta projections are exported in BF16 (they share Megatron's fused ``in_proj`` quantizer and cannot be disabled by name). @@ -71,6 +85,7 @@ Changelog - Fix EAGLE-3 training with context parallelism (``--cp_size > 1`` in ``examples/speculative_decoding``), which failed to start on ``accelerate >= 1.13`` and then raised ``got mixed torch.Tensor and DTensor``. - Polygraphy minimum dependency upgraded to ``0.53.4`` to solve ONNX AutoCast failures when marking optional graph outputs. - Fix ``--kv_cache_free_gpu_memory_fraction`` having no effect on the ``lm_eval`` task of ``examples/hf_ptq/scripts/huggingface_example.sh``, where the KV cache always took TensorRT-LLM's default 90% of free GPU memory and evaluation could run out of memory. ``examples/llm_eval/lm_eval_trtllm.py`` now takes ``kv_cache_free_gpu_memory_fraction`` in ``--model_args``, defaulting to 0.8. +- Fix ``--aux-layers eagle`` failing in the vLLM offline hidden-state dump (``examples/speculative_decoding/collect_hidden_states/compute_hidden_states_vllm.py``). ``eagle`` is the flag's default, but the dump's standalone resolver -- a copy kept so the script runs in a stock vLLM container without ModelOpt -- only handled ``dflash`` and explicit id lists, so the documented invocation aborted with ``invalid literal for int(): 'eagle'`` before any state was written. An unrecognised preset now reports which values are accepted instead of surfacing the raw ``int()`` error. 0.46 (2026-08-17) ^^^^^^^^^^^^^^^^^ diff --git a/docs/source/deployment/1_tensorrt_llm.rst b/docs/source/deployment/1_tensorrt_llm.rst index 43195e3f627..77e4dab649c 100644 --- a/docs/source/deployment/1_tensorrt_llm.rst +++ b/docs/source/deployment/1_tensorrt_llm.rst @@ -2,149 +2,16 @@ TensorRT-LLM ========================== -**Deprecation Notice**: The export_tensorrt_llm_checkpoint API will be deprecated in future releases. Users are encouraged to transition to the :doc:`unified HF export API <3_unified_hf>`, which provides enhanced functionality and flexibility for exporting models to multiple inference frameworks including TensorRT-LLM, vLLM, and SGLang. - -.. note:: - - Please read the `TensorRT-LLM checkpoint workflow `_ - first before going through this section. - - - -ModelOpt toolkit supports automatic conversion of ModelOpt exported LLM to the TensorRT-LLM checkpoint and the engines for accelerated inferencing. - -This conversion is achieved by: - -#. Converting Huggingface, Megatron-Bridge and ModelOpt exported checkpoints to the TensorRT-LLM checkpoint. -#. Building TensorRT-LLM engine from the TensorRT-LLM checkpoint. - - -Export Quantized Model -====================== - -After the model is quantized, the quantized model can be exported to the TensorRT-LLM checkpoint format stored as - -#. A single JSON file recording the model structure and metadata (config.json) -#. A group of safetensors files, each recording the local calibrated model on a single GPU rank (model weights, scaling factors per GPU). - -The export API (:meth:`export_tensorrt_llm_checkpoint `) can be used as follows: - -.. code-block:: python - - from modelopt.torch.export import export_tensorrt_llm_checkpoint - - with torch.inference_mode(): - export_tensorrt_llm_checkpoint( - model, # The quantized model. - decoder_type, # The type of the model as str, e.g gpt, gptj, llama. - dtype, # the weights data type to export the unquantized layers. - export_dir, # The directory where the exported files will be stored. - inference_tensor_parallel, # The number of GPUs used in the inference time for tensor parallelism. - inference_pipeline_parallel, # The number of GPUs used in the inference time for pipeline parallelism. - ) - -If the :meth:`export_tensorrt_llm_checkpoint ` call is successful, the TensorRT-LLM checkpoint will be saved. Otherwise, e.g. the ``decoder_type`` is not supported, a torch state_dict checkpoint will be saved instead. - -.. list-table:: Model support matrix for the TensorRT-LLM checkpoint export - :header-rows: 1 - - * - Model / Quantization - - FP16 / BF16 - - FP8 - - INT8_SQ - - INT4_AWQ - * - GPT2 - - Yes - - Yes - - Yes - - No - * - GPTJ - - Yes - - Yes - - Yes - - Yes - * - LLAMA 2 - - Yes - - Yes - - Yes - - Yes - * - LLAMA 3 - - Yes - - Yes - - No - - Yes - * - Mistral - - Yes - - Yes - - Yes - - Yes - * - Mixtral 8x7B - - Yes - - Yes - - No - - Yes - * - Falcon 40B, 180B - - Yes - - Yes - - Yes - - Yes - * - Falcon 7B - - Yes - - Yes - - Yes - - No - * - MPT 7B, 30B - - Yes - - Yes - - Yes - - Yes - * - Baichuan 1, 2 - - Yes - - Yes - - Yes - - Yes - * - ChatGLM2, 3 6B - - Yes - - No - - No - - Yes - * - Bloom - - Yes - - Yes - - Yes - - Yes - * - Phi-1, 2, 3 - - Yes - - Yes - - Yes - - Yes - * - Nemotron 8 - - Yes - - Yes - - No - - Yes - * - Gemma 2B, 7B - - Yes - - Yes - - No - - Yes - * - Recurrent Gemma - - Yes - - Yes - - Yes - - Yes - * - StarCoder 2 - - Yes - - Yes - - Yes - - Yes - * - Qwen-1, 1.5 - - Yes - - Yes - - Yes - - Yes - -Convert to TensorRT-LLM -======================= - -Once the TensorRT-LLM checkpoint is available, please follow the `TensorRT-LLM build API `_ to build and deploy the quantized LLM. +For current TensorRT-LLM deployments, export quantized models with +:meth:`export_hf_checkpoint ` +and load the exported Hugging Face checkpoint with TensorRT-LLM's PyTorch backend. +See the :doc:`unified HF export guide <3_unified_hf>` for export and deployment +examples, supported models, and quantization formats. This workflow does not require +building a TensorRT engine. + +.. warning:: + + The ``export_tensorrt_llm_checkpoint`` API exports checkpoints for the legacy + TensorRT backend, which current TensorRT-LLM releases no longer support. + The API will be deprecated in a future release. + Use ``export_hf_checkpoint`` instead. diff --git a/docs/source/guides/3_pruning.rst b/docs/source/guides/3_pruning.rst index a65890625ba..779e219f247 100644 --- a/docs/source/guides/3_pruning.rst +++ b/docs/source/guides/3_pruning.rst @@ -7,6 +7,13 @@ Pruning Checkout `Megatron-Bridge Minitron Pruning & Distillation `_ and `ResNet20 on CIFAR-10 Notebook `_ for an end-to-end example of pruning. +.. note:: + + For Minitron (LLM/VLM pruning via Megatron-Bridge/Megatron-LM) and Puzzletron pruning, + the canonical reference is the + `Pruning README `_ + which contains up-to-date API examples, support matrix, guidelines, and distillation + hyperparameters. This guide covers FastNAS pruning for Computer Vision models. ModelOpt provides three main pruning methods (aka ``mode``) - Minitron, Puzzletron, and FastNAS - via a unified API :meth:`mtp.prune `. Given a model, diff --git a/docs/source/guides/_customized_model_quantization.rst b/docs/source/guides/_customized_model_quantization.rst index c8078678ec4..b2fa879ecd4 100644 --- a/docs/source/guides/_customized_model_quantization.rst +++ b/docs/source/guides/_customized_model_quantization.rst @@ -16,7 +16,7 @@ As ModelOpt cannot detect these linear ops out-of-the-box, a HugggingFace plugin #. Rewrite the linear ops (w1, v1 and v2) as a standard ``nn.Linear`` op, and re-implement the ``forward`` method. #. Register the new dynamic ``_QuantDbrxExperts`` to replace the ``DbrxExperts`` from the modeling_dbrx.py in the ``transformers`` library #. Try quantize the DBRX model after the plugin is implemented, feel free to follow the `hf_ptq example `_. -#. TensorRT-LLM is open-sourced. If this customized model is not supported by TensorRT-LLM yet, please modify :meth:`export_tensorrt_llm_checkpoint ` or :meth:`export_hf_checkpoint ` to export the quantized model for deployment with a customized TensorRT-LLM modeling implementation. Feel free to :doc:`contact us <../support/1_contact>` if further support is needed. +#. Export the quantized model with :meth:`export_hf_checkpoint `. If the customized model is not supported by TensorRT-LLM, add support in its PyTorch backend and adapt the HF exporter if needed. See the :doc:`unified HF export guide <../deployment/3_unified_hf>` or :doc:`contact us <../support/1_contact>` for help. The following code snippet is excerpted from ``modelopt/torch/quantization/plugins/huggingface.py`` diff --git a/docs/source/guides/_onnx_quantization.rst b/docs/source/guides/_onnx_quantization.rst index e4d0c2d93d6..22fd9721f44 100644 --- a/docs/source/guides/_onnx_quantization.rst +++ b/docs/source/guides/_onnx_quantization.rst @@ -37,7 +37,7 @@ Requirements Apply Post Training Quantization (PTQ) ====================================== -PTQ should be done with a calibration dataset. If calibration dataset is not provided, ModelOpt will use random scales for the QDQ nodes. +PTQ should be done with a calibration dataset. Random calibration inputs are used when no calibration dataset is provided. Prepare calibration dataset --------------------------- diff --git a/examples/diffusers/quantization/ONNX-TRT-Deployment.md b/examples/diffusers/quantization/ONNX-TRT-Deployment.md index 57448b8a38e..b6933c88f72 100644 --- a/examples/diffusers/quantization/ONNX-TRT-Deployment.md +++ b/examples/diffusers/quantization/ONNX-TRT-Deployment.md @@ -28,12 +28,12 @@ python quantize.py \ #### FLUX-Dev|SDXL|SDXL-Turbo|LTX-Video FP8/FP4 [Script](./quantize.py) -*In our example code, FP4 is only supported for Flux. However, you can modify our script to enable FP4 format support for your own model.* +FP4 ONNX export is supported for Flux and SDXL. ```sh python quantize.py \ --model {flux-dev|sdxl-1.0|sdxl-turbo|ltx-video-dev} --model-dtype {Half|BFloat16} --trt-high-precision-dtype {Half|BFloat16} \ - --format {fp8|fp4} --batch-size 2 --calib-size {128|256} --quantize-mha \ + --format {fp8|fp4} --batch-size 2 --calib-size {128|256} \ --n-steps 20 --quantized-torch-ckpt-save-path ./{MODEL_NAME}.pt --collect-method default \ --onnx-dir {ONNX_DIR} ``` diff --git a/examples/diffusers/quantization/config.py b/examples/diffusers/quantization/config.py index cb8fdf3a5da..3ab8f5db87c 100644 --- a/examples/diffusers/quantization/config.py +++ b/examples/diffusers/quantization/config.py @@ -31,6 +31,9 @@ NVFP4_FP8_MHA_CONFIG = load_config( "configs/ptq/presets/diffusers/nvfp4_fp8_mha", schema_type=QuantizeConfig ).model_dump(exclude_unset=True) +NVFP4_FP8_CONV_CONFIG = load_config( + "configs/ptq/presets/diffusers/nvfp4_fp8_conv", schema_type=QuantizeConfig +).model_dump(exclude_unset=True) def set_quant_config_attr(quant_config, trt_high_precision_dtype, quant_algo, **kwargs): diff --git a/examples/diffusers/quantization/onnx_utils/export.py b/examples/diffusers/quantization/onnx_utils/export.py index 5da795f0f48..246dd1a9881 100644 --- a/examples/diffusers/quantization/onnx_utils/export.py +++ b/examples/diffusers/quantization/onnx_utils/export.py @@ -51,8 +51,6 @@ from modelopt.torch.quantization.export_onnx import configure_linear_module_onnx_quantizers from modelopt.torch.utils import torch_to -from .fp8_onnx_graphsurgeon import convert_zp_fp8 - MODEL_ID_TO_DYNAMIC_AXES = { "sdxl-1.0": { "sample": {0: "batch_size", 1: "num_channels", 2: "height", 3: "width"}, @@ -124,18 +122,6 @@ def flux_convert_rope_weight_type(onnx_graph): return gs.export_onnx(graph) -def generate_fp8_scales(backbone): - # temporary solution due to a known bug in torch.onnx._dynamo_export - for _, module in backbone.named_modules(): - if isinstance(module, (torch.nn.Linear, torch.nn.Conv2d)) and ( - hasattr(module.input_quantizer, "_amax") and module.input_quantizer is not None - ): - module.input_quantizer._num_bits = 8 - module.weight_quantizer._num_bits = 8 - module.input_quantizer._amax = module.input_quantizer._amax * (127 / 448.0) - module.weight_quantizer._amax = module.weight_quantizer._amax * (127 / 448.0) - - def _gen_dummy_inp_and_dyn_shapes_sdxl(backbone, min_bs=1, opt_bs=1): assert isinstance(backbone, UNet2DConditionModel) or isinstance( backbone._orig_mod, UNet2DConditionModel @@ -469,7 +455,6 @@ def modelopt_export_sd(backbone, onnx_dir, model_name, precision): tmp_subfolder = tempfile.mkdtemp(prefix="myapp_") tmp_output = Path(f"{tmp_subfolder}/{model_file_name}") q_output = Path(f"{onnx_dir}/{model_file_name}") - quantizer_context = ( configure_linear_module_onnx_quantizers(backbone) if precision == "fp4" else nullcontext() ) @@ -536,16 +521,8 @@ def modelopt_export_sd(backbone, onnx_dir, model_name, precision): ) print(f"Saved at {tmp_output}") onnx_model = onnx.load(str(tmp_output), load_external_data=True) - if precision == "fp8": - if not model_name.startswith("flux"): - graph = gs.import_onnx(onnx_model) - graph.cleanup().toposort() - onnx_model = gs.export_onnx(graph) - onnx_model = convert_zp_fp8(onnx_model) - graph = gs.import_onnx(onnx_model) - onnx_model = gs.export_onnx(graph.cleanup()) - else: - flux_convert_rope_weight_type(onnx_model) + if precision == "fp8" and model_name.startswith("flux"): + onnx_model = flux_convert_rope_weight_type(onnx_model) if precision == "fp4": onnx_model = NVFP4QuantExporter.process_model(onnx_model) save_onnx(onnx_model, q_output) diff --git a/examples/diffusers/quantization/onnx_utils/fp8_onnx_graphsurgeon.py b/examples/diffusers/quantization/onnx_utils/fp8_onnx_graphsurgeon.py index 7194e672635..7904d27e0a3 100644 --- a/examples/diffusers/quantization/onnx_utils/fp8_onnx_graphsurgeon.py +++ b/examples/diffusers/quantization/onnx_utils/fp8_onnx_graphsurgeon.py @@ -97,29 +97,6 @@ def insert_cast(graph, input_tensor, attrs): next_node.inputs[idx] = output_tensor -def convert_zp_fp8(onnx_graph): - """ - Convert Q/DQ zero datatype from INT8 to FP8. - We use this WAR because FP8 Conv cannot be exported to ONNX directly. - The workaround is to first convert the FP8 QDQs into INT8 QDQs, - then modify the ONNX model afterward to change those INT8 QDQs back into FP8 QDQs. - """ - # Find all zero constant nodes - qdq_zero_nodes = set() - for node in onnx_graph.graph.node: - if node.op_type == "QuantizeLinear" and len(node.input) > 2: - qdq_zero_nodes.add(node.input[2]) - - print(f"[WAR], found {len(qdq_zero_nodes)} INT8 QDQ pairs, you can ignore this message..") - - # Convert zero point datatype from INT8 to FP8. - for node in onnx_graph.graph.node: - if node.output[0] in qdq_zero_nodes: - node.attribute[0].t.data_type = onnx.TensorProto.FLOAT8E4M3FN - - return onnx_graph - - def cast_resize_io(graph): """ After all activations and weights are converted to fp16, we will diff --git a/examples/diffusers/quantization/quantize.py b/examples/diffusers/quantization/quantize.py index 1d71c088652..9262adbaac1 100644 --- a/examples/diffusers/quantization/quantize.py +++ b/examples/diffusers/quantization/quantize.py @@ -27,6 +27,7 @@ FP8_DEFAULT_CONFIG, INT8_DEFAULT_CONFIG, NVFP4_DEFAULT_CONFIG, + NVFP4_FP8_CONV_CONFIG, NVFP4_FP8_MHA_CONFIG, reset_set_int8_config, set_quant_config_attr, @@ -55,6 +56,9 @@ import modelopt.torch.opt as mto import modelopt.torch.quantization as mtq from modelopt.torch.export import export_hf_checkpoint +from modelopt.torch.quantization.nn import TensorQuantizer + +_SDXL_MODEL_TYPES = (ModelType.SDXL_BASE, ModelType.SDXL_TURBO) def setup_logging(verbose: bool = False) -> logging.Logger: @@ -130,7 +134,9 @@ def get_quant_config(self, n_steps: int, backbone: torch.nn.Module) -> Any: elif self.config.format == QuantFormat.FP8: base_cfg = FP8_DEFAULT_CONFIG elif self.config.format == QuantFormat.FP4: - if self.model_config.model_type.value.startswith("flux"): + if self.model_config.model_type in _SDXL_MODEL_TYPES: + base_cfg = NVFP4_FP8_CONV_CONFIG + elif self.model_config.model_type.value.startswith("flux"): base_cfg = NVFP4_FP8_MHA_CONFIG else: base_cfg = NVFP4_DEFAULT_CONFIG @@ -271,23 +277,6 @@ def __init__( self.logger = logger self.pipeline_manager = pipeline_manager - def _has_conv_layers(self, model: torch.nn.Module) -> bool: - """ - Check if the model contains any convolutional layers. - - Args: - model: Model to check - - Returns: - True if model contains Conv layers, False otherwise - """ - for module in model.modules(): - if isinstance(module, (torch.nn.Conv1d, torch.nn.Conv2d, torch.nn.Conv3d)) and ( - module.input_quantizer.is_enabled or module.weight_quantizer.is_enabled - ): - return True - return False - def save_checkpoint( self, backbone: torch.nn.Module, @@ -335,15 +324,10 @@ def export_onnx( # Deferred: the ONNX stack (onnx, onnx_graphsurgeon, ...) is only needed # for --onnx-dir exports; HF-checkpoint-only runs must not require it. - from onnx_utils.export import generate_fp8_scales, modelopt_export_sd + from onnx_utils.export import modelopt_export_sd self.logger.info(f"Starting ONNX export to {self.config.onnx_dir}") - if quant_format == QuantFormat.FP8 and self._has_conv_layers(backbone): - self.logger.info( - "Detected quantizing conv layers in backbone. Generating FP8 scales..." - ) - generate_fp8_scales(backbone) self.logger.info("Preparing models for export...") pipe.to("cpu") torch.cuda.empty_cache() @@ -457,7 +441,7 @@ def create_argument_parser() -> argparse.ArgumentParser: %(prog)s --model ltx-video-dev --format fp8 --batch-size 1 --calib-size 32 --ltx-skip-upsampler # Restore and export a previously quantized model - %(prog)s --model flux-schnell --restore-from checkpoint.pt --onnx-dir ./exports/ + %(prog)s --model flux-schnell --restore-from ./checkpoints/ --onnx-dir ./exports/ """, ) model_group = parser.add_argument_group("Model Configuration") @@ -586,7 +570,9 @@ def create_argument_parser() -> argparse.ArgumentParser: help="Directory for HuggingFace checkpoint export", ) export_group.add_argument( - "--restore-from", type=str, help="Path to restore from previous checkpoint" + "--restore-from", + type=str, + help="Checkpoint directory; quantization format and MHA policy are restored automatically", ) export_group.add_argument( "--trt-high-precision-dtype", @@ -600,6 +586,25 @@ def create_argument_parser() -> argparse.ArgumentParser: return parser +def _infer_restored_quantization_format( + backbones: list[tuple[str, torch.nn.Module]], +) -> QuantFormat: + has_nvfp4 = False + has_fp8 = False + + for _, backbone in backbones: + for module in backbone.modules(): + if isinstance(module, TensorQuantizer) and module.is_enabled: + has_nvfp4 |= module.is_nvfp4_dynamic or module.is_nvfp4_static + has_fp8 |= module.is_fp8 + + if has_nvfp4: + return QuantFormat.FP4 + if has_fp8: + return QuantFormat.FP8 + return QuantFormat.INT8 + + def main() -> None: from diffusers.models.normalization import RMSNorm as DiffuserRMSNorm @@ -674,9 +679,9 @@ def main() -> None: ) logger.info("Validating configurations...") - quant_config.validate() export_config.validate() if not export_config.restore_from: + quant_config.validate() calib_config.validate() pipeline_manager = PipelineManager(model_config, logger) @@ -685,8 +690,12 @@ def main() -> None: export_manager = ExportManager(export_config, logger, pipeline_manager) - if export_config.restore_from and export_config.restore_from.exists(): + if export_config.restore_from: export_manager.restore_checkpoint() + quant_config.format = _infer_restored_quantization_format( + list(pipeline_manager.iter_backbones()) + ) + logger.info(f"Detected restored quantization format: {quant_config.format.value}") else: logger.info("Initializing calibration...") @@ -716,11 +725,12 @@ def forward_loop(mod): mtq.compress(backbone) logger.info(f"{backbone_name} compression completed") - # For VAE backbones, skip check_conv_and_mha — the whole point - # of VAE quantization is to quantize Conv layers. if backbone_name not in ("video_decoder", "vae"): check_conv_and_mha( - backbone, quant_config.format == QuantFormat.FP4, quant_config.quantize_mha + backbone, + quant_config.format == QuantFormat.FP4 + and model_config.model_type not in _SDXL_MODEL_TYPES, + quant_config.quantize_mha, ) export_manager.save_checkpoint(backbone, backbone_name) diff --git a/examples/diffusers/quantization/utils.py b/examples/diffusers/quantization/utils.py index c3cfdcd5cdd..b7a79e49e70 100644 --- a/examples/diffusers/quantization/utils.py +++ b/examples/diffusers/quantization/utils.py @@ -64,11 +64,8 @@ def check_conv_and_mha(backbone, if_fp4, quantize_mha): ): if hasattr(module, attr): getattr(module, attr).disable() - setattr(module, "_disable_fp8_mha", True) print(f"Disabled Attention layer quantization for layer {name}") - else: - setattr(module, "_disable_fp8_mha", False) def filter_func_ltx_video(name: str) -> bool: diff --git a/examples/hf_ptq/hf_ptq.py b/examples/hf_ptq/hf_ptq.py index a56a62b54b5..5cd3fc82fdc 100755 --- a/examples/hf_ptq/hf_ptq.py +++ b/examples/hf_ptq/hf_ptq.py @@ -1020,11 +1020,18 @@ def post_quantize( pass elif model_type != "llama4" and not is_nemotron_vl_model: # Our fake quantizer may not be fully compatible with torch.compile. - generated_ids_after_ptq = full_model.generate( - preview_input_ids, - attention_mask=preview_attention_mask, - max_new_tokens=100, - ) + # 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. + try: + generated_ids_after_ptq = full_model.generate( + preview_input_ids, + attention_mask=preview_attention_mask, + max_new_tokens=100, + ) + except Exception as e: + warnings.warn(f"Post-quantization generation sanity check failed, skipping it: {e}") elif is_nemotron_vl_model and tokenizer is not None: generated_ids_after_ptq = run_nemotron_vl_preview( full_model, diff --git a/examples/llm_sparsity/weight_sparsity/README.md b/examples/llm_sparsity/weight_sparsity/README.md index 97563aff007..42bac030ce9 100644 --- a/examples/llm_sparsity/weight_sparsity/README.md +++ b/examples/llm_sparsity/weight_sparsity/README.md @@ -135,28 +135,39 @@ ROUGE scores: {'rouge1': 42.174, 'rouge2': 19.2724, 'rougeL': 28.6989, 'rougeLsu Please refer to [link](../llm_eval/README.md#Evaluation-scripts-for-LLM-tasks) for more details of how to evaluate the sparsified models on other benchmarks, such as MMLU and HumanEval. -## Export TensorRT-LLM Checkpoint +## Export and Serve a Hugging Face Checkpoint -To export the PTS pytorch model to a TensorRT-LLM checkpoint, run the following command: +Export the PTS model to a Hugging Face checkpoint: ```sh -python export_trtllm_ckpt.py --model_name_or_path meta-llama/Llama-2-7b-hf \ +python export_hf_ckpt.py --model_name_or_path meta-llama/Llama-2-7b-hf \ --model_max_length 1024 \ --dtype fp16 \ - --modelopt_restore_path saved_models_Llama-2-7b-hf_sparsegpt_tp1_pp1/pts_modelopt_state.pth \ - --output_dir saved_models_Llama-2-7b-hf_sparsegpt_tp1_pp1/trtllm/ckpt_pts + --modelopt_restore_path saved_models_Llama-2-7b-hf_sparsegpt_tp1_pp1/pts/pts_modelopt_state.pth \ + --output_dir saved_models_Llama-2-7b-hf_sparsegpt_tp1_pp1/hf/pts ``` -To export the finetuned pytorch model to a TensorRT-LLM checkpoint, run the following command: +Export the finetuned model in the same way: ```sh -python export_trtllm_ckpt.py --model_name_or_path meta-llama/Llama-2-7b-hf \ +python export_hf_ckpt.py --model_name_or_path meta-llama/Llama-2-7b-hf \ --model_max_length 1024 \ --dtype fp16 \ --modelopt_restore_path saved_models_Llama-2-7b-hf_sparsegpt_tp1_pp1/finetuned/finetuned_modelopt_state.pth \ - --output_dir saved_models_Llama-2-7b-hf_sparsegpt_tp1_pp1/trtllm/ckpt_finetuned + --output_dir saved_models_Llama-2-7b-hf_sparsegpt_tp1_pp1/hf/finetuned ``` -## Build TensorRT-LLM Engine +The output is a standard Hugging Face checkpoint, including the tokenizer and sparse model weights. +TensorRT-LLM loads it directly; no checkpoint conversion or engine build is required. -For guidance on how to build TensorRT-LLM engines, please refer to [link](https://nvidia.github.io/TensorRT-LLM/commands/trtllm-build.html#trtllm-build) and use the `--weight_sparsity` flag. +Serve the exported checkpoint and set the deployment parallelism at startup: + +```sh +trtllm-serve saved_models_Llama-2-7b-hf_sparsegpt_tp1_pp1/hf/pts \ + --tp_size 1 \ + --pp_size 1 \ + --host 0.0.0.0 \ + --port 8000 +``` + +After the server starts, send requests to the OpenAI-compatible `/v1/chat/completions` endpoint. diff --git a/examples/llm_sparsity/weight_sparsity/export_trtllm_ckpt.py b/examples/llm_sparsity/weight_sparsity/export_hf_ckpt.py similarity index 84% rename from examples/llm_sparsity/weight_sparsity/export_trtllm_ckpt.py rename to examples/llm_sparsity/weight_sparsity/export_hf_ckpt.py index 2cf7ca3a7ac..7ef22c52eca 100644 --- a/examples/llm_sparsity/weight_sparsity/export_trtllm_ckpt.py +++ b/examples/llm_sparsity/weight_sparsity/export_hf_ckpt.py @@ -23,7 +23,7 @@ import modelopt.torch.opt as mto import modelopt.torch.sparsity as mts -from modelopt.torch.export import export_tensorrt_llm_checkpoint, get_model_type +from modelopt.torch.export import export_hf_checkpoint DEFAULT_PAD_TOKEN = "[PAD]" @@ -103,8 +103,6 @@ def main(args): model=model, ) - # Export the sparse model to trt-llm checkpoint - model_type = get_model_type(model) if args.modelopt_restore_path: print(f"Loading sparsity state from {args.modelopt_restore_path}") if not os.path.isfile(args.modelopt_restore_path): @@ -112,17 +110,11 @@ def main(args): mto.restore(model, args.modelopt_restore_path) - print(f"Exporting trt-llm checkpoint to {args.output_dir}") + print(f"Exporting Hugging Face checkpoint to {args.output_dir}") with torch.inference_mode(): model = mts.export(model) - export_tensorrt_llm_checkpoint( - model, - model_type, - torch.float16, - export_dir=args.output_dir, - inference_tensor_parallel=args.inference_tensor_parallel, - inference_pipeline_parallel=args.inference_pipeline_parallel, - ) + export_hf_checkpoint(model, export_dir=args.output_dir) + tokenizer.save_pretrained(args.output_dir) if __name__ == "__main__": @@ -140,22 +132,11 @@ def main(args): parser.add_argument("--dtype", help="Model data type.", default="fp16") parser.add_argument( "--model_max_length", + type=int, default=2048, help="Maximum sequence length. Sequences will be right padded (and possibly truncated).", ) parser.add_argument("--output_dir", default="output_dir") - parser.add_argument( - "--inference_tensor_parallel", - help="Number of tensor parallel groups for inference.", - type=int, - default=1, - ) - parser.add_argument( - "--inference_pipeline_parallel", - help="Number of pipeline parallel groups for inference.", - type=int, - default=1, - ) parser.add_argument( "--trust_remote_code", help="Set trust_remote_code for Huggingface models and tokenizers", diff --git a/examples/onnx_ptq/README.md b/examples/onnx_ptq/README.md index de983b409bd..65c62f0655a 100644 --- a/examples/onnx_ptq/README.md +++ b/examples/onnx_ptq/README.md @@ -111,7 +111,7 @@ The model can be quantized as an FP8, INT8 or INT4 model using either the CLI or > *For NVFP4 and MXFP8 ONNX, see the [PyTorch to ONNX example](../torch_onnx/).* -> *Minimum opset requirements: int8 (13+), fp8 (21+), int4 (21+). ModelOpt will automatically upgrade lower opset versions to meet these requirements.* +> *Minimum opset requirements: int8 (19+), fp8 (19+), int4 (21+). ModelOpt will automatically upgrade lower opset versions to meet these requirements.* #### Option 1: Command-line interface @@ -119,7 +119,7 @@ The model can be quantized as an FP8, INT8 or INT4 model using either the CLI or python -m modelopt.onnx.quantization \ --onnx_path=vit_base_patch16_224.onnx \ --quantize_mode= \ - --calibration_data=calib.npy \ + --calibration_data_path=calib.npy \ --calibration_method= \ --output_path=vit_base_patch16_224.quant.onnx ``` @@ -127,12 +127,14 @@ python -m modelopt.onnx.quantization \ #### Option 2: Python API ```python +import numpy as np + from modelopt.onnx.quantization import quantize quantize( onnx_path="vit_base_patch16_224.onnx", quantize_mode="int8", # fp8, int8, int4 etc. - calibration_data="calib.npy", + calibration_data=np.load("calib.npy"), calibration_method="max", # max, entropy, awq_clip, rtn_dq etc. output_path="vit_base_patch16_224.quant.onnx", ) @@ -198,7 +200,7 @@ To enable per node calibration, add the `--calibrate_per_node` flag to your quan python -m modelopt.onnx.quantization \ --onnx_path=vit_base_patch16_224.onnx \ --quantize_mode= \ - --calibration_data=calib.npy \ + --calibration_data_path=calib.npy \ --calibrate_per_node \ --output_path=vit_base_patch16_224.quant.onnx ``` @@ -253,9 +255,9 @@ To access this feature in the ONNX quantization workflow, simply add `--autotune ```bash python -m modelopt.onnx.quantization \ --onnx_path=vit_base_patch16_224.onnx \ - --quantize_mode= \ - --calibration_data=calib.npy \ - --calibration_method= \ + --quantize_mode= \ + --calibration_data_path=calib.npy \ + --calibration_method= \ --output_path=vit_base_patch16_224.quant.onnx \ --autotune= ``` diff --git a/examples/speculative_decoding/collect_hidden_states/compute_hidden_states_vllm.py b/examples/speculative_decoding/collect_hidden_states/compute_hidden_states_vllm.py index 77441f8f858..c232852a4c8 100644 --- a/examples/speculative_decoding/collect_hidden_states/compute_hidden_states_vllm.py +++ b/examples/speculative_decoding/collect_hidden_states/compute_hidden_states_vllm.py @@ -55,9 +55,10 @@ def _resolve_aux_layers_standalone( This dump runs in a stock vLLM container. ``common.resolve_aux_layers`` resolves the 'dflash'/'eagle' presets by importing ``modelopt.torch.speculative.plugins`` — which pulls in the full ``modelopt.torch`` init chain (omegaconf, etc.) that the vLLM - container does not have, so the import fails. Resolve the 'dflash' preset inline - (mirroring ``modeling_dflash.build_target_layer_ids`` for ``num_draft`` draft layers) - and accept an explicit comma-separated int list. ``num_draft`` MUST match the recipe's + container does not have, so the import fails. Resolve both presets inline -- 'eagle' + mirrors ``hf_eagle.default_eagle_aux_layer_ids``, 'dflash' mirrors + ``modeling_dflash.build_target_layer_ids`` for ``num_draft`` draft layers -- and + accept an explicit comma-separated int list. ``num_draft`` MUST match the recipe's ``dflash.dflash_architecture_config.num_hidden_layers`` (pass --num-draft-layers) or the dumped aux layers silently mis-align with what the draft consumes at training time. Keep in sync with modelopt. @@ -65,15 +66,39 @@ def _resolve_aux_layers_standalone( TODO: drop this once ``common.resolve_aux_layers`` is decoupled from the heavy ``modelopt.torch`` import chain so it can be reused directly in a vLLM container. """ + # One literal: a future preset means editing this message once, not twice. The bug + # this guards against was precisely a copy drifting from its original. + unsupported = ( + f"--aux-layers={aux_layers!r}: in the stock vLLM container (no modelopt) only the " + "'eagle' / 'dflash' presets or an explicit comma-separated layer-id list " + "are supported." + ) spec = aux_layers.strip().lower() + if spec == "eagle": + # Mirrors hf_eagle.default_eagle_aux_layer_ids: three layers near the start, + # middle, and end of the stack. This is add_aux_layers_args' default, so the + # dump must resolve it without modelopt or the documented invocation fails. + return sorted({1, max(0, num_hidden_layers // 2 - 1), max(0, num_hidden_layers - 4)}) if spec == "dflash": + # build_target_layer_ids raises here; without the same guard the interpolation + # below silently collapses to a short deduplicated list (e.g. [1] for a 4-layer + # target), dumping fewer aux layers than the draft consumes. + if num_hidden_layers < num_draft: + raise ValueError( + f"num_target_layers ({num_hidden_layers}) must be >= num_draft_layers ({num_draft})" + ) if num_draft == 1: return [num_hidden_layers // 2] start = min(1, num_hidden_layers - 1) end = max(start, num_hidden_layers - 3) span = end - start return sorted({round(start + (i * span) / (num_draft - 1)) for i in range(num_draft)}) - ids = sorted({int(t) for t in aux_layers.split(",") if t.strip()}) + try: + ids = sorted({int(t) for t in aux_layers.split(",") if t.strip()}) + except ValueError as exc: + # An unrecognised preset name would otherwise surface as a bare + # "invalid literal for int()", which hides what the caller should pass. + raise ValueError(unsupported) from exc # Match the shared helper's contract: ids must be valid layer indices. out_of_range = [i for i in ids if not 0 <= i < num_hidden_layers] if out_of_range: @@ -82,10 +107,7 @@ def _resolve_aux_layers_standalone( f"for a {num_hidden_layers}-layer model." ) if not ids: - raise ValueError( - f"--aux-layers={aux_layers!r}: in the stock vLLM container (no modelopt) only the " - "'dflash' preset or an explicit comma-separated layer-id list are supported." - ) + raise ValueError(unsupported) return ids diff --git a/examples/speculative_decoding/main.py b/examples/speculative_decoding/main.py index 46848203606..8e5acd91544 100644 --- a/examples/speculative_decoding/main.py +++ b/examples/speculative_decoding/main.py @@ -217,7 +217,10 @@ def train(): assert checkpoint is not None # guaranteed by checkpoint_is_hf with patch_transformers5_params_loading(): model = load_vlm_or_llm( - checkpoint, dtype="auto", trust_remote_code=recipe.model.trust_remote_code + checkpoint, + dtype="auto", + trust_remote_code=recipe.model.trust_remote_code, + config_overrides=recipe.model.config_overrides, ) tokenizer = transformers.AutoTokenizer.from_pretrained( checkpoint, trust_remote_code=recipe.model.trust_remote_code @@ -242,6 +245,7 @@ def train(): dtype="auto", device_map="cpu", trust_remote_code=recipe.model.trust_remote_code, + config_overrides=recipe.model.config_overrides, ) tokenizer = transformers.AutoTokenizer.from_pretrained( model_name_or_path, diff --git a/examples/speculative_decoding/scripts/ar_validate.py b/examples/speculative_decoding/scripts/ar_validate.py index 5699c480b7b..ccfb47dc085 100644 --- a/examples/speculative_decoding/scripts/ar_validate.py +++ b/examples/speculative_decoding/scripts/ar_validate.py @@ -28,7 +28,11 @@ import modelopt.torch.opt as mto from modelopt.torch.speculative.plugins.hf_eagle import HFARValidation -from modelopt.torch.speculative.utils import load_vlm_or_llm +from modelopt.torch.speculative.utils import ( + CONFIG_OVERRIDES_HELP, + load_vlm_or_llm, + parse_config_overrides, +) mto.enable_huggingface_checkpointing() @@ -100,11 +104,22 @@ def main(): default=None, help="Error if AR is below this threshold.", ) + parser.add_argument( + "--config_overrides", + type=str, + default=None, + help=CONFIG_OVERRIDES_HELP, + ) args = parser.parse_args() + config_overrides = parse_config_overrides(args.config_overrides) + accelerator = Accelerator() model = load_vlm_or_llm( - args.model_path, device_map="auto", trust_remote_code=args.trust_remote_code + args.model_path, + device_map="auto", + trust_remote_code=args.trust_remote_code, + config_overrides=config_overrides, ) tokenizer = AutoTokenizer.from_pretrained( args.model_path, trust_remote_code=args.trust_remote_code diff --git a/examples/speculative_decoding/scripts/export_hf_checkpoint.py b/examples/speculative_decoding/scripts/export_hf_checkpoint.py index cee2e45d0eb..2b86a47004b 100644 --- a/examples/speculative_decoding/scripts/export_hf_checkpoint.py +++ b/examples/speculative_decoding/scripts/export_hf_checkpoint.py @@ -21,7 +21,11 @@ import modelopt.torch.opt as mto from modelopt.torch.export import export_speculative_decoding -from modelopt.torch.speculative.utils import load_vlm_or_llm +from modelopt.torch.speculative.utils import ( + CONFIG_OVERRIDES_HELP, + load_vlm_or_llm, + parse_config_overrides, +) def parse_args(): @@ -33,13 +37,24 @@ def parse_args(): parser.add_argument( "--export_path", type=str, default="Destination directory for exported files." ) + parser.add_argument( + "--config_overrides", + type=str, + default=None, + help=CONFIG_OVERRIDES_HELP, + ) return parser.parse_args() mto.enable_huggingface_checkpointing() args = parse_args() -model = load_vlm_or_llm(args.model_path, dtype="auto", trust_remote_code=args.trust_remote_code) +model = load_vlm_or_llm( + args.model_path, + dtype="auto", + trust_remote_code=args.trust_remote_code, + config_overrides=parse_config_overrides(args.config_overrides), +) model.eval() with torch.inference_mode(): export_speculative_decoding( diff --git a/examples/speculative_decoding/scripts/merge_lora.py b/examples/speculative_decoding/scripts/merge_lora.py index 25d311393b5..4580e166832 100644 --- a/examples/speculative_decoding/scripts/merge_lora.py +++ b/examples/speculative_decoding/scripts/merge_lora.py @@ -31,7 +31,13 @@ from pathlib import Path from safetensors.torch import load_file -from transformers import AutoModelForCausalLM, AutoTokenizer +from transformers import AutoTokenizer + +from modelopt.torch.speculative.utils import ( + CONFIG_OVERRIDES_HELP, + load_vlm_or_llm, + parse_config_overrides, +) def parse_args(): @@ -61,11 +67,18 @@ def parse_args(): action="store_true", help="Allow loading models that define custom code on the HF Hub. Off by default.", ) + parser.add_argument( + "--config_overrides", + type=str, + default=None, + help=CONFIG_OVERRIDES_HELP, + ) return parser.parse_args() def main(): args = parse_args() + config_overrides = parse_config_overrides(args.config_overrides) lora_dir = Path(args.exported_lora_dir) # Verify exported files exist (standard peft naming) @@ -81,13 +94,21 @@ def main(): print(f"Loaded {len(lora_sd)} LoRA tensors from {lora_dir}") print(f" Sample keys: {list(lora_sd.keys())[:4]}") - # Load the original base model + # Load the original base model. + # + # Use load_vlm_or_llm rather than AutoModelForCausalLM directly: it falls back to + # AutoModelForCausalLM for plain LLMs (same dtype/device_map, so unchanged behavior), but also + # handles VLMs and registers/loads architectures the Auto* maps don't cover. Cosmos3 is the + # motivating case -- the transformers-cosmos3 plugin registers only the `cosmos3_omni` config, + # never a model under Auto*, so AutoModelForCausalLM raises KeyError('cosmos3_omni') no matter + # what is imported. print(f"Loading base model from {args.base_model_path}...") - model = AutoModelForCausalLM.from_pretrained( + model = load_vlm_or_llm( args.base_model_path, - torch_dtype="auto", + dtype="auto", device_map="cpu", trust_remote_code=args.trust_remote_code, + config_overrides=config_overrides, ) tokenizer = AutoTokenizer.from_pretrained( args.base_model_path, trust_remote_code=args.trust_remote_code @@ -135,9 +156,18 @@ def main(): # Since LoRA only changes weights — not architecture — the original config is correct. import shutil + # ...but only when the loaded config still matches the base's. With --config_overrides the + # weights were built from corrected dims, so copying the uncorrected base config back would + # leave config.json disagreeing with model.safetensors and force every downstream reader to + # re-supply the same overrides. base_config = Path(args.base_model_path) / "config.json" output_config = Path(args.output_path) / "config.json" - if base_config.exists(): + if config_overrides: + print( + " Keeping the saved config.json (config_overrides were applied, so the original " + "base config would not match the merged weights)" + ) + elif base_config.exists(): shutil.copy2(str(base_config), str(output_config)) print(f" Restored original config.json from {base_config}") diff --git a/examples/vllm_serve/Dockerfile b/examples/vllm_serve/Dockerfile index 5fe0799c427..a4c26ab1a7b 100644 --- a/examples/vllm_serve/Dockerfile +++ b/examples/vllm_serve/Dockerfile @@ -1,4 +1,5 @@ -FROM vllm/vllm-openai:v0.26.0 +ARG VLLM_VERSION=0.28.0 +FROM vllm/vllm-openai:v${VLLM_VERSION} # Set environment variables ENV PIP_NO_CACHE_DIR=off \ diff --git a/examples/vllm_serve/README.md b/examples/vllm_serve/README.md index fc4e8a0ebcc..bc2b0fb53f3 100644 --- a/examples/vllm_serve/README.md +++ b/examples/vllm_serve/README.md @@ -4,17 +4,35 @@ This is a simple example to demonstrate calibrating and serving ModelOpt fakequa Compared with realquant, fakequant is 2-5x slower, but doesn't require dedicated kernel support and facilitates research. -The general fakequant example is tested with vLLM 0.9.0, 0.19.1, and 0.26.0. The compact -NVFP4 attention worker documented below requires vLLM 0.15.0 or newer. +The general fakequant example is tested with vLLM 0.9.0, 0.19.1, 0.26.0, and 0.28.0. The +compact NVFP4 attention worker documented below requires vLLM 0.15.0 or newer. ## Prepare environment -Follow the following instruction to build a docker environment, or install vllm with pip. +Use the Dockerfile to build an environment with vLLM 0.28.0: ```bash -docker build -f examples/vllm_serve/Dockerfile -t vllm-modelopt . +docker build -f examples/vllm_serve/Dockerfile -t vllm-modelopt:v0.28.0 . ``` +To build the same environment with another tested vLLM release, override `VLLM_VERSION`: + +```bash +docker build --build-arg VLLM_VERSION=0.26.0 \ + -f examples/vllm_serve/Dockerfile -t vllm-modelopt:v0.26.0 . +``` + +For a direct installation from the ModelOpt repository root, install the tested vLLM +release and the ModelOpt extras used by this example: + +```bash +python3 -m pip install "vllm==0.28.0" +python3 -m pip install -e ".[all,mlflow]" +``` + +See the [ModelOpt installation guide](../../docs/source/getting_started/_installation_for_Linux.rst) +for details about installing partial dependency sets. + ## Calibrate and serve fake quant model in vLLM Step 1: Configure quantization settings. @@ -39,6 +57,18 @@ Step 2: Run the following command, with all supported flag as `vllm serve`: python vllm_serve_fakequant.py -tp 8 --host 0.0.0.0 --port 8000 ``` +Hybrid attention/Mamba models such as Nemotron 3 Nano are supported on vLLM 0.26.0 and +0.28.0. For example, calibrate and serve with NVFP4 KV-cache fakequant as follows: + +```bash +KV_QUANT_CFG=NVFP4_KV_CFG QUANT_CALIB_SIZE=512 \ + python vllm_serve_fakequant.py -tp 8 \ + --max-model-len 8192 --enforce-eager --host 0.0.0.0 --port 8000 +``` + +Calibration uses dedicated scratch KV-cache blocks, so reducing `--max-num-batched-tokens` +is not required to avoid NaNs. + For vLLM versions that expose `--moe-backend`, this launcher defaults to `--moe-backend triton`. ModelOpt expert fakequant needs a decomposed MoE backend so both expert GEMMs are visible during calibration. diff --git a/examples/vllm_serve/vllm_ptq_utils.py b/examples/vllm_serve/vllm_ptq_utils.py index 709d6532fb3..ca9e276b9b4 100644 --- a/examples/vllm_serve/vllm_ptq_utils.py +++ b/examples/vllm_serve/vllm_ptq_utils.py @@ -36,8 +36,131 @@ def _create_new_data_cls(data_cls, **kwargs): return data_cls(**filtered_kwargs) +def _get_calibration_block_count( + model_runner: Any, +) -> Callable[[int, Any], int] | None: + """Return the block reservation policy supported by the installed vLLM.""" + vllm_config = model_runner.vllm_config + + try: + from vllm.v1.worker.gpu.warmup import _reserved_block_count + except ImportError: + try: + from vllm.utils.math_utils import cdiv + from vllm.v1.kv_cache_interface import CrossAttentionSpec, MambaSpec + except ImportError: + return None + + def block_count(num_tokens: int, kv_cache_spec: Any) -> int: + """Calculate the vLLM 0.26 warmup block reservation.""" + # vLLM 0.26's warmup reservation policy. + if isinstance(kv_cache_spec, CrossAttentionSpec): + num_tokens = 0 + num_blocks = cdiv(num_tokens, kv_cache_spec.block_size) + if isinstance(kv_cache_spec, MambaSpec) and kv_cache_spec.mamba_cache_mode == "align": + num_blocks += kv_cache_spec.num_speculative_blocks + return num_blocks + + else: + + def block_count(num_tokens: int, kv_cache_spec: Any) -> int: + """Calculate the current vLLM warmup block reservation.""" + # Calibration runs before model_state is initialized, so call the + # underlying reservation policy rather than _warmup_block_counter. + return _reserved_block_count( + num_tokens, + kv_cache_spec, + num_lookahead_tokens=vllm_config.num_lookahead_tokens, + max_model_len=model_runner.max_model_len, + max_encoder_len=0, + ) + + return block_count + + +def _allocate_calibration_blocks( + self: Any, sequence_lengths: list[int] +) -> tuple[list[tuple[list[int], ...]], list[int] | None]: + """Allocate scheduler-compatible scratch blocks for calibration requests. + + vLLM 0.28 treats block 0 as the null block. Its GPU runner expects real block + tables for hybrid attention/Mamba models, even for one-shot prefill requests. + Use vLLM's warmup reservation policy so this stays aligned with each cache + group's KVCacheSpec. + """ + kv_cache_config = self.model_runner.kv_cache_config + kv_cache_groups = kv_cache_config.kv_cache_groups + block_count = _get_calibration_block_count(self.model_runner) + + if block_count is None: + warnings.warn( + "vLLM warmup block reservation helpers were not found; falling back to " + "empty block tables. Hybrid attention/Mamba models may produce NaNs.", + stacklevel=2, + ) + return [tuple([] for _ in kv_cache_groups) for _ in sequence_lengths], None + + next_block_id = 1 # Block 0 is reserved as the null block. + block_ids_batch: list[tuple[list[int], ...]] = [] + allocated_block_ids: list[int] = [] + + for sequence_length in sequence_lengths: + request_block_ids = [] + for group in kv_cache_groups: + num_blocks = block_count(sequence_length, group.kv_cache_spec) + block_ids = list(range(next_block_id, next_block_id + num_blocks)) + next_block_id += num_blocks + allocated_block_ids.extend(block_ids) + request_block_ids.append(block_ids) + block_ids_batch.append(tuple(request_block_ids)) + + if next_block_id > kv_cache_config.num_blocks: + raise RuntimeError( + "Calibration batch requires " + f"{next_block_id - 1} KV cache blocks, but only " + f"{kv_cache_config.num_blocks - 1} non-null blocks are available." + ) + + scheduler_fields = {field.name for field in dataclasses.fields(SchedulerOutput)} + if "new_block_ids_to_zero" in scheduler_fields: + blocks_to_zero = ( + allocated_block_ids if getattr(kv_cache_config, "needs_kv_cache_zeroing", False) else [] + ) + else: + blocks_to_zero = None + return block_ids_batch, blocks_to_zero + + +def _cleanup_calibration_requests( + self: Any, + cleanup_output: SchedulerOutput, + calibration_error: BaseException | None, +) -> None: + """Clean request state without hiding an active calibration error.""" + try: + # Zero-token steps return before forward/sampling, so no sample_tokens call is needed. + self.execute_model(cleanup_output) + except Exception as execute_error: + finish_requests = getattr(self.model_runner, "finish_requests", None) + if finish_requests is None: + if calibration_error is not None: + raise calibration_error from execute_error + raise + + try: + finish_requests(cleanup_output) + except Exception as finish_error: + if calibration_error is not None: + finish_error.__cause__ = execute_error + raise calibration_error from finish_error + raise finish_error from execute_error + + def calibrate_fun(calib_dataloader: DataLoader, self: Any) -> Callable[[Any], None]: + """Create a calibration loop backed by the vLLM worker scheduler.""" + def calibrate_loop(model: Any) -> None: + """Calibrate the model with batches submitted through the scheduler.""" for batch_idx, batch in tqdm(enumerate(calib_dataloader)): input_ids_batch = batch["input_ids"] @@ -56,7 +179,9 @@ def calibrate_loop(model: Any) -> None: input_ids_list_batch = [input_ids_list_batch] num_groups = len(self.model_runner.kv_cache_config.kv_cache_groups) - empty_block_ids = tuple([] for _ in range(num_groups)) + block_ids_batch, new_block_ids_to_zero = _allocate_calibration_blocks( + self, [len(input_ids) for input_ids in input_ids_list_batch] + ) scheduled_new_reqs = [] num_scheduled_tokens = {} @@ -74,7 +199,7 @@ def calibrate_loop(model: Any) -> None: mm_features=[], sampling_params=SamplingParams(max_tokens=1), pooling_params=None, - block_ids=empty_block_ids, + block_ids=block_ids_batch[seq_idx], num_computed_tokens=0, lora_request=None, ) @@ -96,40 +221,36 @@ def calibrate_loop(model: Any) -> None: kv_connector_metadata=None, structured_output_request_ids={}, grammar_bitmask=None, + new_block_ids_to_zero=new_block_ids_to_zero, + ) + # Submit a zero-token scheduler step after the request has been + # registered. This is the vLLM 0.28 cleanup path and removes + # request-scoped attention/Mamba state from the persistent batch. + cleanup_output = _create_new_data_cls( + type(scheduler_output), + scheduled_new_reqs=[], + scheduled_cached_reqs=CachedRequestData.make_empty(), + num_scheduled_tokens={}, + total_num_scheduled_tokens=0, + scheduled_spec_decode_tokens={}, + scheduled_encoder_inputs={}, + num_common_prefix_blocks=[0] * num_groups, + finished_req_ids=set(num_scheduled_tokens), + free_encoder_mm_hashes=[], + kv_connector_metadata=None, + structured_output_request_ids={}, + grammar_bitmask=None, ) try: output = self.execute_model(scheduler_output) if hasattr(self, "sample_tokens"): if output is None: # TODO: make this default when vllm <= 0.11 is outdated self.sample_tokens(None) - finally: - # finish_requests runs before add_requests inside execute_model, so - # req IDs aren't registered yet at that point — call it directly after. - # Wrap in try/except so a cleanup error never masks the original exception. - try: - if hasattr(self.model_runner, "finish_requests"): - cleanup_output = _create_new_data_cls( - type(scheduler_output), - scheduled_new_reqs=[], - scheduled_cached_reqs=scheduler_output.scheduled_cached_reqs, - num_scheduled_tokens={}, - total_num_scheduled_tokens=0, - scheduled_spec_decode_tokens={}, - scheduled_encoder_inputs={}, - num_common_prefix_blocks=scheduler_output.num_common_prefix_blocks, - finished_req_ids=set(num_scheduled_tokens.keys()), - free_encoder_mm_hashes=[], - kv_connector_metadata=None, - structured_output_request_ids={}, - grammar_bitmask=None, - ) - self.model_runner.finish_requests(cleanup_output) - else: - warnings.warn( - "model_runner.finish_requests not found; request state may leak during calibration." - ) - except Exception: - warnings.warn("Failed to clean up request state after calibration batch.") + except BaseException as calibration_error: + _cleanup_calibration_requests(self, cleanup_output, calibration_error) + raise + + _cleanup_calibration_requests(self, cleanup_output, calibration_error=None) return calibrate_loop @@ -171,6 +292,7 @@ def update_kv_cfg_for_mla(model: torch.nn.Module, kv_quant_cfg: list) -> list: def get_quant_config(quant_config: dict[str, Any], model: Any) -> dict[str, Any]: + """Resolve and merge model and KV-cache quantization configuration.""" import copy if quant_config["recipe_path"]: diff --git a/modelopt/onnx/autocast/convert.py b/modelopt/onnx/autocast/convert.py index d9bc8d68c99..e259235a911 100644 --- a/modelopt/onnx/autocast/convert.py +++ b/modelopt/onnx/autocast/convert.py @@ -23,6 +23,7 @@ nodes. """ +import os from copy import deepcopy import numpy as np @@ -108,7 +109,7 @@ def convert_to_mixed_precision( onnx.ModelProto: The converted mixed precision model. """ # Load and process model - model = onnx.load(onnx_path, load_external_data=True) + model = onnx.load(onnx_path, load_external_data=False) assert low_precision_type in ["fp16", "bf16"], "low_precision_type must be either fp16 or bf16" original_network_io_metadata = _capture_network_io_metadata(model, keep_io_types) @@ -146,6 +147,7 @@ def convert_to_mixed_precision( trt_plugins=trt_plugins, trt_plugins_precision=trt_plugins_precision, max_ir_version=LATEST_IR_VERSION_SUPPORTED_BY_ORT, + onnx_path=onnx_path, ) graph_sanitizer.sanitize() model = graph_sanitizer.model @@ -154,6 +156,9 @@ def convert_to_mixed_precision( # as an exception (triggering infer_types' standalone type-inference fallback) instead of # silently leaving tensors untyped, which would break later type lookups. model = onnx_utils.infer_types(model, use_standalone_type_inference, strict_mode=True) + onnx.external_data_helper.load_external_data_for_model( + model, os.path.dirname(os.path.abspath(onnx_path)) + ) value_info_map, initializer_map, node_to_init_map = utils.setup_mappings(model) # Automatically add 'trt' to list of providers if custom ops are detected @@ -187,6 +192,7 @@ def convert_to_mixed_precision( custom_ops=graph_sanitizer.custom_ops, use_standalone_type_inference=use_standalone_type_inference, original_network_io_metadata=original_network_io_metadata, + sanitize_model=False, ) # Obtain reference data diff --git a/modelopt/onnx/autocast/graphsanitizer.py b/modelopt/onnx/autocast/graphsanitizer.py index 2154a42568e..e9be72c0c27 100644 --- a/modelopt/onnx/autocast/graphsanitizer.py +++ b/modelopt/onnx/autocast/graphsanitizer.py @@ -15,6 +15,8 @@ """Graph sanitization and optimization for ONNX models.""" +import os + import numpy as np import onnx import onnx_graphsurgeon as gs @@ -25,7 +27,12 @@ import modelopt.onnx.utils as onnx_utils from modelopt.onnx.autocast.logging_config import logger from modelopt.onnx.quantization.graph_utils import cast_custom_ops -from modelopt.onnx.trt_utils import interpret_trt_plugins_precision_flag +from modelopt.onnx.trt_utils import ( + get_custom_layers, + infer_types_shapes_tensorrt, + interpret_trt_plugins_precision_flag, + set_trt_plugin_domain, +) class GraphSanitizer: @@ -38,6 +45,7 @@ def __init__( max_ir_version: int | None = None, trt_plugins: list[str] | None = [], trt_plugins_precision: list[str] | None = [], + onnx_path: str | None = None, ) -> None: """Initialize GraphSanitizer. @@ -46,6 +54,7 @@ def __init__( min_opset: minimum opset version to use max_ir_version: maximum IR version supported by ORT trt_plugins: list of TensorRT plugin library paths in .so format (compiled shared library). + onnx_path: path to the source ONNX model, used to resolve external data. """ self.model = model self.min_opset = min_opset @@ -55,6 +64,8 @@ def __init__( self.custom_ops_low_precision_nodes = [] self.trt_plugins = trt_plugins self.trt_plugins_precision = trt_plugins_precision or [] + self.onnx_path = os.path.abspath(onnx_path) if onnx_path is not None else None + self.external_data_dir = os.path.dirname(self.onnx_path) if self.onnx_path else "" def sanitize(self) -> None: """Sanitize the model graph. @@ -118,13 +129,14 @@ def find_custom_nodes(self) -> None: node.op_type for node in self.model.graph.node if node.op_type not in self.standard_ops } if self.custom_ops: - from modelopt.onnx.trt_utils import infer_types_shapes_tensorrt, set_trt_plugin_domain - # Set TensorRT plugin domain info in the graph for ORT compatibility self.model = set_trt_plugin_domain(self.model, self.custom_ops) # Infer types and shapes in the graph for ORT compatibility - self.model = infer_types_shapes_tensorrt(self.model, self.trt_plugins) + _, all_tensor_info = get_custom_layers(self.onnx_path or self.model, self.trt_plugins) + self.model = infer_types_shapes_tensorrt( + self.model, self.trt_plugins, all_tensor_info=all_tensor_info + ) def remove_disconnected_outputs(self) -> None: """Remove disconnected outputs from the model.""" @@ -501,7 +513,7 @@ def _get_initializer_value(self, name: str, return_array: bool = False) -> np.nd """Get value from an initializer by name.""" for init in self.model.graph.initializer: if init.name == name: - value = numpy_helper.to_array(init) + value = numpy_helper.to_array(init, base_dir=self.external_data_dir) return value if return_array else value.item() return None @@ -516,7 +528,7 @@ def _convert_fp64_initializers(self) -> bool: for initializer in self.model.graph.initializer: if initializer.data_type == onnx.TensorProto.DOUBLE: # Convert the data to FP32 - fp64_data = numpy_helper.to_array(initializer) + fp64_data = numpy_helper.to_array(initializer, base_dir=self.external_data_dir) fp32_data = fp64_data.astype(np.float32) # Create new initializer with FP32 data @@ -575,7 +587,7 @@ def _convert_fp64_nodes(self) -> bool: for attr in node.attribute: if attr.name == "value" and attr.t.data_type == onnx.TensorProto.DOUBLE: # Convert the tensor value to FP32 - fp64_data = numpy_helper.to_array(attr.t) + fp64_data = numpy_helper.to_array(attr.t, base_dir=self.external_data_dir) fp32_data = fp64_data.astype(np.float32) new_tensor = numpy_helper.from_array(fp32_data) attr.t.CopyFrom(new_tensor) diff --git a/modelopt/onnx/autocast/precisionconverter.py b/modelopt/onnx/autocast/precisionconverter.py index bd5c0fc13f3..a78249fdca4 100644 --- a/modelopt/onnx/autocast/precisionconverter.py +++ b/modelopt/onnx/autocast/precisionconverter.py @@ -99,6 +99,7 @@ def __init__( tensor_block_dict: dict[str, dict[str, list[int]]] = {}, use_standalone_type_inference: bool = False, original_network_io_metadata: dict[str, list[onnx.ValueInfoProto]] | None = None, + sanitize_model: bool = True, ) -> None: """Initialize PrecisionConverter. @@ -118,11 +119,18 @@ def __init__( tensor_block_dict: Dictionary of tensors (operation type and I/O indices) that should remain in FP32. use_standalone_type_inference: Use standalone type inference instead of ONNX's infer_shapes. original_network_io_metadata: Original public input/output metadata captured at the API boundary. + sanitize_model: Whether to sanitize the model before precision conversion. """ self.model = deepcopy(model) - self.value_info_map = value_info_map - self.initializer_map = initializer_map - self.node_to_init_map = node_to_init_map + self.sanitize_model = sanitize_model + if sanitize_model: + self.value_info_map = value_info_map + self.initializer_map = initializer_map + self.node_to_init_map = node_to_init_map + else: + self.value_info_map, self.initializer_map, self.node_to_init_map = utils.setup_mappings( + self.model + ) self.keep_io_types = keep_io_types self.init_conversion_max_bytes = ( np.inf if init_conversion_max_bytes is None else init_conversion_max_bytes @@ -134,13 +142,6 @@ def __init__( self.low_precision_type = PRECISION_MAP[low_precision_type] self.high_precision_type = PRECISION_MAP["fp32"] - # Preserve original network inputs and outputs for sanity checks - self.original_network_io = { - io.name: io.type.tensor_type.elem_type for io in self.model.graph.input - } - self.original_network_io.update( - {io.name: io.type.tensor_type.elem_type for io in self.model.graph.output} - ) self.original_network_io_metadata = ( { "input": [deepcopy(io) for io in self.model.graph.input], @@ -152,6 +153,13 @@ def __init__( for field, values in original_network_io_metadata.items() } ) + # Preserve the public I/O types captured at the API boundary. Type inference may have + # changed the working model's declarations before the converter is initialized. + self.original_network_io = { + io.name: io.type.tensor_type.elem_type + for values in self.original_network_io_metadata.values() + for io in values + } self.min_opset = min_opset self.max_ir_version = max_ir_version self.trt_plugins = trt_plugins @@ -195,7 +203,8 @@ def convert( "AutoCast can only operate on valid ONNX models, but the input model is invalid. See log for details." ) - self._sanitize_model() + if self.sanitize_model: + self._sanitize_model() # Filter out nodes that are not allowed to be in low precision # This is done here and not in NodeClassifier because it is required for the model to be valid @@ -1431,7 +1440,10 @@ def _add_cast( # Update network output for output in self.model.graph.output: if output.name == tensor_name and ( - (self.keep_io_types and cast_to.onnx_type == output.type.tensor_type.elem_type) + ( + self.keep_io_types + and cast_to.onnx_type == self.original_network_io.get(tensor_name) + ) or ( not self.keep_io_types and cast_to.onnx_type == self.low_precision_type.onnx_type diff --git a/modelopt/onnx/autocast/referencerunner.py b/modelopt/onnx/autocast/referencerunner.py index 6dcf7137689..4db54b3b366 100644 --- a/modelopt/onnx/autocast/referencerunner.py +++ b/modelopt/onnx/autocast/referencerunner.py @@ -26,6 +26,7 @@ import copy import io +import os import sys import tempfile from collections import OrderedDict @@ -112,8 +113,6 @@ def _load_inputs_from_npz(self, input_data_path): Returns: List of input dictionaries, one per batch. """ - import os - if os.path.isdir(input_data_path): # Load all NPZ files in the directory as multiple batches npz_files = sorted([f for f in os.listdir(input_data_path) if f.endswith(".npz")]) @@ -135,7 +134,7 @@ def _load_inputs_from_npz(self, input_data_path): def _validate_inputs(self, data_loader): """Validate that input names and shapes match the model.""" if isinstance(data_loader, list) and ( - isinstance(data_loader[0], (dict, np.lib.npyio.NpzFile)) + isinstance(data_loader[0], dict | np.lib.npyio.NpzFile) ): if sorted(self.input_names) != sorted(data_loader[0].keys()): raise ValueError("Input names from ONNX model do not match provided input names.") @@ -165,8 +164,6 @@ def _load_inputs(self, inputs): # If no inputs are provided, use random inputs data_loader = DataLoader(val_range={"": (-1, 1)}) - import os - if inputs is not None: if isinstance(inputs, str): if inputs.endswith(".json"): @@ -178,7 +175,7 @@ def _load_inputs(self, inputs): f"Invalid input file: {inputs}. Supported input types: .json (Polygraphy JSON format), " ".npz (Numpy), or a directory containing .npz files" ) - elif isinstance(inputs, (dict, OrderedDict)): + elif isinstance(inputs, dict | OrderedDict): data_loader = [inputs] else: raise ValueError( @@ -193,32 +190,32 @@ def _get_ort_runner(self, model): from polygraphy.backend.onnx import BytesFromOnnx from polygraphy.backend.onnxrt import OnnxrtRunner, SessionFromOnnx - # Check if model has external data by checking: - # 1. If any initializer has data_location set to EXTERNAL (even if data is loaded) - # 2. If model size would exceed 2GB (indicating need for external data) - needs_external_data = onnx_utils.check_model_uses_external_data( - self.model - ) or self.model.ByteSize() > 2 * (1024**3) - if needs_external_data: - logger.debug("Model has external data, using file-based approach") - # Get the actual ONNX ModelProto from ModifyOutputs wrapper - modified_model = model() - - # Use a persistent temp file, because we need the file to be present in an broader context - tmp_file = tempfile.NamedTemporaryFile(suffix=".onnx", delete=False) - tmp_file.close() - tmp_file_path = tmp_file.name - onnx_utils.save_onnx(modified_model, tmp_file_path, save_as_external_data=True) - logger.debug(f"Model with all outputs saved to {tmp_file_path}") - build_onnxrt_session = SessionFromOnnx(tmp_file_path, providers=self.providers) + # Get the actual ONNX ModelProto from ModifyOutputs wrapper + modified_model = model() - else: - # For models without external data, use the original BytesFromOnnx approach (no tmp files) - logger.debug("Model has no external data, using BytesFromOnnx approach") - serialize_onnx = BytesFromOnnx(model) - build_onnxrt_session = SessionFromOnnx(serialize_onnx, providers=self.providers) - runners = [OnnxrtRunner(build_onnxrt_session)] - return runners + needs_file_backed_model = onnx_utils.check_model_uses_external_data( + modified_model + ) or onnx_utils.is_model_too_large_for_protobuf(modified_model) + model_temp_dir = None + try: + if needs_file_backed_model: + logger.debug("Model has external data, using file-based approach") + model_temp_dir = tempfile.TemporaryDirectory() + tmp_file_path = os.path.join(model_temp_dir.name, "model.onnx") + onnx_utils.save_onnx(modified_model, tmp_file_path, save_as_external_data=True) + logger.debug(f"Model with all outputs saved to {tmp_file_path}") + build_onnxrt_session = SessionFromOnnx(tmp_file_path, providers=self.providers) + else: + # For models without external data, use the original BytesFromOnnx approach (no tmp files) + logger.debug("Model has no external data, using BytesFromOnnx approach") + serialize_onnx = BytesFromOnnx(modified_model) + build_onnxrt_session = SessionFromOnnx(serialize_onnx, providers=self.providers) + runners = [OnnxrtRunner(build_onnxrt_session)] + except Exception: + if model_temp_dir is not None: + model_temp_dir.cleanup() + raise + return runners, model_temp_dir def _aggregate_tensor_stats(self, all_batch_data: list[OrderedDict]) -> OrderedDict: """Aggregate tensor statistics across multiple batches. @@ -300,22 +297,25 @@ def run(self, inputs=None): modify_outputs = ModifyOnnxOutputs(model_copy, outputs=constants.MARK_ALL) # Load the modified model and create an inference session - runners = self._get_ort_runner(modify_outputs) - - # Comparator is used despite the fact that we are using ONNXRuntime - # because it provides the ability to generate random inputs using DataLoader - data_loader = self._load_inputs(inputs) - - # Temporarily redirect stdout to suppress Comparator.run() output - stdout = sys.stdout - string_buffer = io.StringIO() - sys.stdout = string_buffer + runners, model_temp_dir = self._get_ort_runner(modify_outputs) try: - results = Comparator.run(runners, data_loader=data_loader) + # Comparator is used despite the fact that we are using ONNXRuntime + # because it provides the ability to generate random inputs using DataLoader + data_loader = self._load_inputs(inputs) + + # Temporarily redirect stdout to suppress Comparator.run() output + stdout = sys.stdout + string_buffer = io.StringIO() + sys.stdout = string_buffer + try: + results = Comparator.run(runners, data_loader=data_loader) + finally: + # Capture the output before restoring stdout + captured_output = string_buffer.getvalue() + sys.stdout = stdout finally: - # Capture the output before restoring stdout - captured_output = string_buffer.getvalue() - sys.stdout = stdout + if model_temp_dir is not None: + model_temp_dir.cleanup() if not results: logger.error(f"ONNXRuntime execution failed with output:\n{captured_output}") diff --git a/modelopt/onnx/export/fp8_exporter.py b/modelopt/onnx/export/fp8_exporter.py index 427a7791f3b..8e94f40f332 100644 --- a/modelopt/onnx/export/fp8_exporter.py +++ b/modelopt/onnx/export/fp8_exporter.py @@ -17,6 +17,7 @@ import time +import ml_dtypes import numpy as np import onnx import onnx_graphsurgeon as gs @@ -33,6 +34,13 @@ _FP8_E4M3_SOFTMAX_SCALE = 1.0 / _FP8_E4M3_MAX +def _torch_from_numpy_for_fp8(array: np.ndarray) -> torch.Tensor: + """Convert a NumPy array to the PyTorch dtype used for FP8 normalization.""" + if array.dtype == ml_dtypes.bfloat16: + return torch.from_numpy(array.view(np.int16)).view(torch.bfloat16).float() + return torch.from_numpy(array) + + class FP8QuantExporter(ONNXQuantExporter): """Exporter for FP8 quantization.""" @@ -48,7 +56,7 @@ def compute_scales(onnx_model: onnx.ModelProto) -> onnx.ModelProto: @staticmethod def compress_weights(onnx_model: onnx.ModelProto) -> onnx.ModelProto: - """Compresses FP32/FP16 weights to FP8 by folding QDQ nodes to DQ only. + """Compresses FP32/FP16/BF16 weights to FP8 by folding QDQ nodes to DQ only. Even though modelopt supports FP8 onnx export, the weights are represented in fp32 + QDQ. The storage is therefore very bad. In this function, @@ -56,7 +64,7 @@ def compress_weights(onnx_model: onnx.ModelProto) -> onnx.ModelProto: weights in the output model. TRT custom ops are converted to native ONNX DequantizeLinear. Parameters: - onnx_model: ONNX model with FP32/FP16 weights and TRT_FP8 QDQ nodes. + onnx_model: ONNX model with FP32/FP16/BF16 weights and TRT_FP8 QDQ nodes. Returns: ONNX model with FP8 weights and native ONNX DQ nodes for weights (QDQ preserved for activations). @@ -78,8 +86,8 @@ def compress_weights(onnx_model: onnx.ModelProto) -> onnx.ModelProto: weights = node.inputs[0] scale = node.inputs[1] - torch_weights = torch.from_numpy(weights.values) - torch_scale = torch.from_numpy(scale.values) + torch_weights = _torch_from_numpy_for_fp8(weights.values) + torch_scale = _torch_from_numpy_for_fp8(scale.values) quantizer_name = scale.name.rsplit("/", 1)[0] dq_op = node.outputs[0].outputs[0] if dq_op.op != "TRT_FP8DequantizeLinear": @@ -194,14 +202,25 @@ def _quantize_conv_weights_to_fp8(graph: gs.Graph) -> int: if any(out.op == "DequantizeLinear" for out in weight_input.outputs): continue - torch_weights = torch.from_numpy(weight_input.values.copy()) + torch_weights = _torch_from_numpy_for_fp8(weight_input.values.copy()) amax = torch_weights.abs().max().float() if amax == 0: continue scale_val = (amax / _FP8_E4M3_MAX).item() + scale_data = np.array(scale_val, dtype=weight_input.values.dtype) + # Round up so normalizing by the serialized scale stays within the FP8 range. + if scale_data < scale_val: + np.nextafter( + scale_data, + np.array(np.inf, dtype=scale_data.dtype), + out=scale_data, + ) + torch_scale = _torch_from_numpy_for_fp8(scale_data) # Quantize weights to FP8 (WAR: numpy doesn't support fp8) - fp8_data = (torch_weights / scale_val).to(torch.float8_e4m3fn).view(torch.uint8).numpy() + fp8_data = ( + (torch_weights / torch_scale).to(torch.float8_e4m3fn).view(torch.uint8).numpy() + ) fp8_tensor = onnx.TensorProto() fp8_tensor.data_type = onnx.TensorProto.FLOAT8E4M3FN fp8_tensor.dims.extend(fp8_data.shape) @@ -210,13 +229,16 @@ def _quantize_conv_weights_to_fp8(graph: gs.Graph) -> int: node.name + "/weight_quantizer/fp8_weights", LazyValues(fp8_tensor) ) - # Scale in FP16 — DQ output type matches scale dtype, must match activation type scale_constant = gs.Constant( node.name + "/weight_quantizer/scale", - np.array(scale_val, dtype=np.float16), + scale_data, ) - dq_output = gs.Variable(node.name + "/weight_quantizer/dq_output") + dq_output = gs.Variable( + node.name + "/weight_quantizer/dq_output", + scale_data.dtype, + weight_input.values.shape, + ) dq_node = gs.Node( op="DequantizeLinear", name=node.name + "/weight_quantizer/DequantizeLinear", diff --git a/modelopt/onnx/export/nvfp4_exporter.py b/modelopt/onnx/export/nvfp4_exporter.py index 338e2725b14..42598af5475 100644 --- a/modelopt/onnx/export/nvfp4_exporter.py +++ b/modelopt/onnx/export/nvfp4_exporter.py @@ -430,4 +430,10 @@ def _cast_input_dtypes(node: onnx.NodeProto, precision_dtype: str): utils.topologically_sort_graph_nodes(graph) + if fp4_qdq_nodes: + default_opset = next( + opset for opset in onnx_model.opset_import if opset.domain in {"", "ai.onnx"} + ) + default_opset.version = max(default_opset.version, 23) + return onnx_model diff --git a/modelopt/onnx/quantization/gs_patching.py b/modelopt/onnx/quantization/gs_patching.py index a0eea84951e..503a91f9721 100644 --- a/modelopt/onnx/quantization/gs_patching.py +++ b/modelopt/onnx/quantization/gs_patching.py @@ -101,9 +101,15 @@ def _export_value_info_proto(tensor: gs.Variable, do_type_check: bool) -> onnx.V ) if tensor.dtype is not None: - dtype = getattr( - tensor, "explicit_dtype", onnx.helper.np_dtype_to_tensor_dtype(np.dtype(tensor.dtype)) - ) + dtype = getattr(tensor, "explicit_dtype", None) + if dtype is None: + dtype = tensor.dtype + if isinstance(dtype, (int, np.integer)): + dtype = int(dtype) + if dtype not in onnx.TensorProto.DataType.values(): + raise ValueError(f"Unknown ONNX tensor dtype for {tensor.name}: {dtype}") + else: + dtype = onnx.helper.np_dtype_to_tensor_dtype(np.dtype(dtype)) onnx_tensor = onnx.helper.make_tensor_value_info(tensor.name, dtype, tensor.shape) else: onnx_tensor = onnx.helper.make_empty_tensor_value_info(tensor.name) diff --git a/modelopt/onnx/quantization/ort_patching.py b/modelopt/onnx/quantization/ort_patching.py index f10b97d13d1..d2957ec183f 100755 --- a/modelopt/onnx/quantization/ort_patching.py +++ b/modelopt/onnx/quantization/ort_patching.py @@ -52,7 +52,7 @@ import onnxruntime as ort import pynvml from onnx import onnx_pb -from onnxruntime.quantization import calibrate +from onnxruntime.quantization import calibrate, qdq_quantizer from onnxruntime.quantization.base_quantizer import BaseQuantizer from onnxruntime.quantization.calibrate import ( CalibraterBase, @@ -74,6 +74,7 @@ QuantType, add_infer_metadata, ) +from onnxruntime.quantization.quant_utils import compute_scale_zp as _ort_compute_scale_zp from onnxruntime.quantization.quantize import check_static_quant_arguments from onnxruntime.quantization.registry import QDQRegistry, QLinearOpsRegistry from onnxruntime.tools.symbolic_shape_infer import SymbolicShapeInference @@ -94,6 +95,53 @@ def load_model_with_shape_infer(model_path: Path) -> onnx.ModelProto: return model +def _compute_scale_zp(rmin, rmax, qmin, qmax, symmetric=False, min_real_range=None): + """Retry FP16 scale calculation in FP32 when range subtraction overflows.""" + range_dtype = np.asarray(rmax).dtype + if range_dtype != np.float16: + return _ort_compute_scale_zp(rmin, rmax, qmin, qmax, symmetric, min_real_range) + + with np.errstate(over="ignore", invalid="ignore"): + zero_point, scale = _ort_compute_scale_zp(rmin, rmax, qmin, qmax, symmetric, min_real_range) + if np.all(np.isfinite(scale)): + return zero_point, scale + + zero_point, scale = _ort_compute_scale_zp( + np.asarray(rmin, dtype=np.float32), + np.asarray(rmax, dtype=np.float32), + qmin, + qmax, + symmetric, + min_real_range, + ) + return zero_point, np.asarray(scale, dtype=range_dtype) + + +def _prepare_histogram_data(histogram_collector, tensor, data_arr): + """Use FP32 for histogram math while remembering the source dtype.""" + if data_arr.dtype != np.float16: + return data_arr + + original_dtypes = getattr(histogram_collector, "_modelopt_original_dtypes", {}) + original_dtypes[tensor] = data_arr.dtype + histogram_collector._modelopt_original_dtypes = original_dtypes + return data_arr.astype(np.float32) + + +def _restore_histogram_calibration_dtypes(histogram_collector, tensors_range): + """Restore source dtypes at the calibration-to-quantization boundary.""" + original_dtypes = getattr(histogram_collector, "_modelopt_original_dtypes", {}) + for tensor, dtype in original_dtypes.items(): + if tensor not in tensors_range: + continue + tensor_data = tensors_range[tensor] + dtype_limits = np.finfo(dtype) + for attribute in ("lowest", "highest", "avg", "std"): + if hasattr(tensor_data, attribute): + value = np.clip(getattr(tensor_data, attribute), dtype_limits.min, dtype_limits.max) + setattr(tensor_data, attribute, np.asarray(value, dtype=dtype)) + + def _collect_value(histogram_collector, name_to_arr): """Collect histogram on real value.""" for tensor, data_arr in tqdm(name_to_arr.items()): @@ -105,6 +153,7 @@ def _collect_value(histogram_collector, name_to_arr): curr_data_arr = curr_data_arr.flatten() concat_data_arr = np.concatenate((concat_data_arr, curr_data_arr)) + concat_data_arr = _prepare_histogram_data(histogram_collector, tensor, concat_data_arr) data_arr = concat_data_arr # ========================================================== if data_arr.size > 0: @@ -130,9 +179,6 @@ def _collect_value(histogram_collector, name_to_arr): old_histogram, data_arr, min_value, max_value, threshold ) else: - # Cast range endpoints to Python float so numpy computes bin edges in - # float64. A fp16 threshold here can underflow the 128-bin linspace - # and trip "Too many bins for data range" on numpy >= 2.0. range_max = float(threshold) hist, hist_edges = np.histogram( data_arr, histogram_collector.num_bins, range=(-range_max, range_max) @@ -1126,6 +1172,7 @@ def _collect_value_histogram_collector_single_node_calibration(histogram_collect """Collect histogram on real value.""" for tensor, data_arr in name_to_arr.items(): data_arr = np.asarray(data_arr).flatten() + data_arr = _prepare_histogram_data(histogram_collector, tensor, data_arr) min_value, max_value = (np.min(data_arr), np.max(data_arr)) if data_arr.size > 0 else (0, 0) # Replace inf/nan with float32 min/max @@ -1147,9 +1194,6 @@ def _collect_value_histogram_collector_single_node_calibration(histogram_collect threshold, ) else: - # Cast range endpoints to Python float so numpy computes bin edges in - # float64. A fp16 threshold here can underflow the 128-bin linspace - # and trip "Too many bins for data range" on numpy >= 2.0. range_max = float(threshold) hist, hist_edges = np.histogram( data_arr, histogram_collector.num_bins, range=(-range_max, range_max) @@ -1685,6 +1729,8 @@ def _quantize_static( raise TypeError( f"Unexpected type {type(tensors_range)} for tensors_range and calibrator={type(calibrator)}." ) + if isinstance(calibrator, HistogramCalibrater): + _restore_histogram_calibration_dtypes(calibrator.collector, tensors_range) del calibrator check_static_quant_arguments(quant_format, activation_type, weight_type) @@ -1795,4 +1841,5 @@ def patch_ort_modules(calibrate_per_node: bool = False): CalibraterBase.select_tensors_to_calibrate = _select_tensors_to_calibrate QDQQuantizer.check_opset_version = _check_opset_version BaseQuantizer.adjust_tensor_ranges = _adjust_tensor_ranges + qdq_quantizer.compute_scale_zp = _compute_scale_zp CalibraterBase.__init__ = _init_calibrater_base diff --git a/modelopt/onnx/quantization/quantize.py b/modelopt/onnx/quantization/quantize.py index fb6b662916d..ef16390b90e 100755 --- a/modelopt/onnx/quantization/quantize.py +++ b/modelopt/onnx/quantization/quantize.py @@ -779,7 +779,6 @@ def quantize( ) if calibration_data_reader is None: - # Use random scales if calibration data is not supplied if calibration_data is None: calibration_data_reader = RandomDataProvider(onnx_path, calibration_shapes) else: diff --git a/modelopt/onnx/utils.py b/modelopt/onnx/utils.py index f8b5a41a41a..aa7d27bc8fd 100644 --- a/modelopt/onnx/utils.py +++ b/modelopt/onnx/utils.py @@ -18,9 +18,9 @@ import copy import io import os +import sys import tempfile import uuid -from collections import defaultdict from typing import Any import numpy as np @@ -527,50 +527,114 @@ def name_onnx_nodes(graph: onnx.GraphProto) -> bool: def duplicate_shared_constants(onnx_model: onnx.ModelProto) -> tuple[onnx.ModelProto, bool]: """Duplicate constant tensors if they are shared.""" - graph = gs.import_onnx(onnx_model) - name_dict = defaultdict(lambda: 0) - - def _get_unique_name(old_name): - name_dict[old_name] += 1 - return old_name + "_" + str(name_dict[old_name]) - - # Get tensors with shared constant inputs - tensors = [] - for node in graph.nodes: - for inp_idx, tensor in enumerate(node.inputs): - # constant is shared across multiple nodes - if isinstance(tensor, Constant) and len(tensor.outputs) > 1: - tensors.append({"tensor": tensor, "inp_node": node, "inp_idx": inp_idx}) - - # Duplicate shared tensors - for tensor_dict in tensors: - tensor = tensor_dict["tensor"] - new_tensor = Constant( - name=_get_unique_name(tensor.name), - values=tensor.values, + graph = onnx_model.graph + initializers = {initializer.name: initializer for initializer in graph.initializer} + use_counts = dict.fromkeys(initializers, 0) + for node in graph.node: + for input_name in node.input: + if input_name in use_counts: + use_counts[input_name] += 1 + + shared_names = {name for name, count in use_counts.items() if count > 1} + if not shared_names: + return onnx_model, False + + used_names = set(initializers) + used_names.update( + sparse_initializer.values.name for sparse_initializer in graph.sparse_initializer + ) + for value_info in list(graph.input) + list(graph.output) + list(graph.value_info): + used_names.add(value_info.name) + for node in graph.node: + used_names.update(name for name in node.input if name) + used_names.update(name for name in node.output if name) + + def _get_nested_graphs(node: onnx.NodeProto): + for attribute in node.attribute: + if attribute.type == onnx.AttributeProto.GRAPH: + yield attribute.g + elif attribute.type == onnx.AttributeProto.GRAPHS: + yield from attribute.graphs + + def _find_captured_names(nested_graph: onnx.GraphProto, outer_names: set[str]) -> set[str]: + local_names = {value_info.name for value_info in nested_graph.input} + local_names.update(initializer.name for initializer in nested_graph.initializer) + local_names.update( + sparse_initializer.values.name for sparse_initializer in nested_graph.sparse_initializer ) - tensor_dict["inp_node"].inputs[tensor_dict["inp_idx"]] = new_tensor + for nested_node in nested_graph.node: + local_names.update(name for name in nested_node.output if name) - onnx_model = gs.export_onnx(graph) - is_modified = bool(tensors) - return onnx_model, is_modified + visible_outer_names = outer_names - local_names + captured_names = { + output.name for output in nested_graph.output if output.name in visible_outer_names + } + for nested_node in nested_graph.node: + captured_names.update( + input_name for input_name in nested_node.input if input_name in visible_outer_names + ) + for child_graph in _get_nested_graphs(nested_node): + captured_names.update(_find_captured_names(child_graph, visible_outer_names)) + return captured_names + captured_names = set() + for node in graph.node: + for nested_graph in _get_nested_graphs(node): + captured_names.update(_find_captured_names(nested_graph, shared_names)) -def check_model(model: onnx.ModelProto) -> None: - """Checks if the given model is valid.""" - save_as_external_data = False + next_suffix: dict[str, int] = {} + + def _get_unique_name(old_name: str) -> str: + suffix = next_suffix.get(old_name, 1) + new_name = f"{old_name}_{suffix}" + while new_name in used_names: + suffix += 1 + new_name = f"{old_name}_{suffix}" + next_suffix[old_name] = suffix + 1 + used_names.add(new_name) + return new_name + + for node in graph.node: + for input_index, input_name in enumerate(node.input): + if input_name not in shared_names: + continue + + duplicated_initializer = graph.initializer.add() + duplicated_initializer.CopyFrom(initializers[input_name]) + duplicated_initializer.name = _get_unique_name(input_name) + node.input[input_index] = duplicated_initializer.name + + retained_names = captured_names | { + output.name for output in graph.output if output.name in shared_names + } + removed_initializer_names = shared_names - retained_names + for initializer in initializers.values(): + if initializer.name in removed_initializer_names: + graph.initializer.remove(initializer) + for graph_input in list(graph.input): + if graph_input.name in removed_initializer_names: + graph.input.remove(graph_input) + + return onnx_model, True + + +def is_model_too_large_for_protobuf(model: onnx.ModelProto) -> bool: + """Return whether a model cannot safely use an in-memory protobuf API.""" try: model_size = model.ByteSize() except Exception as e: logger.warning( "Failed to compute model size with ByteSize (%s). Using external data path.", e ) - save_as_external_data = True - else: - if model_size <= 0 or model_size > (2 * (1024**3)): - save_as_external_data = True + return True - if save_as_external_data: + max_model_size = onnx.checker.MAXIMUM_PROTOBUF - sys.getsizeof(b"") + return model_size <= 0 or model_size > max_model_size + + +def check_model(model: onnx.ModelProto) -> None: + """Checks if the given model is valid.""" + if is_model_too_large_for_protobuf(model): with tempfile.TemporaryDirectory() as temp_dir: # ONNX also looks in CWD, so we need to use a unique id unique_id = str(uuid.uuid4())[:8] @@ -1171,19 +1235,7 @@ def infer_types_verification(model: onnx.ModelProto) -> onnx.ModelProto: def infer_shapes(model: onnx.ModelProto, **kwargs): """Infers shapes of the onnx graph, handles large models.""" - save_as_external_data = False - try: - model_size = model.ByteSize() - except Exception as e: - logger.warning( - "Failed to compute model size with ByteSize (%s). Using external data path.", e - ) - save_as_external_data = True - else: - if model_size <= 0 or model_size > (2 * (1024**3)): - save_as_external_data = True - - if save_as_external_data: + if is_model_too_large_for_protobuf(model): with tempfile.TemporaryDirectory() as temp_dir: # ONNX also looks in CWD, so we need to use a unique id unique_id = str(uuid.uuid4())[:8] diff --git a/modelopt/torch/_deploy/utils/onnx_optimizer.py b/modelopt/torch/_deploy/utils/onnx_optimizer.py index 08c4c40ca97..c520983f7c0 100644 --- a/modelopt/torch/_deploy/utils/onnx_optimizer.py +++ b/modelopt/torch/_deploy/utils/onnx_optimizer.py @@ -21,6 +21,8 @@ import onnx import onnx_graphsurgeon as gs +from modelopt.onnx.utils import is_model_too_large_for_protobuf + class Optimizer: """Optimizer for onnx graphs.""" @@ -62,7 +64,7 @@ def select_outputs(self, keep, names=None): def infer_shapes(self, return_onnx=False): """Infers shapes of the onnx graph.""" onnx_graph = gs.export_onnx(self.graph) - if onnx_graph.ByteSize() > (2 * (1024**3)): # 2GB limit + if is_model_too_large_for_protobuf(onnx_graph): # 2GB limit temp_dir = tempfile.TemporaryDirectory().name os.makedirs(temp_dir, exist_ok=True) onnx_orig_path = os.path.join(temp_dir, "model.onnx") diff --git a/modelopt/torch/_deploy/utils/torch_onnx.py b/modelopt/torch/_deploy/utils/torch_onnx.py index 01fb754bbae..b217f1188b3 100644 --- a/modelopt/torch/_deploy/utils/torch_onnx.py +++ b/modelopt/torch/_deploy/utils/torch_onnx.py @@ -56,6 +56,7 @@ get_output_names, get_output_shapes, infer_shapes, + is_model_too_large_for_protobuf, remove_node_training_mode, remove_redundant_casts, ) @@ -96,7 +97,6 @@ def _patched_remove_unnecessary_cast_node(graph): # a few constants... DEFAULT_ONNX_OPSET = 20 ONNX_EXPORT_OUT_PREFIX = "out" -TWO_GB = 2 * 1024 * 1024 * 1024 class OnnxBytes: @@ -507,14 +507,17 @@ def get_onnx_bytes_and_metadata( `torch.onnx.export `_. onnx_opset: The onnx opset version to use for exporting the model. dq_only: If True, the exported onnx model is converted to a dq_only model. - weights_dtype: The dtype of the weights in the onnx model. + weights_dtype: Requested high-precision dtype for exported weights. For an FP8 model, + ``"bf16"`` is accepted only when every floating parameter is already BF16. This is + a weight-focused no-op, not a graph-wide conversion: floating buffers are not + considered for eligibility and may preserve higher-precision regions. Returns: bytes: Onnx model in bytes. ModelMetadata: The model's meta data. Raises: - ValueError: If nn.Module is not passed as model. + ValueError: If model is not an nn.Module or the requested precision conversion is unsupported. """ if not isinstance(model, nn.Module): raise ValueError("Only PyTorch model compilation is supported.") @@ -527,6 +530,22 @@ def get_onnx_bytes_and_metadata( if isinstance(model, (DataParallel, DistributedDataParallel)): model = model.module + source_parameter_dtypes = { + parameter.dtype for parameter in model.parameters() if parameter.is_floating_point() + } + source_parameter_dtype_names = ", ".join(sorted(map(str, source_parameter_dtypes))) or "none" + uses_fp4 = is_fp4_quantized(model) + uses_mxfp8 = is_mxfp8_quantized(model) + uses_fp8 = is_fp8_quantized(model) + uses_int8 = is_int8_quantized(model) + uses_other_unsupported_quantizer = is_int4_quantized(model) or uses_mxfp8 or uses_int8 + is_bf16_fp8_noop = ( + weights_dtype == "bf16" + and source_parameter_dtypes == {torch.bfloat16} + and uses_fp8 + and not (uses_fp4 or uses_other_unsupported_quantizer) + ) + # Standardize model args and also tensorize them so they also appear in the onnx graph! # Floats/ints are tensorized when they are provided, but not tensorized when they are not # provided which is somewhat inconsistent (we always tensorize them!) @@ -549,11 +568,7 @@ def get_onnx_bytes_and_metadata( input_none_names = list(set(tree_spec_input.names) - set(input_names)) use_torch_autocast = not ( - is_fp4_quantized(model) - or is_mxfp8_quantized(model) - or is_fp8_quantized(model) - or is_int8_quantized(model) - or weights_dtype == "fp32" + uses_fp4 or uses_mxfp8 or uses_fp8 or uses_int8 or weights_dtype == "fp32" ) autocast = torch.autocast("cuda") if use_torch_autocast else nullcontext() @@ -575,6 +590,22 @@ def get_onnx_bytes_and_metadata( ) return onnx_model.to_bytes(), model_metadata + if weights_dtype == "fp16" and uses_fp8 and torch.bfloat16 in source_parameter_dtypes: + raise ValueError( + "Converting a BF16 FP8 ONNX graph to FP16 is not supported yet " + f"(source parameter dtypes: {source_parameter_dtype_names})" + ) + + if ( + weights_dtype == "bf16" + and (uses_fp8 or uses_other_unsupported_quantizer) + and not is_bf16_fp8_noop + ): + raise ValueError( + "Converting a quantized ONNX graph to BF16 is not supported yet " + f"(source parameter dtypes: {source_parameter_dtype_names})" + ) + # Export onnx model from pytorch model # As the maximum size of protobuf is 2GB, we cannot use io.BytesIO() buffer during export. model_name = model_name or model.__class__.__name__ @@ -583,16 +614,12 @@ def get_onnx_bytes_and_metadata( # Configure quantizers if the model is quantized in NVFP4 or MXFP8 mode quantizer_context = ( - configure_linear_module_onnx_quantizers(model) - if is_fp4_quantized(model) or is_mxfp8_quantized(model) - else nullcontext() + configure_linear_module_onnx_quantizers(model) if uses_fp4 or uses_mxfp8 else nullcontext() ) # Disable FP8 Conv weight quantizers: TorchScript custom ops produce outputs with # unknown shapes, causing _convolution symbolic to fail. Conv weights are quantized # to FP8 in post-processing by FP8QuantExporter instead. - conv_wq_context = ( - _disable_fp8_conv_weight_quantizers(model) if is_fp8_quantized(model) else nullcontext() - ) + conv_wq_context = _disable_fp8_conv_weight_quantizers(model) if uses_fp8 else nullcontext() with torch.inference_mode(), autocast, quantizer_context, conv_wq_context: additional_kwargs = {} if not dynamo_export: @@ -634,14 +661,8 @@ def get_onnx_bytes_and_metadata( if dq_only: onnx_opt_graph = qdq_to_dq(onnx_opt_graph) - if weights_dtype in ["fp16", "bf16"]: - if ( - is_int4_quantized(model) - or is_mxfp8_quantized(model) - or is_fp8_quantized(model) - or is_int8_quantized(model) - ): - assert weights_dtype == "fp16", "BF16 + MXFP8/INT4 mixed precision is not supported yet" + if weights_dtype in ["fp16", "bf16"] and not is_bf16_fp8_noop: + if uses_other_unsupported_quantizer or uses_fp8: onnx_opt_graph = convert_float_to_float16( onnx_opt_graph, keep_io_types=False, @@ -665,7 +686,7 @@ def get_onnx_bytes_and_metadata( onnx_opt_graph = remove_redundant_casts(onnx_opt_graph) # Remove Cast nodes around Q/DQ for optimal TRT fusion - if is_fp8_quantized(model): + if uses_fp8: onnx_opt_graph = fold_q_fp16_to_fp32_casts(onnx_opt_graph) onnx_opt_graph = fold_dq_fp32_to_fp16_casts(onnx_opt_graph) @@ -744,6 +765,6 @@ def create_model_metadata( "output_onnx_names": get_output_names(onnx_graph), "signature": inspect.signature(model.forward), "onnx_node_names": get_node_names(onnx_graph), - "is_bytes_pickled": onnx_graph.ByteSize() > TWO_GB, + "is_bytes_pickled": is_model_too_large_for_protobuf(onnx_graph), "config": model.config if hasattr(model, "config") else None, } diff --git a/modelopt/torch/export/hf_export_handlers.py b/modelopt/torch/export/hf_export_handlers.py index 21a8a3fe246..8fb8b61e027 100644 --- a/modelopt/torch/export/hf_export_handlers.py +++ b/modelopt/torch/export/hf_export_handlers.py @@ -118,9 +118,30 @@ def _prepare_iterable_experts(name: str, moe_module: nn.Module, ctx: ExportConte # Export handlers are registered in the same precedence as the legacy model walk. -@ExportModuleRegistry.register( - "QuantMoELinear", predicate=lambda module: hasattr(module, "experts") -) +def _is_quant_moe_linear(module: nn.Module) -> bool: + """Whether ``module`` is an expert-indexed ``MoELinear`` expanded by ``_QuantMoELinear``. + + Matched by wrapper type first, not only by the dynamically generated class name + (``Quant`` + the model's own class name): the wrapper is registered structurally, so a + compatible remote-code class under any other name — or a second one, whose generated + name gets uniquified — would bypass this handler and export without the input-amax + fallback. The name check is kept as a fallback so stand-in modules match too. + + The wrapper lives in the optional transformers plugin, hence the lazy import. + """ + if not hasattr(module, "experts"): + return False + try: + from modelopt.torch.quantization.plugins.huggingface import _QuantMoELinear + + if isinstance(module, _QuantMoELinear): + return True + except ImportError: + pass + return any(cls.__name__ == "QuantMoELinear" for cls in type(module).__mro__) + + +@ExportModuleRegistry.register(predicate=_is_quant_moe_linear) def _export_moe_linear(name: str, module: nn.Module, ctx: ExportContext) -> None: """Fill missing input amax before child expert QuantLinears are exported.""" set_expert_quantizer_amax(list(module.experts), quantizer_attrs="input_quantizer") diff --git a/modelopt/torch/nas/plugins/megatron.py b/modelopt/torch/nas/plugins/megatron.py index 2376422a998..c2d31b6f3de 100644 --- a/modelopt/torch/nas/plugins/megatron.py +++ b/modelopt/torch/nas/plugins/megatron.py @@ -15,7 +15,6 @@ """Plugin to add NAS/Pruning support for megatron-core Language models like GPT and Mamba.""" -import copy import types from abc import ABC from collections.abc import Callable, Sequence @@ -35,10 +34,6 @@ ) from megatron.core.models.common.embeddings.language_model_embedding import LanguageModelEmbedding from megatron.core.models.gpt import GPTModel -from megatron.core.models.gpt.moe_module_specs import get_moe_module_spec -from megatron.core.models.hybrid.hybrid_layer_specs import ( - hybrid_stack_spec as _te_hybrid_stack_spec, -) from megatron.core.models.hybrid.hybrid_model import HybridModel from megatron.core.parallel_state import is_pipeline_first_stage, is_pipeline_last_stage from megatron.core.ssm.gated_delta_net import GatedDeltaNet @@ -56,7 +51,6 @@ from megatron.core.transformer.moe.router import TopKRouter from megatron.core.transformer.moe.shared_experts import SharedExpertMLP from megatron.core.transformer.multi_latent_attention import MLASelfAttention -from megatron.core.transformer.spec_utils import ModuleSpec from megatron.core.transformer.transformer_layer import TransformerLayer from modelopt.torch.nas.modules import DynamicModuleList @@ -91,21 +85,8 @@ # Attention module types that _DynamicTransformerLayer converts. _ATTENTION_TYPES: tuple[type, ...] = (SelfAttention, MLASelfAttention, GatedDeltaNet) -__all__ = ["get_te_hybrid_stack_spec"] - - -def get_te_hybrid_stack_spec(moe_grouped_gemm: bool = False) -> ModuleSpec: - """Return the TE Hybrid stack spec.""" - if moe_grouped_gemm: - return _te_hybrid_stack_spec - - # The upstream TE hybrid stack spec hardcodes TEGroupedMLP for MoE. - # Replace it with SequentialMLP (TE linear layers, no grouped gemm dependency). - te_hybrid_stack_spec = copy.deepcopy(_te_hybrid_stack_spec) - te_hybrid_stack_spec.submodules.moe_layer.submodules.mlp = get_moe_module_spec( - use_te=True, num_experts=8, moe_grouped_gemm=False - ) - return te_hybrid_stack_spec +# This module only registers DynamicModules; it exports no public API. +__all__ = [] # Local Parallel Linear DynamicModules ########################################################################## diff --git a/modelopt/torch/opt/plugins/mcore_dist_checkpointing.py b/modelopt/torch/opt/plugins/mcore_dist_checkpointing.py index 2bae53fa6f1..a4796162f6d 100644 --- a/modelopt/torch/opt/plugins/mcore_dist_checkpointing.py +++ b/modelopt/torch/opt/plugins/mcore_dist_checkpointing.py @@ -181,6 +181,10 @@ def _load_extra_state_from_sharded_checkpoint( module, "modelopt_set_extra_state_callbacks" ): module.set_extra_state(extra_state_dict_no_prefix[key]) + for module in model.modules(): + post_load_extra_state = getattr(module, "modelopt_post_load_extra_state", None) + if callable(post_load_extra_state): + post_load_extra_state() def restore_sharded_modelopt_state( diff --git a/modelopt/torch/quantization/algorithms.py b/modelopt/torch/quantization/algorithms.py index 7beeef6ad7f..889290231e3 100644 --- a/modelopt/torch/quantization/algorithms.py +++ b/modelopt/torch/quantization/algorithms.py @@ -23,7 +23,7 @@ from abc import ABC, abstractmethod from collections import defaultdict from collections.abc import Callable, Sequence -from contextlib import nullcontext +from contextlib import ExitStack, nullcontext from typing import Any import regex as re @@ -413,8 +413,10 @@ def __init__( self.allow_no_quant = allow_no_quant self.is_fixed = fixed_recipe is not None - self.quant_modules = list(set(quant_modules or [])) - self.score_modules = list(set(score_modules or self.quant_modules)) + # Module hashes depend on object identity, so sets can produce different orders per rank. + self.quant_modules = list(dict.fromkeys(quant_modules or [])) + self.score_modules = list(dict.fromkeys(score_modules or self.quant_modules)) + self._warned_parallel_state_fallbacks: set[nn.Module] = set() fixed_quantizers = ( { @@ -467,11 +469,11 @@ def __init__( quant_recipe: dict.fromkeys(self.score_modules) for quant_recipe in self.choices } - # Attach this hparam to each score_module's set of hparams it scores + # Registration order follows the rank-stable runtime-group construction order. for score_module in self.score_modules: if not hasattr(score_module, "_hparams_for_scoring"): - score_module._hparams_for_scoring = set() - score_module._hparams_for_scoring.add(self) + score_module._hparams_for_scoring = [] + score_module._hparams_for_scoring.append(self) @property def active(self) -> HPType: @@ -532,6 +534,31 @@ def get_score(self, recipe: QuantRecipe) -> float: continue parallel_state = getattr(score_module, "parallel_state", None) + if parallel_state is None: + # TODO: Prefer parallel_state owned by the score module; this temporary fallback + # inherits the first quantized child's state and assumes all grouped quant modules + # share the same parallel groups. + parallel_state_source = next( + ( + (module, state) + for module in self.quant_modules + if (state := getattr(module, "parallel_state", None)) is not None + ), + None, + ) + if parallel_state_source is not None: + quant_module, parallel_state = parallel_state_source + if ( + torch.distributed.is_initialized() + and score_module not in self._warned_parallel_state_fallbacks + ): + warnings.warn( + "Distributed training is initialized but no parallel_state is set for " + f"score module {type(score_module)}. Using parallel_state from its first " + f"quantized child {type(quant_module)}. All grouped quant modules must " + "share the same parallel groups." + ) + self._warned_parallel_state_fallbacks.add(score_module) if parallel_state is None: total_score += importance.cpu().item() @@ -1438,7 +1465,232 @@ def _add_auto_quantize_score(grad_output, output_diff, score_tensor): score_tensor += _get_auto_quantize_score(grad_output, output_diff) -class AutoQuantizeGradientSearcher(_AutoQuantizeBaseSearcher): +class _AutoQuantizeBackwardScoringSession(ABC): + """Manage temporary model state used by activation-backward scoring.""" + + def __init__( + self, + model: nn.Module, + score_modules: Sequence[nn.Module], + is_param_grad_enabled: Callable, + verbose: bool = False, + ) -> None: + self.model = model + self.score_modules = tuple(score_modules) + self.is_param_grad_enabled = is_param_grad_enabled + self.verbose = verbose + self._stack = ExitStack() + self._original_forwards: dict[nn.Module, Callable] = {} + self._output_grad_hook_handles: set[Any] = set() + self._grad_accumulators: list[Any] = [] + + def __enter__(self): + """Install scoring hooks and parameter settings.""" + try: + hparams = list( + dict.fromkeys( + hparam + for module in self.score_modules + for hparam in module._hparams_for_scoring + ) + ) + for hparam in hparams: + self._stack.callback(setattr, hparam, "active", hparam.active) + + def patched_forward(module, *args, **kwargs): + return self.forward(module, *args, **kwargs) + + for module in self.score_modules: + original_forward = module.forward + self._original_forwards[module] = original_forward + had_instance_forward = "forward" in module.__dict__ + instance_forward = module.__dict__.get("forward") + module.forward = types.MethodType(patched_forward, module) + if had_instance_forward: + self._stack.callback(setattr, module, "forward", instance_forward) + else: + self._stack.callback(module.__dict__.pop, "forward", None) + + for name, param in self.model.named_parameters(): + requires_grad = param.requires_grad + enable_grad = self.is_param_grad_enabled(name, self.model) + param.requires_grad = enable_grad + self._stack.callback(setattr, param, "requires_grad", requires_grad) + if not enable_grad: + continue + if self.verbose: + print_rank_0(f"AutoQuantize: Enabling gradient for param {name}.") + accumulator, hook = create_param_grad_clear_hook(param) + self._grad_accumulators.append(accumulator) + self._stack.callback(hook.remove) + except Exception: + self._stack.close() + raise + return self + + def __exit__(self, exc_type, exc_value, traceback) -> None: + """Restore all model state changed for scoring.""" + self._clear_output_grad_hooks() + self._stack.close() + self._original_forwards.clear() + self._grad_accumulators.clear() + + def original_forward(self, module: nn.Module) -> Callable: + """Return the forward method saved before scoring.""" + return self._original_forwards[module] + + def _clear_output_grad_hooks(self) -> None: + """Remove output hooks whose backward pass has not run.""" + for handle in self._output_grad_hook_handles: + handle.remove() + self._output_grad_hook_handles.clear() + + def _register_output_grad_hook(self, output: torch.Tensor, hook: Callable) -> None: + """Attach an invocation-specific output-gradient hook for this session.""" + + def run_once(grad): + try: + return hook(grad) + finally: + handle.remove() + self._output_grad_hook_handles.discard(handle) + + handle = output.register_hook(run_once) + self._output_grad_hook_handles.add(handle) + + @abstractmethod + def forward(self, module: nn.Module, *args, **kwargs): + """Run a score module forward pass and collect method-specific state.""" + + +class _AutoQuantizeGradientScoringSession(_AutoQuantizeBackwardScoringSession): + """Collect gradient-based scores while candidate recipes are replayed.""" + + def forward(self, module: nn.Module, *args, **kwargs): + """Run the reference forward and cache each recipe's output perturbation.""" + no_quant_recipe = QuantRecipe(quant_cfg=None) + for hparam in module._hparams_for_scoring: + if hparam.is_configurable: + hparam.active = no_quant_recipe + + output = self.original_forward(module)(*args, **kwargs) + + # Checkpointed modules recompute with gradients enabled during backward. + base_output = output[0] if isinstance(output, tuple) else output + if not torch.is_grad_enabled() or not base_output.requires_grad: + return output + + output_diffs = {hparam: {} for hparam in module._hparams_for_scoring} + with torch.no_grad(): + for hparam in module._hparams_for_scoring: + if not hparam.is_configurable: + continue + for recipe in hparam.choices: + if recipe == no_quant_recipe: + continue + hparam.active = recipe + replay = self.original_forward(module)(*args, **kwargs) + output_diff = ( + replay[0] - output[0] if isinstance(replay, tuple) else replay - output + ) + output_diffs[hparam][recipe] = output_diff.detach() + hparam.active = no_quant_recipe + + self._register_output_grad_hook( + base_output, + lambda grad_output: self._accumulate_scores(module, output_diffs, grad_output), + ) + return output + + def _accumulate_scores(self, module, invocation_diffs, grad_output) -> None: + """Accumulate scores for the invocation that produced ``grad_output``.""" + if not torch.isfinite(grad_output).all(): + module_name = next( + name for name, child in self.model.named_modules() if child is module + ) + raise RuntimeError( + f"AutoQuantize: Non-finite output gradients in module '{module_name or ''}'. " + "Cannot compute reliable sensitivity scores. Check the model, data, and loss. " + "cuDNN SDPA backward on fully masked attention rows is one possible cause; " + "try torch.backends.cuda.enable_cudnn_sdp(False) before rerunning auto_quantize." + ) + for hparam, output_diffs in invocation_diffs.items(): + for recipe, output_diff in output_diffs.items(): + importance = hparam._importance_dict[recipe][module] + if importance is None: + hparam._importance_dict[recipe][module] = _get_auto_quantize_score( + grad_output, output_diff + ) + else: + _add_auto_quantize_score(grad_output, output_diff, importance) + + +class _AutoQuantizeBackwardScoringSearcher(_AutoQuantizeBaseSearcher): + """Share orchestration used by activation-backward scoring methods.""" + + score_module_rules = [ + # Score MoE projections together at their enclosing MLP or mixer output. + r"^(.*?\.mlp)\.experts\.\d+\.(gate_proj|up_proj|down_proj)$", + r"^(.*?\.mixer)\.experts\.\d+\.(up_proj|down_proj)$", + r"^(.*?)\.(\d+\.(w1|w2|w3))$", + r"^(.*?)\.((w1_linear|w2_linear|w3_linear)\.\d+)$", + ] + + _custom_support: list[tuple[Callable, Callable, Callable]] = [] + + @classmethod + def register_custom_support( + cls, + is_supported_checker: Callable, + grad_ckpt_context: Callable, + is_param_grad_enabled: Callable, + ) -> None: + """Register optional hooks for memory-efficient backward scoring. + + `is_supported_checker` selects models that use these hooks. + `grad_ckpt_context` enables their gradient-checkpointing context, and + `is_param_grad_enabled` selects the minimum parameters needed to propagate + activation gradients. + """ + cls._custom_support.append((is_supported_checker, grad_ckpt_context, is_param_grad_enabled)) + + def _configurable_score_modules(self) -> list[nn.Module]: + return [ + module + for module in self.model.modules() + if hasattr(module, "_hparams_for_scoring") + and any(hparam.is_configurable for hparam in module._hparams_for_scoring) + ] + + @abstractmethod + def _estimate_auto_quantize_scores(self, is_param_grad_enabled: Callable) -> None: + """Estimate scores while activation gradients are enabled.""" + + def estimate_sensitivity_scores(self) -> None: + """Run backward scoring with the first matching model-specific support hook.""" + self.model.eval() + + def default_is_param_grad_enabled(_name, _model): + return True + + grad_checkpointing_context = None + is_param_grad_enabled = default_is_param_grad_enabled + for is_supported, context_candidate, grad_candidate in self._custom_support: + if is_supported(self.model): + grad_checkpointing_context = context_candidate + is_param_grad_enabled = grad_candidate + break + + context = ( + grad_checkpointing_context(self.model) + if grad_checkpointing_context is not None + else nullcontext() + ) + with context: + self._estimate_auto_quantize_scores(is_param_grad_enabled) + + +class AutoQuantizeGradientSearcher(_AutoQuantizeBackwardScoringSearcher): """A searcher for AutoQuantize algorithm that uses gradient based score estimation. In AutoQuantize, we search for the best per-layer quantization configuration that minimizes the sum of per-layer @@ -1472,17 +1724,6 @@ class AutoQuantizeGradientSearcher(_AutoQuantizeBaseSearcher): method_name = "gradient" - score_module_rules = [ - # Use MLP layer output for gate_proj, up_proj, down_proj for Qwen3 like MoE models (local and shared experts) - r"^(.*?\.mlp)\.experts\.\d+\.(gate_proj|up_proj|down_proj)$", - r"^(.*?\.mixer)\.experts\.\d+\.(up_proj|down_proj)$", # NemotronH MoE experts - r"^(.*?)\.(\d+\.(w1|w2|w3))$", # mixtral experts - r"^(.*?)\.((w1_linear|w2_linear|w3_linear)\.\d+)$", # dbrx experts - ] - - # See `register_custom_support` for details - _custom_support: list[tuple[Callable, Callable, Callable]] = [] - @property def default_search_config(self): """Get the default config for the searcher.""" @@ -1511,30 +1752,6 @@ def sanitize_search_config(self, config: SearchConfig | None) -> SearchConfig: return config - @classmethod - def register_custom_support( - cls, - is_supported_checker: Callable, - grad_ckpt_context: Callable, - is_param_grad_enabled: Callable, - ) -> None: - """(Optional) Register custom support for `AutoQuantize` score estimation. - - This custom support is used to enable memory/compute efficient backward gradient propagation. This involves: - - - `grad_ckpt_context`: backward pass with gradient checkpointing enabled - - `is_param_grad_enabled`: AutoQuantize only needs activation gradients to be computed (not weight - gradients). `is_param_grad_enabled` is used to select which parameters should have gradients enabled, - limiting gradient computation to only what's needed for activation gradients. For LLMs, to trigger all - activation gradient computation, just enabling the embedding layer weight gradient is sufficient. This will - enable gradient computation for all the activation gradients downstream. - - If the `is_supported_checker(model)` returns True, the `grad_ckpt_context(model)` will be - used to enable gradient checkpointing and `is_param_grad_enabled(pname, model)` - will be used to select which parameters have gradients enabled to minimize gradient computation. - """ - cls._custom_support.append((is_supported_checker, grad_ckpt_context, is_param_grad_enabled)) - def _get_default_forward_backward_step(self): def forward_backward_step(model, data): output = self.config["forward_step"](model, data) @@ -1552,143 +1769,35 @@ def forward_backward_step(model, data): @torch.enable_grad() def _estimate_auto_quantize_scores(self, is_param_grad_enabled): - # TODO: remove the no-quant recipe - def auto_quantize_score_estimate_forward(module, input, *args, **kwargs): - for hparam in module._hparams_for_scoring: - if hparam.is_configurable: - hparam.active = QuantRecipe(quant_cfg=None) - - output = module._forward_original(input, *args, **kwargs) - - # If gradient checkpointing is enabled, gradient will not be enabled in the global forward pass. - # With gradient checkpointing, gradients are computed in the local forward pass during backward pass - - # Lets compute the output_diff and save it in memory only if gradient is enabled to be memory efficient - if not torch.is_grad_enabled(): - return output - - module.output_diff_dict = {hparam: {} for hparam in module._hparams_for_scoring} - with torch.no_grad(): - for hparam in module._hparams_for_scoring: - if not hparam.is_configurable: - continue - for recipe in hparam.choices: - if recipe == QuantRecipe(quant_cfg=None): - continue - hparam.active = recipe - output_diff = module._forward_original(input, *args, **kwargs) - - if isinstance(output_diff, tuple): - output_diff = output_diff[0] - output[0] - else: - output_diff -= output - module.output_diff_dict[hparam][recipe] = output_diff.detach() - - # Disable the configurable hparam now that we have computed the diff - hparam.active = QuantRecipe(quant_cfg=None) - - return output - - def backward_hook(module, grad_input, grad_output): - for hparam, output_diff_dict in module.output_diff_dict.items(): - for recipe, output_diff in output_diff_dict.items(): - if hparam._importance_dict[recipe][module] is None: - hparam._importance_dict[recipe][module] = _get_auto_quantize_score( - grad_output[0], output_diff - ) - else: - _add_auto_quantize_score( - grad_output[0], output_diff, hparam._importance_dict[recipe][module] - ) - - def setup_params_for_score_estimation(name, param, params_metadata, enable_grad=True): - # Let us delete the gradient as soon as they are computed to save memory - params_metadata[name] = {"requires_grad": param.requires_grad} - param.requires_grad = enable_grad - if not enable_grad: - return - if self.config.get("verbose", False): - print_rank_0(f"AutoQuantize: Enabling gradient for param {name}.") - accum_grad, handle = create_param_grad_clear_hook(param) - params_metadata[name]["accum_grad"] = accum_grad # We need to keep the accum_grad alive - params_metadata[name]["handle"] = handle - - def setup_module_for_score_estimation(module): - module._forward_original = module.forward - module.forward = types.MethodType(auto_quantize_score_estimate_forward, module) - module._backward_hook_handle = module.register_full_backward_hook(backward_hook) - - def cleanup_module_after_score_estimation(module): - module.forward = module._forward_original - del module._forward_original - - module._backward_hook_handle.remove() - - def cleanup_params_after_score_estimation(name, param, params_metadata): - param.requires_grad = params_metadata[name]["requires_grad"] - handle = params_metadata[name].get("handle", None) - if handle is not None: - handle.remove() - - score_modules = set() - for name, module in self.model.named_modules(): - if ( - hasattr(module, "_hparams_for_scoring") - and any(hparam.is_configurable for hparam in module._hparams_for_scoring) - and module not in score_modules - ): - # Monkey patch the forward methods to cache (Q(Y) - Y) - setup_module_for_score_estimation(module) - score_modules.add(module) - - params_metadata = {} - for name, param in self.model.named_parameters(): - setup_params_for_score_estimation( - name, param, params_metadata, is_param_grad_enabled(name, self.model) + score_modules = self._configurable_score_modules() + with _AutoQuantizeGradientScoringSession( + self.model, + score_modules, + is_param_grad_enabled, + verbose=self.config.get("verbose", False), + ) as scoring_session: + gc.collect() + if torch.cuda.is_available(): + torch.cuda.reset_peak_memory_stats() + report_memory("AutoQuantize: starting score estimation, ") + + def score_step(model, data): + try: + return self.config["forward_backward_step"](model, data) + finally: + scoring_session._clear_output_grad_hooks() + + self._run_func( + score_step, + num_iters=self.config["num_score_steps"], + desc="Estimating auto_quantize scores", ) - gc.collect() - if torch.cuda.is_available(): - torch.cuda.reset_peak_memory_stats() - report_memory("AutoQuantize: starting score estimation, ") - - self._run_func( - self.config["forward_backward_step"], - num_iters=self.config["num_score_steps"], - desc="Estimating auto_quantize scores", - ) - - if torch.cuda.is_available(): - report_memory("AutoQuantize: After score estimation") - - for module in score_modules: - cleanup_module_after_score_estimation(module) + if torch.cuda.is_available(): + report_memory("AutoQuantize: After score estimation") - for name, param in self.model.named_parameters(): - cleanup_params_after_score_estimation(name, param, params_metadata) - - # Delete the params_metadata - del params_metadata gc.collect() - def estimate_sensitivity_scores(self) -> None: - """Estimate sensitivity scores using hessian approximation.""" - self.model.eval() - - def _default_is_param_grad_enabled(pname, model): - return True - - grad_checkpointing_ctxt = None - is_param_grad_enabled = _default_is_param_grad_enabled - for is_supported_checker, ctxt_candidate, grad_enabled_candidate in self._custom_support: - if is_supported_checker(self.model): - grad_checkpointing_ctxt = ctxt_candidate - is_param_grad_enabled = grad_enabled_candidate - break - - with grad_checkpointing_ctxt(self.model) if grad_checkpointing_ctxt else nullcontext(): - self._estimate_auto_quantize_scores(is_param_grad_enabled) - def run_search_with_stats(self, max_weight_size, verbose=False): """Linear Programming Solve for gradient based auto_quantize. diff --git a/modelopt/torch/quantization/export_onnx.py b/modelopt/torch/quantization/export_onnx.py index e5778c3c96b..2ad3383e572 100644 --- a/modelopt/torch/quantization/export_onnx.py +++ b/modelopt/torch/quantization/export_onnx.py @@ -225,9 +225,12 @@ def _fp8_quantize( "Constant", value_t=torch.tensor(scale_inv).to(torch_dtype_map[inputs.type().scalarType()]), ) - return g.op("trt::TRT_FP8QuantizeLinear", inputs, scale).setType( + quantized = g.op("trt::TRT_FP8QuantizeLinear", inputs, scale).setType( inputs.type().with_dtype(torch.uint8).with_sizes(output_shape) ) + # PyTorch runs shape inference before setType for custom ops, so refresh its reliability state. + torch._C._jit_pass_onnx_node_shape_type_inference(quantized.node(), g.params_dict, g.opset) + return quantized def _fp8_dequantize( diff --git a/modelopt/torch/quantization/model_calib.py b/modelopt/torch/quantization/model_calib.py index c85e97a104d..2a504597c3f 100644 --- a/modelopt/torch/quantization/model_calib.py +++ b/modelopt/torch/quantization/model_calib.py @@ -73,6 +73,7 @@ "local_hessian_calibrate", "lsq", "max_calibrate", + "mse_calibrate", "nvfp4_act_headroom_calibrate", "smoothquant", "svdquant", diff --git a/modelopt/torch/quantization/model_quant.py b/modelopt/torch/quantization/model_quant.py index 3f6040fd4ef..4588c6ac54b 100644 --- a/modelopt/torch/quantization/model_quant.py +++ b/modelopt/torch/quantization/model_quant.py @@ -144,6 +144,85 @@ def postprocess_amax(model: nn.Module, key: str, post_process_fn) -> nn.Module: return model +_SKIP_WEIGHT_QUANT_CHECK_ENV = "MODELOPT_SKIP_WEIGHT_QUANT_CHECK" + + +def _check_weight_quantization_took_effect(model: nn.Module, config: QuantizeConfig) -> None: + """Raise when a config asks for weight quantization but no weight quantizer is enabled. + + A config whose module patterns do not match the model is not an error to + :func:`set_quantizer_by_cfg` — every pattern simply matches nothing — so the run + proceeds through calibration and export and produces a checkpoint that is silently + unquantized (``"quant_algo": null`` with an empty ``quantized_layers``). That has + bitten several MoE architectures whose module naming differs from the wildcards in + the general recipes, and it is only noticed when someone reads the exported config. + + By the time this runs, :func:`set_quantizer_by_cfg` (or the ``apply_mode`` conversion + that calls it) has already applied ``config`` to ``model``, so each quantizer's + ``is_enabled`` *is* the true outcome of that application — checking it directly cannot + diverge from what the config actually did. An earlier version of this check instead + re-derived "did this pattern match anything?" via a separate matcher call, which missed + the case of two *different* overlapping patterns (e.g. ``*weight_quantizer`` enabling + something a later, broader ``*`` then disables): the narrower pattern registered as + "matched" even though the quantizer it matched ended up disabled. + + A config that never asks for weight quantization (activation-only or KV-cache-only) + must not raise, so the check first looks at the config's own intent — via each + pattern's *final* entry, since entries apply in order and the last one for a pattern + wins — before looking at the model at all. + """ + if os.environ.get(_SKIP_WEIGHT_QUANT_CHECK_ENV) == "1": + return + + # Later entries override earlier ones, so only each pattern's final state states intent. + # A pattern naming ``weight_quantizer`` explicitly (the common case, e.g. + # ``*weight_quantizer``, ``*.experts.*weight_quantizer``) is caught by the substring + # check. A broad wildcard that never mentions "weight" -- a bare ``"*"`` catch-all, or + # ``"*_quantizer"`` -- can still match weight quantizers at runtime, so it must count + # too, or a config built only from patterns like that would never trip the guard + # regardless of what the model contains. ``fnmatch`` against the literal probe string + # ``"weight_quantizer"`` catches those (a pattern matching that bare name is, by + # construction, asking for one) without replacing the substring check: the probe alone + # would miss ``*.experts.*weight_quantizer`` (there is no ``.experts.`` in the probe + # string), which is what recognizes model-scoped patterns like the Step / MoE recipes use. + last_entry_per_pattern = {entry.quantizer_name: entry for entry in config.quant_cfg} + weight_patterns = [ + pattern + for pattern, entry in last_entry_per_pattern.items() + if entry.enable + and ("weight_quantizer" in pattern or fnmatch.fnmatch("weight_quantizer", pattern)) + ] + if not weight_patterns: + return + + # `SequentialQuantizer.is_enabled` delegates to its first member, so a list-valued `cfg`'s + # quantizers are already covered here without naming `SequentialQuantizer` explicitly: + # `named_modules()` recurses into the container and yields those children too, individually, + # named `...weight_quantizer.0` / `.1` (the substring match below still applies to them). + if any( + module.is_enabled + for name, module in model.named_modules() + if isinstance(module, TensorQuantizer) and "weight_quantizer" in name + ): + return + + patterns = "\n ".join(sorted(weight_patterns)) + raise RuntimeError( + "The quantization config asks for weight quantization but no weight quantizer is " + f"enabled, so nothing would be quantized. These patterns asked for it:\n {patterns}\n" + "Either the patterns do not match this architecture's module names (check the " + "model-specific recipes under modelopt_recipes/huggingface//), or the " + "modules holding the weights were never converted to quantized modules (an " + "unsupported custom module, e.g. a trust_remote_code MoE layout).\n" + "Under pipeline parallelism, a rank whose local stage genuinely has none of the " + "targeted modules (e.g. a pure-attention stage under an experts-only recipe) hits " + "this too, while other ranks proceed into calibration -- a collective hang, not " + f"just a wrong per-rank verdict. Set {_SKIP_WEIGHT_QUANT_CHECK_ENV}=1 to bypass this " + "check in that situation -- note this is process-global, so it silences the check " + "on every rank, not only the one with the legitimately empty stage." + ) + + def quantize( model: nn.Module, config: dict[str, Any | QuantizeConfig], @@ -241,12 +320,14 @@ def forward_loop(model) -> None: Returns: A pytorch model which has been quantized and calibrated. """ + quantize_config = QuantizeConfig(**dict(config)) if not is_quantized(model): model = apply_mode(model, mode=[("quantize", dict(config))], registry=QuantizeModeRegistry) else: # Already quantized, so lets apply the quant_cfg from the config - quant_cfg = QuantizeConfig(**dict(config)).quant_cfg - set_quantizer_by_cfg(model, quant_cfg) + set_quantizer_by_cfg(model, quantize_config.quant_cfg) + # Fail before calibration rather than after exporting an unquantized checkpoint. + _check_weight_quantization_took_effect(model, quantize_config) return calibrate(model, config.get("algorithm"), forward_loop=forward_loop) diff --git a/modelopt/torch/quantization/nn/modules/quant_module.py b/modelopt/torch/quantization/nn/modules/quant_module.py index 1b5fb09062e..522f2d0db09 100644 --- a/modelopt/torch/quantization/nn/modules/quant_module.py +++ b/modelopt/torch/quantization/nn/modules/quant_module.py @@ -128,6 +128,17 @@ def iter_weights_for_calibration(self): weight_quantizer = getattr(self, quantizer_attr_names(weight_name).weight_quantizer) yield getattr(self, weight_name), weight_quantizer + def iter_weight_quantizers_for_calibration(self): + """Yield just the weight quantizers, without materializing the weight views. + + Callers that only inspect quantizer state must use this rather than discarding the + weight from :meth:`iter_weights_for_calibration`. Reading a weight is free here, but a + subclass whose weight view costs something -- the fused-MoE modules slice a 3-D DTensor + per expert, one redistribute collective each under FSDP2 -- overrides this to skip it. + """ + for _, weight_quantizer in self.iter_weights_for_calibration(): + yield weight_quantizer + @staticmethod @torch.no_grad() def _fold_weight_quantizer( diff --git a/modelopt/torch/quantization/nn/modules/tensor_quantizer.py b/modelopt/torch/quantization/nn/modules/tensor_quantizer.py index 18b97ac2774..a2349c6eb3f 100644 --- a/modelopt/torch/quantization/nn/modules/tensor_quantizer.py +++ b/modelopt/torch/quantization/nn/modules/tensor_quantizer.py @@ -1084,10 +1084,38 @@ def _check_per_channel_block_sizes(block_sizes): # remove block_sizes self._block_sizes = None + def _sanitize_export_amax(self, amax: torch.Tensor) -> torch.Tensor: + """Replace zero/NaN amax entries with ``maxbound`` so exported scales stay positive. + + A zero amax means calibration never activated the layer; downstream exporters divide by + the exported amax, so a zero would fail export or produce inf at inference. Kept + branch-free because export flows may pass a meta ``amax``; only the warning reads values. + """ + sanitized = torch.nan_to_num( + torch.where(amax == 0, torch.full_like(amax, self.maxbound), amax), + nan=self.maxbound, + ) + + if not amax.is_meta: + num_invalid = int((torch.isnan(amax) | (amax == 0)).sum()) + if num_invalid: + warnings.warn( + f"{num_invalid}/{amax.numel()} amax entries of this " + f"{type(self).__name__} are zero or NaN at export time, which means " + "calibration never activated the corresponding layer or expert (or saw NaN " + "activations). Substituting maxbound so the exported scaling factor stays " + "positive. Consider increasing the calibration size if the layer is expected " + "to be active.", + stacklevel=3, + ) + return sanitized + def export_amax(self) -> torch.Tensor | None: """Export correctly formatted/shaped amax.""" if self.block_sizes is not None and self.block_sizes.get("type", None) == "dynamic": - return self.amax + # Dynamic block quantizers keep a per-tensor amax (the NVFP4 second-level scale) that + # needs no reshaping, but it still has to be positive for the exporters. + return None if self.amax is None else self._sanitize_export_amax(self.amax) if self.amax is None: return None @@ -1096,8 +1124,7 @@ def export_amax(self) -> torch.Tensor | None: amax = self.amax else: amax = self.amax.reshape(self._amax_shape_for_export) - amax[amax == 0] = self.maxbound - amax = torch.nan_to_num(amax, nan=self.maxbound) + amax = self._sanitize_export_amax(amax) clamp_min, clamp_max = torch.finfo(amax.dtype).tiny, torch.finfo(amax.dtype).max amax = amax.clamp(min=clamp_min, max=clamp_max) diff --git a/modelopt/torch/quantization/plugins/diffusion/diffusers.py b/modelopt/torch/quantization/plugins/diffusion/diffusers.py index f2f6a702479..e92c775c05b 100644 --- a/modelopt/torch/quantization/plugins/diffusion/diffusers.py +++ b/modelopt/torch/quantization/plugins/diffusion/diffusers.py @@ -141,6 +141,14 @@ def _quantized_sdpa(self, *args, **kwargs): q_quantized_scale = self.q_bmm_quantizer._get_amax(query) k_quantized_scale = self.k_bmm_quantizer._get_amax(key) v_quantized_scale = self.v_bmm_quantizer._get_amax(value) + disable_fp8_mha = not all( + quantizer.is_enabled and quantizer.is_fp8 + for quantizer in ( + self.q_bmm_quantizer, + self.k_bmm_quantizer, + self.v_bmm_quantizer, + ) + ) # We don't need to calibrate the output of softmax return self.bmm2_output_quantizer( @@ -155,7 +163,7 @@ def _quantized_sdpa(self, *args, **kwargs): self.q_bmm_quantizer.trt_high_precision_dtype if hasattr(self.q_bmm_quantizer, "trt_high_precision_dtype") else "Half", - self._disable_fp8_mha if hasattr(self, "_disable_fp8_mha") else True, + disable_fp8_mha, ) ) diff --git a/modelopt/torch/quantization/plugins/huggingface.py b/modelopt/torch/quantization/plugins/huggingface.py index 4acb4d30dfa..12e7c49b687 100644 --- a/modelopt/torch/quantization/plugins/huggingface.py +++ b/modelopt/torch/quantization/plugins/huggingface.py @@ -17,6 +17,7 @@ import inspect import logging +import re import warnings from contextlib import contextmanager from functools import partial @@ -246,8 +247,17 @@ def _quantized_attention( return self._eager_p_qdq_attention( original_attention_interface, query_states, key_states, value_states, **kwargs ) - return self._triton_qdq_attention( - p_qdq, query_states, key_states, value_states, **kwargs + if self.p_bmm_quantizer._if_quant: + return self._triton_qdq_attention( + p_qdq, query_states, key_states, value_states, **kwargs + ) + + # Fused P-QDQ paths bypass TensorQuantizer.forward(). + if not self.p_bmm_quantizer.is_enabled or not self.p_bmm_quantizer._if_quant: + if args: + kwargs.pop("attention_mask", None) + return original_attention_interface( + self, query_states, key_states, value_states, *args, **kwargs ) if kitchen is not None and self.kitchen_attn_fn is None: @@ -1100,6 +1110,24 @@ def iter_weights_for_calibration(self): for idx, q in enumerate(quantizers): yield weight[idx], q + def iter_weight_quantizers_for_calibration(self): + """Yield the per-expert weight quantizers without slicing the fused 3-D weight. + + Overrides the base, which would go through :meth:`iter_weights_for_calibration` and + evaluate ``weight[idx]``. Under FSDP2 that weight is a DTensor, so every expert slice + dispatches a redistribute collective to produce a value a quantizer-only caller discards. + """ + for weight_name, quantizers_name in ( + (self._first_proj_attr, self._first_proj_weight_quantizers_attr), + ("down_proj", "down_proj_weight_quantizers"), + ): + # Same skip condition as iter_weights_for_calibration, so the two stay in lockstep. + # Fetching the attribute is free; only indexing into it would collective. + quantizers = getattr(self, quantizers_name, None) + if getattr(self, weight_name, None) is None or quantizers is None: + continue + yield from quantizers + def fold_weight(self, keep_attrs: bool = False): """Bake each per-expert weight quantizer into its slice of the fused 3-D weight. @@ -1877,7 +1905,7 @@ def _is_param_grad_enabled_for_auto_quantize(pname, model): class _QuantMoELinear(QuantModule): - """Quantization wrapper for Step3p5 MoELinear modules (fused expert weights). + """Quantization wrapper for expert-indexed MoELinear modules (fused expert weights). MoELinear has weight shape [num_experts, out_features, in_features] with forward(x, expert_id). We expand it into per-expert nn.Linear modules so @@ -1895,6 +1923,19 @@ class _QuantMoELinear(QuantModule): def _setup(self): from accelerate import init_empty_weights + # Accelerate's CPU/disk offload (`device_map="auto"`, `--offload_folder`) leaves + # `weight` as a meta tensor and keeps the real value in the module's offload hook, + # keyed on the original `weight` name. Expanding that would copy meta storage into + # every expert and then delete the key the hook restores into, silently producing a + # checkpoint of zeros. Refuse instead of corrupting. + if self.weight.is_meta or getattr(getattr(self, "_hf_hook", None), "offload", False): + raise NotImplementedError( + f"{type(self).__name__}: expert-indexed MoELinear weights cannot be quantized " + "while offloaded by Accelerate (the weight is a meta tensor whose value lives " + "in the offload hook). Load the model without CPU/disk offload — more GPUs, or " + "a device_map that keeps the MoE layers resident — and re-run." + ) + dtype, device = self.weight.dtype, self.weight.device with init_empty_weights(): @@ -1914,38 +1955,143 @@ def _setup(self): self.experts = experts def forward(self, x, expert_id): - # experts[expert_id] is a _QuantLinear after quantization wrapping, - # providing per-expert input_quantizer and weight_quantizer. - # Cast input to match expert weight dtype before linear operation, - # then cast output to float32 to match original MoELinear forward behavior. + # experts[expert_id] is a _QuantLinear after quantization wrapping, providing + # per-expert input_quantizer and weight_quantizer. + # + # MoELinear.forward always promotes to fp32 for the matmul regardless of storage + # dtype (`F.linear(x.float(), self.weight[expert_id].float())`), so leaving the + # expert's weight at its native storage dtype (e.g. bf16) and downcasting x to + # match before the matmul would compute in bf16 and change the model's output even + # with every quantizer disabled -- the bf16 rounding this class exists to quantize + # *past*, not to reintroduce as a side effect of conversion. + # + # A prior version of this fix instead expanded every expert's weight in fp32 + # permanently in `_setup`. That reproduces Step's fp32 compute but turns a per-call + # transient promotion into persistent model state: on Step-3.7's full routed-expert + # set (42 layers x 3 projections x 288 experts x 4096 x 1280), doubling from bf16 to + # fp32 adds roughly 354 GiB held throughout calibration, on top of device placement + # already sized for bf16 -- a model that loaded successfully can then OOM. It also + # left disabled/unquantized experts reconstructed at fp32 in the exported checkpoint. + # + # Instead, only the one expert actually being called is promoted, transiently, for + # the duration of this one call -- matching Step's own per-call `.float()` memory + # profile instead of Step-3.7's full expert set. `expert.weight` is read here + # outside any `quantize_weight()` context, so `_get_quantized_weight` passes it + # through unchanged and this is the real underlying nn.Parameter (the same pattern + # `_setup` above uses), not a value computed by the quantizer -- so reassigning its + # `.data` genuinely mutates the persisted storage, not a transient wrapper. + # + # This must keep calling `expert(x)` (`__call__`, not `.forward()`) rather than + # reimplementing the input/weight-quantize/output-quantize sequence inline: some + # calibration algorithms (e.g. `local_hessian_calibrate`) register a + # `forward_pre_hook` directly on the quantized Linear module, which only fires + # through standard `nn.Module.__call__` dispatch. expert = self.experts[expert_id] - x = x.to(expert.weight.dtype) - return expert(x).float() + original_weight = expert.weight.data + with torch.no_grad(): + expert.weight.data = original_weight.float() + try: + out = expert(x.float()) + finally: + with torch.no_grad(): + expert.weight.data = original_weight + return out.float() + + +def _is_expert_indexed_moe_linear(module: nn.Module) -> bool: + """Whether ``module`` packs one projection's experts into an expert-indexed 3-D weight. + + The Step family (``stepfun-ai/Step-3.5-Flash``, ``stepfun-ai/Step-3.7-Flash``) ships a + custom ``MoELinear`` via ``trust_remote_code``: a plain ``nn.Module`` holding a single + ``weight`` of shape ``[num_experts, out_features, in_features]``, whose + ``forward(x, expert_id)`` runs ``F.linear`` against the selected expert's slice. The + weights therefore live on the projection submodule rather than on the expert container, + which is what :func:`_fused_experts_wrapper_class` looks for, and the module is not an + ``nn.Linear``, so neither the fused-experts path nor the plain linear path claims it. + + Detection is structural rather than keyed on class names so new Step revisions are picked + up without another hardcoded name, but the shape alone is not a sufficient contract: the + replacement forward indexes ``self.experts[expert_id]``, so it only works for callers that + pass a **scalar expert index**. Grouped-GEMM MoE layers share the exact same 3-D weight and + attribute set while passing a per-expert token-count *tensor* instead (e.g. Moondream3's + ``MoeFusedLinear.forward(input, m_sizes)``, which would raise ``TypeError: only integer + tensors of a single element can be converted to an index`` on the first calibration + forward). The second parameter must therefore be named ``expert_id``, which is the + scalar-index contract both Step revisions declare. + """ + weight = getattr(module, "weight", None) + if not isinstance(weight, (nn.Parameter, Tensor)) or weight.dim() != 3: + return False + if not all(hasattr(module, attr) for attr in ("num_experts", "in_features", "out_features")): + return False + # The wrapper rebuilds the weight as `num_experts` Linears of (out_features, in_features), + # so a 3-D weight laid out any other way would silently copy the wrong slices. + if tuple(weight.shape) != (module.num_experts, module.out_features, module.in_features): + return False + try: + params = list(inspect.signature(type(module).forward).parameters.values())[1:] + except (TypeError, ValueError): + return False + # The replacement forward is exactly `(x, expert_id)`, so anything the caller could pass + # beyond those two — a keyword-only `router_state`, *args, **kwargs — would raise once + # converted. Require the signature to match what the wrapper can honour. + return ( + len(params) == 2 + and all(p.kind is p.POSITIONAL_OR_KEYWORD and p.default is p.empty for p in params) + and params[1].name == "expert_id" + ) + + +_STEP_FAMILY_RE = re.compile(r"(?i)^step\d") + + +def _is_step_family_model(model: nn.Module) -> bool: + """Whether ``model`` is a Step-family root model (Step-3.5, Step-3.7, or a future revision). + + Matched against the ``step`` convention shared by ``model_type`` (``"step3p5"``, + ``"step3p7"``) and the remote-code class name (``Step3p5ForCausalLM``, + ``Step3p7ForConditionalGeneration``), not an exact revision, so a new Step release is + still picked up without another hardcoded name. This is deliberately narrower than the + shape/signature check in :func:`_is_expert_indexed_moe_linear` alone: that check accepts + any module with a matching 3-D weight and an ``(x, expert_id)`` forward, which is a + coincidence risk on its own -- an unrelated architecture happening to reuse the parameter + name ``expert_id`` with different semantics (a per-expert bias or post-scale, say) would + be claimed and have that behavior silently dropped by the replacement wrapper. Gating on + the model family keeps the shape check doing what it is actually good at: telling + Step revisions apart without a class-name allowlist, rather than distinguishing Step + from arbitrary third-party MoE code. + """ + model_type = str(getattr(getattr(model, "config", None), "model_type", "") or "") + return bool(_STEP_FAMILY_RE.match(model_type) or _STEP_FAMILY_RE.match(type(model).__name__)) -def register_step3p5_moe_on_the_fly(model): - """Register Step3p5 MoELinear for quantization. +def register_moe_linear_on_the_fly(model): + """Register expert-indexed ``MoELinear`` modules (Step-3.5 / Step-3.7) for quantization. - Step3p5 uses a custom MoELinear class (loaded via trust_remote_code) with - weight shape [num_experts, out_features, in_features] and forward(x, expert_id). - We detect it by model class name, then grab the type from the first MoE layer. + Without this the routed experts carry no quantizer at all: an experts-only recipe matches + nothing and the export writes a checkpoint with ``quant_algo: null``. """ - if type(model).__name__ not in ("Step3p5ForCausalLM", "Step3p5Model"): + if not _is_step_family_model(model): return - for module in model.modules(): - if type(module).__name__ == "Step3p5MoEMLP": - moe_linear_type = type(module.up_proj) - if QuantModuleRegistry.get(moe_linear_type) is None: - QuantModuleRegistry.register({moe_linear_type: f"hf.{moe_linear_type.__name__}"})( - _QuantMoELinear - ) - break + visited_types = set() + for name, module in model.named_modules(): + mod_type = type(module) + if mod_type in visited_types or QuantModuleRegistry.get(mod_type) is not None: + continue + visited_types.add(mod_type) + + if _is_expert_indexed_moe_linear(module): + print( + f"\033[1mDetected expert-indexed MoE linear '{name}' of type " + f"{mod_type.__name__}, registering with _QuantMoELinear.\033[0m" + ) + QuantModuleRegistry.register({mod_type: f"hf.{mod_type.__name__}"})(_QuantMoELinear) def _reconstruct_fused_moe_linear(model: nn.Module) -> None: - """Reconstruct QuantMoELinear per-expert weights back to original 3D MoELinear format. + """Reconstruct :class:`_QuantMoELinear` per-expert weights back to the 3-D MoELinear format. - After _process_quantized_modules, each expert's nn.Linear inside QuantMoELinear has: + After _process_quantized_modules, each expert's nn.Linear inside the wrapper has: - weight: fp4-quantized tensor [out_features, in_features] - weight_scale, weight_scale_2: per-block / global scales - input_scale: activation scale (if calibrated) @@ -1953,12 +2099,12 @@ def _reconstruct_fused_moe_linear(model: nn.Module) -> None: This stacks them back into the original MoELinear layout so the exported state_dict uses the original key names (e.g. moe.up_proj.weight with shape [N, out, in]). - Note: QuantMoELinear is the dynamically generated class name (Quant + MoELinear), - not _QuantMoELinear which is the implementation class. + Matched by wrapper type rather than by the dynamically generated class name (``Quant`` + + the model's own class name): a model whose class is not spelled ``MoELinear`` would + otherwise quantize normally but export unusable per-expert keys. """ for _name, module in model.named_modules(): - # Match QuantMoELinear (dynamically generated name) not _QuantMoELinear (implementation class) - if type(module).__name__ != "QuantMoELinear": + if not isinstance(module, _QuantMoELinear): continue n = module.num_experts @@ -1988,7 +2134,7 @@ def _reconstruct_fused_moe_linear(model: nn.Module) -> None: [ register_falcon_linears_on_the_fly, register_dbrx_moe_on_the_fly, - register_step3p5_moe_on_the_fly, + register_moe_linear_on_the_fly, register_fused_experts_on_the_fly, force_eager_experts_impl_on_the_fly, register_sparse_moe_on_the_fly, diff --git a/modelopt/torch/quantization/plugins/megatron.py b/modelopt/torch/quantization/plugins/megatron.py index 8c2352b5265..ec1958649dd 100644 --- a/modelopt/torch/quantization/plugins/megatron.py +++ b/modelopt/torch/quantization/plugins/megatron.py @@ -115,6 +115,39 @@ def _check_nvfp4_static_tp_supported(model: torch.nn.Module) -> None: ) +def _initialize_grouped_weight_quantizer_state(module: torch.nn.Module) -> None: + """Create per-expert buffers as destinations for the subsequent checkpoint load.""" + grouped_leaves = [ + quantizer if isinstance(quantizer, torch.nn.Sequential) else [quantizer] + for quantizer in [module.weight_quantizer[idx] for idx in range(module.num_gemms)] + ] + for sibling_leaves in zip(*grouped_leaves): + eligible_leaves = [ + quantizer + for quantizer in sibling_leaves + if ( + quantizer.is_enabled + and not quantizer.is_mx_format + and not getattr(quantizer, "_dynamic", False) + and not getattr(quantizer, "_lsq", False) + ) + ] + for state_name in ("_amax", "_global_amax"): + reference = next( + ( + state + for quantizer in eligible_leaves + if (state := getattr(quantizer, state_name, None)) is not None + ), + None, + ) + if reference is None: + continue + for quantizer in eligible_leaves: + if getattr(quantizer, state_name, None) is None: + quantizer.register_buffer(state_name, torch.zeros_like(reference)) + + def real_quant_module_get_extra_state(self) -> dict: """Populating real_quantizer_state and q_tensor_state.""" extra_state = {} @@ -818,6 +851,9 @@ class _QuantTELayerNormColumnParallelLinear( # Quantized subclasses to support TEGroupedLinear quantization class _QuantMegatronTEGroupedLinear(_QuantTEGroupedLinear, _MegatronParallelLinear): + def modelopt_post_load_extra_state(self): + _initialize_grouped_weight_quantizer_state(self) + def _load_from_state_dict(self, state_dict, prefix, *args, **kwargs): # _sharded_state_dict_grouped adds _extra_state{gemm_idx} for gemm_idx:[1, num_gemms] in # sharded_state_dict which is same as _extra_state. The _extra_state{gemm_idx} is used for diff --git a/modelopt/torch/quantization/utils/core_utils.py b/modelopt/torch/quantization/utils/core_utils.py index 3bd31a03128..0adcb35a534 100644 --- a/modelopt/torch/quantization/utils/core_utils.py +++ b/modelopt/torch/quantization/utils/core_utils.py @@ -17,6 +17,7 @@ import copy import itertools +import warnings from collections import namedtuple from contextlib import ExitStack, contextmanager, nullcontext from typing import TYPE_CHECKING, Any @@ -1107,10 +1108,13 @@ def promote_static_block_weight_quantizers(model: nn.Module) -> int: for quantizer in state._member_quantizers() } converted = 0 + # Quantizer-only iteration: this loop reads quantizer state and never the weight, and asking + # for weights here costs one DTensor collective per expert on a fused-MoE model under FSDP2. + warned_dtensor_amax = False for _name, module in list(model.named_modules()): if not isinstance(module, QuantModule): continue - for _, quantizer in module.iter_weights_for_calibration(): + for quantizer in module.iter_weight_quantizers_for_calibration(): if isinstance(quantizer, SequentialQuantizer): if len(quantizer) == 0: continue @@ -1123,6 +1127,16 @@ def promote_static_block_weight_quantizers(model: nn.Module) -> int: amax = quantizer.amax if amax is None: continue + if isinstance(amax, DTensor) and not warned_dtensor_amax: + # _amax is a buffer, so FSDP2 leaves it replicated and reduce_amax below stays + # local. If that ever stops holding, the reduction becomes a per-quantizer + # collective and this loop needs a batched reduction instead. + warned_dtensor_amax = True + warnings.warn( + "promote_static_block_weight_quantizers: _amax is a DTensor, so the " + "per-quantizer global-amax reduction is a collective. Batch the reduction " + "before running this at scale." + ) if quantizer.is_nvfp4_static: # Grouped siblings share one canonical global_amax (common FP8 grid); otherwise # fall back to this quantizer's own per-block amax. diff --git a/modelopt/torch/speculative/plugins/hf_training_args.py b/modelopt/torch/speculative/plugins/hf_training_args.py index 38d3f483e6b..b121d9f260d 100644 --- a/modelopt/torch/speculative/plugins/hf_training_args.py +++ b/modelopt/torch/speculative/plugins/hf_training_args.py @@ -44,6 +44,11 @@ class ModelArguments(BaseModel): model_name_or_path: str | None = "TinyLlama/TinyLlama-1.1B-Chat-v1.0" use_fake_base_for_offline: bool = False trust_remote_code: bool = False + # Optional config field overrides applied to the loaded model config (and its + # text_config) before instantiation. Needed for checkpoints whose config doesn't + # round-trip cleanly through transformers (e.g. Cosmos3's Qwen3-VL text tower, + # where intermediate_size/num_key_value_heads don't propagate from text_config). + config_overrides: dict | None = None class DataArguments(BaseModel): diff --git a/modelopt/torch/speculative/utils.py b/modelopt/torch/speculative/utils.py index 9fa8fde5e15..c419d24588f 100644 --- a/modelopt/torch/speculative/utils.py +++ b/modelopt/torch/speculative/utils.py @@ -18,6 +18,7 @@ import contextlib import copy import importlib.util +import json import os import sys import warnings @@ -36,6 +37,11 @@ KIMI_K2_PACKAGE_NAME = "kimi_k2_temp" +# Attributes under which a checkpoint may nest its text-tower config. Mirrors +# modelopt.torch.speculative.plugins.modeling_fakebase._VLM_CONFIG_ATTRS. +NESTED_CONFIG_ATTRS = ["text_config", "llm_config"] + + REMOVE_THINK_CHAT_TEMPLATE = ( "{% if '' in content %}{% set content = content.split('')[-1] %}{% endif %}" ) @@ -584,6 +590,39 @@ def enable_cp_ttt_patch(cp_size: int = 1): modelopt.torch.speculative.plugins.hf_eagle.ENABLE_CP_TTT_PATCH = False +CONFIG_OVERRIDES_HELP = ( + "JSON object of config fields to override on the model config and its text_config before " + "instantiation, e.g. '{\"num_hidden_layers\": 36}'. Needed for checkpoints whose nested " + "text_config dims don't propagate to the parent config." +) + + +def parse_config_overrides(raw: str | None) -> dict | None: + """Parse a ``--config_overrides`` CLI value into a dict, or ``None`` if not supplied. + + Rejects malformed JSON and non-object payloads here, with an actionable message, rather than + letting them surface later as a raw ``JSONDecodeError`` or an ``AttributeError`` from deep + inside model loading. + """ + if not raw: + return None + try: + # Reject the JSON5-ish constants Python's json accepts by default: NaN/Infinity + # would sail through as floats and land on a config field as a dimension. + def _reject(const): + raise ValueError(f"--config_overrides contains non-finite value {const!r}") + + parsed = json.loads(raw, parse_constant=_reject) + except json.JSONDecodeError as e: + raise ValueError(f"--config_overrides is not valid JSON: {e}") from e + if not isinstance(parsed, dict): + raise ValueError( + f"--config_overrides must be a JSON object mapping field names to values, got " + f"{type(parsed).__name__}: {raw!r}" + ) + return parsed + + def load_vlm_or_llm( model_name_or_path: str, use_fake_base: bool = False, @@ -591,6 +630,7 @@ def load_vlm_or_llm( dtype: str | torch.dtype | None = None, device_map: str | None = None, trust_remote_code: bool = False, + config_overrides: dict | None = None, ): """Load a VLM or LLM. Returns the model. @@ -605,26 +645,84 @@ def load_vlm_or_llm( dtype: dtype to use when loading the model. device_map: Device map passed to ``from_pretrained``. trust_remote_code: Whether to trust remote code. + config_overrides: Optional config field overrides applied to the model config and + its ``text_config`` before instantiation (e.g. to correct dims that don't + propagate from a checkpoint's nested text config). """ + + def _warn_overrides_on_fake_base(): + # FakeBaseModel.from_source re-reads the checkpoint config itself, so overrides applied + # here never reach it -- but it is not silently wrong: from_source resolves dims from the + # nested text_config/llm_config first (modeling_fakebase._VLM_CONFIG_ATTRS), which is the + # very problem config_overrides exists to work around, so this path is already correct + # without them. Warn rather than raise: main.py forwards config_overrides unconditionally, + # so hard-failing would leave a single recipe unable to run offline at all. + if config_overrides: + warnings.warn( + "config_overrides is ignored on the FakeBaseModel path: from_source rebuilds the " + "config from the checkpoint, reading dims from the nested text_config/llm_config " + "directly, so the overrides are not needed there.", + stacklevel=2, + ) + if use_offline_training and use_fake_base: + _warn_overrides_on_fake_base() from modelopt.torch.speculative.plugins.modeling_fakebase import FakeBaseModel return FakeBaseModel.from_source(model_name_or_path, trust_remote_code=trust_remote_code) + # Import the transformers-cosmos3 plugin if available: it registers the `cosmos3_omni` + # architecture with AutoConfig on import, so the from_pretrained below recognizes it. + with contextlib.suppress(ImportError): + import transformers_cosmos3 # noqa: F401 + model_config = transformers.AutoConfig.from_pretrained( model_name_or_path, trust_remote_code=trust_remote_code, ) + # Apply caller-supplied config corrections to the parent config and every nested config + # (some checkpoints don't propagate the nested dims up to the parent). + if config_overrides: + # Cover every nested attribute VLM detection accepts, not just text_config: an + # llm_config-nesting checkpoint mirrors the fields on the parent as None, so an override + # would land on the parent, count as "applied", and never reach the real text tower. + targets = [ + cfg + for cfg in ( + model_config, + *(getattr(model_config, a, None) for a in NESTED_CONFIG_ATTRS), + ) + if cfg is not None + ] + unmatched = [] + for key, value in config_overrides.items(): + applied = False + for cfg_obj in targets: + if hasattr(cfg_obj, key): + setattr(cfg_obj, key, value) + applied = True + if not applied: + unmatched.append(key) + if unmatched: + # Silently skipping a key would hand back a wrong-shaped model while appearing to + # have applied the override -- the exact failure this option exists to correct. + raise ValueError( + f"config_overrides key(s) {sorted(unmatched)} matched no field on the model " + f"config (model_type={getattr(model_config, 'model_type', None)!r}) or its " + "text_config. Check for typos." + ) + # Detect VLMs: either "vl" in model_type (e.g. "llava") or has a nested text config # (e.g. Mistral3Config with model_type="mistral3" and text_config attribute). _is_vlm = "vl" in model_config.model_type.lower() or any( - getattr(model_config, attr, None) is not None for attr in ["text_config", "llm_config"] + getattr(model_config, attr, None) is not None for attr in NESTED_CONFIG_ATTRS ) if _is_vlm and use_offline_training: # For VLMs in offline training, FakeBaseModel loads only embed_tokens + lm_head # and auto-detects VLM weight key layouts (e.g. "language_model.model.embed_tokens"). + _warn_overrides_on_fake_base() from modelopt.torch.speculative.plugins.modeling_fakebase import FakeBaseModel return FakeBaseModel.from_source(model_name_or_path, trust_remote_code=trust_remote_code) @@ -641,13 +739,41 @@ def load_vlm_or_llm( model_cls = transformers.AutoModelForCausalLM extra = {} + + # Cosmos3 omni checkpoints: the transformers-cosmos3 plugin registers only the config + # (cosmos3_omni) with AutoConfig, never a model under the Auto* maps, so dispatch to its + # Cosmos3ForConditionalGeneration (a Qwen3-VL subclass) directly. The unused vision + # tower has mismatched dims vs the text-only use, so ignore those on load. + if getattr(model_config, "model_type", None) == "cosmos3_omni": + from transformers_cosmos3 import Cosmos3ForConditionalGeneration + + model_cls = Cosmos3ForConditionalGeneration + extra["ignore_mismatched_sizes"] = True + + # Pass our config object only when we had to modify it (overrides) or when the model class + # needs the plugin-built config; otherwise let from_pretrained build its own, exactly as before. + pass_config = bool(config_overrides) or "ignore_mismatched_sizes" in extra + + # Capture the true depth before any zeroing below, since it is restored after load. + orig_num_hidden_layers = getattr(model_config, "num_hidden_layers", None) + if use_offline_training: - extra["num_hidden_layers"] = 0 - if hasattr(model_config, "layer_types"): - extra["layer_types"] = [] + if pass_config: + # from_pretrained only forwards unrecognized kwargs into the config when it builds + # that config itself. Given a PretrainedConfig instance it deep-copies it and leaves + # the rest in model_kwargs, so num_hidden_layers=0 would never reach the config and + # the full model would be materialized. Set the fields on the config directly. + model_config.num_hidden_layers = 0 + if hasattr(model_config, "layer_types"): + model_config.layer_types = [] + else: + extra["num_hidden_layers"] = 0 + if hasattr(model_config, "layer_types"): + extra["layer_types"] = [] model = model_cls.from_pretrained( model_name_or_path, + config=model_config if pass_config else None, trust_remote_code=trust_remote_code, torch_dtype=dtype, device_map=device_map, @@ -656,7 +782,7 @@ def load_vlm_or_llm( if use_offline_training: # Preserve the original layer count since we loaded with num_hidden_layers=0 - model.config.num_orig_hidden_layers = model_config.num_hidden_layers + model.config.num_orig_hidden_layers = orig_num_hidden_layers return model diff --git a/modelopt/torch/utils/plugins/__init__.py b/modelopt/torch/utils/plugins/__init__.py index da40fe9e565..ff251eaf133 100644 --- a/modelopt/torch/utils/plugins/__init__.py +++ b/modelopt/torch/utils/plugins/__init__.py @@ -23,6 +23,9 @@ with import_plugin("megatron_generate"): from .megatron_generate import * +with import_plugin("megatron_layer_specs"): + from .megatron_layer_specs import * + with import_plugin("megatron_mmlu"): from .megatron_mmlu import * @@ -33,6 +36,7 @@ from .prepare_megatron_data_blend import * # NOTE: Dont pre-import megatron bridge plugin here to avoid circular dependency issues. -# We dont register anything so this isnt a problem. +# It registers an instantiate allowlist prefix on import, which only the ModelOpt entrypoints +# that import it need, so leaving it out here is still fine. # with import_plugin("megatron bridge"): # from .mbridge import * diff --git a/modelopt/torch/utils/plugins/mbridge.py b/modelopt/torch/utils/plugins/mbridge.py index 05713c7b282..7b85fb9d4fb 100644 --- a/modelopt/torch/utils/plugins/mbridge.py +++ b/modelopt/torch/utils/plugins/mbridge.py @@ -22,13 +22,17 @@ from megatron.bridge import AutoBridge from megatron.bridge.models.gpt_provider import GPTModelProvider from megatron.bridge.models.hf_pretrained.utils import is_safe_repo -from megatron.bridge.models.hybrid.hybrid_provider import HybridModelProvider +from megatron.bridge.models.hybrid.hybrid_provider import ( + HybridModelProvider, + transformer_engine_hybrid_stack_spec, +) from megatron.bridge.training.checkpointing import _load_model_weights_from_checkpoint from megatron.bridge.training.post_training.checkpointing import ( _get_modelopt_checkpoint_path, has_modelopt_state, load_modelopt_state, ) +from megatron.bridge.utils.instantiate_utils import register_allowed_target_prefix from megatron.core.models.gpt import GPTModel from megatron.core.models.hybrid.hybrid_model import HybridModel from megatron.core.transformer.module import MegatronModule @@ -37,8 +41,16 @@ from transformers import AutoConfig, AutoTokenizer from modelopt.torch.export.plugins.mcore_common import all_mcore_hf_export_mapping -from modelopt.torch.nas.plugins.megatron import get_te_hybrid_stack_spec from modelopt.torch.utils import print_rank_0, warn_rank_0 +from modelopt.torch.utils.plugins.megatron_layer_specs import te_hybrid_stack_spec_sequential_mlp + +# ``set_moe_expert_layout`` records ``te_hybrid_stack_spec_sequential_mlp`` in a SequentialMLP +# checkpoint's ``run_config.yaml``, and rebuilding that config resolves the target against this +# allowlist. Only a process that imports this module registers it, so such a checkpoint must be +# converted through a ModelOpt entrypoint, not stock ``scripts/conversion/convert.sh``. The prefix +# covers all of ``modelopt`` rather than one module, which assumes a ``run_config.yaml`` is trusted +# input -- it comes from a checkpoint the caller is already choosing to load. +register_allowed_target_prefix("modelopt.") __all__ = [ "get_language_model", @@ -114,10 +126,24 @@ def set_moe_expert_layout(provider, moe_grouped_gemm: bool) -> None: Set ``moe_grouped_gemm`` on the provider (the bridge's native, possibly custom/hybrid spec reads it at build time) rather than replacing the whole layer spec -- overwriting it would drop custom layers (e.g. Qwen3.5's GatedDeltaNet or Gemma3's custom spec). A hybrid provider - additionally needs its stack spec rebuilt, since the native one pins ``TEGroupedMLP``. + additionally has its stack spec set, since the native one pins ``TEGroupedMLP``: the bridge's + own factory for grouped GEMM, a ModelOpt one that swaps in SequentialMLP otherwise. + + Assign a *factory function*, never a built ``ModuleSpec``: the provider is serialized into + every checkpoint's ``run_config.yaml``, and Megatron-LM's YAML writer drops the fields of a + dataclass nested inside a ``functools.partial`` keyword, which is how a stack spec holds + ``MLPSubmodules`` / ``MoESubmodules``. Such a checkpoint cannot be reloaded or exported. The + provider calls the factory at build time, so behavior is unchanged. """ if isinstance(provider, HybridModelProvider): - provider.hybrid_stack_spec = get_te_hybrid_stack_spec(moe_grouped_gemm=moe_grouped_gemm) + # The grouped-GEMM factory is Megatron-Bridge's, and returns Megatron-Core's + # ``hybrid_stack_spec`` unchanged -- the layer composition is identical either way. It is + # named from the bridge so stock tooling resolves the target without importing ModelOpt. + provider.hybrid_stack_spec = ( + transformer_engine_hybrid_stack_spec + if moe_grouped_gemm + else te_hybrid_stack_spec_sequential_mlp + ) provider.moe_grouped_gemm = moe_grouped_gemm elif (provider.num_moe_experts or 0) > 0: provider.moe_grouped_gemm = moe_grouped_gemm diff --git a/modelopt/torch/utils/plugins/megatron_layer_specs.py b/modelopt/torch/utils/plugins/megatron_layer_specs.py new file mode 100644 index 00000000000..d9b08b7b134 --- /dev/null +++ b/modelopt/torch/utils/plugins/megatron_layer_specs.py @@ -0,0 +1,46 @@ +# 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. + +"""Megatron-Core layer specs used to build models for ModelOpt workflows.""" + +import copy + +from megatron.core.models.gpt.moe_module_specs import get_moe_module_spec +from megatron.core.models.hybrid.hybrid_layer_specs import ( + hybrid_stack_spec as _te_hybrid_stack_spec, +) +from megatron.core.transformer.spec_utils import ModuleSpec + +__all__ = ["te_hybrid_stack_spec_sequential_mlp"] + + +def te_hybrid_stack_spec_sequential_mlp() -> ModuleSpec: + """Return the TE Hybrid stack spec with SequentialMLP MoE experts. + + Named and zero-argument so a provider can store this function instead of the ModuleSpec it + builds; see ``set_moe_expert_layout`` for why a built spec cannot be serialized. + + Its module path and name are written into ``run_config.yaml`` as a ``_target_``, so moving or + renaming it breaks every SequentialMLP hybrid checkpoint already saved. + """ + # The upstream TE hybrid stack spec hardcodes TEGroupedMLP for MoE. + # Replace it with SequentialMLP (TE linear layers, no grouped gemm dependency). + # num_experts only has to be non-zero to select the MoE branch; the real count comes from the + # model config at build time. + te_hybrid_stack_spec = copy.deepcopy(_te_hybrid_stack_spec) + te_hybrid_stack_spec.submodules.moe_layer.submodules.mlp = get_moe_module_spec( + use_te=True, num_experts=8, moe_grouped_gemm=False + ) + return te_hybrid_stack_spec diff --git a/modelopt_recipes/configs/ptq/presets/diffusers/nvfp4.yaml b/modelopt_recipes/configs/ptq/presets/diffusers/nvfp4.yaml index 691defb2ae2..47071a4ce9c 100644 --- a/modelopt_recipes/configs/ptq/presets/diffusers/nvfp4.yaml +++ b/modelopt_recipes/configs/ptq/presets/diffusers/nvfp4.yaml @@ -35,3 +35,6 @@ quant_cfg: - quantizer_name: '*softmax_quantizer' cfg: $import: fp8 + - parent_class: 'nn.LayerNorm' + quantizer_name: '*' + enable: false diff --git a/modelopt_recipes/configs/ptq/presets/diffusers/nvfp4_fp8_conv.yaml b/modelopt_recipes/configs/ptq/presets/diffusers/nvfp4_fp8_conv.yaml new file mode 100644 index 00000000000..16e2ebbebd2 --- /dev/null +++ b/modelopt_recipes/configs/ptq/presets/diffusers/nvfp4_fp8_conv.yaml @@ -0,0 +1,53 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 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. + +# Diffusers SDXL preset with dynamic NVFP4 Linears and per-tensor FP8 Conv2d layers. + +# modelopt-schema: modelopt.torch.quantization.config.QuantizeConfig +imports: + base_disable_all: configs/ptq/units/base_disable_all + fp8: configs/numerics/fp8 + nvfp4: configs/numerics/nvfp4 + +algorithm: max +quant_cfg: + - $import: base_disable_all + - parent_class: nn.Linear + quantizer_name: '*weight_quantizer' + cfg: + $import: nvfp4 + - parent_class: nn.Linear + quantizer_name: '*input_quantizer' + cfg: + $import: nvfp4 + - parent_class: nn.Linear + quantizer_name: '*to_[qkv].input_quantizer' + enable: false + - parent_class: nn.Linear + quantizer_name: '*to_[qkv].weight_quantizer' + enable: false + - parent_class: nn.Conv2d + quantizer_name: '*weight_quantizer' + cfg: + $import: fp8 + - parent_class: nn.Conv2d + quantizer_name: '*input_quantizer' + cfg: + $import: fp8 + - quantizer_name: '*output_quantizer' + enable: false + - quantizer_name: '*softmax_quantizer' + cfg: + $import: fp8 diff --git a/modelopt_recipes/configs/ptq/presets/diffusers/nvfp4_fp8_mha.yaml b/modelopt_recipes/configs/ptq/presets/diffusers/nvfp4_fp8_mha.yaml index ee8c2704fc5..d93166ff11e 100644 --- a/modelopt_recipes/configs/ptq/presets/diffusers/nvfp4_fp8_mha.yaml +++ b/modelopt_recipes/configs/ptq/presets/diffusers/nvfp4_fp8_mha.yaml @@ -39,3 +39,6 @@ quant_cfg: - quantizer_name: '*bmm2_output_quantizer' cfg: $import: fp8 + - parent_class: 'nn.LayerNorm' + quantizer_name: '*' + enable: false diff --git a/modelopt_recipes/general/speculative_decoding/dflash.yaml b/modelopt_recipes/general/speculative_decoding/dflash.yaml index 021cccd475d..737789dc4ed 100644 --- a/modelopt_recipes/general/speculative_decoding/dflash.yaml +++ b/modelopt_recipes/general/speculative_decoding/dflash.yaml @@ -16,7 +16,8 @@ data: offline_data_path: # Jinja chat template with {% generation %} tags for answer_only_loss. # Required when answer_only_loss=true. Set in per-model launcher YAML. - # Templates are in modelopt_recipes/general/speculative_decoding/chat_templates/ + # Each model keeps its own beside its launcher example, e.g. + # tools/launcher/examples/Qwen/Qwen3-8B/chat_template_train.jinja chat_template: # maps to TrainingArguments (main.py) diff --git a/modelopt_recipes/huggingface/step3p7/ptq/nvfp4_experts_only-kv_fp8_cast.yaml b/modelopt_recipes/huggingface/step3p7/ptq/nvfp4_experts_only-kv_fp8_cast.yaml new file mode 100644 index 00000000000..086bdf4038c --- /dev/null +++ b/modelopt_recipes/huggingface/step3p7/ptq/nvfp4_experts_only-kv_fp8_cast.yaml @@ -0,0 +1,56 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 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. + +# Model-specific PTQ recipe: Step-3.7 routed-experts-only NVFP4 with FP8 KV-cache cast. +# +# Step's remote code names the MoE block `moe` (not `experts` / `block_sparse_moe`), so the +# general experts-only recipe matches nothing on this architecture. The routed experts are +# `moe.{gate,up,down}_proj`, the router is `moe.gate`, and each MoE layer also has a dense +# `share_expert` that stays in BF16. + +imports: + base_disable_all: configs/ptq/units/base_disable_all + default_disabled_quantizers: configs/ptq/units/default_disabled_quantizers + nvfp4: configs/numerics/nvfp4 + kv_fp8_cast: configs/ptq/units/kv_fp8_cast + +metadata: + recipe_type: ptq + description: >- + Applies dynamic NVFP4 to Step-3.7 routed-expert weight and input quantizers, plus FP8 + KV-cache cast mode using constant amax; the router, the shared expert and the dense MLPs + stay unquantized. Uses max calibration. + +quantize: + algorithm: + method: max + layerwise: {enable: false} + quant_cfg: + - $import: base_disable_all + - quantizer_name: '*moe*weight_quantizer' + cfg: + $import: nvfp4 + - quantizer_name: '*moe*input_quantizer' + cfg: + $import: nvfp4 + - $import: kv_fp8_cast + - $import: default_disabled_quantizers + # The router (`moe.gate`) is matched by `*moe*` above, so it must be disabled last + # (later wins). `share_expert` is not matched by any rule above — it is listed as an + # explicit guard so the shared expert stays BF16 if the patterns are ever widened. + - quantizer_name: '*moe.gate.*' + enable: false + - quantizer_name: '*share_expert*' + enable: false diff --git a/modelopt_recipes/huggingface/step3p7/ptq/nvfp4_mlp_only-kv_fp8.yaml b/modelopt_recipes/huggingface/step3p7/ptq/nvfp4_mlp_only-kv_fp8.yaml new file mode 100644 index 00000000000..b858d895117 --- /dev/null +++ b/modelopt_recipes/huggingface/step3p7/ptq/nvfp4_mlp_only-kv_fp8.yaml @@ -0,0 +1,61 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 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. + +# Model-specific PTQ recipe: Step-3.7 MLP/MoE NVFP4 with calibrated FP8 KV cache. +# +# Same module-naming caveat as the experts-only recipe: Step's MoE block is `moe`, so the +# general `*mlp*` / `*.experts.*` patterns reach only the dense layers' `mlp` submodules and +# leave the routed experts — the bulk of the model — in BF16. + +imports: + base_disable_all: configs/ptq/units/base_disable_all + default_disabled_quantizers: configs/ptq/units/default_disabled_quantizers + nvfp4: configs/numerics/nvfp4 + kv_fp8: configs/ptq/units/kv_fp8 + +metadata: + recipe_type: ptq + description: >- + Applies dynamic NVFP4 to Step-3.7 routed-expert and dense-MLP weight and input + quantizers, plus calibrated FP8 KV-cache quantization; the router and the shared expert + stay unquantized. Uses max calibration. + +quantize: + algorithm: + method: max + layerwise: {enable: false} + quant_cfg: + - $import: base_disable_all + - quantizer_name: '*moe*weight_quantizer' + cfg: + $import: nvfp4 + - quantizer_name: '*moe*input_quantizer' + cfg: + $import: nvfp4 + - quantizer_name: '*mlp*weight_quantizer' + cfg: + $import: nvfp4 + - quantizer_name: '*mlp*input_quantizer' + cfg: + $import: nvfp4 + - $import: kv_fp8 + - $import: default_disabled_quantizers + # The router (`moe.gate`) is matched by `*moe*` above, so it must be disabled last + # (later wins). `share_expert` is not matched by any rule above — it is listed as an + # explicit guard so the shared expert stays BF16 if the patterns are ever widened. + - quantizer_name: '*moe.gate.*' + enable: false + - quantizer_name: '*share_expert*' + enable: false diff --git a/modelopt_recipes/ptq.md b/modelopt_recipes/ptq.md index ee2ca1f1e2a..b12c3404d67 100644 --- a/modelopt_recipes/ptq.md +++ b/modelopt_recipes/ptq.md @@ -170,6 +170,31 @@ How the quantization scales are searched. The default (no suffix) is `max`. (amax) calibrated as in the default recipes. Costs more calibration time but recovers accuracy NVFP4 W4A4 can lose under plain max. Reach for it when a `max` recipe regresses. +- **`nvfp4_act_headroom`** (`nvfp4_act_headroom-kv_fp8_cast`) — leaves headroom + on the NVFP4 **activation** global scale. Plain `max` anchors that scale to the + largest per-block amax seen during calibration, so any larger activation at + inference saturates. This variant instead sets `amax = max(rho * anchor, upper)` + from the per-block amaxes at `anchor_percentile` (default 1) and + `upper_percentile` (default 99.99), with headroom factor `rho` (default 16384): + calibrated blocks sit low in the FP8 block-scale range, leaving the rest for + unseen outliers. Costs one calibration pass and exports a standard NVFP4 + checkpoint; coverage matches `nvfp4_default-kv_fp8_cast`, and the nested + `weight_scale_algorithm` still accepts `mse` / `local_hessian`. Affects NVFP4 + **input** quantizers only — a no-op for FP8 and weight-only recipes. + + Reach for it when a W4A4 recipe regresses and the **activations**, not the + weights, are to blame: an A16 ablation of the same scope clears it while `mse` + does not; the symptom is behavioral (verbose or runaway generations, hitting + the generation cap) rather than a flat score drop; inference contexts run + longer than the calibration set; or a few rare blocks dominate the activation + error. MoE experts-only scopes are the common case — one observer covers every + expert in a layer, and `down_proj` inputs clip first. On a GLM-5.3-Flash + experts-only W4A4 study (SciCode, temperature 1.0) it cut the median + generation-length regression from +38% to +19% and the mean from +19% to +4% + with no capped generations: the best strict-W4A4 result there, but still short + of the p50/p75 gate. A strong first lever, not a guaranteed fix — and sweep + `rho`, since headroom above the calibrated range costs resolution inside it. + - **`input_scale1`** (`nvfp4_experts_only_input_scale1-kv_fp8_cast`) — pins the expert **activation** per-tensor amax to a constant `2688.0` (= E2M1_MAX × E4M3_MAX = 6 × 448) via `constant_amax`, so the exported NVFP4 @@ -215,7 +240,8 @@ These can also be **stacked** when a single method isn't enough — e.g. `mse` + target requires, checking accuracy as you go. 3. **Recover accuracy via calibration before backing off the scope.** If a wider-scope recipe regresses, switch its `max` to the `mse` variant before - retreating to a narrower scope. + retreating to a narrower scope. If `mse` doesn't clear it but an A16 ablation + does, the **activations** are the problem: try `nvfp4_act_headroom` next. 4. **Pick KV by deployment.** `kv_fp8_cast` is the safe default (usually as accurate as calibrated `kv_fp8`); use `kv_nvfp4_cast` for maximum KV compression. @@ -301,6 +327,15 @@ A lighter case: **`models/stepfun-ai/Step-3.5-Flash/ptq/nvfp4-mlp-only`** is clo to one released checkpoint and carrying instance-specific disables (`share_expert`, `moe.gate`, the conv1d branches). +**`step3p7/ptq/{nvfp4_experts_only-kv_fp8_cast,nvfp4_mlp_only-kv_fp8}`** are the +Step-3.7 equivalents, and the reason they exist is **module naming**: Step calls +the MoE block `moe` and the dense sibling `share_expert`, so the general +recipes' `*.experts.*`, `*block_sparse_moe*` and `*mlp*` patterns match nothing +on the routed experts — the general recipe would quantize *nothing* and export a +checkpoint with `quant_algo: null`. These select `*moe*` instead and disable the +router (`moe.gate`) and `share_expert` on top. Use them, not the general +recipes, for Step-3.7 checkpoints; Step-3.5 has its own recipe above. + ### Algorithm overrides — `gemma`, `gemma4`, `mpt` These quantize the **same layers** as the general recipes; only the diff --git a/plugins/modelopt/skills/deployment/references/trtllm.md b/plugins/modelopt/skills/deployment/references/trtllm.md index 27185793307..a24835cdcf8 100644 --- a/plugins/modelopt/skills/deployment/references/trtllm.md +++ b/plugins/modelopt/skills/deployment/references/trtllm.md @@ -2,7 +2,7 @@ ## Requirements -- TensorRT-LLM >= 0.17.0 +- TensorRT-LLM >= 1.2.0 (see `docs/source/deployment/3_unified_hf.rst`) - Typically installed via NVIDIA container: `nvcr.io/nvidia/tensorrt-llm/release:` - Or: `pip install tensorrt-llm` @@ -65,7 +65,7 @@ directly together with a ModelOpt-quantized checkpoint. The legacy export path using `export_tensorrt_llm_checkpoint()` is deprecated. Use the unified HF checkpoint format with `export_hf_checkpoint()` instead. -If you encounter a legacy checkpoint (no `hf_quant_config.json`, has `rank*.safetensors` pattern), it needs the TRT-LLM build API to create an engine before deployment. See `docs/source/deployment/1_tensorrt_llm.rst`. +Current TensorRT-LLM releases no longer support the legacy TensorRT backend. Re-export the quantized source model with `export_hf_checkpoint()` and deploy using the PyTorch backend. See `docs/source/deployment/3_unified_hf.rst`. ## Evaluation with TRT-LLM @@ -91,5 +91,3 @@ explicitly for few-shot tasks. |-------|-----| | `No module named tensorrt_llm` | Install via container or pip | | NVFP4 inference fails on Hopper | NVFP4 requires Blackwell GPUs for inference | -| Slow first inference | Engine compilation happens on first run; subsequent runs are cached | -| OOM during engine build | Reduce `--max_batch_size` or increase TP | diff --git a/tests/_test_utils/torch/megatron/models.py b/tests/_test_utils/torch/megatron/models.py index b2bd2e90921..efdc030a9dc 100644 --- a/tests/_test_utils/torch/megatron/models.py +++ b/tests/_test_utils/torch/megatron/models.py @@ -25,6 +25,7 @@ get_gpt_layer_with_transformer_engine_spec, get_gpt_mtp_block_spec, ) +from megatron.core.models.hybrid.hybrid_layer_specs import hybrid_stack_spec as te_hybrid_stack_spec from megatron.core.models.hybrid.hybrid_model import HybridModel from megatron.core.parallel_state import ( get_pipeline_model_parallel_rank, @@ -37,7 +38,7 @@ from megatron.core.transformer.transformer_config import MLATransformerConfig, TransformerConfig from modelopt.torch.export.unified_export_megatron import import_mcore_gpt_from_hf -from modelopt.torch.nas.plugins.megatron import get_te_hybrid_stack_spec +from modelopt.torch.utils.plugins.megatron_layer_specs import te_hybrid_stack_spec_sequential_mlp try: from megatron.core.extensions.transformer_engine import TENorm @@ -444,11 +445,11 @@ def get_mcore_hybrid_model( "share_embeddings_and_output_weights": False, "position_embedding_type": "none", } - spec = ( - get_te_hybrid_stack_spec(moe_grouped_gemm) - if transformer_impl == "transformer_engine" - else get_hybrid_stack_modelopt_spec(remap_te_layernorm=True) - ) + if transformer_impl == "transformer_engine": + # The upstream TE hybrid stack spec hardcodes TEGroupedMLP for MoE. + spec = te_hybrid_stack_spec if moe_grouped_gemm else te_hybrid_stack_spec_sequential_mlp() + else: + spec = get_hybrid_stack_modelopt_spec(remap_te_layernorm=True) model = HybridModel( hybrid_stack_spec=spec, hybrid_layer_pattern=hybrid_layer_pattern, **common_kwargs ) diff --git a/tests/_test_utils/torch/quantization/tensor_quantizer_common.py b/tests/_test_utils/torch/quantization/tensor_quantizer_common.py index dd8d790ee3f..77d3acb7df3 100644 --- a/tests/_test_utils/torch/quantization/tensor_quantizer_common.py +++ b/tests/_test_utils/torch/quantization/tensor_quantizer_common.py @@ -417,6 +417,49 @@ def test_amax_export(self): amax = quantizer.export_amax() assert amax.shape == (1,) + # One config per ``export_amax`` branch. + UNUSABLE_AMAX_CFGS = { + "dynamic_nvfp4": QuantizerAttributeConfig( + num_bits=(2, 1), block_sizes={-1: 16, "type": "dynamic", "scale_bits": (4, 3)} + ), + "static_per_tensor": QuantizerAttributeConfig(num_bits=4), + } + + @pytest.mark.parametrize("cfg_name", [*UNUSABLE_AMAX_CFGS]) + @pytest.mark.parametrize("bad_value", [0.0, float("nan")]) + def test_amax_export_unusable_amax(self, bad_value, cfg_name): + """An unusable amax must export as a positive scale without mutating the quantizer. + + Regression test for NVBug 6768300. The NaN case pins the ``nan_to_num`` half, which the + zero case alone would not catch. + """ + quantizer = TensorQuantizer(self.UNUSABLE_AMAX_CFGS[cfg_name]).to(self.device) + quantizer.amax = torch.full((1,), bad_value).to(self.device) + + amax = quantizer.export_amax() + + assert torch.all(amax > 0), amax + assert torch.all(amax == quantizer.maxbound), amax + # export must leave the calibrated state alone + stored = quantizer.amax + if bad_value == 0.0: + assert torch.all(stored == 0), stored + else: + assert torch.all(torch.isnan(stored)), stored + + @pytest.mark.parametrize("cfg_name", [*UNUSABLE_AMAX_CFGS]) + def test_amax_export_meta_amax(self, cfg_name): + """``export_amax()`` must stay usable when amax is on the meta device.""" + quantizer = TensorQuantizer(self.UNUSABLE_AMAX_CFGS[cfg_name]) + quantizer.amax = torch.zeros(1, device="meta") + + amax = quantizer.export_amax() + + # Shape differs per branch (the per-tensor path unsqueezes), so pin only that it stays + # meta instead of raising. + assert amax.is_meta, amax + assert amax.numel() == 1, amax.shape + def test_save_restore(self): ref_quantizer = TensorQuantizer(QuantizerAttributeConfig(num_bits=4, axis=0)) diff --git a/tests/examples/diffusers/test_diffusers.py b/tests/examples/diffusers/test_diffusers.py index 15c5eb44934..979894819de 100644 --- a/tests/examples/diffusers/test_diffusers.py +++ b/tests/examples/diffusers/test_diffusers.py @@ -117,6 +117,17 @@ def inference(self, tmp_path: Path) -> None: quant_algo="smoothquant", collect_method="min-mean", ), + pytest.param( + DiffuserModel( + name="sd3-medium", + path=SD3_PATH, + dtype="Half", + format_type="fp8", + quant_algo="max", + collect_method="default", + ), + marks=minimum_sm(89), + ), pytest.param( DiffuserModel( name="sdxl-1.0", @@ -128,6 +139,17 @@ def inference(self, tmp_path: Path) -> None: ), marks=minimum_sm(89), ), + pytest.param( + DiffuserModel( + name="sdxl-1.0", + path=SDXL_PATH, + dtype="Half", + format_type="fp4", + quant_algo="max", + collect_method="default", + ), + marks=minimum_sm(100), + ), DiffuserModel( name="sdxl-1.0", path=SDXL_PATH, @@ -140,7 +162,9 @@ def inference(self, tmp_path: Path) -> None: ids=[ "flux_schnell_bf16_int8_smoothquant_3.0_min_mean", "sd3_medium_fp16_int8_smoothquant_3.0_min_mean", + "sd3_medium_fp16_fp8_max_3.0_default", "sdxl_1.0_fp16_fp8_max_3.0_default", + "sdxl_1.0_fp16_fp4_max_3.0_default", "sdxl_1.0_fp16_int8_smoothquant_3.0_min_mean", ], ) diff --git a/tests/examples/hf_ptq/test_hf_ptq_args.py b/tests/examples/hf_ptq/test_hf_ptq_args.py index 6a2c36e4a7c..90bf671abd6 100644 --- a/tests/examples/hf_ptq/test_hf_ptq_args.py +++ b/tests/examples/hf_ptq/test_hf_ptq_args.py @@ -20,6 +20,7 @@ from types import SimpleNamespace import pytest +import torch import yaml from modelopt.recipe import load_recipe @@ -402,3 +403,44 @@ def test_mlflow_checkpoint_tag_is_absolute(monkeypatch, example_utils): ) assert Path(example_utils._mlflow_run_tags(args)["checkpoint_path"]).is_absolute() + + +# --- post-quantization sanity-check generate() must not block export ---------------------------- + + +def test_post_quantize_export_survives_a_failed_sanity_generate(monkeypatch): + """A failed optional sanity generation must not prevent checkpoint export.""" + hf_ptq = _import_hf_ptq(monkeypatch) + + full_model = SimpleNamespace( + generate=lambda *a, **k: (_ for _ in ()).throw(RuntimeError("boom")) + ) + export_calls = [] + monkeypatch.setattr( + hf_ptq, + "export_quantized", + lambda *a, **k: export_calls.append((a, k)), + ) + monkeypatch.setattr(torch.cuda, "empty_cache", lambda: None) + + args = SimpleNamespace(specdec_offline_dataset=None, verbose=False) + + with pytest.warns(UserWarning, match="Post-quantization generation sanity check failed"): + hf_ptq.post_quantize( + args=args, + full_model=full_model, + language_model=full_model, + model_type="llama", + tokenizer=None, + processor=None, + preview_input_ids=torch.zeros(1, 4, dtype=torch.long), + preview_attention_mask=None, + 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, + ) + + assert len(export_calls) == 1 diff --git a/tests/examples/speculative_decoding/test_vllm_hidden_states_aux_layers.py b/tests/examples/speculative_decoding/test_vllm_hidden_states_aux_layers.py new file mode 100644 index 00000000000..52f3c2b96f2 --- /dev/null +++ b/tests/examples/speculative_decoding/test_vllm_hidden_states_aux_layers.py @@ -0,0 +1,110 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 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. + +"""The vLLM dump resolves aux layers without modelopt -- so it can drift from it. + +``compute_hidden_states_vllm.py`` runs in a stock vLLM container where importing +``modelopt.torch`` fails, so it carries its own copy of the preset logic. A copy that +nothing compares against is a copy that silently diverges: the 'eagle' preset -- which +is ``add_aux_layers_args``' *default* -- was missing from it entirely, so the documented +invocation died on ``int('eagle')`` before the dump began. + +These tests pin the standalone resolver to the real ``hf_eagle`` helper. Unlike the +dump script, the unit suite *can* import modelopt, so the comparison is against the +actual source of truth rather than a third transcription of the formula -- a drift +guard written against its own copy would guard nothing. +""" + +import importlib.util +import sys + +import pytest +from _test_utils.examples.run_command import MODELOPT_ROOT + +from modelopt.torch.speculative.plugins.hf_eagle import default_eagle_aux_layer_ids +from modelopt.torch.speculative.plugins.modeling_dflash import build_target_layer_ids + +_COLLECT = MODELOPT_ROOT / "examples" / "speculative_decoding" / "collect_hidden_states" +sys.path.insert(0, str(_COLLECT)) + +_SPEC = importlib.util.spec_from_file_location( + "chs_vllm", _COLLECT / "compute_hidden_states_vllm.py" +) +assert _SPEC is not None and _SPEC.loader is not None +chs_vllm = importlib.util.module_from_spec(_SPEC) +_SPEC.loader.exec_module(chs_vllm) + +resolve = chs_vllm._resolve_aux_layers_standalone + +# Spans tiny models through 80-layer frontier stacks; includes counts small enough that +# the formula's max(0, ...) clamps collapse ids together. +LAYER_COUNTS = [4, 6, 8, 12, 24, 28, 32, 36, 48, 52, 61, 80] + + +@pytest.mark.parametrize("num_layers", LAYER_COUNTS) +def test_eagle_preset_matches_the_shared_helper(num_layers): + """The standalone copy must agree with modelopt's own EAGLE layer selection. + + Disagreement is silent: the dump writes plausible-looking hidden states from the + wrong layers, and only a poor acceptance rate much later reveals it. + """ + assert resolve("eagle", num_layers) == default_eagle_aux_layer_ids(num_layers) + + +def test_eagle_is_resolvable_because_it_is_the_documented_default(): + """Regression for nvbugs/6753684: `--aux-layers eagle` raised ValueError. + + ``add_aux_layers_args`` defaults to 'eagle', so omitting the flag entirely used to + crash the dump before it started. + """ + assert resolve("eagle", 24) == [1, 11, 20] + + +def test_dflash_preset_still_resolves(): + assert resolve("dflash", 32, num_draft=5) == [1, 8, 15, 22, 29] + + +@pytest.mark.parametrize("num_layers", range(1, 40)) +def test_dflash_preset_matches_the_shared_helper(num_layers): + """Same pinning as EAGLE, including where the shared helper *refuses*. + + A target with fewer layers than the draft has no valid assignment, so + ``build_target_layer_ids`` raises. Without the matching guard the standalone + copy silently deduplicated down to a short list -- e.g. ``[1]`` for a 4-layer + target -- and the dump would write fewer aux layers than the draft consumes. + """ + try: + expected = sorted(set(build_target_layer_ids(num_layers, 5))) + except ValueError: + with pytest.raises(ValueError, match="must be >="): + resolve("dflash", num_layers, num_draft=5) + else: + assert resolve("dflash", num_layers, num_draft=5) == expected + + +def test_explicit_id_list_still_resolves(): + assert resolve("2,5,8", 32) == [2, 5, 8] + + +@pytest.mark.parametrize("spec", ["bogus", "EAGLE3", "", "eagle3"]) +def test_unknown_spec_explains_what_is_accepted(spec): + """A bare ``invalid literal for int()`` hides what the caller should have passed.""" + with pytest.raises(ValueError, match="'eagle' / 'dflash' presets"): + resolve(spec, 32) + + +def test_out_of_range_ids_are_rejected(): + with pytest.raises(ValueError, match="out of range"): + resolve("2,99", 32) diff --git a/tests/examples/test_onnx_ptq.sh b/tests/examples/test_onnx_ptq.sh index 2e4805abfe9..f0a10b8eb1c 100755 --- a/tests/examples/test_onnx_ptq.sh +++ b/tests/examples/test_onnx_ptq.sh @@ -152,7 +152,7 @@ for model_path in "${model_paths[@]}"; do python -m modelopt.onnx.quantization \ --onnx_path=$model_dir/fp16/model.onnx \ --quantize_mode=$quant_mode \ - --calibration_data=$calib_data_path \ + --calibration_data_path=$calib_data_path \ --output_path=$model_dir/$quant_mode/model.quant.onnx \ --calibration_eps=cuda done diff --git a/tests/gpu/onnx/test_ort_patching.py b/tests/gpu/onnx/test_ort_patching.py index 26ca49f39b4..84224dcffa0 100644 --- a/tests/gpu/onnx/test_ort_patching.py +++ b/tests/gpu/onnx/test_ort_patching.py @@ -153,21 +153,6 @@ def test_collect_value(self, mock_histogram_collector, sample_tensor_data): assert "tensor1" in mock_histogram_collector.histogram_dict assert "tensor2" in mock_histogram_collector.histogram_dict - def test_collect_value_fp16_narrow_range(self, mock_histogram_collector): - # fp16 activations with a small range (threshold ~1e-5) used to raise - # "Too many bins for data range" on numpy >= 2.0, because the fp16 range - # produced a fp16 linspace where consecutive bin edges rounded together. - activations = np.zeros(1000, dtype=np.float16) - activations[0] = np.float16(1e-5) - name_to_arr = {"narrow_fp16_tensor": [activations]} - - _collect_value(mock_histogram_collector, name_to_arr) - - hist, edges, _, _, _ = mock_histogram_collector.histogram_dict["narrow_fp16_tensor"] - assert hist.sum() == activations.size - assert len(edges) == mock_histogram_collector.num_bins + 1 - assert not np.any(np.diff(edges) == 0), "fp16 bin edges collapsed" - def test_collect_absolute_value(self, mock_histogram_collector, sample_tensor_data): """Test _collect_absolute_value function.""" # Convert to float32 to avoid the float64 assertion error diff --git a/tests/gpu/torch/export/test_fsdp2_export.py b/tests/gpu/torch/export/test_fsdp2_export.py index 65f33abd41c..dc35974c047 100644 --- a/tests/gpu/torch/export/test_fsdp2_export.py +++ b/tests/gpu/torch/export/test_fsdp2_export.py @@ -261,7 +261,9 @@ def test_fsdp2_weight_update_context_for_export(dist_workers): # mtq.W4A8_AWQ_BETA_CFG, #TODO: Fix unit test for this case # mtq.FP8_2D_BLOCKWISE_WEIGHT_ONLY_CFG, #TODO: Fix unit test for this case mtq.W4A8_MXFP4_FP8_CFG, - mtq.NVFP4_MLP_ONLY_CFG, + # NVFP4_MLP_ONLY_CFG is omitted: SmallQKVModel has no MLP, so that config matched + # nothing and the parametrization exercised an unquantized model. NVFP4_OMLP_ONLY_CFG + # covers the same scoped-recipe shape here because it also selects `o_proj`. mtq.NVFP4_OMLP_ONLY_CFG, ], ) @@ -281,7 +283,9 @@ def test_fsdp2_weight_update_context_for_fuse_layers(dist_workers, quant_config, # mtq.W4A8_AWQ_BETA_CFG, #TODO: Fix unit test for this case # mtq.FP8_2D_BLOCKWISE_WEIGHT_ONLY_CFG, #TODO: Fix unit test for this case mtq.W4A8_MXFP4_FP8_CFG, - mtq.NVFP4_MLP_ONLY_CFG, + # NVFP4_MLP_ONLY_CFG is omitted: SmallQKVModel has no MLP, so that config matched + # nothing and the parametrization exercised an unquantized model. NVFP4_OMLP_ONLY_CFG + # covers the same scoped-recipe shape here because it also selects `o_proj`. mtq.NVFP4_OMLP_ONLY_CFG, ], ) diff --git a/tests/gpu/torch/quantization/test_onnx_export_cuda.py b/tests/gpu/torch/quantization/test_onnx_export_cuda.py index 300abc52e9e..39f422c75a9 100644 --- a/tests/gpu/torch/quantization/test_onnx_export_cuda.py +++ b/tests/gpu/torch/quantization/test_onnx_export_cuda.py @@ -17,7 +17,6 @@ import pytest import torch -import torch.nn as nn from _test_utils.torch.quantization.onnx_export import TEST_MODELS, onnx_export_tester @@ -40,7 +39,4 @@ def test_onnx_export_cuda(model_cls, num_bits, per_channel_quantization, constan torch.manual_seed(0) model = model_cls() - for _, module in model.named_modules(): - if isinstance(module, nn.Conv2d) and num_bits == (4, 3): - pytest.skip("Conv2d with FP8 quantization is not supported yet") onnx_export_tester(model, "cuda", num_bits, per_channel_quantization, constant_folding, dtype) diff --git a/tests/gpu_megatron/torch/quantization/plugins/test_megatron.py b/tests/gpu_megatron/torch/quantization/plugins/test_megatron.py index bec404f9c24..041a0bd3fca 100644 --- a/tests/gpu_megatron/torch/quantization/plugins/test_megatron.py +++ b/tests/gpu_megatron/torch/quantization/plugins/test_megatron.py @@ -15,7 +15,6 @@ import copy import math -import re import sys import types from contextlib import nullcontext @@ -38,7 +37,9 @@ get_batch, get_forward, initialize_for_megatron, + load_distributed_checkpoint, run_mcore_inference, + save_distributed_checkpoint, sharded_state_dict_test_helper, ) from _test_utils.torch.misc import set_seed @@ -54,6 +55,7 @@ from megatron.core.parallel_state import ( destroy_model_parallel, get_data_parallel_group, + get_expert_model_parallel_rank, get_tensor_model_parallel_group, ) from megatron.core.tensor_parallel.layers import ColumnParallelLinear, RowParallelLinear @@ -64,10 +66,15 @@ import modelopt import modelopt.torch.opt as mto import modelopt.torch.quantization as mtq +from modelopt.torch.opt.plugins.mcore_dist_checkpointing import ( + restore_sharded_modelopt_state, + save_sharded_modelopt_state, +) from modelopt.torch.quantization.algorithms import QuantRecipe, _AutoQuantizeBaseSearcher from modelopt.torch.quantization.nn import QuantModuleRegistry, SequentialQuantizer from modelopt.torch.quantization.nn.modules.quant_linear import RealQuantLinear from modelopt.torch.quantization.plugins.megatron import ( + _initialize_grouped_weight_quantizer_state, _output_layer_extra_state_has_data, _output_layer_untied, _QuantMegatronTEGroupedLinear, @@ -1094,100 +1101,207 @@ def test_te_grouped_vs_sequential_default_amax(dist_workers_size_1, quant_cfg): ) -def _te_grouped_expert_identity_from_sharded_state(module): - """Return {local_key: (global_expert_idx, num_global_experts)} for per-expert amax shards. +def _set_te_grouped_weight_quantizer_state(model, ep_rank, num_local_experts): + """Give every local expert distinct quantizer state derived from its global index.""" + for linear in model.modules(): + if not isinstance(linear, _QuantMegatronTEGroupedLinear): + continue + for local_expert_idx in range(linear.num_gemms): + quantizer = linear.weight_quantizer[local_expert_idx] + leaves = list(quantizer) if isinstance(quantizer, SequentialQuantizer) else [quantizer] + for leaf in leaves: + amax = getattr(leaf, "_amax", None) + if amax is not None: + amax.fill_(1.0 + ep_rank * num_local_experts + local_expert_idx) + global_amax = getattr(leaf, "_global_amax", None) + if global_amax is not None: + global_amax.fill_(1.0 + ep_rank * num_local_experts + local_expert_idx) + + +def _assert_te_grouped_weight_quantizer_state(model, expected_amax, expect_global_amax): + checked = 0 + for linear in model.modules(): + if not isinstance(linear, _QuantMegatronTEGroupedLinear): + continue + for local_expert_idx in range(linear.num_gemms): + quantizer = linear.weight_quantizer[local_expert_idx] + leaves = list(quantizer) if isinstance(quantizer, SequentialQuantizer) else [quantizer] + for leaf in leaves: + amax = getattr(leaf, "_amax", None) + assert amax is not None, ( + "TEGrouped per-expert weight quantizer amax was not restored" + ) + checked += 1 + assert torch.equal(amax, torch.full_like(amax, expected_amax[local_expert_idx])) + global_amax = getattr(leaf, "_global_amax", None) + if expect_global_amax: + assert global_amax is not None, ( + "TEGrouped per-expert weight quantizer global_amax was not restored" + ) + assert torch.equal( + global_amax, + torch.full_like(global_amax, expected_amax[local_expert_idx]), + ) + assert checked > 0, "no TEGrouped per-expert weight quantizer amax was checked" - The grouped linear must give each fused expert the same global identity the weights use: - the dict key keeps the local expert index (maps to the local buffer on restore) while the - ShardedTensor carries the global expert offset. Called with sharded_offsets=() so the expert - axis is the (only) prepended axis at index 0. - """ - sharded_sd = module.sharded_state_dict(prefix="", sharded_offsets=(), metadata=None) - identity = {} - for key, sh_ten in sharded_sd.items(): - if re.match(r"weight_quantizer\.\d+\..*_amax$", key): - assert sh_ten.prepend_axis_num >= 1, f"{key}: expected a prepended expert axis" - identity[key] = (int(sh_ten.global_offset[0]), int(sh_ten.global_shape[0])) - return identity +def test_initialize_grouped_weight_quantizer_state_for_restore(): + """Missing grouped state inherits the shape and dtype of a populated sibling.""" + source = mtq.nn.StaticBlockScaleQuantizer.from_tensor_quantizer( + mtq.nn.TensorQuantizer(amax=torch.tensor([1.0, 2.0])), global_amax=torch.tensor(2.0) + ) + target = mtq.nn.StaticBlockScaleQuantizer.from_tensor_quantizer(mtq.nn.TensorQuantizer()) + disabled = mtq.nn.StaticBlockScaleQuantizer.from_tensor_quantizer(mtq.nn.TensorQuantizer()) + disabled.disable() + mx = mtq.nn.TensorQuantizer( + mtq.config.QuantizerAttributeConfig( + num_bits=(4, 3), block_sizes={-1: 32, "type": "dynamic", "scale_bits": (8, 0)} + ) + ) -def _test_te_grouped_sharded_state_dict_global_expert_identity_helper( - tp_size, ep_size, quant_cfg, rank, size -): - """Per-expert quantizer amax must persist all num_global_experts across EP. + model = torch.nn.Module() + model.weight = torch.nn.Parameter(torch.empty(1)) + model.num_gemms = 4 + model.weight_quantizer = torch.nn.ModuleList([target, source, disabled, mx]) - With EP>1 the base linear emitted ``weight_quantizer.{local_i}._amax`` at the local index with - no expert offset, so every rank wrote identical keys and torch_dist dedup collapsed them to a - single rank's experts. Assert each rank's fused experts now carry distinct global identities so - the union across ranks covers every global expert. - """ + _initialize_grouped_weight_quantizer_state(model) + + assert torch.equal(target.amax, torch.zeros_like(source.amax)) + assert torch.equal(target.global_amax, torch.zeros_like(source.global_amax)) + assert disabled.amax is None + assert disabled.global_amax is None + assert mx.amax is None + assert not hasattr(mx, "_amax") + + +def _test_te_grouped_sharded_state_dict_reshard_helper( + save_tp_size, + save_ep_size, + load_tp_size, + load_ep_size, + quant_cfg, + expect_global_amax, + checkpoint_path, + rank, + size, +): + """Round-trip TEGroupedMLP amax through a topology change.""" + num_experts = 4 + save_num_local_experts = num_experts // save_ep_size initialize_for_megatron( - tensor_model_parallel_size=tp_size, - expert_model_parallel_size=ep_size, + tensor_model_parallel_size=save_tp_size, + expert_model_parallel_size=save_ep_size, seed=SEED, ) - num_experts = 4 - num_local = num_experts // ep_size - te_grouped = _gpt_model_provider( - tp_size=tp_size, - ep_size=ep_size, + source = _gpt_model_provider( + tp_size=save_tp_size, + ep_size=save_ep_size, hidden_size=32, moe_grouped_gemm=True, transformer_impl="transformer_engine", num_moe_experts=num_experts, ) - forward = get_forward(te_grouped, batch_size=8) - for module in te_grouped.modules(): + forward = get_forward(source, batch_size=8) + for module in source.modules(): if isinstance(module, TopKRouter): module.topk = module.num_experts - mtq.quantize(te_grouped, quant_cfg, forward) + mtq.quantize(source, copy.deepcopy(quant_cfg), forward) + _set_te_grouped_weight_quantizer_state( + source, get_expert_model_parallel_rank(), save_num_local_experts + ) + save_distributed_checkpoint(checkpoint_path, source) + save_sharded_modelopt_state([source], checkpoint_path) + torch.distributed.barrier() + del source + destroy_model_parallel() - grouped_linears = [ - m for m in te_grouped.modules() if isinstance(m, _QuantMegatronTEGroupedLinear) - ] - assert grouped_linears, "No grouped quant linears found" - - expected_global = {rank * num_local + i for i in range(num_local)} - for linear in grouped_linears: - # Give each expert a distinct amax so a value mix-up would also be observable. - for i in range(linear.num_gemms): - wq = linear.weight_quantizer[i] - leaves = list(wq) if isinstance(wq, SequentialQuantizer) else [wq] - for leaf in leaves: - if hasattr(leaf, "_amax") and leaf._amax is not None: - leaf._amax.fill_(1.0 + rank * num_local + i) - - identity = _te_grouped_expert_identity_from_sharded_state(linear) - # One entry per local expert per amax buffer; dict keys keep the LOCAL index. - local_keys = {int(re.search(r"weight_quantizer\.(\d+)\.", k).group(1)) for k in identity} - assert local_keys == set(range(num_local)), ( - f"Expected local expert keys {set(range(num_local))}, got {local_keys}" - ) - # ShardedTensor global identity: this rank owns experts {rank*num_local + i}. - local_global = {gidx for gidx, _ in identity.values()} - assert local_global == expected_global, ( - f"rank {rank}: expected global experts {expected_global}, got {local_global}" - ) - assert all(total == num_experts for _, total in identity.values()), ( - f"num_global_experts should be {num_experts}, got {identity}" + initialize_for_megatron( + tensor_model_parallel_size=load_tp_size, + expert_model_parallel_size=load_ep_size, + seed=SEED, + ) + target = _gpt_model_provider( + tp_size=load_tp_size, + ep_size=load_ep_size, + hidden_size=32, + moe_grouped_gemm=True, + transformer_impl="transformer_engine", + num_moe_experts=num_experts, + ) + target_models = [target] + restore_sharded_modelopt_state(target_models, checkpoint_path) + target = target_models[0] + load_distributed_checkpoint(checkpoint_path, target) + load_num_local_experts = num_experts // load_ep_size + expected_amax = tuple( + range( + 1 + get_expert_model_parallel_rank() * load_num_local_experts, + 1 + (get_expert_model_parallel_rank() + 1) * load_num_local_experts, ) - - # Gather the global expert indices across all EP ranks: the union must cover every expert. - gathered = [None] * size - torch.distributed.all_gather_object(gathered, sorted(expected_global)) - union = set() - for part in gathered: - union.update(part) - assert union == set(range(num_experts)), ( - f"Union of global experts across EP ranks should be {set(range(num_experts))}, got {union}" ) + _assert_te_grouped_weight_quantizer_state(target, expected_amax, expect_global_amax) -@pytest.mark.parametrize("quant_cfg", [mtq.FP8_DEFAULT_CFG, mtq.NVFP4_DEFAULT_CFG]) -def test_te_grouped_sharded_state_dict_global_expert_identity(dist_workers_size_2, quant_cfg): +@pytest.mark.parametrize( + ( + "quant_cfg", + "expect_global_amax", + "save_tp_size", + "save_ep_size", + "load_tp_size", + "load_ep_size", + ), + [ + pytest.param(mtq.FP8_DEFAULT_CFG, False, 1, 2, 1, 1, id="fp8-ep-downsize"), + pytest.param(mtq.FP8_DEFAULT_CFG, False, 1, 1, 1, 2, id="fp8-ep-upsize"), + pytest.param(mtq.FP8_DEFAULT_CFG, False, 1, 1, 2, 1, id="fp8-tp-upsize"), + pytest.param(mtq.FP8_DEFAULT_CFG, False, 2, 1, 1, 1, id="fp8-tp-downsize"), + pytest.param(mtq.NVFP4_DEFAULT_CFG, False, 1, 2, 1, 1, id="nvfp4-ep-downsize"), + pytest.param(mtq.NVFP4_DEFAULT_CFG, False, 1, 1, 1, 2, id="nvfp4-ep-upsize"), + pytest.param(mtq.NVFP4_DEFAULT_CFG, False, 1, 1, 2, 1, id="nvfp4-tp-upsize"), + pytest.param(mtq.NVFP4_DEFAULT_CFG, False, 2, 1, 1, 1, id="nvfp4-tp-downsize"), + pytest.param( + mtq.NVFP4_W4A4_WEIGHT_MSE_FP8_SWEEP_CFG, + True, + 1, + 2, + 1, + 1, + id="nvfp4-mse-ep-downsize", + ), + pytest.param( + mtq.NVFP4_W4A4_WEIGHT_MSE_FP8_SWEEP_CFG, + True, + 1, + 1, + 1, + 2, + id="nvfp4-mse-ep-upsize", + ), + ], +) +def test_te_grouped_sharded_state_dict_reshard( + dist_workers_size_2, + tmp_path, + save_tp_size, + save_ep_size, + load_tp_size, + load_ep_size, + quant_cfg, + expect_global_amax, +): dist_workers_size_2.run( - partial(_test_te_grouped_sharded_state_dict_global_expert_identity_helper, 1, 2, quant_cfg) + partial( + _test_te_grouped_sharded_state_dict_reshard_helper, + save_tp_size, + save_ep_size, + load_tp_size, + load_ep_size, + quant_cfg, + expect_global_amax, + tmp_path, + ) ) diff --git a/tests/gpu_megatron/torch/utils/plugins/test_mbridge.py b/tests/gpu_megatron/torch/utils/plugins/test_mbridge.py new file mode 100644 index 00000000000..3b5e7ea78c8 --- /dev/null +++ b/tests/gpu_megatron/torch/utils/plugins/test_mbridge.py @@ -0,0 +1,74 @@ +# 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 pytest +import yaml +from megatron.bridge.models.hybrid.hybrid_provider import HybridModelProvider +from megatron.bridge.utils.instantiate_utils import instantiate +from megatron.bridge.utils.yaml_utils import dump_dataclass_to_yaml + +from modelopt.torch.utils.plugins.mbridge import set_moe_expert_layout + + +def _round_trip(value): + """Serialize through the writer used for run_config.yaml, then reload.""" + node = yaml.safe_load(dump_dataclass_to_yaml({"spec": value}))["spec"] + return node["_target_"], instantiate(node) + + +@pytest.mark.parametrize( + ("moe_grouped_gemm", "expected_experts", "expected_target"), + [ + ( + True, + "TEGroupedMLP", + "megatron.bridge.models.hybrid.hybrid_provider.transformer_engine_hybrid_stack_spec", + ), + ( + False, + "SequentialMLP", + "modelopt.torch.utils.plugins.megatron_layer_specs.te_hybrid_stack_spec_sequential_mlp", + ), + ], +) +def test_set_moe_expert_layout_survives_run_config_round_trip( + moe_grouped_gemm, expected_experts, expected_target +): + """A provider's stack spec must still build real submodules after a run_config round trip. + + A built ``ModuleSpec`` loses its ``MLPSubmodules`` / ``MoESubmodules`` when written to + ``run_config.yaml``, so ``set_moe_expert_layout`` stores a factory function instead. + """ + provider = HybridModelProvider(num_layers=2, hidden_size=64, num_attention_heads=4) + set_moe_expert_layout(provider, moe_grouped_gemm=moe_grouped_gemm) + assert provider.moe_grouped_gemm == moe_grouped_gemm + + assert callable(provider.hybrid_stack_spec) + + target, factory = _round_trip(provider.hybrid_stack_spec) + # The target is an on-disk contract: renaming or moving the factory breaks saved checkpoints. + assert target == expected_target + + provider.hybrid_stack_spec = factory + spec = provider._resolve_hybrid_stack_spec() + + mlp = spec.submodules.mlp_layer.submodules.mlp.keywords["submodules"] + assert mlp.linear_fc1 is not None + assert mlp.linear_fc2 is not None + + moe = spec.submodules.moe_layer.submodules.mlp.keywords["submodules"] + assert moe.experts is not None + # Experts are built through a partial for the grouped-GEMM layout. + assert getattr(moe.experts, "func", moe.experts).__name__ == expected_experts diff --git a/tests/gpu_vllm/torch/quantization/test_vllm_dynamic_modules.py b/tests/gpu_vllm/torch/quantization/test_vllm_dynamic_modules.py index 038bbd8e978..0f440b93deb 100644 --- a/tests/gpu_vllm/torch/quantization/test_vllm_dynamic_modules.py +++ b/tests/gpu_vllm/torch/quantization/test_vllm_dynamic_modules.py @@ -28,6 +28,7 @@ from __future__ import annotations +import builtins import gc import importlib.util from pathlib import Path @@ -69,6 +70,245 @@ def _load_example_module(name: str): return module +def _calibration_worker( + num_blocks: int, + *, + cache_specs=("attention", "mamba"), + needs_kv_cache_zeroing=True, +): + cache_groups = [SimpleNamespace(kv_cache_spec=spec) for spec in cache_specs] + return SimpleNamespace( + model_runner=SimpleNamespace( + kv_cache_config=SimpleNamespace( + kv_cache_groups=cache_groups, + num_blocks=num_blocks, + needs_kv_cache_zeroing=needs_kv_cache_zeroing, + ) + ) + ) + + +def _patch_vllm_imports(monkeypatch, modules): + real_import = builtins.__import__ + + def fake_import(name, globals=None, locals=None, fromlist=(), level=0): + if name in modules: + imported = modules[name] + if isinstance(imported, BaseException): + raise imported + return imported + return real_import(name, globals, locals, fromlist, level) + + monkeypatch.setattr(builtins, "__import__", fake_import) + + +def test_get_calibration_block_count_uses_vllm_028_reservation_helper(monkeypatch): + """The current vLLM adapter must forward every warmup reservation argument.""" + module = _load_example_module("vllm_ptq_utils") + reserved_block_count = Mock(return_value=4) + _patch_vllm_imports( + monkeypatch, + {"vllm.v1.worker.gpu.warmup": SimpleNamespace(_reserved_block_count=reserved_block_count)}, + ) + model_runner = SimpleNamespace( + vllm_config=SimpleNamespace(num_lookahead_tokens=3), + max_model_len=2048, + ) + kv_cache_spec = object() + + block_count = module._get_calibration_block_count(model_runner) + + assert block_count is not None + assert block_count(128, kv_cache_spec) == 4 + reserved_block_count.assert_called_once_with( + 128, + kv_cache_spec, + num_lookahead_tokens=3, + max_model_len=2048, + max_encoder_len=0, + ) + + +def test_get_calibration_block_count_uses_vllm_026_reservation_policy(monkeypatch): + """The vLLM 0.26 adapter must preserve its cross-attention and Mamba rules.""" + module = _load_example_module("vllm_ptq_utils") + + class CrossAttentionSpec: + block_size = 16 + + class MambaSpec: + block_size = 16 + mamba_cache_mode = "align" + num_speculative_blocks = 2 + + cdiv = Mock( + side_effect=lambda numerator, denominator: (numerator + denominator - 1) // denominator + ) + _patch_vllm_imports( + monkeypatch, + { + "vllm.v1.worker.gpu.warmup": ImportError("0.28 helper unavailable"), + "vllm.utils.math_utils": SimpleNamespace(cdiv=cdiv), + "vllm.v1.kv_cache_interface": SimpleNamespace( + CrossAttentionSpec=CrossAttentionSpec, + MambaSpec=MambaSpec, + ), + }, + ) + model_runner = SimpleNamespace( + vllm_config=SimpleNamespace(), + max_model_len=2048, + ) + + block_count = module._get_calibration_block_count(model_runner) + + assert block_count is not None + assert block_count(33, SimpleNamespace(block_size=16)) == 3 + assert block_count(33, CrossAttentionSpec()) == 0 + assert block_count(33, MambaSpec()) == 5 + assert cdiv.call_args_list == [ + ((33, 16),), + ((0, 16),), + ((33, 16),), + ] + + +def test_allocate_calibration_blocks_assigns_non_null_blocks(monkeypatch): + """Scratch block tables must use unique non-null blocks for every request and group.""" + module = _load_example_module("vllm_ptq_utils") + block_count = Mock(side_effect=[1, 2, 2, 1]) + monkeypatch.setattr( + module, + "_get_calibration_block_count", + Mock(return_value=block_count), + ) + + block_tables, blocks_to_zero = module._allocate_calibration_blocks( + _calibration_worker(num_blocks=7), + sequence_lengths=[8, 16], + ) + + assert block_tables == [ + ([1], [2, 3]), + ([4, 5], [6]), + ] + assert block_count.call_args_list == [ + ((8, "attention"),), + ((8, "mamba"),), + ((16, "attention"),), + ((16, "mamba"),), + ] + + scheduler_fields = {field.name for field in module.dataclasses.fields(module.SchedulerOutput)} + expected_blocks_to_zero = ( + [1, 2, 3, 4, 5, 6] if "new_block_ids_to_zero" in scheduler_fields else None + ) + assert blocks_to_zero == expected_blocks_to_zero + + +def test_allocate_calibration_blocks_skips_zeroing_for_attention_only_cache(monkeypatch): + """Attention-only caches have no block zeroer and must receive an empty zeroing list.""" + module = _load_example_module("vllm_ptq_utils") + monkeypatch.setattr( + module, + "_get_calibration_block_count", + Mock(return_value=Mock(return_value=1)), + ) + + block_tables, blocks_to_zero = module._allocate_calibration_blocks( + _calibration_worker( + num_blocks=4, + cache_specs=("attention",), + needs_kv_cache_zeroing=False, + ), + sequence_lengths=[8], + ) + + assert block_tables == [([1],)] + scheduler_fields = {field.name for field in module.dataclasses.fields(module.SchedulerOutput)} + expected_blocks_to_zero = [] if "new_block_ids_to_zero" in scheduler_fields else None + assert blocks_to_zero == expected_blocks_to_zero + + +def test_allocate_calibration_blocks_rejects_insufficient_capacity(monkeypatch): + """Scratch block allocation must account for block 0 being unavailable.""" + module = _load_example_module("vllm_ptq_utils") + monkeypatch.setattr( + module, + "_get_calibration_block_count", + Mock(return_value=Mock(side_effect=[1, 2, 2, 1])), + ) + + with pytest.raises( + RuntimeError, + match=( + r"Calibration batch requires 6 KV cache blocks, " + r"but only 5 non-null blocks are available\." + ), + ): + module._allocate_calibration_blocks( + _calibration_worker(num_blocks=6), + sequence_lengths=[8, 16], + ) + + +@pytest.mark.parametrize("has_calibration_error", [False, True]) +def test_cleanup_failure_preserves_calibration_error(has_calibration_error): + """Cleanup must fail closed without replacing an active calibration error.""" + module = _load_example_module("vllm_ptq_utils") + execute_error = RuntimeError("scheduler cleanup failed") + finish_error = RuntimeError("legacy cleanup failed") + calibration_error = ValueError("calibration failed") if has_calibration_error else None + worker = SimpleNamespace( + execute_model=Mock(side_effect=execute_error), + model_runner=SimpleNamespace(finish_requests=Mock(side_effect=finish_error)), + ) + + expected_error = calibration_error or finish_error + with pytest.raises(type(expected_error)) as raised: + module._cleanup_calibration_requests(worker, object(), calibration_error) + + assert raised.value is expected_error + if calibration_error is not None: + assert calibration_error.__cause__ is finish_error + assert finish_error.__cause__ is execute_error + + +@pytest.mark.parametrize("has_calibration_error", [False, True]) +def test_cleanup_without_legacy_fallback_preserves_primary_error(has_calibration_error): + """Missing legacy cleanup must preserve the most useful primary error.""" + module = _load_example_module("vllm_ptq_utils") + execute_error = RuntimeError("scheduler cleanup failed") + calibration_error = ValueError("calibration failed") if has_calibration_error else None + worker = SimpleNamespace( + execute_model=Mock(side_effect=execute_error), + model_runner=SimpleNamespace(), + ) + + expected_error = calibration_error or execute_error + with pytest.raises(type(expected_error)) as raised: + module._cleanup_calibration_requests(worker, object(), calibration_error) + + assert raised.value is expected_error + if calibration_error is not None: + assert calibration_error.__cause__ is execute_error + + +def test_cleanup_uses_legacy_fallback(): + """A successful legacy cleanup may recover from an unsupported scheduler step.""" + module = _load_example_module("vllm_ptq_utils") + cleanup_output = object() + finish_requests = Mock() + worker = SimpleNamespace( + execute_model=Mock(side_effect=RuntimeError("unsupported scheduler cleanup")), + model_runner=SimpleNamespace(finish_requests=finish_requests), + ) + + module._cleanup_calibration_requests(worker, cleanup_output, calibration_error=None) + + finish_requests.assert_called_once_with(cleanup_output) + + class _NativeAttention(torch.nn.Module): def forward(self, query, key, value, *args, **kwargs): return query, key, value diff --git a/tests/unit/examples/test_diffusers_fp4.py b/tests/unit/examples/test_diffusers_fp4.py new file mode 100644 index 00000000000..16023e509ad --- /dev/null +++ b/tests/unit/examples/test_diffusers_fp4.py @@ -0,0 +1,234 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 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 importlib.util +import logging +import sys +from pathlib import Path +from unittest.mock import Mock + +import pytest +import torch +from torch import nn + +pytest.importorskip("onnx") +pytest.importorskip("onnx_graphsurgeon") +pytest.importorskip("diffusers") + +import modelopt.torch.quantization as mtq +from examples.diffusers.quantization.onnx_utils import export as diffusion_export +from modelopt.torch.quantization.config import QuantizerAttributeConfig +from modelopt.torch.quantization.nn import TensorQuantizer +from modelopt.torch.quantization.plugins.diffusion import diffusers as diffusers_plugin + +_QUANTIZATION_EXAMPLE = ( + Path(__file__).resolve().parents[3] / "examples" / "diffusers" / "quantization" +) +_LOCAL_IMPORT_NAMES = ( + "calib.plugin_calib", + "calib", + "calibration", + "config", + "models_utils", + "pipeline_manager", + "quantize_config", + "utils", +) + + +def _load_quantize_example(): + spec = importlib.util.spec_from_file_location( + "diffusers_quantize_example", _QUANTIZATION_EXAMPLE / "quantize.py" + ) + assert spec is not None and spec.loader is not None + + original_modules = { + name: sys.modules.pop(name) for name in _LOCAL_IMPORT_NAMES if name in sys.modules + } + sys.path.insert(0, str(_QUANTIZATION_EXAMPLE)) + try: + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + finally: + sys.path.pop(0) + for name in _LOCAL_IMPORT_NAMES: + sys.modules.pop(name, None) + sys.modules.update(original_modules) + return module + + +_quantize = _load_quantize_example() +ModelType = _quantize.ModelType +ModelConfig = _quantize.ModelConfig +QuantFormat = _quantize.QuantFormat +QuantizationConfig = _quantize.QuantizationConfig +Quantizer = _quantize.Quantizer +_infer_restored_quantization_format = _quantize._infer_restored_quantization_format + + +class _RecipeBackbone(nn.Module): + def __init__(self): + super().__init__() + self.linear = nn.Linear(16, 16, bias=False) + self.attn = nn.Module() + self.attn.to_q = nn.Linear(16, 16, bias=False) + self.attn.to_k = nn.Linear(16, 16, bias=False) + self.attn.to_v = nn.Linear(16, 16, bias=False) + self.conv = nn.Conv2d(4, 4, kernel_size=1, bias=False) + + +def _quantizer(*, num_bits, enabled=True, block_sizes=None): + quantizer = TensorQuantizer( + QuantizerAttributeConfig(num_bits=num_bits, axis=None, block_sizes=block_sizes) + ) + quantizer.amax = torch.tensor(448.0) + if not enabled: + quantizer.disable() + return quantizer + + +_FP8_QUANTIZER_CONFIG = {"num_bits": (4, 3)} +_NVFP4_QUANTIZER_CONFIG = { + "num_bits": (2, 1), + "block_sizes": {-1: 16, "type": "dynamic", "scale_bits": (4, 3)}, +} + + +@pytest.mark.parametrize("model_type", [ModelType.SDXL_BASE, ModelType.SDXL_TURBO]) +def test_sdxl_fp4_recipe(model_type): + model = _RecipeBackbone() + config = Quantizer( + QuantizationConfig(format=QuantFormat.FP4), + ModelConfig(model_type=model_type), + logging.getLogger(__name__), + ).get_quant_config(n_steps=1, backbone=model) + + mtq.replace_quant_module(model) + mtq.set_quantizer_by_cfg(model, config["quant_cfg"]) + + for quantizer in (model.linear.input_quantizer, model.linear.weight_quantizer): + assert quantizer.is_enabled + assert quantizer.is_nvfp4_dynamic + assert quantizer.block_sizes[-1] == 16 + for projection in (model.attn.to_q, model.attn.to_k, model.attn.to_v): + assert not projection.input_quantizer.is_enabled + assert not projection.weight_quantizer.is_enabled + for quantizer in (model.conv.input_quantizer, model.conv.weight_quantizer): + assert quantizer.is_enabled + assert quantizer.is_fp8 + + +@pytest.mark.parametrize( + ("format_config", "mha_config", "expected_format", "disable_fp8_mha"), + [ + pytest.param( + _NVFP4_QUANTIZER_CONFIG, + _FP8_QUANTIZER_CONFIG, + QuantFormat.FP4, + False, + id="mixed-fp4", + ), + pytest.param( + _FP8_QUANTIZER_CONFIG, + _FP8_QUANTIZER_CONFIG, + QuantFormat.FP8, + False, + id="fp8", + ), + pytest.param( + {"num_bits": 8}, + {**_FP8_QUANTIZER_CONFIG, "enabled": False}, + QuantFormat.INT8, + True, + id="int8-disabled-fp8", + ), + pytest.param( + {"num_bits": 8}, + {"num_bits": 8}, + QuantFormat.INT8, + True, + id="int8-mha", + ), + ], +) +def test_restored_quantizer_state_drives_format_and_fp8_mha( + monkeypatch, format_config, mha_config, expected_format, disable_fp8_mha +): + backbone = nn.Module() + backbone.quantizer = _quantizer(**format_config) + backbone.attention = nn.Module() + for name in ("q_bmm_quantizer", "k_bmm_quantizer", "v_bmm_quantizer"): + setattr(backbone.attention, name, _quantizer(**mha_config)) + backbone.attention.bmm2_output_quantizer = lambda output: output + + fp8_sdpa = Mock(return_value=torch.empty(0)) + monkeypatch.setattr(diffusers_plugin.FP8SDPA, "apply", fp8_sdpa) + monkeypatch.setattr(torch.onnx, "is_in_onnx_export", lambda: True) + + assert _infer_restored_quantization_format([("transformer", backbone)]) == expected_format + diffusers_plugin._quantized_sdpa(backbone.attention, *(torch.empty(1) for _ in range(3))) + assert fp8_sdpa.call_args.args[-1] is disable_fp8_mha + + +def test_restore_infers_checkpoint_format_for_export(monkeypatch, tmp_path): + backbone = nn.Module() + backbone.quantizer = _quantizer(**_NVFP4_QUANTIZER_CONFIG) + + pipeline_manager = Mock() + pipeline_manager.create_pipeline.return_value = object() + pipeline_manager.iter_backbones.return_value = [("transformer", backbone)] + export_manager = Mock() + monkeypatch.setattr(_quantize, "PipelineManager", lambda *args: pipeline_manager) + monkeypatch.setattr(_quantize, "ExportManager", lambda *args: export_manager) + monkeypatch.setattr( + sys, + "argv", + [ + "quantize.py", + "--model", + "flux-schnell", + "--restore-from", + str(tmp_path), + "--onnx-dir", + str(tmp_path / "onnx"), + ], + ) + + _quantize.main() + + export_manager.restore_checkpoint.assert_called_once_with() + assert export_manager.export_onnx.call_args.args[-1] == QuantFormat.FP4 + export_manager.export_hf_ckpt.assert_called_once() + + +def test_flux_fp8_export_saves_converted_rope_graph(monkeypatch, tmp_path): + original_model = Mock() + converted_model = Mock() + monkeypatch.setattr( + diffusion_export, + "generate_dummy_kwargs_and_dynamic_axes_and_shapes", + lambda *args: ({}, {}, None), + ) + monkeypatch.setattr(diffusion_export, "onnx_export", lambda *args, **kwargs: None) + monkeypatch.setattr(diffusion_export.onnx, "load", lambda *args, **kwargs: original_model) + convert_rope_weight_type = Mock(return_value=converted_model) + monkeypatch.setattr(diffusion_export, "flux_convert_rope_weight_type", convert_rope_weight_type) + save_onnx = Mock() + monkeypatch.setattr(diffusion_export, "save_onnx", save_onnx) + + diffusion_export.modelopt_export_sd(nn.Module(), tmp_path, "flux-dev", "fp8") + + convert_rope_weight_type.assert_called_once_with(original_model) + save_onnx.assert_called_once_with(converted_model, tmp_path / "model.onnx") diff --git a/tests/unit/onnx/autocast/test_autocast.py b/tests/unit/onnx/autocast/test_autocast.py index ccad0c7201b..1377b20af9f 100644 --- a/tests/unit/onnx/autocast/test_autocast.py +++ b/tests/unit/onnx/autocast/test_autocast.py @@ -27,6 +27,7 @@ from modelopt.onnx.autocast import convert_to_mixed_precision from modelopt.onnx.autocast.__main__ import get_parser, main from modelopt.onnx.autocast.convert import convert_to_f16 +from modelopt.onnx.autocast.graphsanitizer import GraphSanitizer from modelopt.onnx.autocast.logging_config import configure_logging configure_logging("DEBUG") @@ -153,6 +154,42 @@ def test_convert_simple_model(temp_model_path, temp_output_path, keep_io_types): onnx.checker.check_model(loaded_model) +def test_convert_external_data_sanitizes_once_and_materializes_initializers( + tmp_path, simple_model, monkeypatch +): + model_path = tmp_path / "external_model.onnx" + onnx.save_model( + simple_model, + model_path, + save_as_external_data=True, + all_tensors_to_one_file=True, + location="external_model.data", + size_threshold=0, + ) + + sanitize_calls = [] + original_sanitize = GraphSanitizer.sanitize + + def record_sanitize(sanitizer): + initializer = sanitizer.model.graph.initializer[0] + sanitize_calls.append( + (sanitizer.onnx_path, initializer.data_location, bool(initializer.raw_data)) + ) + return original_sanitize(sanitizer) + + monkeypatch.setattr(GraphSanitizer, "sanitize", record_sanitize) + + converted_model = convert_to_mixed_precision(onnx_path=str(model_path), data_max=np.inf) + + assert sanitize_calls == [(str(model_path.resolve()), onnx.TensorProto.EXTERNAL, False)] + for initializer in converted_model.graph.initializer: + assert initializer.data_location != onnx.TensorProto.EXTERNAL + assert not initializer.external_data + assert initializer.raw_data + assert initializer.data_type == onnx.TensorProto.FLOAT16 + onnx.checker.check_model(converted_model) + + def assert_input_precision(nodes, dtype="float16"): for node in nodes: for inp in node.inputs: @@ -199,11 +236,6 @@ def test_conv_resize_conversion(tmp_path): # Convert the model converted_model = convert_to_mixed_precision(onnx_path=onnx_path) - # Output model should be produced in the same tmp_path - output_onnx_path = onnx_path.replace(".onnx", ".fp16.onnx") - onnx.save(converted_model, output_onnx_path) - - # Load the output model graph = gs.import_onnx(converted_model) # Check that Resize is correctly converted: @@ -213,6 +245,7 @@ def test_conv_resize_conversion(tmp_path): assert all(inp.dtype == np.float16 for inp in resize_node.inputs[0:2]), ( "Resize data and ROI inputs should be FP16" ) + assert resize_node.inputs[1].name != resize_node.inputs[2].name @pytest.mark.parametrize("target_opset", [13, 17, 19, 21]) diff --git a/tests/unit/onnx/autocast/test_graphsanitizer.py b/tests/unit/onnx/autocast/test_graphsanitizer.py index cb487b56bf0..f42324a6338 100644 --- a/tests/unit/onnx/autocast/test_graphsanitizer.py +++ b/tests/unit/onnx/autocast/test_graphsanitizer.py @@ -13,10 +13,14 @@ # See the License for the specific language governing permissions and # limitations under the License. +from unittest.mock import Mock + import numpy as np +import onnx import pytest from onnx import TensorProto, helper, numpy_helper +import modelopt.onnx.autocast.graphsanitizer as graphsanitizer from modelopt.onnx.autocast.graphsanitizer import GraphSanitizer @@ -407,3 +411,63 @@ def test_convert_fp64_no_changes_needed(): assert sanitizer._convert_fp64_initializers() is False assert sanitizer._convert_fp64_io_types() is False assert sanitizer._convert_fp64_nodes() is False + + +def test_sanitize_large_external_initializer_metadata(): + external_data_bytes = onnx.checker.MAXIMUM_PROTOBUF + 2048 + initializer = TensorProto( + name="weight", + data_type=TensorProto.FLOAT, + dims=[external_data_bytes // 4], + data_location=TensorProto.EXTERNAL, + ) + for key, value in ( + ("location", "weight.bin"), + ("offset", "0"), + ("length", str(external_data_bytes)), + ): + initializer.external_data.add(key=key, value=value) + + output = helper.make_tensor_value_info("output", TensorProto.FLOAT, initializer.dims) + identity = helper.make_node("Identity", [initializer.name], [output.name]) + graph = helper.make_graph( + [identity], "large_external_initializer", [], [output], initializer=[initializer] + ) + model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 22)]) + model.ir_version = 10 + assert model.ByteSize() < 1024 + + sanitizer = GraphSanitizer(model, min_opset=22) + sanitizer.sanitize() + + initializer = sanitizer.model.graph.initializer[0] + assert initializer.data_location == TensorProto.EXTERNAL + assert not initializer.raw_data + assert {entry.key: entry.value for entry in initializer.external_data} == { + "location": "weight.bin", + "offset": "0", + "length": str(external_data_bytes), + } + assert [node.op_type for node in sanitizer.model.graph.node] == ["Identity"] + + +def test_find_custom_nodes_uses_source_model_path(tmp_path, monkeypatch): + x = helper.make_tensor_value_info("X", TensorProto.FLOAT, [1]) + y = helper.make_tensor_value_info("Y", TensorProto.FLOAT, [1]) + custom_node = helper.make_node("CustomOp", [x.name], [y.name], name="custom") + graph = helper.make_graph([custom_node], "custom_graph", [x], [y]) + model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 22)]) + model_path = tmp_path / "custom.onnx" + tensor_info = {"Y": {"dtype": TensorProto.FLOAT, "shape": [1]}} + get_custom_layers = Mock(return_value=([custom_node.name], tensor_info)) + infer_types_shapes = Mock(return_value=model) + + monkeypatch.setattr(graphsanitizer, "set_trt_plugin_domain", Mock(return_value=model)) + monkeypatch.setattr(graphsanitizer, "get_custom_layers", get_custom_layers) + monkeypatch.setattr(graphsanitizer, "infer_types_shapes_tensorrt", infer_types_shapes) + + sanitizer = GraphSanitizer(model, min_opset=22, onnx_path=str(model_path)) + sanitizer.find_custom_nodes() + + get_custom_layers.assert_called_once_with(str(model_path.resolve()), []) + infer_types_shapes.assert_called_once_with(model, [], all_tensor_info=tensor_info) diff --git a/tests/unit/onnx/autocast/test_precisionconverter.py b/tests/unit/onnx/autocast/test_precisionconverter.py index b480bb7c24e..361cef26a5e 100644 --- a/tests/unit/onnx/autocast/test_precisionconverter.py +++ b/tests/unit/onnx/autocast/test_precisionconverter.py @@ -2220,6 +2220,33 @@ def test_convert_to_f16_restores_public_io_metadata_from_entry_boundary(): onnx.checker.check_model(converted, full_check=True) +def test_convert_to_f16_preserves_declared_output_type_after_inference_changes_it(): + graph_input = helper.make_tensor_value_info("X", TensorProto.FLOAT16, [2, 3]) + graph_output = helper.make_tensor_value_info("Y", TensorProto.FLOAT, [2, 3]) + node = helper.make_node("Identity", ["X"], ["Y"], name="Identity_0") + graph = helper.make_graph([node], "inferred_output_type", [graph_input], [graph_output]) + model = helper.make_model( + graph, + producer_name="inferred_output_type", + opset_imports=[helper.make_opsetid("", 19)], + ir_version=10, + ) + + converted = convert_to_f16( + model, keep_io_types=True, op_block_list=[], trt_plugins=[], opset=19 + ) + + output = next(vi for vi in converted.graph.output if vi.name == "Y") + assert output.type.tensor_type.elem_type == TensorProto.FLOAT + output_producers = [node for node in converted.graph.node if "Y" in node.output] + assert len(output_producers) == 1 + assert output_producers[0].op_type == "Cast" + assert next(attr.i for attr in output_producers[0].attribute if attr.name == "to") == ( + TensorProto.FLOAT + ) + onnx.checker.check_model(converted, full_check=True) + + def test_convert_to_f16_refreshes_gathernd_pre_cast_declaration(monkeypatch): def discover_test_plugins_without_trt(self): self.custom_ops = { diff --git a/tests/unit/onnx/autocast/test_referencerunner.py b/tests/unit/onnx/autocast/test_referencerunner.py index 00075714c69..5c5c3c00ab2 100644 --- a/tests/unit/onnx/autocast/test_referencerunner.py +++ b/tests/unit/onnx/autocast/test_referencerunner.py @@ -17,9 +17,12 @@ import os import tempfile from collections import OrderedDict +from unittest.mock import Mock import numpy as np import onnx +import polygraphy.backend.onnx as polygraphy_onnx +import polygraphy.backend.onnxrt as polygraphy_onnxrt import pytest from onnx import TensorProto, helper @@ -67,6 +70,54 @@ def test_init(simple_model): assert isinstance(runner.model, onnx.ModelProto) +def test_get_ort_runner_uses_file_when_modified_byte_size_fails(monkeypatch, reference_runner): + modified_model = Mock() + modified_model.ByteSize.side_effect = ValueError("model size unavailable") + check_external_data = Mock(return_value=False) + save_onnx = Mock() + bytes_from_onnx = Mock(side_effect=AssertionError("Byte serialization should not be used")) + session_from_onnx = Mock(return_value="session") + onnxrt_runner = Mock(return_value="runner") + + monkeypatch.setattr(onnx_utils, "check_model_uses_external_data", check_external_data) + monkeypatch.setattr(onnx_utils, "save_onnx", save_onnx) + monkeypatch.setattr(polygraphy_onnx, "BytesFromOnnx", bytes_from_onnx) + monkeypatch.setattr(polygraphy_onnxrt, "SessionFromOnnx", session_from_onnx) + monkeypatch.setattr(polygraphy_onnxrt, "OnnxrtRunner", onnxrt_runner) + + runners, model_temp_dir = reference_runner._get_ort_runner(lambda: modified_model) + + assert model_temp_dir is not None + with model_temp_dir as model_temp_path: + save_onnx.assert_called_once() + model_path = save_onnx.call_args.args[1] + assert runners == ["runner"] + assert os.path.dirname(model_path) == model_temp_path + check_external_data.assert_called_once_with(modified_model) + modified_model.ByteSize.assert_called_once_with() + save_onnx.assert_called_once_with(modified_model, model_path, save_as_external_data=True) + bytes_from_onnx.assert_not_called() + session_from_onnx.assert_called_once_with(model_path, providers=reference_runner.providers) + onnxrt_runner.assert_called_once_with("session") + + +def test_run_cleans_model_tempdir_when_input_loading_fails(monkeypatch, reference_runner): + model_temp_dir = Mock() + monkeypatch.setattr( + reference_runner, "_get_ort_runner", Mock(return_value=([], model_temp_dir)) + ) + monkeypatch.setattr( + reference_runner, + "_load_inputs", + Mock(side_effect=ValueError("input loading failed")), + ) + + with pytest.raises(ValueError, match="input loading failed"): + reference_runner.run() + + model_temp_dir.cleanup.assert_called_once_with() + + def test_run_with_random_inputs(reference_runner): """Test running inference with random inputs.""" results = reference_runner.run() diff --git a/tests/unit/onnx/quantization/test_ort_patching_histogram.py b/tests/unit/onnx/quantization/test_ort_patching_histogram.py new file mode 100644 index 00000000000..e986454bb57 --- /dev/null +++ b/tests/unit/onnx/quantization/test_ort_patching_histogram.py @@ -0,0 +1,182 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Tests for ONNX Runtime histogram quantization patches.""" + +import numpy as np +import onnx +import onnxruntime as ort +import pytest +from onnx import TensorProto, helper, numpy_helper +from onnxruntime.quantization.calibrate import ( + CalibrationDataReader, + CalibrationMethod, + HistogramCollector, + TensorData, + TensorsData, +) + +from modelopt.onnx.quantization.ort_patching import ( + _collect_value, + _collect_value_histogram_collector_single_node_calibration, + _compute_scale_zp, + _prepare_histogram_data, + _quantize_static, + _restore_histogram_calibration_dtypes, + patch_ort_modules, +) + + +def test_compute_scale_zp_fp16_overflow_fallback(): + zero_point, scale = _compute_scale_zp( + np.array(-65504, dtype=np.float16), + np.array(65504, dtype=np.float16), + np.array(-128, dtype=np.int8), + np.array(127, dtype=np.int8), + symmetric=True, + ) + + assert zero_point.dtype == np.int8 + assert zero_point == 0 + assert scale.dtype == np.float16 + assert scale == np.float16(514) + + +def test_quantize_static_fp16_high_range_scale(tmp_path): + class HighRangeDataReader(CalibrationDataReader): + def __init__(self): + self.rewind() + + def get_next(self): + return next(self.data, None) + + def rewind(self): + values = np.array([[-65504, 65504, -32752, 32752]], dtype=np.float16) + self.data = iter([{"input": values}]) + + model_path = tmp_path / "model.onnx" + output_path = tmp_path / "model.quant.onnx" + graph = helper.make_graph( + [helper.make_node("MatMul", ["input", "weight"], ["output"], name="matmul")], + "fp16_high_range", + [helper.make_tensor_value_info("input", TensorProto.FLOAT16, [1, 4])], + [helper.make_tensor_value_info("output", TensorProto.FLOAT16, [1, 4])], + [numpy_helper.from_array(np.eye(4, dtype=np.float16), name="weight")], + ) + model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 21)]) + model.ir_version = min(model.ir_version, 10) + onnx.save(model, model_path) + + patch_ort_modules(False) + _quantize_static( + model_path, + output_path, + HighRangeDataReader(), + nodes_to_quantize=["matmul"], + op_types_to_quantize=["MatMul"], + calibrate_method=CalibrationMethod.Entropy, + extra_options={ + "ExecutionProviders": ["CPUExecutionProvider"], + "ActivationSymmetric": True, + "AddQDQPairToWeight": True, + }, + ) + + quantized_model = onnx.load(output_path) + initializers = { + initializer.name: initializer for initializer in quantized_model.graph.initializer + } + scale_initializers = [ + initializers[node.input[1]] + for node in quantized_model.graph.node + if node.op_type in {"QuantizeLinear", "DequantizeLinear"} and node.input[1] in initializers + ] + assert scale_initializers + assert {initializer.data_type for initializer in scale_initializers} == {TensorProto.FLOAT16} + assert all( + np.isfinite(numpy_helper.to_array(initializer)).all() for initializer in scale_initializers + ) + ort.InferenceSession(output_path, providers=["CPUExecutionProvider"]) + + +@pytest.mark.parametrize( + "collect_value", + [ + _collect_value, + _collect_value_histogram_collector_single_node_calibration, + ], +) +def test_collect_value_fp16_narrow_range(collect_value): + collector = HistogramCollector( + method="entropy", + symmetric=False, + num_bins=128, + num_quantized_bins=128, + percentile=None, + scenario="same", + ) + activations = np.zeros(1000, dtype=np.float16) + for activation_max in (1e-6, 1e-6, 2e-6): + activations[0] = np.float16(activation_max) + collect_value(collector, {"narrow_fp16_tensor": [activations]}) + + hist, edges, _, _, threshold = collector.histogram_dict["narrow_fp16_tensor"] + assert hist.sum() == 3 * activations.size + assert len(hist) > collector.num_bins + assert edges.dtype == np.float32 + assert np.all(np.diff(edges) > 0), "fp16 bin edges are not strictly increasing" + assert np.asarray(threshold).dtype.itemsize >= np.dtype(np.float32).itemsize + + tensors_range = TensorsData(CalibrationMethod.Entropy, collector.compute_collection_result()) + _restore_histogram_calibration_dtypes(collector, tensors_range) + tensor_range = tensors_range["narrow_fp16_tensor"] + assert tensor_range.lowest.dtype == np.float16 + assert tensor_range.highest.dtype == np.float16 + assert tensor_range.bins.dtype == np.float32 + + +def test_restore_histogram_calibration_dtypes_clamps_fp16(): + collector = HistogramCollector( + method="distribution", + symmetric=False, + num_bins=512, + num_quantized_bins=128, + percentile=None, + scenario="same", + ) + _prepare_histogram_data(collector, "tensor", np.array([], dtype=np.float16)) + _prepare_histogram_data(collector, "missing_tensor", np.array([], dtype=np.float16)) + + fp32_max = np.finfo(np.float32).max + tensor_data = TensorData( + lowest=np.float32(-fp32_max), + highest=np.float32(fp32_max), + avg=np.float32(fp32_max), + std=np.float32(fp32_max), + hist=np.array([1]), + hist_edges=np.array([-1, 1], dtype=np.float32), + ) + tensors_range = TensorsData(CalibrationMethod.Distribution, {"tensor": tensor_data}) + + _restore_histogram_calibration_dtypes(collector, tensors_range) + + fp16_limits = np.finfo(np.float16) + tensor_range = tensors_range["tensor"] + assert tensor_range.lowest == fp16_limits.min + assert tensor_range.highest == fp16_limits.max + assert tensor_range.avg == fp16_limits.max + assert tensor_range.std == fp16_limits.max + assert tensor_range.hist_edges.dtype == np.float32 + assert "missing_tensor" not in tensors_range diff --git a/tests/unit/onnx/quantization/test_qdq_utils.py b/tests/unit/onnx/quantization/test_qdq_utils.py index 4b1e69ec538..95944e1f7ab 100644 --- a/tests/unit/onnx/quantization/test_qdq_utils.py +++ b/tests/unit/onnx/quantization/test_qdq_utils.py @@ -15,14 +15,22 @@ import warnings +import ml_dtypes import numpy as np +import onnx import onnx_graphsurgeon as gs import onnxruntime as ort import pytest from onnx import TensorProto, helper, numpy_helper -from modelopt.onnx.export import INT4QuantExporter, MXFP8QuantExporter, NVFP4QuantExporter +from modelopt.onnx.export import ( + FP8QuantExporter, + INT4QuantExporter, + MXFP8QuantExporter, + NVFP4QuantExporter, +) from modelopt.onnx.export.nvfp4_exporter import _cast_fp4 +from modelopt.onnx.quantization.gs_patching import _export_value_info_proto from modelopt.onnx.quantization.qdq_utils import ( _cast_fp8, apply_column_major_transformation, @@ -34,6 +42,7 @@ replace_zero_scale_with_smallest_nonzero, ) from modelopt.onnx.quantization.quant_utils import pack_float32_to_4bit_cpp_based +from modelopt.onnx.utils import get_opset_version def create_test_model_with_int4_dq_reshape_transpose_matmul(constant_scale: bool = False): @@ -335,8 +344,7 @@ def create_test_model_with_nvfp4_qdq(with_transpose: bool = False): value_info=value_info, ) - model = helper.make_model(graph) - return model + return helper.make_model(graph, opset_imports=[helper.make_opsetid("", 20)]) class TestQuantizeWeightsToInt4: @@ -484,6 +492,57 @@ def test_cast_fp4(self, input_array, expected_array): assert np.all(result == expected_array) +class TestFP8QuantExporter: + """Test suite for FP8QuantExporter.""" + + def test_bf16_weights_and_scale_are_compressed(self): + weight_data = np.array([0.001312255859375], dtype=ml_dtypes.bfloat16) + scale_data = np.array(0.00099945068359375, dtype=ml_dtypes.bfloat16) + weight = gs.Constant("weight", weight_data) + scale = gs.Constant("linear/weight_quantizer/scale", scale_data) + quantized = gs.Variable("quantized", dtype=np.uint8, shape=weight_data.shape) + dequantized = gs.Variable( + "dequantized", dtype=TensorProto.BFLOAT16, shape=weight_data.shape + ) + value_info = _export_value_info_proto(dequantized, do_type_check=True) + assert value_info.type.tensor_type.elem_type == TensorProto.BFLOAT16 + graph = gs.Graph( + nodes=[ + gs.Node( + op="TRT_FP8QuantizeLinear", + inputs=[weight, scale], + outputs=[quantized], + ), + gs.Node( + op="TRT_FP8DequantizeLinear", + inputs=[quantized, scale], + outputs=[dequantized], + ), + ], + outputs=[dequantized], + opset=23, + ) + + converted_model = FP8QuantExporter.compress_weights(gs.export_onnx(graph)) + + onnx.checker.check_model(converted_model) + assert [node.op_type for node in converted_model.graph.node] == ["DequantizeLinear"] + assert converted_model.graph.output[0].type.tensor_type.elem_type == TensorProto.BFLOAT16 + fp8_weight = next( + initializer + for initializer in converted_model.graph.initializer + if initializer.name == "linear/weight_quantizer/fp8_weights" + ) + assert fp8_weight.data_type == TensorProto.FLOAT8E4M3FN + assert fp8_weight.raw_data == b"\x3b" + output_scale = next( + initializer + for initializer in converted_model.graph.initializer + if initializer.name == scale.name + ) + assert output_scale.data_type == TensorProto.BFLOAT16 + + class TestMXFP8QuantExporter: """Test suite for MXFP8QuantExporter.""" @@ -607,6 +666,8 @@ def test_fp4qdq_conversion(self, with_transpose): # Run FP4QDQ to 2DQ conversion converted_model = NVFP4QuantExporter.process_model(model) + assert get_opset_version(converted_model) == 23 + # Verify TRT_FP4QDQ node is removed fp4qdq_nodes = [node for node in converted_model.graph.node if node.op_type == "TRT_FP4QDQ"] assert len(fp4qdq_nodes) == 0 diff --git a/tests/unit/onnx/test_autocast_quantize.py b/tests/unit/onnx/test_autocast_quantize.py index bc123aad9dc..930187980f8 100644 --- a/tests/unit/onnx/test_autocast_quantize.py +++ b/tests/unit/onnx/test_autocast_quantize.py @@ -55,7 +55,18 @@ def test_autocast_quantize_int8(tmp_path, keep_io_types, bias_add): assert os.path.isfile(output_onnx_path) # Load the output model and check QDQ node placements - graph = gs.import_onnx(onnx.load(output_onnx_path)) + quantized_model = onnx.load(output_onnx_path) + graph = gs.import_onnx(quantized_model) + + activation_scale_names = { + node.input[1] for node in quantized_model.graph.node if node.op_type == "QuantizeLinear" + } + activation_scale_types = { + initializer.data_type + for initializer in quantized_model.graph.initializer + if initializer.name in activation_scale_names + } + assert activation_scale_types == {onnx.TensorProto.FLOAT16} # Check that all MatMul nodes are quantized mm_nodes = [n for n in graph.nodes if n.op == "MatMul"] diff --git a/tests/unit/onnx/test_onnx_utils.py b/tests/unit/onnx/test_onnx_utils.py index 36face35b90..590a9f24186 100644 --- a/tests/unit/onnx/test_onnx_utils.py +++ b/tests/unit/onnx/test_onnx_utils.py @@ -14,6 +14,8 @@ # limitations under the License. import os +import sys +from unittest.mock import Mock import numpy as np import onnx @@ -28,6 +30,7 @@ make_tensor_value_info, ) +import modelopt.onnx.utils as onnx_utils from modelopt.onnx.trt_utils import load_onnx_model from modelopt.onnx.utils import ( clear_stale_value_info, @@ -56,6 +59,190 @@ def test_save_onnx(tmp_path): assert os.path.exists(os.path.join(tmp_path, "test.onnx")) +def _make_external_initializer( + name: str, location: str = "missing-shared-data.bin" +) -> onnx.TensorProto: + initializer = onnx.TensorProto(name=name, data_type=onnx.TensorProto.FLOAT, dims=[1]) + initializer.data_location = onnx.TensorProto.EXTERNAL + for key, value in [ + ("location", location), + ("offset", "0"), + ("length", "4"), + ]: + entry = initializer.external_data.add() + entry.key = key + entry.value = value + return initializer + + +def test_duplicate_shared_constants_preserves_and_materializes_external_data(tmp_path): + external_values = np.array([3.25], dtype=np.float32) + external_data_path = tmp_path / "shared.bin" + external_data_path.write_bytes(external_values.tobytes()) + shared = _make_external_initializer("shared", external_data_path.name) + collision = make_tensor("shared_1", onnx.TensorProto.FLOAT, [1], [0.0]) + sparse_collision = onnx.helper.make_sparse_tensor( + make_tensor("shared_2", onnx.TensorProto.FLOAT, [1], [0.0]), + make_tensor("", onnx.TensorProto.INT64, [1], [0]), + [1], + ) + nodes = [ + make_node("Add", ["input", "shared"], ["intermediate"]), + make_node("Add", ["intermediate", "shared"], ["output"]), + ] + graph = make_graph( + nodes, + "shared_external_initializer", + [ + make_tensor_value_info("input", onnx.TensorProto.FLOAT, [1]), + make_tensor_value_info("shared", onnx.TensorProto.FLOAT, [1]), + ], + [make_tensor_value_info("output", onnx.TensorProto.FLOAT, [1])], + initializer=[shared, collision], + ) + graph.sparse_initializer.append(sparse_collision) + model = make_model(graph) + + result, modified = onnx_utils.duplicate_shared_constants(model) + + assert result is model + assert modified + duplicated_names = ("shared_3", "shared_4") + assert [node.input[1] for node in result.graph.node] == list(duplicated_names) + initializers = {initializer.name: initializer for initializer in result.graph.initializer} + assert set(initializers) == {"shared_1", *duplicated_names} + assert result.graph.sparse_initializer[0].values.name == "shared_2" + assert {graph_input.name for graph_input in result.graph.input} == {"input"} + for name in duplicated_names: + initializer = initializers[name] + assert initializer.data_location == onnx.TensorProto.EXTERNAL + assert [(entry.key, entry.value) for entry in initializer.external_data] == [ + ("location", "shared.bin"), + ("offset", "0"), + ("length", "4"), + ] + assert not initializer.HasField("raw_data") + + onnx.external_data_helper.load_external_data_for_model(result, str(tmp_path)) + for name in duplicated_names: + np.testing.assert_array_equal( + onnx.numpy_helper.to_array(initializers[name]), external_values + ) + + +def test_duplicate_shared_constants_fast_path_returns_original_model(): + shared = _make_external_initializer("single_use") + graph = make_graph( + [make_node("Add", ["input", "single_use"], ["output"])], + "single_external_initializer", + [make_tensor_value_info("input", onnx.TensorProto.FLOAT, [1])], + [make_tensor_value_info("output", onnx.TensorProto.FLOAT, [1])], + initializer=[shared], + ) + model = make_model(graph) + serialized_model = model.SerializeToString() + + result, modified = onnx_utils.duplicate_shared_constants(model) + + assert result is model + assert not modified + assert result.SerializeToString() == serialized_model + + +def test_duplicate_shared_constants_retains_initializer_captured_by_subgraphs(): + def make_branch(name): + output = make_tensor_value_info("branch_output", onnx.TensorProto.FLOAT, [1]) + return make_graph([make_node("Identity", ["shared"], [output.name])], name, [], [output]) + + shared = make_tensor("shared", onnx.TensorProto.FLOAT, [1], [1.0]) + nodes = [ + make_node("Add", ["input", "shared"], ["left"]), + make_node("Add", ["input", "shared"], ["right"]), + make_node( + "If", + ["condition"], + ["branch_value"], + then_branch=make_branch("then_branch"), + else_branch=make_branch("else_branch"), + ), + make_node("Sum", ["left", "right", "branch_value"], ["output"]), + ] + graph = make_graph( + nodes, + "captured_initializer", + [ + make_tensor_value_info("input", onnx.TensorProto.FLOAT, [1]), + make_tensor_value_info("condition", onnx.TensorProto.BOOL, []), + make_tensor_value_info("shared", onnx.TensorProto.FLOAT, [1]), + ], + [make_tensor_value_info("output", onnx.TensorProto.FLOAT, [1])], + initializer=[shared], + ) + model = make_model(graph, opset_imports=[make_opsetid("", 17)]) + + result, modified = onnx_utils.duplicate_shared_constants(model) + + assert modified + assert [node.input[1] for node in result.graph.node[:2]] == ["shared_1", "shared_2"] + assert {initializer.name for initializer in result.graph.initializer} == { + "shared", + "shared_1", + "shared_2", + } + assert "shared" in {graph_input.name for graph_input in result.graph.input} + onnx.checker.check_model(result) + + +def test_duplicate_shared_constants_retains_initializer_exposed_as_graph_output(): + shared = make_tensor("shared", onnx.TensorProto.FLOAT, [1], [1.0]) + graph = make_graph( + [ + make_node("Add", ["input", "shared"], ["left"]), + make_node("Add", ["input", "shared"], ["right"]), + ], + "initializer_graph_output", + [make_tensor_value_info("input", onnx.TensorProto.FLOAT, [1])], + [ + make_tensor_value_info("left", onnx.TensorProto.FLOAT, [1]), + make_tensor_value_info("right", onnx.TensorProto.FLOAT, [1]), + make_tensor_value_info("shared", onnx.TensorProto.FLOAT, [1]), + ], + initializer=[shared], + ) + model = make_model(graph) + + result, modified = onnx_utils.duplicate_shared_constants(model) + + assert modified + assert [node.input[1] for node in result.graph.node] == ["shared_1", "shared_2"] + assert {initializer.name for initializer in result.graph.initializer} == { + "shared", + "shared_1", + "shared_2", + } + assert [output.name for output in result.graph.output] == ["left", "right", "shared"] + onnx.checker.check_model(result) + + +@pytest.mark.parametrize( + ("model_size", "expected"), + [ + (1, False), + (onnx.checker.MAXIMUM_PROTOBUF - sys.getsizeof(b""), False), + (onnx.checker.MAXIMUM_PROTOBUF - sys.getsizeof(b"") + 1, True), + (0, True), + (ValueError("model size unavailable"), True), + ], +) +def test_is_model_too_large_for_protobuf(model_size, expected): + model = Mock() + if isinstance(model_size, Exception): + model.ByteSize.side_effect = model_size + else: + model.ByteSize.return_value = model_size + assert onnx_utils.is_model_too_large_for_protobuf(model) is expected + + def make_onnx_model_for_matmul_op(): input_left = np.array([1, 2]) input_right = np.array([1, 3]) diff --git a/tests/unit/recipe/test_step3p7_recipes.py b/tests/unit/recipe/test_step3p7_recipes.py new file mode 100644 index 00000000000..1f62c96da14 --- /dev/null +++ b/tests/unit/recipe/test_step3p7_recipes.py @@ -0,0 +1,177 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 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. + +"""Step-3.7 PTQ recipes: what the `moe` / `share_expert` naming does and does not match.""" + +import types + +import pytest +import torch +import torch.nn as nn +import torch.nn.functional as F + +import modelopt.torch.quantization as mtq +from modelopt.recipe import load_recipe +from modelopt.torch.quantization.nn import QuantModuleRegistry + +HIDDEN_SIZE = 32 +MOE_INTERMEDIATE_SIZE = 16 +NUM_EXPERTS = 2 + + +class _MoELinear(nn.Module): + """Step's expert-indexed projection: one 3-D weight, ``forward(x, expert_id)``.""" + + def __init__(self, num_experts, in_features, out_features): + super().__init__() + self.num_experts = num_experts + self.in_features = in_features + self.out_features = out_features + self.weight = nn.Parameter(torch.randn(num_experts, out_features, in_features) * 0.02) + + def forward(self, x, expert_id): + return F.linear(x.float(), self.weight[expert_id].float()) + + +class _StepMoEMLP(nn.Module): + def __init__(self): + super().__init__() + self.num_experts = NUM_EXPERTS + self.gate = nn.Linear(HIDDEN_SIZE, NUM_EXPERTS, bias=False) # router + self.up_proj = _MoELinear(NUM_EXPERTS, HIDDEN_SIZE, MOE_INTERMEDIATE_SIZE) + self.gate_proj = _MoELinear(NUM_EXPERTS, HIDDEN_SIZE, MOE_INTERMEDIATE_SIZE) + self.down_proj = _MoELinear(NUM_EXPERTS, MOE_INTERMEDIATE_SIZE, HIDDEN_SIZE) + + +class _StepMLP(nn.Module): + """Dense FFN — used both as the non-MoE layers' ``mlp`` and as ``share_expert``.""" + + def __init__(self): + super().__init__() + self.gate_proj = nn.Linear(HIDDEN_SIZE, MOE_INTERMEDIATE_SIZE, bias=False) + self.up_proj = nn.Linear(HIDDEN_SIZE, MOE_INTERMEDIATE_SIZE, bias=False) + self.down_proj = nn.Linear(MOE_INTERMEDIATE_SIZE, HIDDEN_SIZE, bias=False) + + +class _StepMoELayer(nn.Module): + def __init__(self): + super().__init__() + self.self_attn = nn.Module() + self.self_attn.q_proj = nn.Linear(HIDDEN_SIZE, HIDDEN_SIZE, bias=False) + self.moe = _StepMoEMLP() + self.share_expert = _StepMLP() + + +class _StepDenseLayer(nn.Module): + def __init__(self): + super().__init__() + self.self_attn = nn.Module() + self.self_attn.q_proj = nn.Linear(HIDDEN_SIZE, HIDDEN_SIZE, bias=False) + self.mlp = _StepMLP() + + +class _StepModel(nn.Module): + """Mirrors Step3p7ForConditionalGeneration's module paths (one MoE, one dense layer).""" + + def __init__(self): + super().__init__() + # register_moe_linear_on_the_fly gates on the Step-family model_type; real Step + # checkpoints carry this in config.json. + self.config = types.SimpleNamespace(model_type="step3p7") + self.model = nn.Module() + self.model.language_model = nn.Module() + self.model.language_model.layers = nn.ModuleList([_StepMoELayer(), _StepDenseLayer()]) + self.lm_head = nn.Linear(HIDDEN_SIZE, HIDDEN_SIZE, bias=False) + + +@pytest.fixture(autouse=True) +def _unregister_moe_linear(): + yield + if QuantModuleRegistry.get(_MoELinear) is not None: + QuantModuleRegistry.unregister(_MoELinear) + + +def _quantize_with_recipe(name): + """Convert (no calibration) with a built-in recipe and return the model.""" + model = _StepModel() + config = load_recipe(name).quantize.model_dump() + config["algorithm"] = None + # `mtq.quantize` runs the custom-model plugins itself, which is what registers Step's + # `MoELinear` — no explicit registration here, so this also covers that hook firing. + mtq.quantize(model, config) + return model + + +def _enabled(module, quantizer="weight_quantizer"): + return getattr(module, quantizer).is_enabled + + +@pytest.mark.parametrize( + "recipe_name", + [ + "huggingface/step3p7/ptq/nvfp4_experts_only-kv_fp8_cast", + "huggingface/step3p7/ptq/nvfp4_mlp_only-kv_fp8", + ], +) +def test_routed_experts_are_quantized(recipe_name): + """The routed experts — the bulk of the model — must end up quantized, per expert.""" + model = _quantize_with_recipe(recipe_name) + moe = model.model.language_model.layers[0].moe + + for proj in ("up_proj", "gate_proj", "down_proj"): + experts = getattr(moe, proj).experts + assert len(experts) == NUM_EXPERTS + for expert in experts: + assert _enabled(expert) + assert _enabled(expert, "input_quantizer") + # Dynamic NVFP4: 16-element blocks along the input dim. + assert expert.weight_quantizer.block_sizes[-1] == 16 + assert expert.weight_quantizer.num_bits == (2, 1) + + +@pytest.mark.parametrize( + "recipe_name", + [ + "huggingface/step3p7/ptq/nvfp4_experts_only-kv_fp8_cast", + "huggingface/step3p7/ptq/nvfp4_mlp_only-kv_fp8", + ], +) +def test_router_shared_expert_and_head_stay_bf16(recipe_name): + """`*moe*` also matches the router; the shared expert and lm_head stay unquantized too.""" + model = _quantize_with_recipe(recipe_name) + moe_layer = model.model.language_model.layers[0] + + assert not _enabled(moe_layer.moe.gate) + for proj in ("gate_proj", "up_proj", "down_proj"): + assert not _enabled(getattr(moe_layer.share_expert, proj)) + assert not _enabled(model.lm_head) + + +def test_experts_only_leaves_dense_mlp_bf16(): + model = _quantize_with_recipe("huggingface/step3p7/ptq/nvfp4_experts_only-kv_fp8_cast") + dense_mlp = model.model.language_model.layers[1].mlp + + for proj in ("gate_proj", "up_proj", "down_proj"): + assert not _enabled(getattr(dense_mlp, proj)) + + +def test_mlp_only_also_quantizes_dense_mlp(): + model = _quantize_with_recipe("huggingface/step3p7/ptq/nvfp4_mlp_only-kv_fp8") + dense_mlp = model.model.language_model.layers[1].mlp + + for proj in ("gate_proj", "up_proj", "down_proj"): + assert _enabled(getattr(dense_mlp, proj)) + # Attention projections are out of scope for both recipes. + assert not _enabled(model.model.language_model.layers[1].self_attn.q_proj) diff --git a/tests/unit/torch/deploy/utils/test_torch_onnx_utils.py b/tests/unit/torch/deploy/utils/test_torch_onnx_utils.py index 399085ad699..2fd1a1c8504 100644 --- a/tests/unit/torch/deploy/utils/test_torch_onnx_utils.py +++ b/tests/unit/torch/deploy/utils/test_torch_onnx_utils.py @@ -13,7 +13,9 @@ # See the License for the specific language governing permissions and # limitations under the License. +import copy import json +import tempfile from contextlib import nullcontext import numpy as np @@ -24,6 +26,7 @@ import torch.nn as nn from _test_utils.torch.deploy.lib_test_models import BaseDeployModel, get_deploy_models +import modelopt.torch.quantization as mtq from modelopt.onnx.utils import get_batch_size_from_bytes, validate_batch_size from modelopt.torch._deploy.utils import ( OnnxBytes, @@ -57,6 +60,64 @@ } +class _FP8ModelWithBuffer(nn.Sequential): + fp32_buffer: torch.Tensor + + def forward(self, inputs): + return super().forward(inputs) + self.fp32_buffer + + +def _make_fp8_model(source_dtype, kind="fp8"): + if kind == "format": + model = nn.Sequential(*(nn.Linear(128, 128, bias=False) for _ in range(2))) + sample_input = torch.ones(1, 128, dtype=source_dtype) + config = copy.deepcopy(mtq.FP8_DEFAULT_CFG) + config["quant_cfg"].extend( + [ + { + "quantizer_name": "1.weight_quantizer", + "cfg": {"num_bits": 4, "block_sizes": {-1: 128, "type": "static"}}, + }, + {"quantizer_name": "1.input_quantizer", "enable": False}, + ] + ) + else: + model = _FP8ModelWithBuffer( + nn.Conv2d(1, 1, 1, bias=False), + nn.Flatten(), + nn.Linear(4, 4, bias=False), + ) + sample_input = torch.ones(1, 1, 2, 2, dtype=source_dtype) + config = mtq.FP8_DEFAULT_CFG + model = model.eval().to(source_dtype) + if kind != "format": + model.register_buffer("fp32_buffer", torch.ones(4)) + if kind == "parameters": + model.register_parameter("unused_fp32_parameter", nn.Parameter(torch.ones(1))) + quantized_model = mtq.quantize( + model, + config, + forward_loop=lambda quantized_model: quantized_model(sample_input), + ) + if kind != "format": + # Keep calibration nonzero while exercising Conv scale underflow during export. + with torch.no_grad(): + quantized_model[0].weight.fill_(1e-38 if source_dtype == torch.bfloat16 else 1.0) + return quantized_model, sample_input + + +def _export_fp8_model(source_dtype, weights_dtype): + model, sample_input = _make_fp8_model(source_dtype) + onnx_bytes, _ = get_onnx_bytes_and_metadata( + model, + (sample_input,), + weights_dtype=weights_dtype, + onnx_opset=23, + ) + onnx_bytes_obj = OnnxBytes.from_bytes(onnx_bytes) + return onnx.load_model_from_string(onnx_bytes_obj.get_onnx_model_file_bytes()) + + @pytest.mark.parametrize( "model", deploy_benchmark_dynamo.values(), ids=deploy_benchmark_dynamo.keys() ) @@ -157,6 +218,87 @@ def test_onnx_export_and_inputs(model: BaseDeployModel): ) +@pytest.mark.parametrize( + ("source_dtype", "weights_dtype", "expected_onnx_dtype"), + [ + (torch.bfloat16, "bf16", onnx.TensorProto.BFLOAT16), + (torch.float32, "fp16", onnx.TensorProto.FLOAT16), + ], + ids=["bf16-weight-focused-buffer", "fp32-to-fp16"], +) +def test_fp8_export_with_supported_weights_dtype(source_dtype, weights_dtype, expected_onnx_dtype): + exported_model = _export_fp8_model(source_dtype, weights_dtype) + + onnx.checker.check_model(exported_model, full_check=True) + assert {"TRT_FP8QuantizeLinear", "TRT_FP8DequantizeLinear"}.isdisjoint( + node.op_type for node in exported_model.graph.node + ) + initializer_by_name = { + initializer.name: initializer for initializer in exported_model.graph.initializer + } + fp8_weight_dq_nodes = [ + node + for node in exported_model.graph.node + if node.op_type == "DequantizeLinear" + and node.input[0] in initializer_by_name + and initializer_by_name[node.input[0]].data_type == onnx.TensorProto.FLOAT8E4M3FN + ] + assert len(fp8_weight_dq_nodes) == 2 + for node in fp8_weight_dq_nodes: + assert initializer_by_name[node.input[1]].data_type == expected_onnx_dtype + assert {0x7F, 0xFF}.isdisjoint(initializer_by_name[node.input[0]].raw_data) + assert all( + value.type.tensor_type.elem_type == expected_onnx_dtype + for value in exported_model.graph.input + ) + expected_output_dtype = ( + onnx.TensorProto.FLOAT if weights_dtype == "bf16" else expected_onnx_dtype + ) + assert all( + value.type.tensor_type.elem_type == expected_output_dtype + for value in exported_model.graph.output + ) + if weights_dtype == "bf16": + fp32_buffer = next( + initializer + for initializer in exported_model.graph.initializer + if initializer.name.endswith("fp32_buffer") + ) + assert fp32_buffer.data_type == onnx.TensorProto.FLOAT + assert any( + node.op_type == "Add" and fp32_buffer.name in node.input + for node in exported_model.graph.node + ) + + +@pytest.mark.parametrize( + ("kind", "source_dtype", "weights_dtype", "error"), + [ + ("fp8", torch.float32, "bf16", "torch.float32"), + ("fp8", torch.bfloat16, "fp16", "torch.bfloat16"), + ("parameters", torch.bfloat16, "bf16", "torch.bfloat16, torch.float32"), + ("format", torch.bfloat16, "bf16", "torch.bfloat16"), + ], + ids=["fp32-to-bf16", "bf16-to-fp16", "mixed-parameters", "mixed-format"], +) +def test_fp8_export_rejects_unsupported_dtype_conversion( + kind, source_dtype, weights_dtype, error, monkeypatch, tmp_path +): + monkeypatch.setattr(tempfile, "tempdir", str(tmp_path)) + model, sample_input = _make_fp8_model(source_dtype, kind) + with pytest.raises( + ValueError, + match=rf"Converting .* to {weights_dtype.upper()}.*source parameter dtypes: {error}", + ): + get_onnx_bytes_and_metadata( + model, + (sample_input,), + weights_dtype=weights_dtype, + onnx_opset=23, + ) + assert not any(tmp_path.iterdir()) + + class SingleArgModel(nn.Module): def forward(self, x: torch.Tensor): return torch.add(x, x) - x diff --git a/tests/unit/torch/quantization/plugins/test_attention_quant.py b/tests/unit/torch/quantization/plugins/test_attention_quant.py index 702cf3ad1db..e7b8fb69d61 100644 --- a/tests/unit/torch/quantization/plugins/test_attention_quant.py +++ b/tests/unit/torch/quantization/plugins/test_attention_quant.py @@ -185,3 +185,129 @@ def test_p_qdq_mode_detection(): sq.block_sizes = None sq.disable() assert quant_attention._p_qdq_mode() is None + + +@pytest.mark.parametrize("quantization_active", [True, False]) +def test_causal_p_qdq_dispatch_respects_quant_state(monkeypatch, quantization_active): + """Causal attention must dispatch according to the P quantizer runtime state.""" + quant_attention = make_quant_attention() + for name in ("q_bmm_quantizer", "k_bmm_quantizer", "v_bmm_quantizer"): + getattr(quant_attention, name).disable() + + pq = quant_attention.p_bmm_quantizer + pq.num_bits = (4, 3) + pq.block_sizes = None + if quantization_active: + pq.enable_quant() + else: + pq.disable_quant() + + calls = [] + expected = object() + expected_attention_mask = object() + + def triton_attention(*args, **kwargs): + calls.append("triton") + assert kwargs["attention_mask"] is expected_attention_mask + return expected + + monkeypatch.setattr( + quant_attention, + "_triton_qdq_attention", + triton_attention, + ) + monkeypatch.setattr( + quant_attention, + "_init_kitchen_attn_fn", + lambda: pytest.fail("P quantizer dispatch reached Kitchen initialization"), + ) + + def original_attention(_self, _query, _key, _value, attention_mask): + calls.append("original") + assert attention_mask is expected_attention_mask + return expected + + states = torch.zeros(1, 4, 2, 32) + output = quant_attention._quantized_attention( + original_attention, + quant_attention, + states, + states, + states, + expected_attention_mask, + ) + + assert output is expected + assert calls == ["triton" if quantization_active else "original"] + + +@pytest.mark.parametrize( + ("disable_method", "enable_method"), + [("disable", "enable"), ("disable_quant", "enable_quant")], +) +@pytest.mark.parametrize("preinitialized", [False, True]) +def test_kitchen_dispatch_respects_quantizer_runtime_state( + monkeypatch, disable_method, enable_method, preinitialized +): + """Kitchen dispatch must stop while P quantization is inactive and resume afterward.""" + quant_attention = make_quant_attention() + for name in ("q_bmm_quantizer", "k_bmm_quantizer", "v_bmm_quantizer"): + getattr(quant_attention, name).disable() + + pq = quant_attention.p_bmm_quantizer + pq.num_bits = (4, 3) + pq.block_sizes = {-1: 32, "type": "dynamic", "scale_bits": (8, 0)} + + calls = [] + + def kitchen_attention(query, _key, _value): + calls.append("kitchen") + return query.flatten(2) + + def init_kitchen(): + calls.append("init") + quant_attention.use_kitchen = True + quant_attention.kitchen_attn_fn = kitchen_attention + + monkeypatch.setattr( + "modelopt.torch.quantization.plugins.huggingface.kitchen", + object(), + ) + monkeypatch.setattr(quant_attention, "_init_kitchen_attn_fn", init_kitchen) + + if preinitialized: + quant_attention.use_kitchen = True + quant_attention.kitchen_attn_fn = kitchen_attention + + expected = object() + + def original_attention(_self, _query, _key, _value): + calls.append("original") + return expected + + states = torch.zeros(1, 4, 2, 32) + getattr(pq, disable_method)() + output = quant_attention._quantized_attention( + original_attention, + quant_attention, + states, + states, + states, + ) + + assert output is expected + assert calls == ["original"] + + calls.clear() + getattr(pq, enable_method)() + output = quant_attention._quantized_attention( + original_attention, + quant_attention, + states, + states, + states, + ) + + assert output[0].shape == (1, 2, 4, 32) + assert output[1] is None + assert calls == (["kitchen"] if preinitialized else ["init", "kitchen"]) diff --git a/tests/unit/torch/quantization/plugins/test_fused_experts.py b/tests/unit/torch/quantization/plugins/test_fused_experts.py index 07ef2aaff24..49cd999a205 100644 --- a/tests/unit/torch/quantization/plugins/test_fused_experts.py +++ b/tests/unit/torch/quantization/plugins/test_fused_experts.py @@ -385,6 +385,60 @@ def test_fold_weight_keep_attrs_keeps_amax_disables_rotation(self, monkeypatch): # --------------------------------------------------------------------------- # Tests for export # --------------------------------------------------------------------------- +class TestIterWeightQuantizersForCalibration: + """The quantizer-only iterator must agree with the weight iterator and never index a weight. + + Indexing the fused 3-D weight is what dispatches a redistribute collective per expert under + FSDP2, so callers that only read quantizer state go through the quantizer-only path. + """ + + @staticmethod + def _convert(model): + expert_type = type(model.moe.experts) + TestQuantFusedExperts._cleanup_registry(expert_type) + register_fused_experts_on_the_fly(model) + converted = QuantModuleRegistry.convert(model.moe.experts) + TestQuantFusedExperts._cleanup_registry(expert_type) + return converted + + @pytest.mark.parametrize( + "model_cls", [_TinyMoEModel, _TinyNonGatedMoEModel], ids=["gated", "non_gated"] + ) + def test_yields_the_same_quantizers_in_the_same_order(self, model_cls): + experts = self._convert(model_cls()) + + from_weights = [q for _, q in experts.iter_weights_for_calibration()] + quantizers_only = list(experts.iter_weight_quantizers_for_calibration()) + + assert quantizers_only, "expected per-expert weight quantizers" + assert [id(q) for q in quantizers_only] == [id(q) for q in from_weights] + + @pytest.mark.parametrize( + "model_cls", [_TinyMoEModel, _TinyNonGatedMoEModel], ids=["gated", "non_gated"] + ) + def test_does_not_index_the_fused_weight(self, model_cls): + """The point of the override: no ``weight[idx]``, which is the per-expert collective.""" + experts = self._convert(model_cls()) + + class _NoIndexing(torch.Tensor): + @staticmethod + def __new__(cls, data): + return torch.Tensor._make_subclass(cls, data, False) + + def __getitem__(self, item): + raise AssertionError("iter_weight_quantizers_for_calibration indexed the weight") + + for name in (experts._first_proj_attr, "down_proj"): + weight = getattr(experts, name) + setattr(experts, name, nn.Parameter(_NoIndexing(weight.data), requires_grad=False)) + + assert list(experts.iter_weight_quantizers_for_calibration()) + # The weight iterator is the expensive one; it must still slice, or the guard above is + # not actually testing anything. + with pytest.raises(AssertionError, match="indexed the weight"): + list(experts.iter_weights_for_calibration()) + + class TestExportFusedExperts: @staticmethod def _cleanup_registry(mod_type): diff --git a/tests/unit/torch/quantization/plugins/test_moe_linear.py b/tests/unit/torch/quantization/plugins/test_moe_linear.py new file mode 100644 index 00000000000..cceaacd44d5 --- /dev/null +++ b/tests/unit/torch/quantization/plugins/test_moe_linear.py @@ -0,0 +1,395 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Tests for _QuantMoELinear: expert-indexed MoE weights (Step-3.5 / Step-3.7 remote code).""" + +import types + +import pytest +import torch +import torch.nn as nn +import torch.nn.functional as F + +pytest.importorskip("transformers") + +import modelopt.torch.quantization as mtq +from modelopt.torch.export.hf_export_handlers import _export_moe_linear +from modelopt.torch.export.registry import ExportContext, ExportModuleRegistry +from modelopt.torch.quantization.nn import QuantModuleRegistry +from modelopt.torch.quantization.plugins.huggingface import ( + _is_expert_indexed_moe_linear, + _QuantMoELinear, + _reconstruct_fused_moe_linear, + register_moe_linear_on_the_fly, +) + +NUM_EXPERTS = 4 +HIDDEN_SIZE = 32 +MOE_INTERMEDIATE_SIZE = 16 +TOP_K = 2 + + +class _SyntheticMoELinear(nn.Module): + """Mimics Step-3.5 / Step-3.7 ``MoELinear`` (verbatim layout from their remote code).""" + + def __init__(self, num_experts, in_features, out_features): + super().__init__() + self.num_experts = num_experts + self.in_features = in_features + self.out_features = out_features + self.weight = nn.Parameter(torch.randn(num_experts, out_features, in_features) * 0.02) + + def forward(self, x, expert_id): + return F.linear(x.float(), self.weight[expert_id].float()) + + +class _SyntheticStepMoEMLP(nn.Module): + """Mimics ``Step3p7MoEMLP``: a router plus three expert-indexed projections.""" + + def __init__(self): + super().__init__() + self.num_experts = NUM_EXPERTS + self.top_k = TOP_K + self.gate = nn.Linear(HIDDEN_SIZE, NUM_EXPERTS, bias=False) + self.act_fn = nn.SiLU() + self.up_proj = _SyntheticMoELinear(NUM_EXPERTS, HIDDEN_SIZE, MOE_INTERMEDIATE_SIZE) + self.gate_proj = _SyntheticMoELinear(NUM_EXPERTS, HIDDEN_SIZE, MOE_INTERMEDIATE_SIZE) + self.down_proj = _SyntheticMoELinear(NUM_EXPERTS, MOE_INTERMEDIATE_SIZE, HIDDEN_SIZE) + + def forward(self, hidden_states): + tokens = hidden_states.view(-1, HIDDEN_SIZE) + routing = F.softmax(self.gate(tokens).float(), dim=-1) + weights, indices = torch.topk(routing, self.top_k, dim=-1) + out = torch.zeros_like(tokens) + for expert_id in range(self.num_experts): + pos, token_idx = torch.where(indices == expert_id) + if token_idx.numel() == 0: + continue + current = tokens[pos] + gate = self.act_fn(self.gate_proj(current, expert_id)) + up = self.up_proj(current, expert_id) + expert_out = self.down_proj(gate * up, expert_id) + out.index_add_(0, pos, (expert_out * weights[pos, token_idx, None]).to(out.dtype)) + return out.view_as(hidden_states) + + +class _TinyStepModel(nn.Module): + def __init__(self): + super().__init__() + # register_moe_linear_on_the_fly gates on the Step-family model_type; real Step + # checkpoints carry this in config.json. + self.config = types.SimpleNamespace(model_type="step3p7") + self.moe = _SyntheticStepMoEMLP() + + def forward(self, x): + return self.moe(x) + + +@pytest.fixture(autouse=True) +def _unregister_synthetic_moe_linear(): + """Keep the on-the-fly registration from leaking into other tests.""" + yield + if QuantModuleRegistry.get(_SyntheticMoELinear) is not None: + QuantModuleRegistry.unregister(_SyntheticMoELinear) + + +def _moe_quant_cfg(): + """Per-tensor INT8 on the expert projections only — CPU-friendly, no kernels needed.""" + return { + "quant_cfg": [ + {"quantizer_name": "*", "enable": False}, + {"quantizer_name": "*moe*weight_quantizer", "cfg": {"num_bits": 8, "axis": 0}}, + {"quantizer_name": "*moe*input_quantizer", "cfg": {"num_bits": 8, "axis": None}}, + {"quantizer_name": "*moe.gate.*", "enable": False}, + ], + "algorithm": "max", + } + + +def test_expert_indexed_moe_linear_is_detected(): + assert _is_expert_indexed_moe_linear( + _SyntheticMoELinear(NUM_EXPERTS, HIDDEN_SIZE, MOE_INTERMEDIATE_SIZE) + ) + + +@pytest.mark.parametrize( + "module", + [ + pytest.param(nn.Linear(HIDDEN_SIZE, HIDDEN_SIZE), id="plain_linear_2d_weight"), + pytest.param(nn.LayerNorm(HIDDEN_SIZE), id="norm_1d_weight"), + ], +) +def test_unrelated_modules_are_not_claimed(module): + assert not _is_expert_indexed_moe_linear(module) + + +def test_module_with_3d_weight_but_other_forward_is_not_claimed(): + """A 3-D weight alone is not enough — the forward must take ``(x, expert_id)``.""" + + class _NotExpertIndexed(_SyntheticMoELinear): + def forward(self, x, top_k_index, top_k_weights): + return x + + assert not _is_expert_indexed_moe_linear( + _NotExpertIndexed(NUM_EXPERTS, HIDDEN_SIZE, MOE_INTERMEDIATE_SIZE) + ) + + +def test_disabled_quantizers_reproduce_bf16_weight_fp32_compute_parity(): + """Conversion must not change the model's output when every quantizer is disabled, and + must not permanently promote expert storage to fp32 to get there. + + ``MoELinear.forward`` always promotes to fp32 for the matmul regardless of storage + dtype (``F.linear(x.float(), self.weight[expert_id].float())``). A wrapper that + instead downcasts the fp32 activation to the weight's original storage dtype (e.g. + bf16) before the matmul silently changes the model even with quantization off. The + opposite mistake -- promoting every expert's *storage* to fp32 in `_setup` to match -- + reproduces Step's numerics but doubles the model's expert-weight memory footprint for + its entire lifetime (on Step-3.7's full routed-expert set, ~354 GiB); the promotion + must be transient, scoped to the one expert actually being called. + """ + torch.manual_seed(0) + num_experts, in_features, out_features = 2, 4096, 1280 + module = _SyntheticMoELinear(num_experts, in_features, out_features) + module.weight.data = module.weight.data.to(torch.bfloat16) + module.config = types.SimpleNamespace(model_type="step3p7") # satisfy the family gate + x = torch.randn(8, in_features, dtype=torch.bfloat16) + reference = module(x, 0) + + mtq.quantize(module, {"quant_cfg": [{"quantizer_name": "*", "enable": False}]}) + assert isinstance(module, _QuantMoELinear), "conversion did not happen; test is vacuous" + + # Storage stays at the checkpoint's own dtype -- only the matmul promotes, transiently. + for expert in module.experts: + assert expert.weight.dtype == torch.bfloat16 + + converted = module(x, 0) + assert torch.equal(converted, reference) + + # Reconstruction (export) must also see -- and keep -- the original storage dtype, not + # a permanently-promoted one. + _reconstruct_fused_moe_linear(module) + assert module.weight.dtype == torch.bfloat16 + + +def test_grouped_routing_module_is_not_claimed(): + """A grouped-GEMM MoE layer has the identical shape but no scalar-index contract. + + Moondream3's ``MoeFusedLinear`` carries the same 3-D weight and the same three + attributes, but its second argument is a per-expert token-count *tensor*. Claiming it + would make ``_QuantMoELinear.forward`` evaluate ``self.experts[m_sizes]`` and raise + ``TypeError: only integer tensors of a single element can be converted to an index`` + on the first calibration forward. + """ + + class _MoeFusedLinear(nn.Module): + def __init__(self): + super().__init__() + self.num_experts = NUM_EXPERTS + self.in_features = HIDDEN_SIZE + self.out_features = MOE_INTERMEDIATE_SIZE + self.weight = nn.Parameter( + torch.randn(NUM_EXPERTS, MOE_INTERMEDIATE_SIZE, HIDDEN_SIZE) * 0.02 + ) + + def forward(self, input, m_sizes): + return input + + assert not _is_expert_indexed_moe_linear(_MoeFusedLinear()) + + +def test_extra_forward_parameters_are_not_claimed(): + """The replacement forward is exactly ``(x, expert_id)``. + + Anything else the caller could pass — a keyword-only `router_state`, `*args`, + `**kwargs` — would raise `TypeError` once the module is converted. + """ + + class _WithRouterState(_SyntheticMoELinear): + def forward(self, x, expert_id, *, router_state=None): + return x + + class _WithKwargs(_SyntheticMoELinear): + def forward(self, x, expert_id, **kwargs): + return x + + for cls in (_WithRouterState, _WithKwargs): + assert not _is_expert_indexed_moe_linear( + cls(NUM_EXPERTS, HIDDEN_SIZE, MOE_INTERMEDIATE_SIZE) + ), cls.__name__ + + +def test_transposed_weight_layout_is_not_claimed(): + """`[num_experts, in_features, out_features]` would rebuild each expert from wrong slices.""" + + class _TransposedMoELinear(_SyntheticMoELinear): + def __init__(self, num_experts, in_features, out_features): + super().__init__(num_experts, in_features, out_features) + self.weight = nn.Parameter(torch.randn(num_experts, in_features, out_features) * 0.02) + + assert not _is_expert_indexed_moe_linear( + _TransposedMoELinear(NUM_EXPERTS, HIDDEN_SIZE, MOE_INTERMEDIATE_SIZE) + ) + + +def test_offloaded_weights_are_refused_not_silently_corrupted(): + """Accelerate offload leaves `weight` on meta, with the value in the module's hook. + + Expanding that would copy meta storage into each expert and delete the key the hook + restores into, exporting a checkpoint of zeros. Conversion must refuse instead. + """ + pytest.importorskip("accelerate") + from accelerate import cpu_offload # local: accelerate is an optional dependency + + model = _TinyStepModel() + cpu_offload(model.moe.up_proj, execution_device=torch.device("cpu")) + assert model.moe.up_proj.weight.is_meta + + with pytest.raises(NotImplementedError, match="offloaded by Accelerate"): + mtq.quantize(model, _moe_quant_cfg(), forward_loop=None) + + +def test_registration_is_not_gated_on_exact_revision_class_name(): + """Any Step-family root (matched by `model_type`/class-name convention, not an exact + revision) registers its `MoELinear` modules -- Step-3.7 as readily as Step-3.5.""" + model = _TinyStepModel() + assert QuantModuleRegistry.get(_SyntheticMoELinear) is None + + register_moe_linear_on_the_fly(model) + + assert issubclass(QuantModuleRegistry.get(_SyntheticMoELinear), _QuantMoELinear) + + +def test_non_step_model_with_identical_signature_is_not_registered(): + """A same-shape, same-signature module is not enough on its own to be claimed. + + The structural check in `_is_expert_indexed_moe_linear` cannot tell a real Step + `MoELinear` apart from unrelated code that happens to reuse the `(x, expert_id)` + parameter names with different semantics (a per-expert bias or post-scale, say) -- + `_QuantMoELinear` would silently drop that behavior. Registration is therefore also + gated on the model being Step-family; a structurally identical module on a model that + is not must not be registered. + """ + + class _ThirdPartyMoEModel(nn.Module): + def __init__(self): + super().__init__() + self.config = types.SimpleNamespace(model_type="not_step") + self.moe = _SyntheticStepMoEMLP() + + def forward(self, x): + return self.moe(x) + + model = _ThirdPartyMoEModel() + register_moe_linear_on_the_fly(model) + + assert QuantModuleRegistry.get(_SyntheticMoELinear) is None + + +def test_local_hessian_calibration_fires_through_the_transient_weight_swap(): + """`forward` must keep calling `expert(x)` (`__call__`), not `.forward()` directly. + + `local_hessian_calibrate` registers a `forward_pre_hook` on each quantized Linear + module and relies on standard `nn.Module.__call__` dispatch to fire it. A `forward` + that bypassed `__call__` -- e.g. to reimplement the input/weight-quantize/output-quantize + sequence inline instead of transiently swapping the expert's weight storage -- would + silently skip this hook and leave the weight quantizer uncalibrated (amax stays None). + """ + torch.manual_seed(0) + model = _TinyStepModel() + cfg = { + "quant_cfg": [ + {"quantizer_name": "*", "enable": False}, + {"quantizer_name": "*moe*weight_quantizer", "cfg": {"num_bits": 8, "axis": 0}}, + {"quantizer_name": "*moe*input_quantizer", "cfg": {"num_bits": 8, "axis": None}}, + {"quantizer_name": "*moe.gate.*", "enable": False}, + ], + "algorithm": "local_hessian", + } + + def forward_loop(m): + for _ in range(3): + m(torch.randn(2, 8, HIDDEN_SIZE)) + + mtq.quantize(model, cfg, forward_loop=forward_loop) + + for expert in model.moe.up_proj.experts: + assert expert.weight_quantizer.amax is not None + + +def test_expert_indexed_moe_is_quantized_and_reconstructed(): + """Each expert gets its own quantizers, and export folds them back to the 3-D layout.""" + torch.manual_seed(0) + model = _TinyStepModel() + reference_weight = model.moe.up_proj.weight.detach().clone() + + def forward_loop(m): + m(torch.randn(2, 8, HIDDEN_SIZE)) + + mtq.quantize(model, _moe_quant_cfg(), forward_loop=forward_loop) + + # Every expert of every projection carries its own calibrated quantizer pair. + for proj in ("up_proj", "gate_proj", "down_proj"): + experts = getattr(model.moe, proj).experts + assert len(experts) == NUM_EXPERTS + for expert in experts: + assert expert.weight_quantizer.is_enabled + assert expert.weight_quantizer.amax is not None + # The router stays untouched. + assert not model.moe.gate.weight_quantizer.is_enabled + + _reconstruct_fused_moe_linear(model) + + # Back to the original ``[num_experts, out_features, in_features]`` parameter, so the + # exported keys match the hub checkpoint instead of per-expert names. + up_proj = model.moe.up_proj + assert not hasattr(up_proj, "experts") + assert up_proj.weight.shape == reference_weight.shape + assert torch.equal(up_proj.weight, reference_weight) + + +def test_export_handler_matches_a_differently_named_wrapper(): + """Export dispatch must key on the wrapper type, not the generated class name. + + The registration is structural, so a compatible remote-code class can be named + anything; its generated class is then ``Quant``. If the export handler only + matched the literal name ``QuantMoELinear``, such a module would skip + ``_export_moe_linear`` and export without the input-amax fallback for experts that + calibration never routed to. + """ + model = _TinyStepModel() + mtq.quantize(model, _moe_quant_cfg(), forward_loop=lambda m: m(torch.randn(2, 8, HIDDEN_SIZE))) + + converted = model.moe.up_proj + # The generated name is derived from the model's own class, not from `MoELinear`. + assert type(converted).__name__ == "Quant_SyntheticMoELinear" + assert ExportModuleRegistry.match(converted) is _export_moe_linear + + +def test_export_handler_fills_input_amax_for_unrouted_experts(): + """The handler is what gives never-routed experts an input amax before export.""" + torch.manual_seed(0) + model = _TinyStepModel() + mtq.quantize(model, _moe_quant_cfg(), forward_loop=lambda m: m(torch.randn(2, 8, HIDDEN_SIZE))) + + experts = model.moe.up_proj.experts + # Simulate an expert that calibration never routed a token to. + experts[0].input_quantizer.reset_amax() + assert experts[0].input_quantizer.amax is None + assert any(e.input_quantizer.amax is not None for e in experts), "need a donor amax" + + _export_moe_linear("moe.up_proj", model.moe.up_proj, ExportContext(model, torch.float16)) + + assert experts[0].input_quantizer.amax is not None diff --git a/tests/unit/torch/quantization/test_autoquant.py b/tests/unit/torch/quantization/test_autoquant.py index e83f7fa0a70..58ba4c4f627 100644 --- a/tests/unit/torch/quantization/test_autoquant.py +++ b/tests/unit/torch/quantization/test_autoquant.py @@ -15,6 +15,8 @@ import copy import io +import warnings +from contextlib import nullcontext from types import SimpleNamespace import pytest @@ -36,6 +38,7 @@ QuantRecipe, QuantRecipeHparam, _AutoQuantizeBaseSearcher, + _AutoQuantizeGradientScoringSession, _module_search_space_signature, estimate_quant_compression, ) @@ -98,6 +101,47 @@ def get_input(self): return torch.randn(1, 4, 32) +class _ScoredMoeExpert(torch.nn.Module): + def __init__(self): + super().__init__() + self.gate_proj = torch.nn.Linear(8, 8) + self.up_proj = torch.nn.Linear(8, 8) + self.down_proj = torch.nn.Linear(8, 8) + + def forward(self, x): + return self.down_proj(self.gate_proj(x) + self.up_proj(x)) + + +class _ScoredMoeMlp(torch.nn.Module): + def __init__(self): + super().__init__() + self.experts = torch.nn.ModuleList([_ScoredMoeExpert(), _ScoredMoeExpert()]) + + def forward(self, x): + output = torch.zeros_like(x) + for expert in self.experts: + output = output + expert(x) + return output + + +class _ScoredMoeLayer(torch.nn.Module): + def __init__(self): + super().__init__() + self.mlp = _ScoredMoeMlp() + + def forward(self, x): + return self.mlp(x) + + +class _ScoredMoeModel(torch.nn.Module): + def __init__(self): + super().__init__() + self.layers = torch.nn.ModuleList([_ScoredMoeLayer()]) + + def forward(self, x): + return self.layers[0](x) + + @pytest.mark.parametrize( ("quant_cfg", "other_quant_cfg", "is_less_than"), [ @@ -905,6 +949,309 @@ def test_data_parallel_auto_quantize(skip_on_windows): spawn_multiprocess_job(2, _test_data_parallel_auto_quantize, backend="gloo") +def _test_data_parallel_moe_score_module(rank, size): + torch.manual_seed(1234) + model = _ScoredMoeModel() + data_loader = [torch.randn(2, 3, 8) for _ in range(2)] + with warnings.catch_warnings(record=True) as caught_warnings: + warnings.simplefilter("always") + model, search_history = mtq.auto_quantize( + model, + constraints={"effective_bits": 12.0}, + quantization_formats=[mtq.INT8_DEFAULT_CFG], + data_loader=data_loader, + forward_step=lambda model, batch: model(batch), + loss_func=lambda output, data: output.square().mean(), + num_calib_steps=2, + num_score_steps=2, + ) + + hparam = model.layers[0].mlp.experts[0].gate_proj.get_hparam("quant_recipe") + assert hparam.score_modules == [model.layers[0].mlp] + assert isinstance(model.layers[0].mlp._hparams_for_scoring, list) + + fallback_warnings = [ + str(warning.message) + for warning in caught_warnings + if "no parallel_state is set for score module" in str(warning.message) + ] + assert len(fallback_warnings) == len(model.layers[0].mlp._hparams_for_scoring) + assert all(str(type(model.layers[0].mlp)) in message for message in fallback_warnings) + assert all("first quantized child" in message for message in fallback_warnings) + + recipe = QuantRecipe(mtq.INT8_DEFAULT_CFG) + local_score = sum(hparam._importance_dict[recipe][m] for m in hparam.score_modules) + candidate = next( + candidate + for candidate in search_history["candidate_stats"].values() + if "layers.0.mlp.experts.0.gate_proj" in candidate["module_names"] + ) + recipe_idx = candidate["formats"].index(recipe) + torch.testing.assert_close( + local_score * size, + torch.tensor( + candidate["scores"][recipe_idx], + device=local_score.device, + dtype=local_score.dtype, + ), + ) + + scores = { + name: candidate["scores"] for name, candidate in search_history["candidate_stats"].items() + } + rank_zero_scores = DistributedProcessGroup.get_dist_syncd_obj( + scores if rank == 0 else None, + DistributedProcessGroup(None), + lambda values: values[0], + ) + assert scores == rank_zero_scores + + +def test_data_parallel_moe_score_module(skip_on_windows): + spawn_multiprocess_job(2, _test_data_parallel_moe_score_module, backend="gloo") + + +def test_score_hparam_registration_preserves_order(): + quant_modules = [ + mtq.quantize(torch.nn.Linear(4, 4), mtq.INT8_DEFAULT_CFG), + mtq.quantize(torch.nn.Linear(4, 4), mtq.INT8_DEFAULT_CFG), + ] + score_module = torch.nn.Identity() + recipe = QuantRecipe(mtq.INT8_DEFAULT_CFG) + + first = QuantRecipeHparam( + [recipe], + quant_modules=[quant_modules[0], quant_modules[1], quant_modules[0]], + score_modules=[score_module, score_module], + ) + second = QuantRecipeHparam( + [recipe], + quant_modules=[quant_modules[1]], + score_modules=[score_module], + ) + + assert first.quant_modules == quant_modules + assert first.score_modules == [score_module] + assert score_module._hparams_for_scoring == [first, second] + + +def test_gradient_scoring_tracks_reused_module_invocations(): + """Each autograd use of a shared score module retains its own replay difference.""" + no_quant_recipe = QuantRecipe(quant_cfg=None) + quant_recipe = QuantRecipe(mtq.INT8_DEFAULT_CFG) + + class TestHparam: + is_configurable = True + choices = [no_quant_recipe, quant_recipe] + active = no_quant_recipe + + class ScoreModule(torch.nn.Module): + def forward(self, x): + scale = 1.0 if hparam.active == no_quant_recipe else 2.0 + return scale * x + + class ReusedScoreModule(torch.nn.Module): + def __init__(self): + super().__init__() + self.shared = ScoreModule() + + def forward(self, x): + first = self.shared(x) + second = self.shared(3.0 * x) + # This use does not participate in autograd. + self.shared(5.0 * x.detach()) + # This output requires grad but is intentionally unused by the loss. + self.shared(7.0 * x) + return first.sum() + 2.0 * second.sum() + + model = ReusedScoreModule() + score_module = model.shared + hparam = TestHparam() + hparam._importance_dict = {recipe: {score_module: None} for recipe in hparam.choices} + score_module._hparams_for_scoring = [hparam] + inputs = torch.tensor([[1.0, 2.0]], requires_grad=True) + + hparam._importance_dict[quant_recipe][score_module] = None + delayed_session = _AutoQuantizeGradientScoringSession(model, [score_module], lambda *_: True) + with delayed_session: + delayed_loss = model(inputs) + assert delayed_session._output_grad_hook_handles + + assert not delayed_session._output_grad_hook_handles + delayed_loss.backward() + assert hparam._importance_dict[quant_recipe][score_module] is None + + with _AutoQuantizeGradientScoringSession(model, [score_module], lambda *_: True): + model(inputs).backward() + + # First use: sum(x**2) = 5. Second use: sum((2 * 3x)**2) = 180. + importance = hparam._importance_dict[quant_recipe][score_module] + torch.testing.assert_close( + importance, + torch.tensor(185.0, device=importance.device), + rtol=0, + atol=1e-6, + ) + + +def test_gradient_scoring_restores_model_after_failure(): + model = SimpleLinear() + patched_modules = [] + + def fail_during_scoring(model, data): + model(data) + patched_modules.extend( + module + for module in model.modules() + if getattr(module.forward, "__name__", None) == "patched_forward" + ) + raise RuntimeError("stop after scoring forward") + + with pytest.raises(RuntimeError, match="stop after scoring forward"): + mtq.auto_quantize( + model, + constraints={"effective_bits": 12.0}, + quantization_formats=[mtq.INT8_DEFAULT_CFG], + data_loader=[model.get_input()], + forward_step=lambda model, batch: model(batch), + forward_backward_step=fail_during_scoring, + num_calib_steps=1, + num_score_steps=1, + ) + + assert patched_modules + assert all( + getattr(module.forward, "__name__", None) != "patched_forward" for module in patched_modules + ) + assert all("forward" not in module.__dict__ for module in patched_modules) + assert all(param.requires_grad for param in model.parameters()) + for module in model.modules(): + for hparam in getattr(module, "_hparams_for_scoring", []): + assert hparam.active == hparam.original + + +@pytest.mark.parametrize("cudnn_enabled", [False, True]) +@pytest.mark.parametrize("fail", [False, True]) +def test_backward_scoring_session_preserves_sdpa_backends(cudnn_enabled, fail): + original = torch.backends.cuda.cudnn_sdp_enabled() + others = ( + torch.backends.cuda.flash_sdp_enabled(), + torch.backends.cuda.mem_efficient_sdp_enabled(), + torch.backends.cuda.math_sdp_enabled(), + ) + try: + torch.backends.cuda.enable_cudnn_sdp(cudnn_enabled) + with ( + pytest.raises(RuntimeError, match="scoring failed") if fail else nullcontext(), + _AutoQuantizeGradientScoringSession(torch.nn.Identity(), [], lambda *_: False), + ): + assert torch.backends.cuda.cudnn_sdp_enabled() == cudnn_enabled + assert others == ( + torch.backends.cuda.flash_sdp_enabled(), + torch.backends.cuda.mem_efficient_sdp_enabled(), + torch.backends.cuda.math_sdp_enabled(), + ) + if fail: + raise RuntimeError("scoring failed") + finally: + restored = torch.backends.cuda.cudnn_sdp_enabled() + torch.backends.cuda.enable_cudnn_sdp(original) + assert restored == cudnn_enabled + + +@pytest.mark.parametrize("bad_gradient", [float("nan"), float("inf"), -float("inf")]) +def test_auto_quantize_fails_fast_on_nonfinite_gradients(bad_gradient): + model = SimpleLinear() + original_cudnn = torch.backends.cuda.cudnn_sdp_enabled() + + def loss_func(output, _data): + output.register_hook(lambda grad: torch.full_like(grad, bad_gradient)) + return output.sum() + + with pytest.raises(RuntimeError, match="Non-finite output gradients in module") as exc_info: + mtq.auto_quantize( + model, + constraints={"effective_bits": 12.0}, + quantization_formats=[mtq.INT8_DEFAULT_CFG], + data_loader=[model.get_input()], + forward_step=lambda model, batch: model(batch), + loss_func=loss_func, + num_calib_steps=1, + num_score_steps=1, + ) + + assert "torch.backends.cuda.enable_cudnn_sdp(False)" in str(exc_info.value) + assert torch.backends.cuda.cudnn_sdp_enabled() == original_cudnn + assert any( + f"module '{name}'" in str(exc_info.value) + for name, module in model.named_modules() + if getattr(module, "_hparams_for_scoring", []) + ) + for name, module in model.named_modules(): + for hparam in getattr(module, "_hparams_for_scoring", []): + assert "forward" not in module.__dict__ + assert hparam.active == hparam.original + for recipe in hparam.choices: + importance = hparam._importance_dict[recipe][module] + if f"module '{name}'" in str(exc_info.value): + assert importance is None + else: + assert importance is None or torch.isfinite(importance).all() + + +def test_backward_scoring_session_restores_partial_setup(): + model = torch.nn.Sequential(torch.nn.Linear(4, 4)) + score_module = model[0] + score_module._hparams_for_scoring = [] + score_module.weight.requires_grad = False + original_requires_grad = {name: param.requires_grad for name, param in model.named_parameters()} + + def fail_on_second_parameter(name, _model): + if name.endswith("bias"): + raise RuntimeError("stop during scoring setup") + return True + + session = _AutoQuantizeGradientScoringSession( + model, + [score_module], + fail_on_second_parameter, + ) + with pytest.raises(RuntimeError, match="stop during scoring setup"), session: + pytest.fail("scoring setup should not complete") + + assert "forward" not in score_module.__dict__ + assert { + name: param.requires_grad for name, param in model.named_parameters() + } == original_requires_grad + + +@pytest.mark.parametrize("instance_override", [False, True]) +def test_gradient_scoring_restores_forward_attribute_layout(instance_override): + module = torch.nn.Identity() + module._hparams_for_scoring = [] + + def original_forward(x): + return x + 1 + + original_override = original_forward + if instance_override: + module.forward = original_override + + session = _AutoQuantizeGradientScoringSession( + module, + [module], + lambda _name, _model: False, + ) + with pytest.raises(RuntimeError, match="stop during scoring"), session: + assert module.__dict__["forward"] is not original_override + raise RuntimeError("stop during scoring") + + if instance_override: + assert module.__dict__["forward"] is original_override + else: + assert "forward" not in module.__dict__ + + def test_auto_quantize_budget_uses_no_quant_candidate_cost(monkeypatch): class _BudgetCaptureSearcher(AutoQuantizeGradientSearcher): def run_search_with_stats(self, max_weight_size, verbose=False): diff --git a/tests/unit/torch/quantization/test_onnx_export_cpu.py b/tests/unit/torch/quantization/test_onnx_export_cpu.py index ce2ef626d63..9b470f7b019 100644 --- a/tests/unit/torch/quantization/test_onnx_export_cpu.py +++ b/tests/unit/torch/quantization/test_onnx_export_cpu.py @@ -39,6 +39,15 @@ from modelopt.torch.quantization.utils import is_quantized_linear +def _export_to_onnx(model, sample_input, **kwargs): + buffer = io.BytesIO() + if "enable_onnx_checker" in inspect.signature(torch.onnx.export).parameters: + kwargs["enable_onnx_checker"] = False + torch.onnx.export(model, sample_input, buffer, dynamo=False, **kwargs) + buffer.seek(0) + return onnx.load_model_from_string(buffer.read()) + + @pytest.mark.parametrize("model_cls", TEST_MODELS) @pytest.mark.parametrize( ("num_bits", "per_channel_quantization", "constant_folding"), @@ -59,6 +68,41 @@ def test_onnx_export_cpu(model_cls, num_bits, per_channel_quantization, constant ) +def test_fp8_conv_export_preserves_custom_qdq_and_kernel_shape(): + model = torch.nn.Conv2d(3, 4, 3, bias=False).eval() + sample_input = torch.randn(1, 3, 8, 8) + model = mtq.quantize( + model, + mtq.FP8_DEFAULT_CFG, + forward_loop=lambda quantized_model: quantized_model(sample_input), + ) + + exported_model = _export_to_onnx(model, sample_input, opset_version=20) + producers = {output: node for node in exported_model.graph.node for output in node.output} + conv = next(node for node in exported_model.graph.node if node.op_type == "Conv") + + for conv_input in conv.input[:2]: + dequantize = producers[conv_input] + quantize = producers[dequantize.input[0]] + assert dequantize.op_type == "TRT_FP8DequantizeLinear" + assert quantize.op_type == "TRT_FP8QuantizeLinear" + + value_info = {value.name: value for value in exported_model.graph.value_info} + weight_dequantize = producers[conv.input[1]] + weight_quantize = producers[weight_dequantize.input[0]] + for value_name in (*weight_quantize.output, *weight_dequantize.output): + shape = [ + dimension.dim_value for dimension in value_info[value_name].type.tensor_type.shape.dim + ] + assert shape == [4, 3, 3, 3] + + kernel_shape = next( + attribute for attribute in conv.attribute if attribute.name == "kernel_shape" + ) + assert list(kernel_shape.ints) == [3, 3] + onnx.checker.check_model(exported_model) + + def test_nvfp4_exported_onnx_is_topologically_sorted(monkeypatch): def forward_loop(model): model(sample_input) @@ -78,26 +122,14 @@ def cpu_dynamic_block_quantize(inputs, *args): module.input_quantizer.disable() module.weight_quantizer._onnx_quantizer_type = "static" - buffer = io.BytesIO() - if "enable_onnx_checker" in inspect.signature(torch.onnx.export).parameters: - kwargs = {"enable_onnx_checker": False} - else: - kwargs = {} - - torch.onnx.export( + exported_model = _export_to_onnx( model, sample_input, - buffer, input_names=["input"], output_names=["output"], export_params=True, opset_version=21, - dynamo=False, - **kwargs, ) - - buffer.seek(0) - exported_model = onnx.load_model_from_string(buffer.read()) assert any(node.op_type == "TRT_FP4QDQ" for node in exported_model.graph.node) converted_model = NVFP4QuantExporter.process_model(exported_model) diff --git a/tests/unit/torch/quantization/test_quantize_cpu.py b/tests/unit/torch/quantization/test_quantize_cpu.py index 3e4925e7b63..eb62e2fbbcb 100644 --- a/tests/unit/torch/quantization/test_quantize_cpu.py +++ b/tests/unit/torch/quantization/test_quantize_cpu.py @@ -16,6 +16,8 @@ """High-level tests for quantization.""" import copy +import os +from unittest import mock import pytest import torch @@ -441,6 +443,188 @@ def test_enable_only_entry_preserves_attributes(): assert module.axis == 0, "axis should be preserved by enable-only entry" +def test_weight_patterns_matching_nothing_raise(): + """A config whose weight patterns match no module must fail, not quantize nothing. + + Otherwise calibration and export run to completion and produce a checkpoint that is + silently unquantized (``"quant_algo": null``). + """ + model = SimpleLinear() + config = { + "quant_cfg": [ + {"quantizer_name": "*", "enable": False}, + # No module in this model is named `experts`. + {"quantizer_name": "*.experts.*weight_quantizer", "cfg": {"num_bits": 8, "axis": 0}}, + ], + "algorithm": "max", + } + with pytest.raises(RuntimeError, match="no weight quantizer is enabled"): + mtq.quantize(model, config, lambda m: m(m.get_input())) + + +def test_config_without_weight_quantization_is_allowed(): + """Activation-only configs quantize no weight on purpose and must still run.""" + model = SimpleLinear() + config = { + "quant_cfg": [ + {"quantizer_name": "*", "enable": False}, + {"quantizer_name": "*input_quantizer", "cfg": {"num_bits": 8, "axis": None}}, + ], + "algorithm": "max", + } + model = mtq.quantize(model, config, lambda m: m(m.get_input())) + + for name, module in model.named_modules(): + if name.endswith("weight_quantizer"): + assert not module.is_enabled + + +def test_weight_quantizers_disabled_by_a_later_entry_are_allowed(): + """Patterns that match and are then switched off are a choice, not a mismatch.""" + model = SimpleLinear() + config = { + "quant_cfg": [ + {"quantizer_name": "*weight_quantizer", "cfg": {"num_bits": 4, "axis": 0}}, + {"quantizer_name": "*weight_quantizer", "enable": False}, + ], + "algorithm": "max", + } + model = mtq.quantize(model, config, lambda m: m(m.get_input())) + + for name, module in model.named_modules(): + if name.endswith("weight_quantizer"): + assert not module.is_enabled + + +def test_sequential_weight_quantizers_do_not_trip_the_guard(): + """List-valued `cfg` builds `SequentialQuantizer`s, but the guard only ever looks at + `TensorQuantizer` instances -- this pins that it still doesn't wrongly raise for them. + + `SequentialQuantizer` (itself an `nn.Sequential`) is not special-cased: its `TensorQuantizer` + children are reachable directly via `named_modules()`, individually named + `...weight_quantizer.0` / `.1`, so the substring match already finds them. + """ + model = SimpleLinear() + calib_data = [model.get_input() for _ in range(2)] + quantize_model_and_forward(model, copy.deepcopy(WINT4INT8_CFG), calib_data) + + for name, module in model.named_modules(): + if name.endswith("weight_quantizer"): + assert isinstance(module, SequentialQuantizer) + + +def test_fused_experts_quantizer_names_do_not_trip_the_guard(): + """Fused-experts quantizers are named `..._weight_quantizers.N`, plural and indexed, but + still contain `weight_quantizer` as a substring and so are still read by the guard. + """ + model = SimpleLinear() + mtq.quantize(model, mtq.INT8_DEFAULT_CFG, lambda m: m(m.get_input())) + + # Rename as the fused-experts path does; the config's `*weight_quantizer` must still match. + linear = model.net[0] + linear.add_module("gate_up_proj_weight_quantizers", torch.nn.ModuleList([TensorQuantizer()])) + + config = { + "quant_cfg": [ + {"quantizer_name": "*", "enable": False}, + {"quantizer_name": "*gate_up_proj_weight_quantizer", "cfg": {"num_bits": 8, "axis": 0}}, + ], + "algorithm": "max", + } + mtq.quantize(model, config, lambda m: m(m.get_input())) + + +def test_refining_an_already_quantized_model_does_not_raise(): + """A second config whose own patterns match nothing still sees the earlier weight + quantizers as enabled, so this is refining an already-quantized model, not a no-op run. + """ + model = SimpleLinear() + model = mtq.quantize(model, mtq.INT8_DEFAULT_CFG, lambda m: m(m.get_input())) + assert any(m.is_enabled for n, m in model.named_modules() if n.endswith("weight_quantizer")) + + refinement = { + "quant_cfg": [ + {"quantizer_name": "*.experts.*weight_quantizer", "cfg": {"num_bits": 4, "axis": 0}}, + ], + "algorithm": None, + } + mtq.quantize(model, refinement) + + # The earlier weight quantization is untouched. + assert any(m.is_enabled for n, m in model.named_modules() if n.endswith("weight_quantizer")) + + +def test_weight_patterns_enabled_then_retracted_do_not_raise(): + """An unmatched pattern that a later entry disables asks for nothing by the end.""" + model = SimpleLinear() + config = { + "quant_cfg": [ + {"quantizer_name": "*", "enable": False}, + {"quantizer_name": "*.missing.*weight_quantizer", "cfg": {"num_bits": 8, "axis": 0}}, + {"quantizer_name": "*.missing.*weight_quantizer", "enable": False}, + ], + "algorithm": "max", + } + model = mtq.quantize(model, config, lambda m: m(m.get_input())) + + for name, module in model.named_modules(): + if name.endswith("weight_quantizer"): + assert not module.is_enabled + + +def test_bare_wildcard_pattern_matching_nothing_raises(): + """A bare `"*"` (or another pattern never mentioning "weight") still expresses weight + intent if it would match a `weight_quantizer` name -- and must still raise if nothing in + the model actually has one, exactly like an explicit `*weight_quantizer` pattern would. + + A model with zero quantizable modules (e.g. `nn.Module()`, no Linear/Conv anywhere) is + the degenerate case where this matters: nothing in the config's own text says "weight", + so a substring-only check would silently return without ever looking at the model. + """ + model = torch.nn.Module() # no quantizable submodules at all + config = {"quant_cfg": [{"quantizer_name": "*", "cfg": {"num_bits": 8, "axis": 0}}]} + with pytest.raises(RuntimeError, match="no weight quantizer is enabled"): + mtq.quantize(model, config) + + +def test_overlapping_patterns_disabled_by_a_broader_later_one_still_raise(): + """A narrower pattern "matching" is not enough -- the final enabled state is what counts. + + `*weight_quantizer` enables real quantizers, but the later, broader `*` disables + everything again; the config's net effect is still "nothing quantized" and must raise. + A check that asked "did any weight pattern match something" instead of "is anything + actually enabled" would miss this, since the narrower pattern did match. + """ + model = SimpleLinear() + config = { + "quant_cfg": [ + {"quantizer_name": "*weight_quantizer", "cfg": {"num_bits": 8, "axis": 0}}, + {"quantizer_name": "*", "enable": False}, + ], + "algorithm": "max", + } + with pytest.raises(RuntimeError, match="no weight quantizer is enabled"): + mtq.quantize(model, config, lambda m: m(m.get_input())) + + +def test_skip_weight_quant_check_env_var_bypasses_the_guard(): + """Documented escape hatch for pipeline-parallel ranks whose local stage legitimately + has none of the targeted modules (e.g. a pure-attention stage under an experts-only + recipe): raising there while other ranks proceed into calibration is a collective hang, + not just a wrong per-rank verdict. + """ + model = SimpleLinear() + config = { + "quant_cfg": [ + {"quantizer_name": "*", "enable": False}, + {"quantizer_name": "*.experts.*weight_quantizer", "cfg": {"num_bits": 8, "axis": 0}}, + ], + "algorithm": "max", + } + with mock.patch.dict(os.environ, {"MODELOPT_SKIP_WEIGHT_QUANT_CHECK": "1"}): + mtq.quantize(model, config, lambda m: m(m.get_input())) + + def test_atomicity_later_cfg_entry_does_not_inherit_earlier(): """When two cfg-bearing entries match the same quantizer, the second fully replaces the first.""" model = SimpleLinear() diff --git a/tools/launcher/examples/nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-NVFP4/offline_kd_qad.yaml b/tools/launcher/examples/nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16/offline_kd_qad.yaml similarity index 99% rename from tools/launcher/examples/nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-NVFP4/offline_kd_qad.yaml rename to tools/launcher/examples/nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16/offline_kd_qad.yaml index 29cb5280ab8..741235d253f 100644 --- a/tools/launcher/examples/nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-NVFP4/offline_kd_qad.yaml +++ b/tools/launcher/examples/nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16/offline_kd_qad.yaml @@ -34,7 +34,7 @@ # # Usage: # source .env-slurm && cd tools/launcher -# uv run launch.py --yaml examples/nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-NVFP4/offline_kd_qad.yaml --yes +# uv run launch.py --yaml examples/nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16/offline_kd_qad.yaml --yes job_name: NVIDIA-Nemotron-3-Nano-30B-A3B-NVFP4-QAD-OfflineKD pipeline: