diff --git a/.github/instructions/python-model-builder.instructions.md b/.github/instructions/python-model-builder.instructions.md index 8f9ade9e95..45ea329b01 100644 --- a/.github/instructions/python-model-builder.instructions.md +++ b/.github/instructions/python-model-builder.instructions.md @@ -17,3 +17,4 @@ Read both documents to understand the intended usage, supported models, design p 2. Ignore any CodeQL warnings about how an __init__ method calls an overridden method. These warnings are false positives and can be safely ignored. The warning message is: "this call to ABC in an initialization method is overwritten by XYZ". 3. Find ways to reduce code duplication by reusing existing functionality and implementing common patterns. 4. Discover ways to leverage the use of shared code in the base classes to avoid code duplication and improve maintainability. +5. For any new `extra_options` that are added, make sure that they are documented in both `README.md` and `builder.py`. In `README.md`, there should be a description of the option and its purpose. There should be a usage example thereafter showing how to use the option when calling the model builder from wheel or from source. In `builder.py`, there should be a description of the option and its purpose, its default value, and any possible values. Any constraints or limitations should also be documented. Make sure the documentation is consistent across both files. \ No newline at end of file diff --git a/src/python/py/models/README.md b/src/python/py/models/README.md index dfb7d3e959..8246560e59 100644 --- a/src/python/py/models/README.md +++ b/src/python/py/models/README.md @@ -13,18 +13,31 @@ This folder contains the model builder for quickly creating optimized and quanti - [Quantized PyTorch Model](#quantized-pytorch-model) - [GGUF Model](#gguf-model) - [Extra Options](#extra-options) + - [Number of Hidden Layers](#number-of-hidden-layers) + - [Filename](#filename) - [Config Only](#config-only) - [Hugging Face Authentication](#hugging-face-authentication) - [Hugging Face Remote Code](#hugging-face-remote-code) - [Exclude Embedding Layer](#exclude-embedding-layer) - [Exclude Language Modeling Head](#exclude-language-modeling-head) - - [Prune Language Modeling Head](@prune-language-modeling-head) + - [Prune Language Modeling Head](#prune-language-modeling-head) - [Include Last Hidden States Output](#include-last-hidden-states-output) - [Enable Shared Embeddings](#enable-shared-embeddings) + - [Enable CUDA Graph Capture](#enable-cuda-graph-capture) + - [Enable WebGPU Graph Capture](#enable-webgpu-graph-capture) - [Disable QKV Projections Fusion](#disable-qkv-projections-fusion) - - [Enable CUDA Graph](#enable-cuda-graph) - - [Use 8 Bits Quantization in QMoE](#use-8-bits-quantization-in-qmoe) - - [Use QDQ Pattern for Quantization](#use-qdq-pattern-for-quantization) + - [Quantization Options](#quantization-options) + - [Accuracy Level](#accuracy-level) + - [MatMul Block Size](#matmul-block-size) + - [QMoE Block Size](#qmoe-block-size) + - [Is Symmetric](#is-symmetric) + - [Op Types To Quantize](#op-types-to-quantize) + - [Nodes To Exclude](#nodes-to-exclude) + - [Algo Config](#algo-config) + - [Use QDQ Pattern for Quantization](#use-qdq-pattern-for-quantization) + - [Use 8 Bits Quantization in QMoE](#use-8-bits-quantization-in-qmoe) + - [FP32 I/O for WebGPU EP](#fp32-io-for-webgpu-ep) + - [BF16 I/O for CUDA EP](#bf16-io-for-cuda-ep) - [LoRA Models](#lora-models) - [Unit Testing Models](#unit-testing-models) - [Option 1: Use the model builder directly](#option-1-use-the-model-builder-directly) @@ -60,7 +73,7 @@ It is intended for supporting the latest, popular state-of-the-art models. For all available options, please use the `-h/--help` flag. -``` +```bash # From wheel: python -m onnxruntime_genai.models.builder --help @@ -72,7 +85,7 @@ python builder.py --help This scenario is where your PyTorch model is not downloaded locally (either in the default Hugging Face cache directory or in a local folder on disk). -``` +```bash # From wheel: python -m onnxruntime_genai.models.builder -m model_name -o path_to_output_folder -p precision -e execution_provider -c cache_dir_to_save_hf_files @@ -84,7 +97,7 @@ python builder.py -m model_name -o path_to_output_folder -p precision -e executi This scenario is where your PyTorch model is already downloaded locally (either in the default Hugging Face cache directory or in a local folder on disk). -``` +```bash # From wheel: python -m onnxruntime_genai.models.builder -m model_name -o path_to_output_folder -p precision -e execution_provider -c cache_dir_where_hf_files_are_saved @@ -96,7 +109,7 @@ python builder.py -m model_name -o path_to_output_folder -p precision -e executi This scenario is where your PyTorch model has been customized or finetuned for one of the currently supported model architectures and your model can be loaded in Hugging Face. -``` +```bash # From wheel: python -m onnxruntime_genai.models.builder -i path_to_local_folder_on_disk -o path_to_output_folder -p precision -e execution_provider -c cache_dir_to_store_temp_files @@ -108,7 +121,7 @@ python builder.py -i path_to_local_folder_on_disk -o path_to_output_folder -p pr This scenario is where your PyTorch model is one of the currently supported model architectures, has already been quantized to INT4 precision, and your model can be loaded in the Hugging Face style via [AutoGPTQ](https://github.com/AutoGPTQ/AutoGPTQ) or [AutoAWQ](https://github.com/casper-hansen/AutoAWQ). -``` +```bash # From wheel: python -m onnxruntime_genai.models.builder -i path_to_local_folder_on_disk -o path_to_output_folder -p int4 -e execution_provider -c cache_dir_to_store_temp_files @@ -120,7 +133,7 @@ python builder.py -i path_to_local_folder_on_disk -o path_to_output_folder -p in This scenario is where your float16/float32 GGUF model is already on disk. -``` +```bash # From wheel: python -m onnxruntime_genai.models.builder -m model_name -i path_to_gguf_file -o path_to_output_folder -p precision -e execution_provider -c cache_dir_for_hf_files @@ -132,7 +145,7 @@ python builder.py -m model_name -i path_to_gguf_file -o path_to_output_folder -p This scenario is for when you want to have control over some specific settings. The below example shows how you can pass key-value arguments to `--extra_options`. -``` +```bash # From wheel: python -m onnxruntime_genai.models.builder -m model_name -o path_to_output_folder -p precision -e execution_provider -c cache_dir_for_hf_files --extra_options filename=decoder.onnx @@ -142,11 +155,35 @@ python builder.py -m model_name -o path_to_output_folder -p precision -e executi To see all available options through `--extra_options`, please use the `help` commands in the `Full Usage` section above. +#### Number of Hidden Layers + +This scenario is for when you want to manually set the number of hidden layers that the model builder exports. + +```bash +# From wheel: +python -m onnxruntime_genai.models.builder -m model_name -o path_to_output_folder -p precision -e execution_provider --extra_options num_hidden_layers=4 + +# From source: +python builder.py -m model_name -o path_to_output_folder -p precision -e execution_provider --extra_options num_hidden_layers=4 +``` + +#### Filename + +This scenario is for when you want to use a custom ONNX filename instead of the default `model.onnx`. + +```bash +# From wheel: +python -m onnxruntime_genai.models.builder -m model_name -o path_to_output_folder -p precision -e execution_provider --extra_options filename=decoder.onnx + +# From source: +python builder.py -m model_name -o path_to_output_folder -p precision -e execution_provider --extra_options filename=decoder.onnx +``` + #### Config Only This scenario is for when you already have your optimized and/or quantized ONNX model and you need to create the config files to run with ONNX Runtime GenAI. -``` +```bash # From wheel: python -m onnxruntime_genai.models.builder -m model_name -o path_to_output_folder -p precision -e execution_provider -c cache_dir_for_hf_files --extra_options config_only=true @@ -160,7 +197,7 @@ Afterwards, please open the `genai_config.json` file in the output folder and mo This scenario is for when you need to disable the Hugging Face authentication or use a different authentication token than the one stored in [huggingface-cli login](https://huggingface.co/docs/huggingface_hub/main/en/guides/cli#huggingface-cli-login). -``` +```bash # From wheel: python -m onnxruntime_genai.models.builder -m model_name -o path_to_output_folder -p precision -e execution_provider -c cache_dir_for_hf_files --extra_options hf_token=false @@ -172,7 +209,7 @@ python builder.py -m model_name -o path_to_output_folder -p precision -e executi This scenario is for when you need to disable trusting remote code from a Hugging Face repo. -``` +```bash # From wheel: python -m onnxruntime_genai.models.builder -m model_name -o path_to_output_folder -p precision -e execution_provider -c cache_dir_for_hf_files --extra_options hf_remote=false @@ -184,7 +221,7 @@ python builder.py -m model_name -o path_to_output_folder -p precision -e executi This scenario is for when you want to exclude the embedding layer from your ONNX model. -``` +```bash # From wheel: python -m onnxruntime_genai.models.builder -i path_to_local_folder_on_disk -o path_to_output_folder -p precision -e execution_provider -c cache_dir_to_store_temp_files --extra_options exclude_embeds=true @@ -196,7 +233,7 @@ python builder.py -i path_to_local_folder_on_disk -o path_to_output_folder -p pr This scenario is for when you want to exclude the language modeling head from your ONNX model. -``` +```bash # From wheel: python -m onnxruntime_genai.models.builder -i path_to_local_folder_on_disk -o path_to_output_folder -p precision -e execution_provider -c cache_dir_to_store_temp_files --extra_options exclude_lm_head=true @@ -208,7 +245,7 @@ python builder.py -i path_to_local_folder_on_disk -o path_to_output_folder -p pr This scenario is for when you want to prune the language modeling head to only compute the last token's logits. -``` +```bash # From wheel: python -m onnxruntime_genai.models.builder -i path_to_local_folder_on_disk -o path_to_output_folder -p precision -e execution_provider -c cache_dir_to_store_temp_files --extra_options prune_lm_head=true @@ -220,7 +257,7 @@ python builder.py -i path_to_local_folder_on_disk -o path_to_output_folder -p pr This scenario is for when you want to include the last hidden states as an output to your ONNX model. -``` +```bash # From wheel: python -m onnxruntime_genai.models.builder -i path_to_local_folder_on_disk -o path_to_output_folder -p precision -e execution_provider -c cache_dir_to_store_temp_files --extra_options include_hidden_states=true @@ -234,8 +271,9 @@ Note that this is the same as outputting embeddings since the last hidden states This scenario is for when you want to enable weight sharing between the embedding layer and the language modeling head. This reduces model size and can improve memory efficiency, especially useful for models with tied embeddings (where `tie_word_embeddings=true` in config.json). Shared embeddings are automatically enabled if `tie_word_embeddings=true` in the model's config.json (can be overridden with `shared_embeddings=false`), but cannot be used with `exclude_embeds=true` or `exclude_lm_head=true`. -##### Option 1: INT4 (for RTN and K-Quant) -``` +##### Example 1: INT4 weights + INT4 embeddings (for RTN and K-Quant) + +```bash # From wheel: python -m onnxruntime_genai.models.builder -m model_name -o path_to_output_folder -p int4 -e cuda --extra_options shared_embeddings=true int4_algo_config=k_quant @@ -243,8 +281,9 @@ python -m onnxruntime_genai.models.builder -m model_name -o path_to_output_folde python builder.py -m model_name -o path_to_output_folder -p int4 -e cuda --extra_options shared_embeddings=true int4_algo_config=k_quant ``` -##### Option 2: INT4 + INT8 embeddings (for RTN Last and K-Quant Last) -``` +##### Example 2: INT4 weights + INT8 embeddings (for RTN Last and K-Quant Last) + +```bash # From wheel: python -m onnxruntime_genai.models.builder -m model_name -o path_to_output_folder -p int4 -e cuda --extra_options shared_embeddings=true int4_algo_config=k_quant_last @@ -252,8 +291,9 @@ python -m onnxruntime_genai.models.builder -m model_name -o path_to_output_folde python builder.py -m model_name -o path_to_output_folder -p int4 -e cuda --extra_options shared_embeddings=true int4_algo_config=k_quant_last ``` -##### Option 3: INT4 embeddings + FP16 embeddings -``` +##### Example 3: INT4 weights + FP16 embeddings + +```bash # From wheel: python -m onnxruntime_genai.models.builder -m model_name -o path_to_output_folder -p int4 -e cuda --extra_options shared_embeddings=true int4_algo_config=rtn int4_nodes_to_exclude=/lm_head/MatMul @@ -261,8 +301,9 @@ python -m onnxruntime_genai.models.builder -m model_name -o path_to_output_folde python builder.py -m model_name -o path_to_output_folder -p int4 -e cuda --extra_options shared_embeddings=true int4_algo_config=rtn int4_nodes_to_exclude=/lm_head/MatMul ``` -##### Option 4: FP16 embeddings -``` +##### Example 4: FP16 weights + FP16 embeddings + +```bash # From wheel: python -m onnxruntime_genai.models.builder -m model_name -o path_to_output_folder -p fp16 -e cuda --extra_options shared_embeddings=true @@ -270,11 +311,35 @@ python -m onnxruntime_genai.models.builder -m model_name -o path_to_output_folde python builder.py -m model_name -o path_to_output_folder -p fp16 -e cuda --extra_options shared_embeddings=true ``` -#### Disable QKV Projections Fusion +#### Enable CUDA Graph Capture -This scenario is for when you want to keep Q/K/V projections in the attention layer separate instead of fusing them into a single packed MatMul operation. +This scenario is for when you want to enable CUDA graph capture for your ONNX model. +```bash +# From wheel: +python -m onnxruntime_genai.models.builder -i path_to_local_folder_on_disk -o path_to_output_folder -p precision -e execution_provider -c cache_dir_to_store_temp_files --extra_options enable_cuda_graph=true + +# From source: +python builder.py -i path_to_local_folder_on_disk -o path_to_output_folder -p precision -e execution_provider -c cache_dir_to_store_temp_files --extra_options enable_cuda_graph=true ``` + +#### Enable WebGPU Graph Capture + +This scenario is for when you want to enable WebGPU graph capture for your ONNX model. + +```bash +# From wheel: +python -m onnxruntime_genai.models.builder -i path_to_local_folder_on_disk -o path_to_output_folder -p precision -e execution_provider -c cache_dir_to_store_temp_files --extra_options enable_webgpu_graph=true + +# From source: +python builder.py -i path_to_local_folder_on_disk -o path_to_output_folder -p precision -e execution_provider -c cache_dir_to_store_temp_files --extra_options enable_webgpu_graph=true +``` + +#### Disable QKV Projections Fusion + +This scenario is for when you want to keep Q/K/V projections in the attention layer separate instead of fusing them into a single packed MatMul operation. + +```bash # From wheel: python -m onnxruntime_genai.models.builder -i path_to_local_folder_on_disk -o path_to_output_folder -p precision -e execution_provider -c cache_dir_to_store_temp_files --extra_options disable_qkv_fusion=true @@ -282,35 +347,101 @@ python -m onnxruntime_genai.models.builder -i path_to_local_folder_on_disk -o pa python builder.py -i path_to_local_folder_on_disk -o path_to_output_folder -p precision -e execution_provider -c cache_dir_to_store_temp_files --extra_options disable_qkv_fusion=true ``` -#### Enable CUDA Graph +#### Quantization Options -This scenario is for when you want to enable CUDA graph for your ONNX model. +These options apply when exporting quantized models (for example `-p int4`). +##### Accuracy Level + +This scenario is for when you want to control the accuracy level used for MatMul activation handling. + +```bash +# From wheel: +python -m onnxruntime_genai.models.builder -m model_name -o path_to_output_folder -p int4 -e execution_provider --extra_options int4_accuracy_level=4 + +# From source: +python builder.py -m model_name -o path_to_output_folder -p int4 -e execution_provider --extra_options int4_accuracy_level=4 ``` + +##### MatMul Block Size + +This scenario is for when you want to set the block size for MatMul quantization. + +```bash # From wheel: -python -m onnxruntime_genai.models.builder -i path_to_local_folder_on_disk -o path_to_output_folder -p precision -e execution_provider -c cache_dir_to_store_temp_files --extra_options enable_cuda_graph=true +python -m onnxruntime_genai.models.builder -m model_name -o path_to_output_folder -p int4 -e execution_provider --extra_options int4_block_size=32 # From source: -python builder.py -i path_to_local_folder_on_disk -o path_to_output_folder -p precision -e execution_provider -c cache_dir_to_store_temp_files --extra_options enable_cuda_graph=true +python builder.py -m model_name -o path_to_output_folder -p int4 -e execution_provider --extra_options int4_block_size=32 ``` -#### Use 8 Bits Quantization in QMoE +##### QMoE Block Size -This scenario is for when you want to use 8-bit quantization for MoE layers. Default is using 4-bit quantization. +This scenario is for when you want to set the block size for QMoE expert weights. +```bash +# From wheel: +python -m onnxruntime_genai.models.builder -m model_name -o path_to_output_folder -p int4 -e execution_provider --extra_options qmoe_block_size=128 + +# From source: +python builder.py -m model_name -o path_to_output_folder -p int4 -e execution_provider --extra_options qmoe_block_size=128 ``` + +##### Is Symmetric + +This scenario is for when you want to choose symmetric (`int4`) or asymmetric (`uint4`) weight quantization. + +```bash # From wheel: -python -m onnxruntime_genai.models.builder -i path_to_local_folder_on_disk -o path_to_output_folder -p precision -e execution_provider -c cache_dir_to_store_temp_files --extra_options use_8bits_moe=true +python -m onnxruntime_genai.models.builder -m model_name -o path_to_output_folder -p int4 -e execution_provider --extra_options int4_is_symmetric=false # From source: -python builder.py -i path_to_local_folder_on_disk -o path_to_output_folder -p precision -e execution_provider -c cache_dir_to_store_temp_files --extra_options use_8bits_moe=true +python builder.py -m model_name -o path_to_output_folder -p int4 -e execution_provider --extra_options int4_is_symmetric=false ``` -#### Use QDQ Pattern for Quantization +##### Op Types To Quantize -This scenario is for when you want to use the QDQ pattern when quantizing the model to 4 bits. +This scenario is for when you want to target specific operator types for quantization. +```bash +# From wheel: +python -m onnxruntime_genai.models.builder -m model_name -o path_to_output_folder -p int4 -e execution_provider --extra_options int4_op_types_to_quantize=MatMul/Gather + +# From source: +python builder.py -m model_name -o path_to_output_folder -p int4 -e execution_provider --extra_options int4_op_types_to_quantize=MatMul/Gather ``` + +##### Nodes To Exclude + +This scenario is for when you want to skip quantizing specific nodes. + +```bash +# From wheel: +python -m onnxruntime_genai.models.builder -m model_name -o path_to_output_folder -p int4 -e execution_provider --extra_options int4_nodes_to_exclude=/lm_head/MatMul,/model/embed_tokens/Gather + +# From source: +python builder.py -m model_name -o path_to_output_folder -p int4 -e execution_provider --extra_options int4_nodes_to_exclude=/lm_head/MatMul,/model/embed_tokens/Gather +``` + +##### Algo Config + +This scenario is for when you want to select the quantization algorithm mode. + +```bash +# From wheel: +python -m onnxruntime_genai.models.builder -m model_name -o path_to_output_folder -p int4 -e execution_provider --extra_options int4_algo_config=default + +# From source: +python builder.py -m model_name -o path_to_output_folder -p int4 -e execution_provider --extra_options int4_algo_config=default +``` + +Supported values are: `default`, `rtn`, `rtn_last`, `k_quant`, `k_quant_mixed`, `k_quant_last`, `k_quant_linear`. + +##### Use QDQ Pattern for Quantization + +This scenario is for when you want to use the QDQ pattern when quantizing the model. + +```bash # From wheel: python -m onnxruntime_genai.models.builder -i path_to_local_folder_on_disk -o path_to_output_folder -p precision -e execution_provider -c cache_dir_to_store_temp_files --extra_options use_qdq=true @@ -318,11 +449,47 @@ python -m onnxruntime_genai.models.builder -i path_to_local_folder_on_disk -o pa python builder.py -i path_to_local_folder_on_disk -o path_to_output_folder -p precision -e execution_provider -c cache_dir_to_store_temp_files --extra_options use_qdq=true ``` +##### Use 8 Bits Quantization in QMoE + +This scenario is for when you want to use 8-bit quantization for MoE layers. Default is using 4-bit quantization. + +```bash +# From wheel: +python -m onnxruntime_genai.models.builder -i path_to_local_folder_on_disk -o path_to_output_folder -p precision -e execution_provider -c cache_dir_to_store_temp_files --extra_options use_8bits_moe=true + +# From source: +python builder.py -i path_to_local_folder_on_disk -o path_to_output_folder -p precision -e execution_provider -c cache_dir_to_store_temp_files --extra_options use_8bits_moe=true +``` + +#### FP32 I/O for WebGPU EP + +This scenario is for when you want to force FP32 model I/O for WebGPU (useful for GPUs without FP16 support on WebGPU, such as GTX 10xx). + +```bash +# From wheel: +python -m onnxruntime_genai.models.builder -i path_to_local_folder_on_disk -o path_to_output_folder -p int4 -e webgpu -c cache_dir_to_store_temp_files --extra_options use_webgpu_fp32=true + +# From source: +python builder.py -i path_to_local_folder_on_disk -o path_to_output_folder -p int4 -e webgpu -c cache_dir_to_store_temp_files --extra_options use_webgpu_fp32=true +``` + +#### BF16 I/O for CUDA EP + +This scenario is for when you want to use BF16 I/O precision in quantized ONNX models for CUDA / TRT-RTX. + +```bash +# From wheel: +python -m onnxruntime_genai.models.builder -i path_to_local_folder_on_disk -o path_to_output_folder -p int4 -e cuda -c cache_dir_to_store_temp_files --extra_options use_cuda_bf16=true + +# From source: +python builder.py -i path_to_local_folder_on_disk -o path_to_output_folder -p int4 -e cuda -c cache_dir_to_store_temp_files --extra_options use_cuda_bf16=true +``` + #### LoRA Models This scenario is where you have a finetuned model with LoRA adapters and your model can be loaded in the Hugging Face style via [PEFT](https://github.com/huggingface/peft). -``` +```bash # From wheel: python -m onnxruntime_genai.models.builder -i path_to_local_folder_on_disk -o path_to_output_folder -p fp16 -e execution_provider -c cache_dir_to_store_temp_files --extra_options adapter_path=path_to_adapter_files @@ -336,7 +503,7 @@ Base weights should be located in `path_to_local_folder_on_disk` and adapter wei This scenario is where your PyTorch model is already downloaded locally (either in the default Hugging Face cache directory or in a local folder on disk). If it is not already downloaded locally, here is an example of how you can download it. -``` +```py from transformers import AutoModelForCausalLM, AutoTokenizer model_name = "your_model_name" @@ -353,7 +520,7 @@ tokenizer.save_pretrained(cache_dir) This option is the simplest but it will download another copy of the PyTorch model onto disk to accommodate the change in the number of hidden layers. -``` +```bash # From wheel: python -m onnxruntime_genai.models.builder -m model_name -o path_to_output_folder -p precision -e execution_provider --extra_options num_hidden_layers=4 @@ -367,7 +534,7 @@ python builder.py -m model_name -o path_to_output_folder -p precision -e executi 2. Modify `num_hidden_layers` in `config.json` to your desired target (e.g. 4 layers). 3. Run the below command for the model builder. -``` +```bash # From wheel: python -m onnxruntime_genai.models.builder -m model_name -o path_to_output_folder -p precision -e execution_provider -c cache_dir_where_hf_files_are_saved diff --git a/src/python/py/models/builder.py b/src/python/py/models/builder.py index 04c1d30dff..b22a9bc67c 100644 --- a/src/python/py/models/builder.py +++ b/src/python/py/models/builder.py @@ -86,6 +86,9 @@ def check_extra_options(kv_pairs, execution_provider): else: raise ValueError(f"{key} must be false/False/0 or true/True/1.") + if "hf_token" in kv_pairs: + kv_pairs["hf_token"] = parse_hf_token(kv_pairs["hf_token"]) + if "int4_op_types_to_quantize" in kv_pairs: op_types_to_quantize = () for op_type in kv_pairs["int4_op_types_to_quantize"].split("/"): @@ -192,7 +195,7 @@ def create_model( # Load model config extra_kwargs = {} if os.path.isdir(input_path) else {"cache_dir": cache_dir} hf_name = input_path if os.path.isdir(input_path) else model_name - hf_token = parse_hf_token(extra_options.get("hf_token", "true")) + hf_token = extra_options.get("hf_token", True) hf_remote = extra_options.get("hf_remote", True) config = AutoConfig.from_pretrained(hf_name, token=hf_token, trust_remote_code=hf_remote, **extra_kwargs) @@ -427,17 +430,13 @@ def get_args(): Separate the node names with a ',' when passing them here (e.g. int4_nodes_to_exclude=/lm_head/MatMul,/model/embed_tokens/Gather) int4_algo_config = Method for int4 quantization. Default is 'default'. Currently supported options are: 'default', 'rtn', 'rtn_last', 'k_quant', 'k_quant_mixed', 'k_quant_last', 'k_quant_linear'. - default = algo_config passed to MatMulNBitsQuantizer is None. Quantizer uses default RTN algorithm. All MatMuls are quantized as int4.(different node naming conventions to `rtn`) + default = algo_config passed to MatMulNBitsQuantizer is None. Quantizer uses default RTN algorithm. All MatMuls are quantized as int4. Uses different node naming conventions to `rtn`. rtn = RTN algorithm for int4 quantization. rtn_last = RTN algorithm where only the last MatMul (/lm_head/MatMul) is quantized as int8. Other MatMuls are quantized as int4. k_quant = k_quant algorithm for int4 quantization. k_quant_mixed = k_quant algorithm with mixed precision (int4 + int8). k_quant_last = k_quant algorithm where only the last MatMul (/lm_head/MatMul) is quantized as int8. Other MatMuls are quantized as int4. k_quant_linear = k_quant algorithm with linear attention layer projections and MLPs promoted to int8 (for hybrid attention models like Qwen3.5). - shared_embeddings = Enable weight sharing between embedding and LM head layers. Default is false. - Use this option to share weights and reduce model size by eliminating duplicate weights. - For quantized models (INT4/UINT4): Shares quantized weights using GatherBlockQuantized. Only works with rtn and k_quant algorithms, and cannot be used if LM head is excluded. - For float models (FP16/FP32/BF16): Shares float weights using Gather. Works for pure FP models or INT4 models where LM head is excluded from quantization. num_hidden_layers = Manually specify the number of layers in your ONNX model. Used for unit testing purposes. filename = Filename for ONNX model (default is 'model.onnx'). @@ -457,28 +456,33 @@ def get_args(): exclude_lm_head = Remove language modeling head from your ONNX model. Use this option when you want to remove the language modeling head from within your ONNX model. Instead of `logits`, you will have `hidden_states` as the output to your ONNX model. + prune_lm_head = Prune the LM head to only compute last-token logits during prefill. Default is false. + Inserts Gather+Unsqueeze before the LM head so the MatMul input is [B,1,H] instead of [B,S,H], + eliminating ~(S-1)/S of the compute. Cannot be combined with exclude_lm_head. include_hidden_states = Include hidden states as output from your ONNX model. Use this option when you want to have the hidden states as an output from your ONNX model. In addition to `logits`, you will have `hidden_states` as an output to your ONNX model. + shared_embeddings = Enable weight sharing between embedding and LM head layers. Default is false. + Use this option to share weights and reduce model size by eliminating duplicate weights. + Shares quantized weights using GatherBlockQuantized and shares unquantized weights using Gather. enable_cuda_graph = Enable CUDA graph capture during inference. Default is false. If enabled, all nodes being placed on the CUDA EP is the prerequisite for the CUDA graph to be used correctly. It is not guaranteed that CUDA graph be enabled as it depends on the model and the graph structure. enable_webgpu_graph = Enable WebGPU graph capture during inference. Default is false. If enabled, the model structure will be optimized for WebGPU graph execution. This affects attention mask reformatting and position IDs handling. - use_8bits_moe = Use 8-bit quantization for MoE layers. Default is false. - If true, the QMoE op will use 8-bit quantization. If false, the QMoE op will use 4-bit quantization. use_qdq = Use the QDQ decomposition for ops. Use this option when you want to use quantize-dequantize ops. For example, you will have a quantized MatMul op instead of the MatMulNBits op. + use_8bits_moe = Use 8-bit quantization for MoE layers. Default is false. + If true, the QMoE op will use 8-bit quantization. If false, the QMoE op will use 4-bit quantization. + disable_qkv_fusion = Disable QKV fusion in the model. Default is false. + If true, the model will not fuse the Q, K, and V projections. Automatically assumed for certain EPs. use_webgpu_fp32 = Use FP32 I/O precision for WebGPU EP. Use this option to enable GPUs that do not support FP16 on WebGPU (e.g. GTX 10xx). use_cuda_bf16 = Use BF16 I/O precision in quantized ONNX models for CUDA EP. Use this option to create quantized ONNX models that use BF16 precision. adapter_path = Path to folder on disk containing the adapter files (adapter_config.json and adapter model weights). Use this option for LoRA models. - prune_lm_head = Prune the LM head to only compute last-token logits during prefill. Default is false. - Inserts Gather+Unsqueeze before the LM head so the MatMul input is [B,1,H] instead of [B,S,H], - eliminating ~(S-1)/S of the compute. Cannot be combined with exclude_lm_head. """), ) diff --git a/src/python/py/models/builders/base.py b/src/python/py/models/builders/base.py index de416c18d1..da22e0655a 100644 --- a/src/python/py/models/builders/base.py +++ b/src/python/py/models/builders/base.py @@ -33,23 +33,6 @@ ) -def parse_hf_token(hf_token): - """ - Returns the authentication token needed for Hugging Face. - Token is obtained either from the user or the environment. - """ - if hf_token.lower() in {"false", "0"}: - # Default is `None` for disabling authentication - return None - - if hf_token.lower() in {"true", "1"}: - # Return token in environment - return True - - # Return user-provided token as string - return hf_token - - class Model: def __init__(self, config, io_dtype, onnx_dtype, ep, cache_dir, extra_options): # Model attributes from config @@ -102,7 +85,7 @@ def __init__(self, config, io_dtype, onnx_dtype, ep, cache_dir, extra_options): self.cache_dir = cache_dir self.filename = extra_options.get("filename", "model.onnx") - self.hf_token = parse_hf_token(extra_options.get("hf_token", "true")) + self.hf_token = extra_options.get("hf_token", True) self.hf_remote = extra_options.get("hf_remote", True) self.extra_options = extra_options @@ -328,7 +311,7 @@ def __init__(self, config, io_dtype, onnx_dtype, ep, cache_dir, extra_options): self.make_lm_head_init(config) # Quantization-specific variables (INT4, INT8, etc.) - algo_config = self.make_algo_config(extra_options.get("int4_algo_config", "default")) + self.algo_config_name = extra_options.get("int4_algo_config", "default") self.matmul_block_size = int(extra_options.get("int4_block_size", 32)) self.qmoe_block_size = int(extra_options.get("qmoe_block_size", 128 if self.ep in {"cuda", "trt-rtx"} else 32)) self.quant_attrs = { @@ -338,7 +321,7 @@ def __init__(self, config, io_dtype, onnx_dtype, ep, cache_dir, extra_options): "is_symmetric": extra_options.get("int4_is_symmetric", True), "op_types_to_quantize": extra_options.get("int4_op_types_to_quantize", ("MatMul",)), "nodes_to_exclude": extra_options.get("int4_nodes_to_exclude", []), - "algo_config": algo_config, + "algo_config": self.make_algo_config(), "use_qdq": extra_options.get("use_qdq", False), } self.make_quant_init(config) @@ -519,17 +502,78 @@ def make_quant_init(self, config): ) def make_tied_embeddings_init(self, config): - # Determine if lm_head is unquantized. int4/8 can have options to int4_nodes_to_exclude. FP models are always unquantized. - self.unquantized_lm_head = "/lm_head/MatMul" in self.quant_attrs["nodes_to_exclude"] or self.onnx_dtype in {ir.DataType.FLOAT, ir.DataType.FLOAT16, ir.DataType.BFLOAT16} - self.shared_embeddings = self.extra_options.get("shared_embeddings", config.tie_word_embeddings if hasattr(config, "tie_word_embeddings") and config.tie_word_embeddings is not None else False) - self.int8_lm_head = self.extra_options.get("int4_algo_config", "default") in {"k_quant_mixed", "k_quant_last", "k_quant_linear", "rtn_last"} - - # shared_embeddings conflicts with exclude_embeds and exclude_lm_head - if self.exclude_embeds or self.exclude_lm_head: - self.shared_embeddings = False - elif self.shared_embeddings and not self.unquantized_lm_head: - # matmul_nbits_quantizer.py has a different naming for default quantization, so lm_head.MatMul.weight_Q{}G{} does not match. - self.shared_embeddings = self.int8_lm_head or self.extra_options.get("int4_algo_config", "default") in {"rtn", "k_quant"} + # Determine if tied embeddings is even possible on the graph + shared_embeddings = ( + self.extra_options.get("shared_embeddings", config.tie_word_embeddings if hasattr(config, "tie_word_embeddings") and config.tie_word_embeddings is not None else False) + and not self.exclude_embeds + and not self.exclude_lm_head + and not self.prune_lm_head + ) + + # Determine if embeddings and lm_head will be quantized or not + quantized_embeds = ( + self.onnx_dtype in {ir.DataType.INT4, ir.DataType.UINT4} + and "Gather" in self.quant_attrs["op_types_to_quantize"] + and "/model/embed_tokens/Gather" not in self.quant_attrs["nodes_to_exclude"] + ) + quantized_lm_head = ( + self.onnx_dtype in {ir.DataType.INT4, ir.DataType.UINT4} + and "MatMul" in self.quant_attrs["op_types_to_quantize"] + and "/lm_head/MatMul" not in self.quant_attrs["nodes_to_exclude"] + ) + + if shared_embeddings: + self.tied_quantized_embeddings = quantized_embeds and quantized_lm_head + self.tied_unquantized_embeddings = not quantized_embeds and not quantized_lm_head + else: + self.tied_quantized_embeddings = False + self.tied_unquantized_embeddings = False + + def make_tied_quantized_embedding_input_names(self): + # Quantized tied embeddings in make_embedding() consume lm_head weights using + # algorithm-specific naming. + # + # Reference for quantized input names that will be produced: + # +------------+-------------+--------------------------------------+--------------------------+-------------------------------+ + # | Symmetry | Algorithm | Weight Input Name Format | Scales Input Name Format | Zero-Points Input Name Format | + # +------------+-------------+--------------------------------------+--------------------------+-------------------------------+ + # | asymmetric | default | *.MatMul.weight | N/A | N/A | + # | symmetric | default | *.MatMul.weight_Q4 | *.MatMul.weight_scales | N/A | + # | asymmetric | rtn* | *.MatMul.weight_Q4G32 or _Q8G32 | *.MatMul.weight_scale | *.MatMul.weight_zp | + # | symmetric | rtn* | *.MatMul.weight_Q4G32 or _Q8G32 | *.MatMul.weight_scale | N/A | + # | asymmetric | k_quant* | *.MatMul.weight_Q4G32 or _Q8G32 | *.MatMul.weight_scale | *.MatMul.weight_zp | + # | symmetric | k_quant* | *.MatMul.weight_Q4G32 or _Q8G32 | *.MatMul.weight_scale | *.MatMul.weight_zp | + # +------------+-------------+--------------------------------------+--------------------------+-------------------------------+ + # where rtn* = rtn, rtn_last + # k_quant* = k_quant, k_quant_last, k_quant_linear, k_quant_mixed + + bits = 8 if self.algo_config_name in {"k_quant_mixed", "k_quant_last", "k_quant_linear", "rtn_last"} else 4 + is_symmetric = self.quant_attrs["is_symmetric"] + + if self.algo_config_name in {"rtn", "rtn_last"}: + return ( + bits, + f"lm_head.MatMul.weight_Q{bits}G{self.matmul_block_size}", + "lm_head.MatMul.weight_scale", + "lm_head.MatMul.weight_zp" if not is_symmetric else "", + ) + + if self.algo_config_name in {"k_quant", "k_quant_mixed", "k_quant_last", "k_quant_linear"}: + return ( + bits, + f"lm_head.MatMul.weight_Q{bits}G{self.matmul_block_size}", + "lm_head.MatMul.weight_scale", + "lm_head.MatMul.weight_zp", + ) + + # Fallback to default convention for unknown values. + assert self.algo_config_name == "default", "Unknown quantization algo config name detected" + return ( + bits, + f"lm_head.MatMul.weight_Q{bits}" if is_symmetric else "lm_head.MatMul.weight", + "lm_head.MatMul.weight_scales" if is_symmetric else "", + "", + ) def make_genai_config(self, model_name_or_path, extra_kwargs, out_dir): # Create config with attributes from config.json and generation_config.json (if latter file exists) @@ -690,20 +734,21 @@ def save_processing(self, model_name_or_path, extra_kwargs, out_dir): print(f"Saving processing files in {out_dir} for GenAI") tokenizer.save_pretrained(out_dir) - def make_algo_config(self, quant_method: str): + def make_algo_config(self): + # Create a quantization configuration based on the algorithm name. customized_weight_config = {} algo_config = None - if quant_method in {"rtn", "rtn_last"}: - if quant_method == "rtn_last": + if self.algo_config_name in {"rtn", "rtn_last"}: + if self.algo_config_name == "rtn_last": customized_weight_config["/lm_head/MatMul"] = {"bits": 8} algo_config = RTNWeightOnlyQuantConfig(customized_weight_config=customized_weight_config) - elif quant_method in {"k_quant", "k_quant_mixed", "k_quant_last", "k_quant_linear"}: - if quant_method != "k_quant": + elif self.algo_config_name in {"k_quant", "k_quant_mixed", "k_quant_last", "k_quant_linear"}: + if self.algo_config_name != "k_quant": customized_weight_config["/lm_head/MatMul"] = {"bits": 8} - if quant_method == "k_quant_mixed": + if self.algo_config_name == "k_quant_mixed": # k_quant_mixed is from llama.cpp. # Reference: https://github.com/ggml-org/llama.cpp/blob/36667c8edcded08063ed51c7d57e9e086bbfc903/src/llama-quant.cpp#L136 # We also consider some MatMuls are more senstive to quantization than other MatMuls. @@ -719,7 +764,7 @@ def make_algo_config(self, quant_method: str): customized_weight_config["/model/layers." + str(i) + "/attn/v_proj/MatMul"] = {"bits": 8} customized_weight_config["/model/layers." + str(i) + "/mlp/down_proj/MatMul"] = {"bits": 8} - if quant_method == "k_quant_linear" and hasattr(self, "layer_types"): + if self.algo_config_name == "k_quant_linear" and hasattr(self, "layer_types"): # Promote linear attention projections and their MLPs to INT8. # Linear attention recurrence accumulates quantization errors across # the full sequence (no softmax normalization). @@ -730,7 +775,6 @@ def make_algo_config(self, quant_method: str): for proj in ("gate_proj", "up_proj", "down_proj"): customized_weight_config[f"/model/layers.{i}/mlp/{proj}/MatMul"] = {"bits": 8} - customized_weight_config["/lm_head/MatMul"] = {"bits": 8} algo_config = KQuantWeightOnlyQuantConfig(customized_weight_config=customized_weight_config) return algo_config @@ -752,9 +796,8 @@ def to_int4(self) -> ir.Model: def save_model(self, out_dir): print(f"Saving ONNX model in {out_dir}") - already_quantized_in_qdq_format = ( - self.quant_type is not None and self.quant_attrs["use_qdq"] - ) # Skip quantizing `MatMul` in `DequantizeLinear --> Transpose --> MatMul` path + # Skip quantizing `MatMul` in `DequantizeLinear --> Transpose --> MatMul` path + already_quantized_in_qdq_format = self.quant_type is not None and self.quant_attrs["use_qdq"] if self.onnx_dtype in {ir.DataType.INT4, ir.DataType.UINT4} and not already_quantized_in_qdq_format: model = self.to_int4() else: @@ -1413,24 +1456,28 @@ def make_embedding(self, embedding): basename = "/model/embed_tokens" # Use GatherBlockQuantized if and only if tied embeddings are enabled and export model is quantized. quantized d_type in set_onnx_dtype is INT4/UINT4 - if self.shared_embeddings and self.onnx_dtype in {ir.DataType.INT4, ir.DataType.UINT4}: + if self.tied_quantized_embeddings: + bits, tied_weight_name, tied_weight_scale_name, tied_weight_zp_name = self.make_tied_quantized_embedding_input_names() + gather_name = f"{basename}/GatherBlockQuantized" gather_output = f"{gather_name}/output_0" weight_reshape_name = f"{basename}/Reshape" - bits = 8 if self.int8_lm_head else 4 flat_dim = self.hidden_size * bits // 8 weight_reshape_inputs = [ - f"lm_head.MatMul.weight_Q{bits}G{self.matmul_block_size}", + tied_weight_name, f"/model/constants/INT64/[{self.vocab_size}, {flat_dim}]", ] weight_reshape_output = f"{weight_reshape_name}/output_0" + # Quantized weight dtype is uint8. See here for more info: # https://github.com/microsoft/onnxruntime/blob/0c9356cb986fd4cd2c5d510909d31186010ba226/onnxruntime/python/tools/quantization/neural_compressor/weight_only.py#L73 self.make_reshape(weight_reshape_name, weight_reshape_inputs, dtype=ir.DataType.UINT8, shape=[self.vocab_size, flat_dim]) - input_names = [weight_reshape_output, self.input_names["input_ids"], "lm_head.MatMul.weight_scale"]; - if not self.quant_attrs["is_symmetric"]: - input_names.append("lm_head.MatMul.weight_zp") + input_names = [weight_reshape_output, self.input_names["input_ids"]] + if tied_weight_scale_name: + input_names.append(tied_weight_scale_name) + if tied_weight_zp_name: + input_names.append(tied_weight_zp_name) self.make_node( "GatherBlockQuantized", @@ -1445,7 +1492,7 @@ def make_embedding(self, embedding): ) # Use Transpose + Gather for tied embeddings for float embedding layers - elif self.shared_embeddings and self.unquantized_lm_head: + elif self.tied_unquantized_embeddings: transpose_name = f"{basename}/Transpose" transpose_output = f"{transpose_name}/output_0" self.make_transpose( diff --git a/src/python/py/models/quantized_model.py b/src/python/py/models/quantized_model.py index 3f5a5753c5..a3823175fb 100644 --- a/src/python/py/models/quantized_model.py +++ b/src/python/py/models/quantized_model.py @@ -1258,6 +1258,7 @@ def __init__(self, quant_type, input_path, quant_attrs, q_size, kv_size, interme for i, layer in enumerate(self.layers): if i >= self.num_layers: break + print(f"Unpacking and repacking layer {i}") # Unpack and repack all `QuantizedTensorModule` classes in attention self_attn = getattr(layer, "self_attn", None) or getattr(layer, "self_attention", None) diff --git a/test/python/builder/test_tied_embeddings.py b/test/python/builder/test_tied_embeddings.py new file mode 100644 index 0000000000..250ceb78d2 --- /dev/null +++ b/test/python/builder/test_tied_embeddings.py @@ -0,0 +1,493 @@ +from __future__ import annotations + +import importlib.util +import sys +import types +from pathlib import Path + +import onnx_ir as ir +import pytest + +BUILDERS_DIR = Path(__file__).parents[3] / "src" / "python" / "py" / "models" / "builders" +sys.path.insert(0, str(BUILDERS_DIR.parents[1])) + + +def _load_builder_module(module_name): + spec = importlib.util.spec_from_file_location(f"models.builders.{module_name}", BUILDERS_DIR / f"{module_name}.py") + module = importlib.util.module_from_spec(spec) + sys.modules[f"models.builders.{module_name}"] = module + spec.loader.exec_module(module) + return module + + +sys.modules.setdefault("models", types.ModuleType("models")) +builders_package = sys.modules.setdefault("models.builders", types.ModuleType("models.builders")) +builders_package.__path__ = [str(BUILDERS_DIR)] + +base_module = _load_builder_module("base") +Model = base_module.Model + + +def _make_model_for_tied_embeddings( + *, + shared_embeddings=None, + tie_word_embeddings=None, + onnx_dtype=ir.DataType.FLOAT16, + op_types=("MatMul", "Gather"), + nodes_to_exclude=(), + exclude_embeds=False, + exclude_lm_head=False, + prune_lm_head=False, + int4_algo_config="default", +): + model = Model.__new__(Model) + model.extra_options = {"int4_algo_config": int4_algo_config} + if shared_embeddings is not None: + model.extra_options["shared_embeddings"] = shared_embeddings + model.onnx_dtype = onnx_dtype + model.quant_attrs = { + "op_types_to_quantize": op_types, + "nodes_to_exclude": list(nodes_to_exclude), + } + model.exclude_embeds = exclude_embeds + model.exclude_lm_head = exclude_lm_head + model.prune_lm_head = prune_lm_head + + config = types.SimpleNamespace(tie_word_embeddings=tie_word_embeddings) + model.make_tied_embeddings_init(config) + return model + + +def test_shared_embeddings_option_overrides_config_tie_word_embeddings(): + model = _make_model_for_tied_embeddings( + shared_embeddings=False, + tie_word_embeddings=True, + onnx_dtype=ir.DataType.INT4, + ) + + assert model.tied_quantized_embeddings is False + assert model.tied_unquantized_embeddings is False + + +def test_tie_word_embeddings_defaults_to_false_when_unset_or_none(): + model_unset = _make_model_for_tied_embeddings() + model_none = _make_model_for_tied_embeddings(tie_word_embeddings=None) + + assert model_unset.tied_quantized_embeddings is False + assert model_unset.tied_unquantized_embeddings is False + assert model_none.tied_quantized_embeddings is False + assert model_none.tied_unquantized_embeddings is False + + +@pytest.mark.parametrize( + "exclude_embeds, exclude_lm_head, prune_lm_head", + [ + (True, False, False), + (False, True, False), + (False, False, True), + ], +) +def test_shared_embeddings_are_disabled_when_embeddings_or_lm_head_are_excluded( + exclude_embeds, + exclude_lm_head, + prune_lm_head, +): + model = _make_model_for_tied_embeddings( + shared_embeddings=True, + tie_word_embeddings=False, + onnx_dtype=ir.DataType.INT4, + exclude_embeds=exclude_embeds, + exclude_lm_head=exclude_lm_head, + prune_lm_head=prune_lm_head, + ) + + assert model.tied_quantized_embeddings is False + assert model.tied_unquantized_embeddings is False + + +@pytest.mark.parametrize( + "onnx_dtype, op_types, nodes_to_exclude, exclude_embeds, exclude_lm_head, prune_lm_head, int4_algo_config, expected_tied_quantized, expected_tied_unquantized", + [ + (ir.DataType.INT4, ("MatMul", "Gather"), (), False, False, False, "default", True, False), + (ir.DataType.INT4, ("MatMul", "Gather"), (), False, False, False, "rtn", True, False), + (ir.DataType.UINT4, ("MatMul", "Gather"), (), False, False, False, "default", True, False), + (ir.DataType.UINT4, ("MatMul", "Gather"), (), False, False, False, "k_quant", True, False), + (ir.DataType.FLOAT16, ("MatMul", "Gather"), (), False, False, False, "default", False, True), + (ir.DataType.INT4, ("MatMul",), (), False, False, False, "rtn", False, False), + (ir.DataType.INT4, ("Gather",), (), False, False, False, "rtn", False, False), + (ir.DataType.INT4, ("MatMul", "Gather"), ("/model/embed_tokens/Gather",), False, False, False, "rtn", False, False), + (ir.DataType.INT4, ("MatMul", "Gather"), ("/lm_head/MatMul",), False, False, False, "rtn", False, False), + (ir.DataType.INT4, ("MatMul", "Gather"), (), True, False, False, "rtn", False, False), + (ir.DataType.INT4, ("MatMul", "Gather"), (), False, True, False, "rtn", False, False), + (ir.DataType.INT4, ("MatMul", "Gather"), (), False, False, True, "rtn", False, False), + ], +) +def test_tied_embedding_path_selection_matches_current_base_logic( + onnx_dtype, + op_types, + nodes_to_exclude, + exclude_embeds, + exclude_lm_head, + prune_lm_head, + int4_algo_config, + expected_tied_quantized, + expected_tied_unquantized, +): + model = _make_model_for_tied_embeddings( + shared_embeddings=True, + tie_word_embeddings=False, + onnx_dtype=onnx_dtype, + op_types=op_types, + nodes_to_exclude=nodes_to_exclude, + exclude_embeds=exclude_embeds, + exclude_lm_head=exclude_lm_head, + prune_lm_head=prune_lm_head, + int4_algo_config=int4_algo_config, + ) + + assert model.tied_quantized_embeddings is expected_tied_quantized + assert model.tied_unquantized_embeddings is expected_tied_unquantized + + +@pytest.mark.parametrize( + "onnx_dtype, op_types, nodes_to_exclude", + [ + (ir.DataType.INT4, ("MatMul",), ("/lm_head/MatMul",)), + (ir.DataType.UINT4, ("MatMul",), ("/lm_head/MatMul",)), + (ir.DataType.INT4, ("MatMul", "Gather"), ("/lm_head/MatMul", "/model/embed_tokens/Gather")), + (ir.DataType.UINT4, ("MatMul", "Gather"), ("/lm_head/MatMul", "/model/embed_tokens/Gather")), + ], +) +def test_tied_unquantized_embeddings_can_be_true_in_int4_mode_when_both_quant_paths_are_disabled( + onnx_dtype, + op_types, + nodes_to_exclude, +): + model = _make_model_for_tied_embeddings( + shared_embeddings=True, + tie_word_embeddings=False, + onnx_dtype=onnx_dtype, + op_types=op_types, + nodes_to_exclude=nodes_to_exclude, + int4_algo_config="rtn", + ) + + assert model.tied_quantized_embeddings is False + assert model.tied_unquantized_embeddings is True + + +@pytest.mark.parametrize( + "quantized_embeds, quantized_lm_head, int4_algo_config, expected_tied_quantized, expected_tied_unquantized", + [ + (True, True, "default", True, False), + (True, True, "rtn", True, False), + (True, True, "rtn_last", True, False), + (True, True, "k_quant", True, False), + (True, True, "k_quant_last", True, False), + (True, True, "k_quant_mixed", True, False), + (True, True, "k_quant_linear", True, False), + (True, False, "rtn", False, False), + (False, True, "rtn", False, False), + (False, False, "rtn", False, True), + ], +) +def test_shared_embeddings_prefers_quantized_path_only_when_both_layers_are_quantized( + quantized_embeds, + quantized_lm_head, + int4_algo_config, + expected_tied_quantized, + expected_tied_unquantized, +): + op_types = tuple(op for enabled, op in ((quantized_embeds, "Gather"), (quantized_lm_head, "MatMul")) if enabled) + model = _make_model_for_tied_embeddings( + shared_embeddings=True, + tie_word_embeddings=False, + onnx_dtype=ir.DataType.INT4, + op_types=op_types, + int4_algo_config=int4_algo_config, + ) + + assert model.tied_quantized_embeddings is expected_tied_quantized + assert model.tied_unquantized_embeddings is expected_tied_unquantized + + +@pytest.mark.parametrize( + "int4_algo_config, matmul_block_size, is_symmetric, expected_bits, expected_weight, expected_scale, expected_zp", + [ + ("default", 32, True, 4, "lm_head.MatMul.weight_Q4", "lm_head.MatMul.weight_scales", ""), + ("default", 32, False, 4, "lm_head.MatMul.weight", "", ""), + ("rtn", 32, True, 4, "lm_head.MatMul.weight_Q4G32", "lm_head.MatMul.weight_scale", ""), + ("rtn", 32, False, 4, "lm_head.MatMul.weight_Q4G32", "lm_head.MatMul.weight_scale", "lm_head.MatMul.weight_zp"), + ("rtn_last", 32, True, 8, "lm_head.MatMul.weight_Q8G32", "lm_head.MatMul.weight_scale", ""), + ("rtn_last", 32, False, 8, "lm_head.MatMul.weight_Q8G32", "lm_head.MatMul.weight_scale", "lm_head.MatMul.weight_zp"), + ("k_quant", 32, True, 4, "lm_head.MatMul.weight_Q4G32", "lm_head.MatMul.weight_scale", "lm_head.MatMul.weight_zp"), + ("k_quant", 32, False, 4, "lm_head.MatMul.weight_Q4G32", "lm_head.MatMul.weight_scale", "lm_head.MatMul.weight_zp"), + ("k_quant_last", 32, True, 8, "lm_head.MatMul.weight_Q8G32", "lm_head.MatMul.weight_scale", "lm_head.MatMul.weight_zp"), + ("k_quant_last", 32, False, 8, "lm_head.MatMul.weight_Q8G32", "lm_head.MatMul.weight_scale", "lm_head.MatMul.weight_zp"), + ("k_quant_mixed", 32, True, 8, "lm_head.MatMul.weight_Q8G32", "lm_head.MatMul.weight_scale", "lm_head.MatMul.weight_zp"), + ("k_quant_mixed", 32, False, 8, "lm_head.MatMul.weight_Q8G32", "lm_head.MatMul.weight_scale", "lm_head.MatMul.weight_zp"), + ("k_quant_linear", 32, True, 8, "lm_head.MatMul.weight_Q8G32", "lm_head.MatMul.weight_scale", "lm_head.MatMul.weight_zp"), + ("k_quant_linear", 32, False, 8, "lm_head.MatMul.weight_Q8G32", "lm_head.MatMul.weight_scale", "lm_head.MatMul.weight_zp"), + ], +) +def test_tied_quantized_embedding_weight_names_cover_all_supported_algorithms( + int4_algo_config, + matmul_block_size, + is_symmetric, + expected_bits, + expected_weight, + expected_scale, + expected_zp, +): + model = Model.__new__(Model) + model.extra_options = {"int4_algo_config": int4_algo_config} + model.algo_config_name = int4_algo_config + model.matmul_block_size = matmul_block_size + model.quant_attrs = {"is_symmetric": is_symmetric} + + bits, weight_name, scale_name, zp_name = model.make_tied_quantized_embedding_input_names() + + assert bits == expected_bits + assert weight_name == expected_weight + assert scale_name == expected_scale + assert zp_name == expected_zp + + +def test_tied_quantized_embedding_weight_names_raise_for_unknown_algorithm(): + model = Model.__new__(Model) + model.extra_options = {"int4_algo_config": "unexpected"} + model.algo_config_name = "unexpected" + model.matmul_block_size = 32 + model.quant_attrs = {"is_symmetric": True} + + with pytest.raises(AssertionError, match="Unknown quantization algo config name detected"): + model.make_tied_quantized_embedding_input_names() + + +def _make_minimal_model_for_quantized_tied_embedding(*, int4_algo_config, is_symmetric=True, quant_type=None): + model = Model.__new__(Model) + model.extra_options = {"int4_algo_config": int4_algo_config} + model.algo_config_name = int4_algo_config + model.matmul_block_size = 32 + model.hidden_size = 64 + model.vocab_size = 32000 + model.io_dtype = ir.DataType.FLOAT16 + model.quant_attrs = {"is_symmetric": is_symmetric} + model.quant_type = quant_type + model.input_names = {"input_ids": "input_ids"} + model.embed_attrs = {"scale": 1} + model.layernorm_attrs = { + "cast": {"use_fp32": False}, + "root_input": "", + "skip_input": "", + } + model.tied_quantized_embeddings = True + model.tied_unquantized_embeddings = False + + model._reshape_calls = [] + model._node_calls = [] + + def _make_reshape(name, inputs, dtype, shape): + model._reshape_calls.append((name, inputs, dtype, shape)) + + def _make_node(op_type, inputs, outputs, name, **kwargs): + model._node_calls.append((op_type, inputs, outputs, name, kwargs)) + + def _make_value(_name, _dtype, shape=None): + return shape + + model.make_reshape = _make_reshape + model.make_node = _make_node + model.make_value = _make_value + + return model + + +@pytest.mark.parametrize( + "int4_algo_config, is_symmetric, quant_type, expected_weight_name, expected_scale_name, expected_zp_name, expect_zp_input", + [ + ("default", True, None, "lm_head.MatMul.weight_Q4", "lm_head.MatMul.weight_scales", None, False), + ("default", False, None, "lm_head.MatMul.weight", "", None, False), + ("rtn", True, None, "lm_head.MatMul.weight_Q4G32", "lm_head.MatMul.weight_scale", None, False), + ("rtn", False, None, "lm_head.MatMul.weight_Q4G32", "lm_head.MatMul.weight_scale", "lm_head.MatMul.weight_zp", True), + ("k_quant", True, None, "lm_head.MatMul.weight_Q4G32", "lm_head.MatMul.weight_scale", "lm_head.MatMul.weight_zp", True), + ("k_quant", False, None, "lm_head.MatMul.weight_Q4G32", "lm_head.MatMul.weight_scale", "lm_head.MatMul.weight_zp", True), + ], +) +def test_make_embedding_uses_algo_specific_lm_head_initializer_names_for_tied_quantized_embeddings( + int4_algo_config, + is_symmetric, + quant_type, + expected_weight_name, + expected_scale_name, + expected_zp_name, + expect_zp_input, +): + model = _make_minimal_model_for_quantized_tied_embedding( + int4_algo_config=int4_algo_config, + is_symmetric=is_symmetric, + quant_type=quant_type, + ) + + model.make_embedding(embedding=None) + + assert model._reshape_calls[0][1][0] == expected_weight_name + + gather_calls = [call for call in model._node_calls if call[0] == "GatherBlockQuantized"] + assert len(gather_calls) == 1 + gather_inputs = gather_calls[0][1] + if expected_scale_name: # Only check scale name if it's not empty + assert expected_scale_name in gather_inputs + if expected_zp_name is not None: + assert (expected_zp_name in gather_inputs) is expect_zp_input + else: + assert "lm_head.MatMul.weight_zp" not in gather_inputs + assert "lm_head.MatMul.weight_zero_points" not in gather_inputs + + +def _make_minimal_model_for_embedding_branches(*, tied_quantized_embeddings=False, tied_unquantized_embeddings=False): + model = Model.__new__(Model) + model.hidden_size = 64 + model.vocab_size = 32000 + model.io_dtype = ir.DataType.FLOAT16 + model.input_names = {"input_ids": "input_ids"} + model.embed_attrs = {"scale": 1} + model.layernorm_attrs = { + "cast": {"use_fp32": False}, + "root_input": "", + "skip_input": "", + } + model.tied_quantized_embeddings = tied_quantized_embeddings + model.tied_unquantized_embeddings = tied_unquantized_embeddings + + model._transpose_calls = [] + model._initializer_calls = [] + model._node_calls = [] + model._value_calls = [] + + def _make_transpose(name, root_input, dtype, shape, perm): + model._transpose_calls.append((name, root_input, dtype, shape, perm)) + + def _make_initializer(tensor, name, to=None): + model._initializer_calls.append((tensor, name, to)) + + def _make_node(op_type, **kwargs): + model._node_calls.append((op_type, kwargs)) + + def _make_value(name, dtype, shape): + model._value_calls.append((name, dtype, shape)) + + model.make_transpose = _make_transpose + model.make_initializer = _make_initializer + model.make_node = _make_node + model.make_value = _make_value + return model + + +def test_make_embedding_unquantized_tied_path_emits_transpose_and_gather(): + model = _make_minimal_model_for_embedding_branches( + tied_quantized_embeddings=False, + tied_unquantized_embeddings=True, + ) + + model.make_embedding(embedding=None) + + assert len(model._transpose_calls) == 1 + transpose_call = model._transpose_calls[0] + assert transpose_call[0] == "/model/embed_tokens/Transpose" + assert transpose_call[1] == "lm_head.MatMul.weight" + + gather_calls = [call for call in model._node_calls if call[0] == "Gather"] + assert len(gather_calls) == 1 + gather_inputs = gather_calls[0][1]["inputs"] + assert gather_inputs[0] == "/model/embed_tokens/Transpose/output_0" + assert gather_inputs[1] == "input_ids" + + assert model._initializer_calls == [] + + +def test_make_embedding_non_tied_path_uses_embed_tokens_initializer_and_gather(): + model = _make_minimal_model_for_embedding_branches( + tied_quantized_embeddings=False, + tied_unquantized_embeddings=False, + ) + embedding = object() + + model.make_embedding(embedding=embedding) + + assert len(model._initializer_calls) == 1 + initializer_call = model._initializer_calls[0] + assert initializer_call[0] is embedding + assert initializer_call[1] == "model.embed_tokens.weight" + assert initializer_call[2] == ir.DataType.FLOAT16 + + gather_calls = [call for call in model._node_calls if call[0] == "Gather"] + assert len(gather_calls) == 1 + gather_inputs = gather_calls[0][1]["inputs"] + assert gather_inputs[0] == "model.embed_tokens.weight" + assert gather_inputs[1] == "input_ids" + + assert model._transpose_calls == [] + + +def _make_minimal_model_for_int4_matmul(): + model = Model.__new__(Model) + model.io_dtype = ir.DataType.FLOAT16 + model.quant_attrs = {"accuracy_level": 0} + + model._float_called = False + model._initializers = [] + model._nodes = [] + model._values = [] + + def _make_matmul_float(_matmul, _basename, _root_input, **_kwargs): + model._float_called = True + return "float_fallback" + + def _make_initializer(tensor, name, to=None): + model._initializers.append((name, to, tensor)) + + def _make_node(op_type, **kwargs): + model._nodes.append((op_type, kwargs)) + + def _make_value(name, dtype, shape): + model._values.append((name, dtype, shape)) + + model.make_matmul_float = _make_matmul_float + model.make_initializer = _make_initializer + model.make_node = _make_node + model.make_value = _make_value + return model + + +def test_int4_matmul_uses_float_fallback_when_model_not_already_quantized(): + model = _make_minimal_model_for_int4_matmul() + + matmul = types.SimpleNamespace(weight=object()) + result = model.make_matmul_int4(matmul, "/lm_head/MatMul", "hidden_states") + + assert result == "float_fallback" + assert model._float_called is True + assert model._nodes == [] + + +def test_int4_matmul_emits_matmul_nbits_when_model_already_quantized(): + model = _make_minimal_model_for_int4_matmul() + + matmul = types.SimpleNamespace( + qweight=object(), + scales=object(), + qzeros=object(), + g_idx=object(), + bits=4, + group_size=32, + in_features=64, + out_features=128, + ) + result = model.make_matmul_int4(matmul, "/lm_head/MatMul", "hidden_states") + + assert result == "/lm_head/MatMulNBits" + assert model._float_called is False + assert any(op_type == "MatMulNBits" for op_type, _ in model._nodes) + assert any(name == "lm_head.MatMulNBits.qweight" for name, _, _ in model._initializers) + assert any(name == "lm_head.MatMulNBits.scales" for name, _, _ in model._initializers) diff --git a/test/python/models/test_qwen_2.5_vl.py b/test/python/models/test_qwen_2_5_vl.py similarity index 99% rename from test/python/models/test_qwen_2.5_vl.py rename to test/python/models/test_qwen_2_5_vl.py index 3613434766..7933b19904 100644 --- a/test/python/models/test_qwen_2.5_vl.py +++ b/test/python/models/test_qwen_2_5_vl.py @@ -120,7 +120,7 @@ def ort_io_binding_helper( sess.run_with_iobinding(bind) -def test_parity( +def run_parity( hf_model_name: str, cache_dir: str, onnx_model_path: str, @@ -405,7 +405,7 @@ def test_parity( print("Warning: CPU testing with IOBinding is not set up. Forcing GPU.") # This script is now GPU-only - test_parity( + run_parity( hf_model_name=args.hf_model, cache_dir=args.cache_dir, onnx_model_path=args.onnx_model, diff --git a/test/python/models/test_qwen_3.5_text_only.py b/test/python/models/test_qwen_3_5_text_only.py similarity index 97% rename from test/python/models/test_qwen_3.5_text_only.py rename to test/python/models/test_qwen_3_5_text_only.py index b62fd308f4..ecea33d30b 100644 --- a/test/python/models/test_qwen_3.5_text_only.py +++ b/test/python/models/test_qwen_3_5_text_only.py @@ -9,7 +9,7 @@ "qwen3_5_text"), which uses 2D position_ids and hybrid KV/recurrent state. Usage: - pytest test_qwen_3.5_text_only.py --test_models=test/models + pytest test_qwen_3_5_text_only.py --test_models=test/models """ import os diff --git a/test/python/models/test_yarn_rope_parity.py b/test/python/models/test_yarn_rope_parity.py index 88a98f6420..2cf7c97f61 100644 --- a/test/python/models/test_yarn_rope_parity.py +++ b/test/python/models/test_yarn_rope_parity.py @@ -34,7 +34,7 @@ # Import Model from the source tree so tests always run against the working copy. # The installed onnxruntime_genai package may be out of date during development. -sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "..", "src", "python", "py")) +sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "..", "..", "src", "python", "py")) from models.builders.base import Model diff --git a/test/python/test_onnxruntime_genai.py b/test/python/test_onnxruntime_genai.py index 8dc5ba2065..3ad93cbd15 100644 --- a/test/python/test_onnxruntime_genai.py +++ b/test/python/test_onnxruntime_genai.py @@ -9,8 +9,6 @@ import onnxruntime_genai as og from _test_utils import download_models, is_webgpu_ep_available, run_subprocess -from models.test_gemma4_models import run_gemma4_vision_tests -from models.test_qwen_fara_models import run_qwen_fara_vision_tests logging.basicConfig(format="%(asctime)s %(name)s [%(levelname)s] - %(message)s", level=logging.DEBUG) log = logging.getLogger("onnxruntime-genai-tests") @@ -21,7 +19,7 @@ def run_onnxruntime_genai_api_tests( log: logging.Logger, test_models: str | bytes | os.PathLike, ): - log.debug("Running: ONNX Runtime GenAI API Tests") + log.debug("Running: ONNX Runtime GenAI API, builder, and model tests") command = [ sys.executable, @@ -29,6 +27,8 @@ def run_onnxruntime_genai_api_tests( "pytest", "-sv", "test_onnxruntime_genai_api.py", + "builder", + "models", "--test_models", test_models, ] @@ -108,10 +108,6 @@ def main(): # Run ONNX Runtime GenAI tests run_onnxruntime_genai_api_tests(os.path.abspath(args.cwd), log, os.path.abspath(args.test_models)) - # Run vision model tests (tests auto-skip if models are not present) - run_gemma4_vision_tests(os.path.abspath(args.cwd), log, os.path.abspath(args.test_models)) - run_qwen_fara_vision_tests(os.path.abspath(args.cwd), log, os.path.abspath(args.test_models)) - if args.e2e: run_onnxruntime_genai_e2e_tests(os.path.abspath(args.cwd), log, output_paths)