diff --git a/CHANGELOG.rst b/CHANGELOG.rst index aca584633e1..09e2a89660a 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -81,6 +81,8 @@ Changelog **Bug Fixes** +- Fix ``ShapeInferenceError`` during ONNX INT8 + FP16 quantization (``--high_precision_dtype fp16``) of weakly-typed models (e.g. TensorFlow exports) that carry stale rank-0 ``graph.output`` shapes or ops such as ``TopK`` that ONNX's static shape inference cannot resolve. ``clear_stale_value_info`` now reconciles stale output shapes via symbolic shape inference (keeping every output's shape field populated), and AutoCast runs ONNX shape inference in strict mode and falls back to schema-based standalone type inference when it fails, so unresolved ops no longer leave tensors untyped. +- Always list unquantized MoE routers/gates in the exported ``exclude_modules`` (NVBug 5718750). ``get_quant_config`` only recorded modules that carry a quantizer, but on ``transformers>=5.0`` MoE routers are no longer ``nn.Linear`` (e.g. ``TopKRouter``) and never receive one, so the BF16 router weight was written to the checkpoint yet omitted from ``exclude_modules``. vLLM / SGLang then treated it as quantized and failed to load (e.g. Qwen3-30B-A3B NVFP4: ``AssertionError: Tried to load weights of size [128, 2048] to a parameter of size [128, 1024]``). Routers are now detected structurally (an MoE block with an ``experts`` container plus a weight-bearing ``gate`` / ``router`` / ``shared_expert_gate`` submodule) and recorded as unquantized regardless of quantizer attachment. - In Megatron-Core only do EP amax sync for routed expert weights if ``sync_expert_weight_amax=True``. Previously EP amax sync would sync routed expert weights across EP ranks even when ``sync_expert_weight_amax`` was False. - Fix Megatron-Core HF importer to load fused ``TELayerNormColumnParallelLinear.layer_norm_weight`` from HF for GPT-family models (Qwen3 etc.) under ``--export-default-te-spec``. Importer now prefers per-context keys ``fused_input_layernorm`` / ``fused_pre_mlp_layernorm`` (fallback ``fused_norm`` for Nemotron-H backward compatibility); ``mcore_qwen.py`` provides the new rules. Without this fix, post-prune MMLU sat at chance. - Fix ONNX AutoCast ``keep_io_types=True`` sanity-check failure (``Unexpected type in I/O tensor ...``) when a network input/output is an empty tensor (a dimension of size 0). Such tensors were "fake-cast" (retyped in place) to the low precision type; because the value-info map aliases the ``graph.input``/``graph.output`` ``ValueInfoProto``, this silently changed the model's I/O type. AutoCast now inserts a real ``Cast`` for protected I/O tensors instead. @@ -88,6 +90,7 @@ Changelog - Fix the GPT-OSS MXFP4 → NVFP4 PTQ path in ``examples/llm_ptq/hf_ptq.py`` (used with ``--cast_mxfp4_to_nvfp4``). ``get_model`` now loads native MXFP4 checkpoints (``openai/gpt-oss-*``) dequantized to BF16 ``GptOssExperts`` via ``Mxfp4Config(dequantize=True)`` on a sequential device map. This fixes a CUDA illegal-memory access during the multi-GPU dequant load and the ``NotImplementedError`` for experts type ``Mxfp4GptOssExperts`` during unified HF export (the packed-kernel experts wrapper, used when the optional ``kernels`` package is installed, is unsupported by export); ``kernels`` is no longer required. The ``--cast_mxfp4_to_nvfp4`` step now also resolves a HF Hub ID ``--pyt_ckpt_path`` to its local snapshot directory instead of failing with ``FileNotFoundError``. - Fix ``_QuantGptOssExperts`` / ``_QuantLlama4TextExperts`` static-block NVFP4 weight calibration raising ``ValueError: Input shape has changed`` during the calibration forward. These experts quantize their weights transposed (``_transposed_quantize``); ``iter_weights_for_calibration`` now yields the same transposed view so weight-only calibration and the forward agree on the block-quant shape (and the export ``_amax`` orientation). - Fix unified HF checkpoint export for Llama4 MoE models. The uncalibrated-experts input-quantizer ``amax`` fallback in ``_export_transformers_checkpoint`` special-cased only ``QuantGptOssExperts``; ``QuantLlama4TextExperts`` uses the same fused ``gate_up_proj`` / ``down_proj`` layout and is now handled by the same branch, fixing the export failure. +- Fix ``NotImplementedError: "max_all_cuda" not implemented for 'Float8_e4m3fn'`` during quantization calibration of models with natively FP8 (``float8_e4m3fn`` / ``float8_e5m2``) weights, such as DeepSeek-V3. FP8 dtypes implement no reduction (``max``/``amax``), ``abs``, or elementwise ``maximum`` kernels, so ``reduce_amax`` now upcasts FP8 inputs to the default float dtype before reducing; the upcast is lossless and only affects the FP8 path. 0.44 (2026-05-14) ^^^^^^^^^^^^^^^^^ diff --git a/examples/llm_eval/README.md b/examples/llm_eval/README.md index 79f1b85d7e0..560aa17ec87 100644 --- a/examples/llm_eval/README.md +++ b/examples/llm_eval/README.md @@ -16,18 +16,24 @@ The supported eval tasks are [here](https://github.com/EleutherAI/lm-evaluation- ### Baseline +Both standard HuggingFace models and heterogeneous pruned checkpoints produced by Puzzletron are supported. + - For models which fit on a single GPU: ```sh python lm_eval_hf.py --model hf --model_args pretrained= --tasks --batch_size 4 ``` +For a quick smoke test, add `--limit 10` to any of the above commands to evaluate on only 10 samples. + - With model-sharding (for models which require multiple GPUs): ```sh python lm_eval_hf.py --model hf --model_args pretrained=,parallelize=True --tasks --batch_size 4 ``` +> **Note (Slurm interactive nodes):** On Slurm interactive nodes, `WORLD_SIZE` is set to the number of available GPUs in the shell environment. Running `python` directly causes `lm_eval` to hang waiting for peer ranks that were never spawned. Prepend `WORLD_SIZE=1` to the `python` commands above to fix this. This does not limit GPU usage — `parallelize=True` independently enables model parallelism across all available GPUs within the single process. The `accelerate launch` command manages `WORLD_SIZE` itself and does not require this workaround. + - For data-parallel evaluation with model-sharding: With the following command, the model will be sharded across `total_num_of_available_gpus/num_copies_of_your_model` with a data-parallelism of `num_copies_of_your_model` @@ -40,22 +46,6 @@ accelerate launch --multi_gpu --num_processes \ --batch_size 4 ``` -### Heterogeneous Pruned Checkpoints (Puzzletron) - -Heterogeneous pruned checkpoints produced by Puzzletron are automatically detected and loaded with the appropriate model patcher. No additional flags are needed beyond specifying the checkpoint path: - -```sh -python lm_eval_hf.py --model hf \ - --model_args pretrained=path/to/anymodel/checkpoint,dtype=bfloat16,parallelize=True \ - --tasks mmlu \ - --num_fewshot 5 \ - --batch_size 4 -``` - -For a quick smoke test, add `--limit 10`. - -> **Note:** Requires the `puzzletron` extra to be installed (`pip install -e ".[puzzletron]"`). - ### Quantized (simulated) - For simulated quantization with any of the default quantization formats: diff --git a/examples/llm_ptq/example_utils.py b/examples/llm_ptq/example_utils.py index 65363a1460e..6ffaa06a3cb 100755 --- a/examples/llm_ptq/example_utils.py +++ b/examples/llm_ptq/example_utils.py @@ -552,6 +552,44 @@ def get_original_hf_quant_method(config) -> str | None: return None +def _resolve_init_config(hf_config, auto_model_module, ckpt_path, config_kwargs): + """Re-derive a built-in config when a remote-code config is used with a built-in model + class, so it matches the model definition's version; fall back to hf_config otherwise. + """ + if auto_model_module in [AutoModelForCausalLM, AutoModel]: + return hf_config + if not type(hf_config).__module__.startswith("transformers_modules"): + return hf_config + builtin_config_kwargs = {k: v for k, v in config_kwargs.items() if k != "trust_remote_code"} + try: + return AutoConfig.from_pretrained(ckpt_path, **builtin_config_kwargs) + except Exception as e: + warnings.warn( + f"Could not re-derive a built-in config for {ckpt_path} ({e}); using the " + "remote-code config for device-map inference." + ) + return hf_config + + +def _get_config_dtype(config): + config_dtype = ( + getattr(config, "dtype", None) or getattr(config, "torch_dtype", None) or torch.bfloat16 + ) + if isinstance(config_dtype, str): + config_dtype = getattr(torch, config_dtype) + return config_dtype + + +def _apply_dtype_to_config(model_kwargs, config_dtype, architecture, apply_config_dtype=False): + model_kwargs = model_kwargs.copy() + if "DeciLM" in architecture: + model_kwargs["torch_dtype"] = config_dtype + model_kwargs.pop("dtype", None) + elif apply_config_dtype: + model_kwargs["dtype"] = config_dtype + return model_kwargs + + def get_model( ckpt_path, device="cuda", @@ -701,16 +739,21 @@ def has_pack_quantized_config(config): auto_model_module = getattr(transformers, architecture) from_config = auto_model_module._from_config + config_for_init = _resolve_init_config( + hf_config, auto_model_module, ckpt_path, config_kwargs + ) + with init_empty_weights(include_buffers=True): # When computing the device_map, assuming bfloat16 precision by default, # unless specified by the hf_config. - torch_dtype = getattr(hf_config, "torch_dtype", torch.bfloat16) - model_kwargs2 = model_kwargs.copy() + config_dtype = _get_config_dtype(config_for_init) + model_kwargs2 = _apply_dtype_to_config( + model_kwargs, config_dtype, architecture, apply_config_dtype=True + ) if auto_model_module not in [AutoModelForCausalLM, AutoModel]: model_kwargs2.pop("trust_remote_code", None) - model_kwargs2["dtype"] = torch_dtype model_kwargs2.pop("max_memory", None) - model = from_config(hf_config, **model_kwargs2) + model = from_config(config_for_init, **model_kwargs2) max_memory = get_max_memory() inferred_device_map = infer_auto_device_map(model, max_memory=max_memory) @@ -730,10 +773,11 @@ def has_pack_quantized_config(config): ) model_kwargs["max_memory"] = max_memory + model_kwargs2 = _apply_dtype_to_config(model_kwargs, config_dtype, architecture) model = auto_model_module.from_pretrained( ckpt_path, device_map=device_map, - **model_kwargs, + **model_kwargs2, ) model.eval() if has_pack_quantized_config(hf_config): diff --git a/examples/puzzletron/README.md b/examples/puzzletron/README.md index 497650ebdea..6a5b4c03b77 100644 --- a/examples/puzzletron/README.md +++ b/examples/puzzletron/README.md @@ -256,18 +256,7 @@ The plot shows how token accuracy changes with different compression rates. High ## Evaluation -Evaluate AnyModel checkpoints using [lm-eval](https://github.com/EleutherAI/lm-evaluation-harness) directly. - -```bash -python examples/llm_eval/lm_eval_hf.py \ - --model hf \ - --model_args pretrained=path/to/checkpoint,dtype=bfloat16,parallelize=True \ - --tasks mmlu \ - --num_fewshot 5 \ - --batch_size 4 -``` - -For a quick smoke test, add `--limit 10`. +Evaluate AnyModel checkpoints using lm-eval. See the [LM-Eval-Harness section](../llm_eval/README.md#lm-eval-harness) in `examples/llm_eval/README.md` for full instructions, including multi-GPU and Slurm setup. > **Alternative:** For server-based evaluation via an OpenAI-compatible endpoint, > see [evaluation/nemo_evaluator_instructions.md](./evaluation/nemo_evaluator_instructions.md). diff --git a/examples/puzzletron/configs/gptoss-20b_remove_experts_memory/validate_model_defaults.yaml b/examples/puzzletron/configs/gptoss-20b_remove_experts_memory/validate_model_defaults.yaml index b80faea5f5f..4ed2529e4ff 100644 --- a/examples/puzzletron/configs/gptoss-20b_remove_experts_memory/validate_model_defaults.yaml +++ b/examples/puzzletron/configs/gptoss-20b_remove_experts_memory/validate_model_defaults.yaml @@ -3,7 +3,7 @@ autocast_dtype: torch.bfloat16 # dtype for torch.autocast for validate_model block_size: 8192 bos_rate: 0.5 data_column: messages -val_dataset_name: valid +val_dataset_name: validation shuffle_seed: 81436 seed: 42 fim_rate: 0 diff --git a/examples/puzzletron/configs/llama-3_2-3B_pruneffn_memory/validate_model_defaults.yaml b/examples/puzzletron/configs/llama-3_2-3B_pruneffn_memory/validate_model_defaults.yaml index b80faea5f5f..4ed2529e4ff 100644 --- a/examples/puzzletron/configs/llama-3_2-3B_pruneffn_memory/validate_model_defaults.yaml +++ b/examples/puzzletron/configs/llama-3_2-3B_pruneffn_memory/validate_model_defaults.yaml @@ -3,7 +3,7 @@ autocast_dtype: torch.bfloat16 # dtype for torch.autocast for validate_model block_size: 8192 bos_rate: 0.5 data_column: messages -val_dataset_name: valid +val_dataset_name: validation shuffle_seed: 81436 seed: 42 fim_rate: 0 diff --git a/examples/puzzletron/configs/mistral-small-24b-instruct-2501_pruneffn_memory/validate_model_defaults.yaml b/examples/puzzletron/configs/mistral-small-24b-instruct-2501_pruneffn_memory/validate_model_defaults.yaml index ce1749d9698..6b36142a3a8 100644 --- a/examples/puzzletron/configs/mistral-small-24b-instruct-2501_pruneffn_memory/validate_model_defaults.yaml +++ b/examples/puzzletron/configs/mistral-small-24b-instruct-2501_pruneffn_memory/validate_model_defaults.yaml @@ -3,7 +3,7 @@ autocast_dtype: torch.bfloat16 # dtype for torch.autocast for validate_model block_size: 8192 bos_rate: 0.5 data_column: messages -val_dataset_name: valid +val_dataset_name: validation shuffle_seed: 81436 seed: 42 fim_rate: 0 diff --git a/examples/puzzletron/configs/nemotron-nano-12b-v2/validate_model_defaults.yaml b/examples/puzzletron/configs/nemotron-nano-12b-v2/validate_model_defaults.yaml index ce1749d9698..6b36142a3a8 100644 --- a/examples/puzzletron/configs/nemotron-nano-12b-v2/validate_model_defaults.yaml +++ b/examples/puzzletron/configs/nemotron-nano-12b-v2/validate_model_defaults.yaml @@ -3,7 +3,7 @@ autocast_dtype: torch.bfloat16 # dtype for torch.autocast for validate_model block_size: 8192 bos_rate: 0.5 data_column: messages -val_dataset_name: valid +val_dataset_name: validation shuffle_seed: 81436 seed: 42 fim_rate: 0 diff --git a/examples/puzzletron/configs/qwen2_5_7b_instruct_pruneffn_memory/validate_model_defaults.yaml b/examples/puzzletron/configs/qwen2_5_7b_instruct_pruneffn_memory/validate_model_defaults.yaml index ce1749d9698..6b36142a3a8 100644 --- a/examples/puzzletron/configs/qwen2_5_7b_instruct_pruneffn_memory/validate_model_defaults.yaml +++ b/examples/puzzletron/configs/qwen2_5_7b_instruct_pruneffn_memory/validate_model_defaults.yaml @@ -3,7 +3,7 @@ autocast_dtype: torch.bfloat16 # dtype for torch.autocast for validate_model block_size: 8192 bos_rate: 0.5 data_column: messages -val_dataset_name: valid +val_dataset_name: validation shuffle_seed: 81436 seed: 42 fim_rate: 0 diff --git a/examples/puzzletron/configs/qwen3-8b_pruneffn_memory/validate_model_defaults.yaml b/examples/puzzletron/configs/qwen3-8b_pruneffn_memory/validate_model_defaults.yaml index ce1749d9698..6b36142a3a8 100644 --- a/examples/puzzletron/configs/qwen3-8b_pruneffn_memory/validate_model_defaults.yaml +++ b/examples/puzzletron/configs/qwen3-8b_pruneffn_memory/validate_model_defaults.yaml @@ -3,7 +3,7 @@ autocast_dtype: torch.bfloat16 # dtype for torch.autocast for validate_model block_size: 8192 bos_rate: 0.5 data_column: messages -val_dataset_name: valid +val_dataset_name: validation shuffle_seed: 81436 seed: 42 fim_rate: 0 diff --git a/examples/torch_onnx/README.md b/examples/torch_onnx/README.md index cfd4dc380ce..c378c1eea40 100644 --- a/examples/torch_onnx/README.md +++ b/examples/torch_onnx/README.md @@ -122,8 +122,8 @@ source venv/bin/activate pip3 install . # Verify installation -tensorrt-edgellm-quantize-llm --help -tensorrt-edgellm-export-llm --help +tensorrt-edgellm-quantize --help +tensorrt-edgellm-export --help ``` **System requirements:** @@ -137,11 +137,8 @@ tensorrt-edgellm-export-llm --help | Tool | Purpose | | :--- | :--- | -| `tensorrt-edgellm-quantize-llm` | Quantize LLM models using ModelOpt (FP8, INT4 AWQ, NVFP4) | -| `tensorrt-edgellm-export-llm` | Export LLM to ONNX with precision-specific optimizations | -| `tensorrt-edgellm-export-visual` | Export visual encoders for multimodal VLM models | -| `tensorrt-edgellm-quantize-draft` | Quantize EAGLE draft models for speculative decoding | -| `tensorrt-edgellm-export-draft` | Export EAGLE draft models to ONNX | +| `tensorrt-edgellm-quantize` | Quantize models using ModelOpt (FP8, INT4 AWQ, NVFP4); subcommands: `llm`, `draft` | +| `tensorrt-edgellm-export` | Export quantized or FP16/BF16 checkpoint to ONNX; auto-detects VLM and audio components | | `tensorrt-edgellm-insert-lora` | Insert LoRA patterns into existing ONNX models | | `tensorrt-edgellm-process-lora` | Process LoRA adapter weights for runtime loading | @@ -149,64 +146,58 @@ tensorrt-edgellm-export-llm --help ```bash # Step 1: Quantize with ModelOpt -tensorrt-edgellm-quantize-llm \ +tensorrt-edgellm-quantize llm \ --model_dir Qwen/Qwen2.5-3B-Instruct \ --quantization fp8 \ --output_dir quantized/qwen2.5-3b-fp8 # Step 2: Export to ONNX -tensorrt-edgellm-export-llm \ - --model_dir quantized/qwen2.5-3b-fp8 \ - --output_dir onnx_models/qwen2.5-3b +tensorrt-edgellm-export \ + quantized/qwen2.5-3b-fp8 \ + onnx_models/qwen2.5-3b ``` ### Example: Quantize and Export a VLM ```bash -# Quantize the language model component -tensorrt-edgellm-quantize-llm \ +# Quantize with ModelOpt (handles both LLM and visual components) +tensorrt-edgellm-quantize llm \ --model_dir Qwen/Qwen2.5-VL-3B-Instruct \ --quantization fp8 \ --output_dir quantized/qwen2.5-vl-3b -# Export the language model -tensorrt-edgellm-export-llm \ - --model_dir quantized/qwen2.5-vl-3b \ - --output_dir onnx_models/qwen2.5-vl-3b/llm - -# Export the visual encoder -tensorrt-edgellm-export-visual \ - --model_dir Qwen/Qwen2.5-VL-3B-Instruct \ - --output_dir onnx_models/qwen2.5-vl-3b/visual +# Export to ONNX (auto-detects VLM and exports LLM + visual encoder to separate subdirs) +tensorrt-edgellm-export \ + quantized/qwen2.5-vl-3b \ + onnx_models/qwen2.5-vl-3b ``` ### Example: EAGLE Speculative Decoding ```bash # Quantize base model -tensorrt-edgellm-quantize-llm \ +tensorrt-edgellm-quantize llm \ --model_dir meta-llama/Llama-3.1-8B-Instruct \ --quantization fp8 \ --output_dir quantized/llama3.1-8b-base # Export base model with EAGLE flag -tensorrt-edgellm-export-llm \ - --model_dir quantized/llama3.1-8b-base \ - --output_dir onnx_models/llama3.1-8b/base \ - --is_eagle_base +tensorrt-edgellm-export \ + quantized/llama3.1-8b-base \ + onnx_models/llama3.1-8b/base \ + --eagle-base # Quantize EAGLE draft model -tensorrt-edgellm-quantize-draft \ +tensorrt-edgellm-quantize draft \ --base_model_dir meta-llama/Llama-3.1-8B-Instruct \ --draft_model_dir EAGLE3-LLaMA3.1-Instruct-8B \ --quantization fp8 \ --output_dir quantized/llama3.1-8b-draft # Export draft model -tensorrt-edgellm-export-draft \ - --draft_model_dir quantized/llama3.1-8b-draft \ - --base_model_dir meta-llama/Llama-3.1-8B-Instruct \ - --output_dir onnx_models/llama3.1-8b/draft +tensorrt-edgellm-export \ + quantized/llama3.1-8b-draft \ + onnx_models/llama3.1-8b/draft ``` ### Quantization Methods diff --git a/modelopt/onnx/autocast/convert.py b/modelopt/onnx/autocast/convert.py index ec66ec11edd..0288c729dd4 100644 --- a/modelopt/onnx/autocast/convert.py +++ b/modelopt/onnx/autocast/convert.py @@ -136,8 +136,10 @@ def convert_to_mixed_precision( graph_sanitizer.sanitize() model = graph_sanitizer.model - # Setup internal mappings - model = onnx_utils.infer_types(model, use_standalone_type_inference) + # Setup internal mappings. Use strict shape inference so an op ONNX cannot resolve surfaces + # 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) 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 @@ -267,8 +269,10 @@ def convert_to_f16( sanitizer.convert_fp64_to_fp32() model = sanitizer.model - # Setup internal mappings - model = onnx_utils.infer_types(model, use_standalone_type_inference) + # Setup internal mappings. Use strict shape inference so an op ONNX cannot resolve surfaces + # 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) value_info_map, initializer_map, node_to_init_map = utils.setup_mappings(model) precision_converter = PrecisionConverter( diff --git a/modelopt/onnx/autocast/precisionconverter.py b/modelopt/onnx/autocast/precisionconverter.py index 1603a05a2f2..853fa3d1cb3 100644 --- a/modelopt/onnx/autocast/precisionconverter.py +++ b/modelopt/onnx/autocast/precisionconverter.py @@ -255,7 +255,9 @@ def convert( self.model = self._propagate_types_shapes_custom_ops(self.model) else: # Clear type/shape information for intermediates and outputs (including subgraphs) - self._clear_types_and_shapes_recursive(self.model.graph) + utils.clear_types_and_shapes_recursive( + self.model.graph, clear_shapes=not self.use_standalone_type_inference + ) # Populate type information with inferred types self.model = onnx_utils.infer_types( self.model, self.use_standalone_type_inference, strict_mode=True, check_type=False @@ -284,66 +286,6 @@ def _ensure_types_are_defined(self): if vi.type.tensor_type.elem_type == onnx.TensorProto.UNDEFINED: vi.type.tensor_type.elem_type = self.low_precision_type.onnx_type - def _clear_types_and_shapes_recursive( - self, graph: onnx.GraphProto, is_subgraph: bool = False - ) -> None: - """Recursively clear type/shape information for a graph and all its subgraphs. - - If use_standalone_type_inference is True, we clear only types, not shapes. - For subgraphs, input types/shapes are cleared, so that the input types/shapes are propagated - from the main graph. - - Args: - graph: The ONNX graph to clear types and shapes for. - is_subgraph: Whether this is a subgraph (True) or the main graph (False). - """ - - def _clear_callback(g: onnx.GraphProto, parent: onnx.NodeProto, is_sub: bool) -> None: - logger.debug( - f"Clearing types/shapes in {'subgraph' if is_sub else 'main graph'}: {g.name}" - ) - - # Clear type/shape information for inputs (only for subgraphs, not main graph inputs) - if is_sub: - for inp in g.input: - if inp.type.HasField("tensor_type"): - inp.type.tensor_type.elem_type = onnx.TensorProto.UNDEFINED - if not self.use_standalone_type_inference: - for idx, d in enumerate(inp.type.tensor_type.shape.dim): - if d.dim_value: - inp.type.tensor_type.shape.dim[idx].dim_param = "unk" - - if is_sub: - # Identify which tensors are produced by nodes in this subgraph - subgraph_outputs = set() - for node in g.node: - subgraph_outputs.update(node.output) - - # Clear value_info only for intermediates produced by nodes in this subgraph - for vi in g.value_info: - if vi.name in subgraph_outputs: - vi.type.tensor_type.elem_type = onnx.TensorProto.UNDEFINED - if not self.use_standalone_type_inference: - for idx, d in enumerate(vi.type.tensor_type.shape.dim): - if d.dim_value: - vi.type.tensor_type.shape.dim[idx].dim_param = "unk" - else: - for vi in g.value_info: - vi.type.tensor_type.elem_type = onnx.TensorProto.UNDEFINED - for idx, d in enumerate(vi.type.tensor_type.shape.dim): - if d.dim_value: - vi.type.tensor_type.shape.dim[idx].dim_param = "unk" - - # Clear outputs for both main graph and subgraphs - for out in g.output: - out.type.tensor_type.elem_type = onnx.TensorProto.UNDEFINED - if not self.use_standalone_type_inference: - for idx, d in enumerate(out.type.tensor_type.shape.dim): - if d.dim_value: - out.type.tensor_type.shape.dim[idx].dim_param = "unk" - - utils.walk_subgraphs_recursive(graph, _clear_callback, is_subgraph=is_subgraph) - def _propagate_types_shapes_custom_ops(self, model): """Propagate types and shapes after insertion of 'Cast' nodes or other graph modifications.""" logger.info("Propagating tensor shapes and types in model with custom ops.") diff --git a/modelopt/onnx/autocast/utils.py b/modelopt/onnx/autocast/utils.py index a9ad5484c0c..218034a6413 100644 --- a/modelopt/onnx/autocast/utils.py +++ b/modelopt/onnx/autocast/utils.py @@ -28,6 +28,7 @@ import onnx import modelopt.onnx.utils as onnx_utils +from modelopt.onnx.autocast.logging_config import logger from modelopt.onnx.utils import get_opset_version @@ -115,6 +116,53 @@ def walk_subgraphs_recursive( walk_subgraphs_recursive(subgraph, callback, parent_node=node, is_subgraph=True) +def clear_types_and_shapes_recursive( + graph: onnx.GraphProto, clear_shapes: bool = True, is_subgraph: bool = False +) -> None: + """Recursively clear type/shape information for a graph and all its subgraphs. + + Resets intermediate (``value_info``) and output tensor types to ``UNDEFINED`` and, when + ``clear_shapes`` is True, replaces concrete dims with a symbolic ``"unk"`` so a subsequent + :func:`modelopt.onnx.utils.infer_types` re-derives them from the operator graph. For subgraphs, + input types/shapes are cleared too so they propagate from the parent graph. This does not change + tensor *rank*, so it cannot repair a stale rank (see ``_reconcile_stale_output_shapes``). + + Args: + graph: The ONNX graph to clear types and shapes for. + clear_shapes: If True, also clear shapes (False keeps shapes for type-only inference). + is_subgraph: Whether this is a subgraph (True) or the main graph (False). + """ + + def _clear(value_info: onnx.ValueInfoProto, clear_shape: bool) -> None: + value_info.type.tensor_type.elem_type = onnx.TensorProto.UNDEFINED + if clear_shape: + for dim in value_info.type.tensor_type.shape.dim: + if dim.dim_value: + dim.dim_param = "unk" + + def _clear_callback(g: onnx.GraphProto, parent: onnx.NodeProto, is_sub: bool) -> None: + logger.debug(f"Clearing types/shapes in {'subgraph' if is_sub else 'main graph'}: {g.name}") + + if is_sub: + # Subgraph inputs are cleared so they propagate from the parent graph. + for inp in g.input: + if inp.type.HasField("tensor_type"): + _clear(inp, clear_shapes) + # Only clear value_info for intermediates produced within this subgraph. + subgraph_outputs = {out for node in g.node for out in node.output} + for vi in g.value_info: + if vi.name in subgraph_outputs: + _clear(vi, clear_shapes) + else: + for vi in g.value_info: + _clear(vi, clear_shape=True) + + for out in g.output: + _clear(out, clear_shapes) + + walk_subgraphs_recursive(graph, _clear_callback, is_subgraph=is_subgraph) + + def get_op_types_not_supported_in_low_precision( model: onnx.ModelProto, min_opset: int, diff --git a/modelopt/onnx/utils.py b/modelopt/onnx/utils.py index 64b1fd1c414..3b8edf76a84 100644 --- a/modelopt/onnx/utils.py +++ b/modelopt/onnx/utils.py @@ -1205,19 +1205,32 @@ def infer_types( When use_standalone_type_inference is True, uses a standalone type inference implementation that only infers types. Otherwise, uses ONNX's infer_shapes which infers both types and shapes. + ONNX's ``infer_shapes`` can fail on weakly-typed models -- with ``strict_mode=True`` it raises + on an op it cannot resolve (e.g. a ``TopK`` whose axis it resolves to a stale dimension) + instead of silently leaving that node's outputs untyped. On any shape-inference failure this + falls back to the standalone type inferencer, which derives types from operator schemas + regardless of shapes, so downstream type lookups (e.g. in AutoCast) do not fail. Callers that + need a fully typed graph should pass ``strict_mode=True`` so incomplete inference surfaces as + an exception that triggers the fallback. + Args: model: ONNX model to infer types/shapes for. use_standalone_type_inference: If True, use standalone type inference (_infer_types_only). If False, use ONNX's shape inference (infer_shapes). - **kwargs: Additional arguments passed to infer_shapes when not using standalone type inference. + **kwargs: Additional arguments passed to infer_shapes when not using standalone type + inference (e.g. ``strict_mode``, ``check_type``, ``data_prop``). Returns: onnx.ModelProto: Model with inferred types (and shapes if not using standalone type inference). """ if use_standalone_type_inference: return _infer_types_only(model) - else: + + try: return infer_shapes(model, **kwargs) + except Exception as e: + logger.debug("ONNX shape inference failed (%s); using standalone type inference.", e) + return _infer_types_only(model) def onnx_type_str_to_enum(dtype: str) -> int: @@ -1862,21 +1875,121 @@ def change_casts_to_fp16(model: onnx.ModelProto, target_op_types: list[str]) -> return model +def _reconcile_stale_output_shapes(model: onnx.ModelProto) -> int: + """Re-derive stale ``graph.output`` shapes from the operator graph. + + Weakly-typed models (e.g. exported from TensorFlow) can declare an output rank + that conflicts with the graph topology -- most commonly a leftover rank-0 + (scalar) annotation on a tensor that is really rank-2+. Such a stale rank poisons + downstream shape inference: ORT fails while augmenting the model for INT8 + calibration (``axis must be in [-rank, rank-1]. Input rank was 0``), and + ``onnx.shape_inference`` with ``strict_mode=True`` raises ``Inferred shape and + existing shape differ in rank`` during fp16 autocast. + + Strategy: snapshot the declared output shapes, clear them, and re-derive them from + the operator graph -- preferring ORT's symbolic shape inference (it resolves ops + such as ``TopK`` that ONNX's static inference gives up on) and falling back to the + size-aware ``infer_shapes`` wrapper. A declared shape is only overwritten when it is + genuinely stale -- a rank mismatch (the rank-0-vs-rank-N bug) or a conflicting + concrete dimension. Outputs that merely differ in symbolic ``dim_param`` names (e.g. + a re-derived ``unk__0`` vs a declared ``batch``) keep their original declaration, so + healthy models -- including dynamic batch/sequence dims -- are left untouched. A + graph output is never left without a shape (``onnx.checker`` requires the field). + + Args: + model: Loaded in-memory onnx ModelProto, ideally with ``value_info`` already + cleared so re-inference derives shapes from the operator graph. + + Returns: + Number of graph outputs whose shape was changed. + """ + outputs = model.graph.output + if not outputs: + return 0 + + def _outputs_with_shapes(m: onnx.ModelProto) -> dict[str, onnx.TensorShapeProto]: + return { + o.name: o.type.tensor_type.shape + for o in m.graph.output + if o.type.tensor_type.HasField("shape") + } + + def _is_stale(declared: onnx.TensorShapeProto | None, inferred: onnx.TensorShapeProto | None): + # Only treat a declaration as stale when inference contradicts it: a different + # rank, or a concrete dim that disagrees with an inferred concrete dim. A missing + # declaration is "stale" (adopt whatever was inferred); a missing inference is not + # (keep the declaration). Symbolic dim_param renames are intentionally ignored. + if inferred is None: + return False + if declared is None: + return True + if len(declared.dim) != len(inferred.dim): + return True + return any( + d.HasField("dim_value") and i.HasField("dim_value") and d.dim_value != i.dim_value + for d, i in zip(declared.dim, inferred.dim) + ) + + # Snapshot declared shapes, then clear them so re-inference starts from the + # topology instead of being biased by the stale annotations. + declared: dict[str, onnx.TensorShapeProto | None] = {} + for o in outputs: + tt = o.type.tensor_type + if tt.HasField("shape"): + snapshot = onnx.TensorShapeProto() + snapshot.CopyFrom(tt.shape) + declared[o.name] = snapshot + else: + declared[o.name] = None + tt.ClearField("shape") + + # Re-derive output shapes from the cleared model (neither inference call mutates it): + # prefer ORT symbolic shape inference, then fall back to the size-aware infer_shapes + # wrapper if it is unavailable or yields nothing. + inferred: dict[str, onnx.TensorShapeProto] = {} + try: + from onnxruntime.tools.symbolic_shape_infer import SymbolicShapeInference + + inferred = _outputs_with_shapes(SymbolicShapeInference.infer_shapes(model, auto_merge=True)) + except Exception as e: + logger.debug("Symbolic shape inference unavailable/failed: %s", e) + if not inferred: + try: + inferred = _outputs_with_shapes(infer_shapes(model, strict_mode=False, data_prop=True)) + except Exception as e: + logger.debug("ONNX shape inference for output reconciliation failed: %s", e) + + changed = 0 + for o in outputs: + decl = declared[o.name] + inf = inferred.get(o.name) + # Adopt the inferred shape only when the declaration is genuinely stale; otherwise + # restore the declared shape (never leaving a graph output shapeless). + if _is_stale(decl, inf): + o.type.tensor_type.shape.CopyFrom(inf) + changed += 1 + elif decl is not None: + o.type.tensor_type.shape.CopyFrom(decl) + return changed + + def clear_stale_value_info(model: onnx.ModelProto) -> int: - """Clear stale type metadata that would otherwise trip ORT's type checker. + """Clear stale type/shape metadata that would otherwise trip ORT's type checker. - Walks every ``Cast`` node and forces the ``elem_type`` of any - ``graph.output`` entry produced by that Cast to match the Cast's ``to`` - attribute (the spec-defined contract for a Cast's output dtype). Then - clears ``value_info`` so ORT/shape-inference re-derives intermediate-tensor - types from the operator graph during session setup -- except entries for - outputs of ``trt.plugins`` custom-op nodes, whose types ORT cannot infer. + Walks every ``Cast`` node and forces the ``elem_type`` of any ``graph.output`` + entry produced by that Cast to match the Cast's ``to`` attribute (the spec-defined + contract for a Cast's output dtype). Clears ``value_info`` so ORT/shape-inference + re-derives intermediate-tensor types from the operator graph during session setup + -- except entries for outputs of ``trt.plugins`` custom-op nodes, whose types ORT + cannot infer. Finally, reconciles stale ``graph.output`` *shapes* (e.g. a leftover + rank-0 scalar on a tensor that is really rank-2+) which would otherwise propagate a + wrong rank into downstream shape inference. Args: model: Loaded in-memory onnx ModelProto. Returns: - Number of Cast outputs reconciled plus value_info entries cleared. + Total number of entries reconciled or cleared. """ cast_to_by_output = { node.output[0]: get_cast_to_type(node) @@ -1901,4 +2014,9 @@ def clear_stale_value_info(model: onnx.ModelProto) -> int: if n_cleared: del model.graph.value_info[:] model.graph.value_info.extend(preserved) - return fixed_outputs + n_cleared + + # Reconcile output shapes after value_info is cleared so the re-inference inside + # the helper derives shapes cleanly from the operator graph. + fixed_shapes = _reconcile_stale_output_shapes(model) + + return fixed_outputs + fixed_shapes + n_cleared diff --git a/modelopt/torch/export/quant_utils.py b/modelopt/torch/export/quant_utils.py index d8ddf442924..8d5828c2d1e 100755 --- a/modelopt/torch/export/quant_utils.py +++ b/modelopt/torch/export/quant_utils.py @@ -1349,6 +1349,39 @@ def preprocess_linear_fusion(modules: list[torch.nn.Module], resmooth_only=False module.weight_quantizer.amax = weight_amax +def _get_unquantized_moe_router_names(model: nn.Module) -> list[str]: + """Return the names of MoE router/gate submodules left in original precision. + + A module is added to ``exclude_modules`` during unified HF export only if it carries a + quantizer (even a disabled one) -- see :func:`get_quant_config`. MoE routers are kept + unquantized on purpose, but on ``transformers>=5.0`` they are no longer ``nn.Linear`` + modules (e.g. ``TopKRouter``), so ``mtq.quantize`` never attaches a quantizer to them. + Without this, the BF16 router weight is written to the checkpoint but omitted from + ``exclude_modules``, and deployment frameworks (vLLM / SGLang) then try to load it as a + quantized weight -- e.g. ``AssertionError: Tried to load weights of size [E, H] to a + parameter of size [E, H/2]`` for Qwen3-MoE. + + Routers are detected structurally: an MoE block exposes an ``experts`` container plus a + ``gate`` / ``router`` (or ``shared_expert_gate``) submodule that owns a weight tensor. + Routers that the user opted to quantize (a non-NONE format) are skipped. + """ + router_attrs = ("gate", "router", "shared_expert_gate") + router_names = [] + for name, module in model.named_modules(): + if not hasattr(module, "experts"): + continue + for attr in router_attrs: + router = getattr(module, attr, None) + if not isinstance(router, nn.Module): + continue + if not isinstance(getattr(router, "weight", None), torch.Tensor): + continue + if get_quantization_format(router) != QUANTIZATION_NONE: + continue + router_names.append(f"{name + '.' if name else ''}{attr}") + return router_names + + def get_quant_config( model: nn.Module, is_modelopt_qlora: bool = False, @@ -1457,6 +1490,14 @@ def get_quant_config( "Do not support mixed precision kv cache quantization" ) + # MoE routers/gates are intentionally kept in original precision. On transformers>=5.0 they + # are not nn.Linear modules (e.g. TopKRouter), never receive a quantizer, and would otherwise + # be missing from exclude_modules even though their BF16 weight is exported -- causing + # deployment frameworks to load them as quantized weights. Record them explicitly as + # unquantized so they land in exclude_modules. + for router_name in _get_unquantized_moe_router_names(model): + layer_config_dict.setdefault(router_name + ".quantization", QUANTIZATION_NONE) + # Process per layer quantization config dict quant_config["quantization"].update(process_layer_quant_config(layer_config_dict)) diff --git a/modelopt/torch/opt/plugins/mcore_dist_checkpointing.py b/modelopt/torch/opt/plugins/mcore_dist_checkpointing.py index d564023f7db..e3bc6e49d8e 100644 --- a/modelopt/torch/opt/plugins/mcore_dist_checkpointing.py +++ b/modelopt/torch/opt/plugins/mcore_dist_checkpointing.py @@ -23,8 +23,8 @@ import torch from megatron.core import dist_checkpointing, mpu -from megatron.core.dist_checkpointing.serialization import get_default_load_sharded_strategy from megatron.core.dist_checkpointing.strategies.common import COMMON_STATE_FNAME +from megatron.core.dist_checkpointing.strategies.torch import TorchDistLoadShardedStrategy from megatron.core.dist_checkpointing.validation import StrictHandling from megatron.core.transformer.module import Float16Module @@ -162,7 +162,7 @@ def _load_extra_state_from_sharded_checkpoint( extra_state_dict = dist_checkpointing.load( extra_sharded_state_dict, checkpoint_name, - get_default_load_sharded_strategy(checkpoint_name), + TorchDistLoadShardedStrategy(), strict=StrictHandling.LOG_UNEXPECTED, ) extra_state_dict_no_prefix = {} diff --git a/modelopt/torch/puzzletron/dataset/prepare_dataset.py b/modelopt/torch/puzzletron/dataset/prepare_dataset.py index 3d80062ae0f..74bb9622860 100644 --- a/modelopt/torch/puzzletron/dataset/prepare_dataset.py +++ b/modelopt/torch/puzzletron/dataset/prepare_dataset.py @@ -63,7 +63,7 @@ def process_and_save_dataset( ds_dict = datasets.DatasetDict( { "train": ds_split["train"], - "valid": ds_split["test"], + "validation": ds_split["test"], } ) # Save locally diff --git a/modelopt/torch/quantization/backends/fp8_per_tensor_gemm.py b/modelopt/torch/quantization/backends/fp8_per_tensor_gemm.py index d89ed35c6ca..b2b8f92b1fa 100644 --- a/modelopt/torch/quantization/backends/fp8_per_tensor_gemm.py +++ b/modelopt/torch/quantization/backends/fp8_per_tensor_gemm.py @@ -20,7 +20,6 @@ from modelopt.torch.quantization.backends.gemm_registry import gemm_registry from modelopt.torch.quantization.config import FP8_DEFAULT_CFG, find_quant_cfg_entry_by_path -from modelopt.torch.quantization.nn.modules.quant_linear import RealQuantLinear from modelopt.torch.quantization.qtensor import FP8QTensor, QTensorWrapper from modelopt.torch.quantization.utils import reduce_amax @@ -112,6 +111,9 @@ def _fp8_availability_check(module, input, args, kwargs): if not torch.cuda.is_available() or not fp8_compatible(): return False + # Import lazily because backend registration can run while quant_linear is initializing. + from modelopt.torch.quantization.nn.modules.quant_linear import RealQuantLinear + # Check module type if not isinstance(module, RealQuantLinear): return False diff --git a/modelopt/torch/quantization/backends/nvfp4_gemm.py b/modelopt/torch/quantization/backends/nvfp4_gemm.py index fdf6babb695..b62834c8935 100644 --- a/modelopt/torch/quantization/backends/nvfp4_gemm.py +++ b/modelopt/torch/quantization/backends/nvfp4_gemm.py @@ -21,7 +21,6 @@ import modelopt.torch.quantization as mtq from modelopt.torch.quantization.backends.gemm_registry import gemm_registry from modelopt.torch.quantization.backends.utils import fp4_compatible -from modelopt.torch.quantization.nn.modules.quant_linear import RealQuantLinear from modelopt.torch.quantization.qtensor import NVFP4QTensor, QTensorWrapper from modelopt.torch.quantization.utils import reduce_amax @@ -203,6 +202,9 @@ def _nvfp4_availability_check(module, input, args, kwargs): if not torch.cuda.is_available() or not fp4_compatible(): return False + # Import lazily because backend registration can run while quant_linear is initializing. + from modelopt.torch.quantization.nn.modules.quant_linear import RealQuantLinear + # Check module type if not isinstance(module, RealQuantLinear): return False diff --git a/modelopt/torch/quantization/utils/core_utils.py b/modelopt/torch/quantization/utils/core_utils.py index e95e175b71d..2faaaa4d91f 100644 --- a/modelopt/torch/quantization/utils/core_utils.py +++ b/modelopt/torch/quantization/utils/core_utils.py @@ -33,6 +33,10 @@ if TYPE_CHECKING: from collections.abc import Generator +# FP8 dtypes do not implement reduction kernels (e.g. ``max_all_cuda``), ``abs``, or +# elementwise ``maximum``, so tensors of these dtypes must be upcast before amax reduction. +_FP8_DTYPES = (torch.float8_e4m3fn, torch.float8_e5m2) + def reduce_block_amax(input_tensor: torch.Tensor, block_sizes: dict): """Computes the amax of the input tensor using block-based reduction for each dimension. @@ -157,6 +161,10 @@ def reduce_amax(input, axis=None, keepdims=True, squeeze_scalar=True): Returns: The reduced tensor. """ + # FP8 dtypes lack reduction/abs kernels (e.g. ``max_all_cuda``); upcast to the default + # float dtype, which represents every FP8 value exactly so the amax is computed losslessly. + if input.dtype in _FP8_DTYPES: + input = input.to(torch.get_default_dtype()) # A memory-efficient implementation that avoids copying input tensor if axis is None: max_val = torch.max(input) diff --git a/modelopt_recipes/configs/ptq/units/default_disabled_quantizers.yaml b/modelopt_recipes/configs/ptq/units/default_disabled_quantizers.yaml index 87fd67300fa..5e48dc73b7e 100644 --- a/modelopt_recipes/configs/ptq/units/default_disabled_quantizers.yaml +++ b/modelopt_recipes/configs/ptq/units/default_disabled_quantizers.yaml @@ -43,10 +43,13 @@ # Multimodal vision branch: keep the vision encoder (SigLIP / ViT) and any # multimodal embedding projection in BF16 by default. Recipes that enable bare # `*weight_quantizer` / `*input_quantizer` or `*mlp*` wildcards otherwise also - # match the vision tower (`model.vision_tower.*`, `model.visual.*`) and the - # embedding projection (`model.embed_vision.*`); quantizing the vision branch - # crashes export / produces garbage image embeddings on VL models (gemma-4, - # Qwen3.5-VL — NVBugs 6293731, 6293762, 6294017). A recipe that intentionally + # match the vision tower (`model.vision_tower.*`, `model.visual.*`, + # Llama-4 `vision_model.*`) and the embedding projection (`model.embed_vision.*`, + # Llama-4 `multi_modal_projector.*`); quantizing the vision branch crashes + # export / produces garbage image embeddings on VL models (gemma-4, + # Qwen3.5-VL — NVBugs 6293731, 6293762, 6294017; Llama-4-Scout — NVBugs + # 6359097, where `vision_model.patch_embedding.linear` has in_features=588, + # not divisible by the NVFP4 block size). A recipe that intentionally # quantizes vision must re-enable these after importing this unit. - quantizer_name: '*embed_vision*' enable: false @@ -54,6 +57,10 @@ enable: false - quantizer_name: '*visual*' enable: false + - quantizer_name: '*vision_model*' + enable: false + - quantizer_name: '*multi_modal_projector*' + enable: false - parent_class: 'nn.BatchNorm1d' quantizer_name: '*' enable: false diff --git a/modelopt_recipes/general/ptq/nvfp4_mlp_only-novit-kv_fp8.yaml b/modelopt_recipes/general/ptq/nvfp4_mlp_only-novit-kv_fp8.yaml new file mode 100644 index 00000000000..7bbf21393f8 --- /dev/null +++ b/modelopt_recipes/general/ptq/nvfp4_mlp_only-novit-kv_fp8.yaml @@ -0,0 +1,75 @@ +# 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. + +# Composed PTQ recipe for MLP/MoE-only dynamic NVFP4 quantization with FP8 KV-cache +# quantization, with the multimodal vision tower EXCLUDED from quantization. +# +# Identical to `nvfp4_mlp_only-kv_fp8` except for the trailing `*visual*` / +# `*vision_tower*` disable rules. The bare `*mlp*` enable rules below match BOTH +# the language-model MLPs (`model.language_model.layers.*.mlp`) AND the vision +# tower's block MLPs (e.g. Kimi-K2.5 `vision_tower.encoder.blocks.*.mlp.*`, +# Qwen3.5-VL `model.visual.blocks.*.mlp`). Quantizing the ViT FFNs to NVFP4 is +# both quality-harmful (degenerate image embeddings) and can fail export: the +# Kimi-K2.5 MoonViT `vt_intermediate_size=4304` is not divisible by the NVFP4 +# packing constraint (2 x group_size = 32), so the compressed-tensors export +# raises `tensor column shape must be divisible by the given group_size 32`. +# The vision branch is small and quant-sensitive, so it is kept in BF16, +# mirroring `nvfp4_mlp_only_mse-kv_fp8_cast-novit` and NVIDIA's reference +# `nvidia/Kimi-K2.5-NVFP4` checkpoint (which excludes the vision tower and +# multimodal projector). + +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 only to MLP/MoE weight and input quantizers, plus FP8 KV-cache + quantization; uses max calibration. Vision tower (model.visual.* / vision_tower.*) + excluded for VL models. +quantize: + algorithm: max + quant_cfg: + - $import: base_disable_all + - quantizer_name: '*mlp*weight_quantizer' + cfg: + $import: nvfp4 + - quantizer_name: '*mlp*input_quantizer' + cfg: + $import: nvfp4 + - quantizer_name: '*block_sparse_moe*weight_quantizer' + cfg: + $import: nvfp4 + - quantizer_name: '*block_sparse_moe*input_quantizer' + cfg: + $import: nvfp4 + - quantizer_name: '*.experts.*weight_quantizer' + cfg: + $import: nvfp4 + - quantizer_name: '*.experts.*input_quantizer' + cfg: + $import: nvfp4 + - $import: kv_fp8 + - $import: default_disabled_quantizers + # VL vision-branch exclusion: keep the ViT (and its block MLPs) in BF16. + # Must come after the `*mlp*` enables above so the disable wins (later + # entries override earlier ones). + - quantizer_name: '*visual*' + enable: false + - quantizer_name: '*vision_tower*' + enable: false diff --git a/modelopt_recipes/ptq.md b/modelopt_recipes/ptq.md index b6c00d151ab..c9975ea7911 100644 --- a/modelopt_recipes/ptq.md +++ b/modelopt_recipes/ptq.md @@ -28,7 +28,7 @@ supported combinations. ### The shipped recipes
-All 18 general/ptq/ recipes (click to expand) +All 19 general/ptq/ recipes (click to expand) | Recipe | Model body | KV cache | Calibration | |--------|-----------|----------|-------------| @@ -39,6 +39,7 @@ supported combinations. | `nvfp4_default-kv_nvfp4_cast` | NVFP4 W4A4, all linears | NVFP4 (constant amax) | max | | `nvfp4_default-kv_none-gptq` | NVFP4 W4A4 (static W), all linears | none | GPTQ (layerwise) | | `nvfp4_mlp_only-kv_fp8` | NVFP4 W4A4, MLP + MoE experts | FP8 (calibrated) | max | +| `nvfp4_mlp_only-novit-kv_fp8` | NVFP4 W4A4, MLP + MoE experts (VL vision tower excluded) | FP8 (calibrated) | max | | `nvfp4_mlp_only-kv_fp8_cast` | NVFP4 W4A4, MLP + MoE experts | FP8 (constant amax) | max | | `nvfp4_mlp_only_mse-kv_fp8_cast` | NVFP4 W4A4, MLP + MoE experts | FP8 (constant amax) | MSE + FP8 sweep | | `nvfp4_experts_only-kv_fp8` | NVFP4 W4A4, MoE experts only | FP8 (calibrated) | max | diff --git a/pyproject.toml b/pyproject.toml index 4b633c6610a..5835c759b73 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -103,7 +103,7 @@ dev-lint = [ dev-docs = [ "autodoc_pydantic>=2.1.0", "sphinx~=8.1.0", - "sphinx-argparse>=0.5.2", + "sphinx-argparse>=0.5.2,<0.6.0", "sphinx-autobuild>=2024.10.3", "sphinx-copybutton>=0.5.2", "sphinx-inline-tabs>=2023.4.21", diff --git a/tests/examples/llm_ptq/test_example_utils.py b/tests/examples/llm_ptq/test_example_utils.py index 0ed392340ad..e9ebdaef117 100644 --- a/tests/examples/llm_ptq/test_example_utils.py +++ b/tests/examples/llm_ptq/test_example_utils.py @@ -19,8 +19,11 @@ """ import json +from contextlib import nullcontext from types import SimpleNamespace +from unittest.mock import patch +import pytest import torch from _test_utils.examples.llm_ptq_example_utils import example_utils from safetensors.torch import save_file @@ -194,3 +197,122 @@ def test_get_original_hf_quant_method_none_for_unquantized(): example_utils.get_original_hf_quant_method(SimpleNamespace(quantization_config=None)) is None ) + + +# ---------- _resolve_init_config --------------------------------------------- + + +def _remote_config(): + # Config whose class module lives under "transformers_modules" (remote code). + cls = type("_RemoteConfig", (), {"__module__": "transformers_modules.ckpt.config"}) + return cls() + + +def test_resolve_init_config_rederives_for_remote_config(): + builtin_cfg = SimpleNamespace() + with patch.object( + example_utils.AutoConfig, "from_pretrained", return_value=builtin_cfg + ) as mock: + out = example_utils._resolve_init_config( + _remote_config(), object, "/ckpt", {"trust_remote_code": True} + ) + assert out is builtin_cfg + mock.assert_called_once_with("/ckpt") # trust_remote_code stripped + + +def test_resolve_init_config_keeps_non_remote_config(): + cfg = SimpleNamespace() # module is "types", not remote + with patch.object(example_utils.AutoConfig, "from_pretrained") as mock: + assert example_utils._resolve_init_config(cfg, object, "/ckpt", {}) is cfg + mock.assert_not_called() + + +def test_resolve_init_config_falls_back_when_rederive_raises(): + cfg = _remote_config() + with patch.object(example_utils.AutoConfig, "from_pretrained", side_effect=ValueError()): + assert example_utils._resolve_init_config(cfg, object, "/ckpt", {}) is cfg + + +@pytest.mark.parametrize( + ( + "architecture", + "model_class_name", + "expected_config_dtype_kwarg", + "unexpected_config_dtype_kwarg", + ), + [ + ("DeciLMForCausalLM", "AutoModelForCausalLM", "torch_dtype", "dtype"), + ("LlamaForCausalLM", "LlamaForCausalLM", "dtype", "torch_dtype"), + ], +) +def test_get_model_uses_expected_dtype_kwarg( + monkeypatch, + architecture, + model_class_name, + expected_config_dtype_kwarg, + unexpected_config_dtype_kwarg, +): + calls = {} + hf_config = SimpleNamespace( + architectures=[architecture], + dtype=torch.float16, + model_type="llama", + torch_dtype=torch.bfloat16, + ) + + class FakeModel: + def eval(self): + calls["eval"] = True + + class FakeAutoModelForCausalLM: + @staticmethod + def from_config(config, **kwargs): + calls["from_config"] = kwargs + assert config is hf_config + assert kwargs[expected_config_dtype_kwarg] is torch.float16 + assert unexpected_config_dtype_kwarg not in kwargs + assert "max_memory" not in kwargs + return FakeModel() + + @staticmethod + def from_pretrained(*args, **kwargs): + calls["from_pretrained"] = kwargs + assert "dtype" not in kwargs + assert kwargs["torch_dtype"] is torch.float16 + return FakeModel() + + class FakeLlamaForCausalLM(FakeAutoModelForCausalLM): + _from_config = FakeAutoModelForCausalLM.from_config + + @staticmethod + def from_pretrained(*args, **kwargs): + calls["from_pretrained"] = kwargs + assert kwargs["dtype"] == "auto" + assert "torch_dtype" not in kwargs + return FakeModel() + + monkeypatch.setattr( + example_utils.AutoConfig, + "from_pretrained", + lambda *args, **kwargs: hf_config, + ) + if model_class_name == "AutoModelForCausalLM": + monkeypatch.setattr(example_utils, "AutoModelForCausalLM", FakeAutoModelForCausalLM) + monkeypatch.delattr(example_utils.transformers, architecture, raising=False) + else: + monkeypatch.setattr(example_utils.transformers, model_class_name, FakeLlamaForCausalLM) + monkeypatch.setattr(example_utils, "is_nemotron_vl", lambda config: False) + monkeypatch.setattr(example_utils, "is_speculative", lambda config: False) + monkeypatch.setattr(example_utils, "init_empty_weights", lambda include_buffers: nullcontext()) + monkeypatch.setattr(example_utils, "get_max_memory", lambda: {0: 1024}) + monkeypatch.setattr(example_utils, "infer_auto_device_map", lambda model, max_memory: {"": 0}) + + model = example_utils.get_model("checkpoint", device="cpu", trust_remote_code=True) + + assert isinstance(model, FakeModel) + assert calls["eval"] + if expected_config_dtype_kwarg == "torch_dtype": + assert calls["from_config"]["trust_remote_code"] is True + else: + assert "trust_remote_code" not in calls["from_config"] + assert calls["from_pretrained"]["trust_remote_code"] is True diff --git a/tests/unit/onnx/autocast/test_autocast.py b/tests/unit/onnx/autocast/test_autocast.py index f761e1e9e3c..ccad0c7201b 100644 --- a/tests/unit/onnx/autocast/test_autocast.py +++ b/tests/unit/onnx/autocast/test_autocast.py @@ -26,6 +26,7 @@ import modelopt.onnx.utils as onnx_utils 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.logging_config import configure_logging configure_logging("DEBUG") @@ -321,3 +322,36 @@ def test_opset_parser_argument(): # Test parsing without opset (should be None) args = parser.parse_args(["--onnx_path", "test.onnx"]) assert args.opset is None + + +@pytest.fixture +def weakly_typed_topk_model(): + # TopK k (5) exceeds the static axis size (3), so ONNX shape inference cannot resolve it. + x = onnx.helper.make_tensor_value_info("X", onnx.TensorProto.FLOAT, [1, 3]) + out = onnx.helper.make_tensor_value_info("out", onnx.TensorProto.FLOAT, [1, 5]) + weight = onnx.numpy_helper.from_array(np.ones((1, 3), dtype=np.float32), name="weight") + k = onnx.numpy_helper.from_array(np.array([5], dtype=np.int64), name="k") + nodes = [ + onnx.helper.make_node("Add", ["X", "weight"], ["a"], name="add"), + onnx.helper.make_node("TopK", ["a", "k"], ["vals", "inds"], axis=1, name="topk"), + onnx.helper.make_node("Cast", ["inds"], ["out"], to=onnx.TensorProto.FLOAT, name="cast"), + ] + graph = onnx.helper.make_graph(nodes, "weakly_typed_topk", [x], [out], [weight, k]) + model = onnx.helper.make_model(graph, opset_imports=[onnx.helper.make_opsetid("", 17)]) + model.ir_version = 10 + return model + + +def test_convert_to_f16_falls_back_on_unresolvable_op(weakly_typed_topk_model): + """A weakly-typed graph ONNX shape inference cannot resolve must still convert. + + The TopK ``k`` (5) exceeds the static axis size (3), so ONNX shape inference raises + in strict mode (the same failure class as NVBug 6058907). ``convert_to_f16`` -- the + path used by INT8 + ``--high_precision_dtype fp16`` quantization -- runs infer_types + in strict mode and must fall back to standalone type inference instead of crashing, + typing the TopK's int64 indices output that feeds the downstream Cast. + """ + converted_model = convert_to_f16(weakly_typed_topk_model, keep_io_types=True) + + onnx.checker.check_model(converted_model) + assert any(n.op_type == "TopK" for n in converted_model.graph.node) diff --git a/tests/unit/onnx/test_onnx_utils.py b/tests/unit/onnx/test_onnx_utils.py index ec1c5ef75a5..36face35b90 100644 --- a/tests/unit/onnx/test_onnx_utils.py +++ b/tests/unit/onnx/test_onnx_utils.py @@ -33,6 +33,7 @@ clear_stale_value_info, get_input_names_from_bytes, get_output_names_from_bytes, + infer_types, randomize_weights_onnx_bytes, remove_node_training_mode, remove_weights_data, @@ -364,3 +365,94 @@ def test_clear_stale_value_info(output_elem_type, with_value_info, expected_coun assert model.graph.output[0].type.tensor_type.elem_type == onnx.TensorProto.FLOAT assert len(model.graph.value_info) == 0 assert count == expected_count + + +def _make_matmul_model(output_shape): + """Build an X[3,4] @ W[4,5] -> Y model with Y declared using ``output_shape``.""" + weights = make_tensor("W", onnx.TensorProto.FLOAT, [4, 5], np.zeros(20, dtype=np.float32)) + nodes = [make_node("MatMul", ["X", "W"], ["Y"], name="matmul")] + inputs = [make_tensor_value_info("X", onnx.TensorProto.FLOAT, [3, 4])] + outputs = [make_tensor_value_info("Y", onnx.TensorProto.FLOAT, output_shape)] + graph = make_graph(nodes, "matmul_graph", inputs, outputs, initializer=[weights]) + return make_model(graph, producer_name="modelopt test", opset_imports=[make_opsetid("", 17)]) + + +def test_clear_stale_value_info_reconciles_stale_rank0_output(): + # Y is really rank-2 [3, 5] but the model declares it as a rank-0 scalar (stale + # metadata typical of weakly-typed exports). This is the rank-(N)-vs-(0) class of + # conflict that crashes downstream shape inference (NVBug 6058907). + model = _make_matmul_model(output_shape=[]) + assert len(model.graph.output[0].type.tensor_type.shape.dim) == 0 # stale rank-0 + + clear_stale_value_info(model) + + out_type = model.graph.output[0].type.tensor_type + assert out_type.HasField("shape") # shape field must remain (onnx.checker requires it) + assert [d.dim_value for d in out_type.shape.dim] == [3, 5] # reconciled to the real shape + onnx.checker.check_model(model) + + +def test_clear_stale_value_info_preserves_valid_output_shape(): + # A correct output shape must be left untouched (no-op for healthy models). + model = _make_matmul_model(output_shape=[3, 5]) + + clear_stale_value_info(model) + + out_type = model.graph.output[0].type.tensor_type + assert [d.dim_value for d in out_type.shape.dim] == [3, 5] + + +def _make_dynamic_dim_model(): + """Build an X[batch,4] -> Relu -> Y[my_batch,4] model (output declares a different dim_param).""" + nodes = [make_node("Relu", ["X"], ["Y"], name="relu")] + inputs = [make_tensor_value_info("X", onnx.TensorProto.FLOAT, ["batch", 4])] + outputs = [make_tensor_value_info("Y", onnx.TensorProto.FLOAT, ["my_batch", 4])] + graph = make_graph(nodes, "dyn_graph", inputs, outputs) + return make_model(graph, producer_name="modelopt test", opset_imports=[make_opsetid("", 17)]) + + +def test_clear_stale_value_info_preserves_dynamic_dim_names(): + # A healthy output with a named dynamic dim must not be rewritten just because + # symbolic shape inference re-derives a different dim_param. Y is declared with + # "my_batch" while the graph would infer "batch" from the input: same rank, no + # concrete-dim conflict, so the declaration (incl. its dim_param) must be preserved. + model = _make_dynamic_dim_model() + + clear_stale_value_info(model) + + out_dims = model.graph.output[0].type.tensor_type.shape.dim + assert [d.dim_param or d.dim_value for d in out_dims] == ["my_batch", 4] + onnx.checker.check_model(model) + + +def _make_topk_overflow_model(): + """Build a model whose TopK ``k`` (5) exceeds the static axis dim (3). + + ONNX shape inference raises "Axis has less than the requested k elements" on this + model (the same failure class seen in NVBug 6058907), while standalone type + inference can still derive the output types (values float, indices int64). + """ + k = make_tensor("k", onnx.TensorProto.INT64, [1], [5]) + nodes = [ + make_node("TopK", ["X", "k"], ["vals", "inds"], axis=1, name="topk"), + make_node("Cast", ["inds"], ["out"], to=onnx.TensorProto.FLOAT, name="cast_inds"), + ] + inputs = [make_tensor_value_info("X", onnx.TensorProto.FLOAT, [1, 3])] + outputs = [make_tensor_value_info("out", onnx.TensorProto.FLOAT, [1, 5])] + graph = make_graph(nodes, "topk_overflow", inputs, outputs, initializer=[k]) + return make_model(graph, producer_name="modelopt test", opset_imports=[make_opsetid("", 17)]) + + +def test_infer_types_falls_back_to_standalone_when_onnx_fails(): + # ONNX shape inference cannot resolve this model's TopK. With strict_mode=True it raises + # (instead of silently leaving the TopK outputs untyped), so infer_types catches the + # error and falls back to standalone type inference, which still types every tensor. + model = _make_topk_overflow_model() + + inferred = infer_types(model, strict_mode=True) + + value_info_types = {vi.name: vi.type.tensor_type.elem_type for vi in inferred.graph.value_info} + output_types = {o.name: o.type.tensor_type.elem_type for o in inferred.graph.output} + assert value_info_types.get("vals") == onnx.TensorProto.FLOAT + assert value_info_types.get("inds") == onnx.TensorProto.INT64 # TopK indices + assert output_types.get("out") == onnx.TensorProto.FLOAT diff --git a/tests/unit/recipe/test_loader.py b/tests/unit/recipe/test_loader.py index b1325e3a69e..18f7648bc59 100644 --- a/tests/unit/recipe/test_loader.py +++ b/tests/unit/recipe/test_loader.py @@ -164,6 +164,7 @@ def test_load_recipe_builtin_description(): "general/ptq/nvfp4_experts_only-kv_fp8_cast", "general/ptq/nvfp4_experts_only-kv_fp8_layerwise", "general/ptq/nvfp4_mlp_only-kv_fp8", + "general/ptq/nvfp4_mlp_only-novit-kv_fp8", "general/ptq/nvfp4_mlp_only-kv_fp8_cast", "general/ptq/nvfp4_omlp_only-kv_fp8", "general/ptq/nvfp4_omlp_only-kv_fp8_cast", @@ -197,6 +198,17 @@ def test_nvfp4_weight_only_recipe_disables_vllm_marlin_incompatible_projections( } <= disabled_quantizers +def test_nvfp4_mlp_only_novit_recipe_disables_vision_quantizers(): + recipe = load_recipe("general/ptq/nvfp4_mlp_only-novit-kv_fp8") + disabled_quantizers = { + entry["quantizer_name"] + for entry in recipe.quantize.model_dump()["quant_cfg"] + if entry.get("enable") is False + } + + assert {"*visual*", "*vision_tower*"} <= disabled_quantizers + + # --------------------------------------------------------------------------- # load_recipe — error cases # --------------------------------------------------------------------------- diff --git a/tests/unit/torch/export/test_get_quantization.py b/tests/unit/torch/export/test_get_quantization.py index 1199f4c7cf0..e7ca68d0b69 100644 --- a/tests/unit/torch/export/test_get_quantization.py +++ b/tests/unit/torch/export/test_get_quantization.py @@ -13,6 +13,8 @@ # See the License for the specific language governing permissions and # limitations under the License. +import fnmatch + import pytest import torch from _test_utils.torch.export.utils import ( @@ -60,3 +62,108 @@ def test_nvfp4_static_quantizer_export(): quant_config = get_quant_config(model) assert quant_config["quantization"]["quant_algo"] == "NVFP4" assert quant_config["quantization"]["group_size"] == 16 + + +class _FakeTopKRouter(torch.nn.Module): + """Mimics a transformers>=5.0 MoE router: owns a ``weight`` but is NOT an ``nn.Linear``. + + ``mtq.quantize`` only attaches quantizers to registered modules (e.g. ``nn.Linear``), so a + router like this never receives one -- reproducing the condition behind NVBug 5718750. + """ + + def __init__(self, hidden: int, num_experts: int): + super().__init__() + self.weight = torch.nn.Parameter(torch.randn(num_experts, hidden)) + self.top_k = 2 + self.num_experts = num_experts + + def forward(self, x): + return torch.nn.functional.linear(x, self.weight) + + +class _FakeMoEBlock(torch.nn.Module): + def __init__(self, hidden: int = 16, num_experts: int = 4): + super().__init__() + self.gate = _FakeTopKRouter(hidden, num_experts) + self.experts = torch.nn.ModuleList( + torch.nn.Linear(hidden, hidden, bias=False) for _ in range(num_experts) + ) + + def forward(self, x): + self.gate(x) # exercise the router so it is reachable + out = x + for expert in self.experts: + out = expert(out) + return out + + +class _FakeMoEModel(torch.nn.Module): + def __init__(self, hidden: int = 16, num_experts: int = 4): + super().__init__() + self.block = _FakeMoEBlock(hidden, num_experts) + + def forward(self, x): + return self.block(x) + + +_nvfp4_all_linears_config = { + "quant_cfg": [ + {"quantizer_name": "*", "enable": False}, + { + "quantizer_name": "*weight_quantizer", + "cfg": { + "num_bits": (2, 1), + "block_sizes": {-1: 16, "type": "dynamic", "scale_bits": (4, 3)}, + "axis": None, + }, + "enable": True, + }, + { + "quantizer_name": "*input_quantizer", + "cfg": { + "num_bits": (2, 1), + "block_sizes": {-1: 16, "type": "dynamic", "scale_bits": (4, 3)}, + "axis": None, + }, + "enable": True, + }, + ], + "algorithm": "max", +} + + +def test_moe_router_excluded_when_not_quantized(): + """NVBug 5718750: a non-Linear MoE router (transformers>=5.0 TopKRouter) gets no quantizer. + + Its BF16 weight is still exported, so it must be listed in ``exclude_modules``; otherwise + deployment frameworks treat it as a quantized weight and fail to load the checkpoint. + """ + hidden = 16 + model = _FakeMoEModel(hidden=hidden) + mtq.quantize(model, _nvfp4_all_linears_config, lambda m: m(torch.randn(2, hidden))) + + # The router is not an nn.Linear, so quantize attached no quantizer to it. + assert not hasattr(model.block.gate, "weight_quantizer") + # The experts are quantized to NVFP4. + assert get_quantization_format(model.block.experts[0]) == QUANTIZATION_NVFP4 + + quant_config = get_quant_config(model) + assert quant_config["quantization"]["quant_algo"] == "NVFP4" + + exclude_modules = quant_config["quantization"]["exclude_modules"] + assert any(fnmatch.fnmatch("block.gate", pattern) for pattern in exclude_modules), ( + f"MoE router 'block.gate' missing from exclude_modules: {exclude_modules}" + ) + # The quantized experts must NOT be excluded. + assert not any(fnmatch.fnmatch("block.experts.0", pattern) for pattern in exclude_modules), ( + f"Quantized expert wrongly excluded: {exclude_modules}" + ) + + +def test_moe_router_names_handle_root_module(): + """When the MoE block itself is the root module, router names have no leading dot.""" + from modelopt.torch.export.quant_utils import _get_unquantized_moe_router_names + + block = _FakeMoEBlock(hidden=16) + # name == "" for the root module; the router must be "gate", not ".gate". + assert _get_unquantized_moe_router_names(block) == ["gate"] diff --git a/tests/unit/torch/quantization/test_utils.py b/tests/unit/torch/quantization/test_utils.py index 73d3423ba55..933c553ce12 100644 --- a/tests/unit/torch/quantization/test_utils.py +++ b/tests/unit/torch/quantization/test_utils.py @@ -18,6 +18,7 @@ from modelopt.torch.quantization.utils import ( convert_quantization_axis_to_reduce_axis, + reduce_amax, reduce_block_amax, ) from modelopt.torch.quantization.utils.layerwise_calib import LayerActivationCollector @@ -58,6 +59,25 @@ def test_reduce_block_amax(block_sizes, test_input, expected_scales): torch.allclose(scales, expected_scales) +@pytest.mark.parametrize("fp8_dtype", [torch.float8_e4m3fn, torch.float8_e5m2]) +@pytest.mark.parametrize("axis", [None, 0, 1, (0, 1)]) +def test_reduce_amax_fp8(fp8_dtype, axis): + """FP8 tensors have no reduction/abs kernels; reduce_amax must upcast them. + + Regression test for ``NotImplementedError: "max_all_cuda" not implemented for + 'Float8_e4m3fn'`` when calibrating models with natively FP8 weights (e.g. DeepSeek-V3). + """ + # Values chosen to be exactly representable in both FP8 formats so the upcast is lossless. + ref = torch.tensor([[1.0, -3.0, 2.0], [0.5, -0.25, 4.0]]) + x_fp8 = ref.to(fp8_dtype) + + out = reduce_amax(x_fp8, axis=axis) + expected = reduce_amax(ref, axis=axis) + + assert out.dtype == torch.get_default_dtype() + assert torch.equal(out, expected) + + @pytest.mark.parametrize( ("shape", "quant_axis", "expected_reduce_axis"), [