diff --git a/.pipelines/stages/jobs/py-validation-job.yml b/.pipelines/stages/jobs/py-validation-job.yml
index 2afb279a75..b83f7827d4 100644
--- a/.pipelines/stages/jobs/py-validation-job.yml
+++ b/.pipelines/stages/jobs/py-validation-job.yml
@@ -208,7 +208,9 @@ jobs:
displayName: 'Download CUDA $(cuda_version)'
workingDirectory: '$(Build.Repository.LocalPath)'
- powershell: |
- python -m pip install -r test/python/requirements.txt
+ if ("$(arch)" -ne "arm64") {
+ python -m pip install -r test/python/requirements.txt
+ }
if ("$(ep)" -eq "cuda") {
$env:CUDA_PATH = '$(Build.Repository.LocalPath)\cuda_sdk\v$(cuda_version)'
$env:PATH = "$env:CUDA_PATH\bin;$env:CUDA_PATH\extras\CUPTI\lib64;$env:PATH"
@@ -220,6 +222,10 @@ jobs:
python -m pip install -r test/python/directml/torch/requirements.txt
python -m pip install -r test/python/directml/ort/requirements.txt
}
+ elseif ("$(arch)" -eq "arm64") {
+ python -m pip install numpy<2
+ python -m pip install onnxruntime-qnn==1.20.0
+ }
else {
python -m pip install -r test/python/cpu/torch/requirements.txt
python -m pip install -r test/python/cpu/ort/requirements.txt
@@ -227,7 +233,7 @@ jobs:
cd examples\python
python -m pip install --no-index --find-links=$(Build.BinariesDirectory)/wheel $(pip_package_name)
- python model-generate.py -m .\models\$(prebuild_phi3_mini_model_folder) --min_length 25 --max_length 50 --verbose
+ python model-generate.py -m .\models\$(prebuild_phi3_mini_model_folder) --min_length 25 --max_length 50 --batch_size_for_cuda_graph 3 --verbose
displayName: 'Run Example With Artifact'
workingDirectory: '$(Build.Repository.LocalPath)'
diff --git a/.pipelines/stages/jobs/steps/utils/download-huggingface-model.yml b/.pipelines/stages/jobs/steps/utils/download-huggingface-model.yml
index 0dc5a8fc1d..3bb0caa222 100644
--- a/.pipelines/stages/jobs/steps/utils/download-huggingface-model.yml
+++ b/.pipelines/stages/jobs/steps/utils/download-huggingface-model.yml
@@ -28,7 +28,8 @@ steps:
- powershell: |
python -m pip install "huggingface_hub[cli]"
huggingface-cli login --token $env:HF_TOKEN
- huggingface-cli download ${{ parameters.HuggingFaceRepo }} --include ${{ parameters.RepoFolder }}/* --local-dir ${{ parameters.LocalFolder }} --local-dir-use-symlinks False
+ # Use maximum path length for Windows... otherwises hits the path character limit
+ huggingface-cli download ${{ parameters.HuggingFaceRepo }} --include ${{ parameters.RepoFolder }}/* --local-dir "\\?\${{ parameters.WorkingDirectory }}\\${{ parameters.LocalFolder }}" --local-dir-use-symlinks False
displayName: ${{ parameters.StepName }}
workingDirectory: ${{ parameters.WorkingDirectory }}
env:
diff --git a/VERSION_INFO b/VERSION_INFO
index 0a6a4023a1..79a2734bbf 100644
--- a/VERSION_INFO
+++ b/VERSION_INFO
@@ -1 +1 @@
-0.5.0-rc1
\ No newline at end of file
+0.5.0
\ No newline at end of file
diff --git a/documents/Runtime_option.md b/documents/Runtime_option.md
new file mode 100644
index 0000000000..d164f0ff39
--- /dev/null
+++ b/documents/Runtime_option.md
@@ -0,0 +1,15 @@
+# Runtime Options
+
+This file will provide details on the usage of SetRuntimeOption API. It will list all the current key value pairs which can be used as an input for this API.
+
+## Set Terminate
+
+Set Terminate is a runtime option to terminate the current session or continue/restart an already terminated session. There are two valid ways to call Set Terminate.
+
+To enable terminate, the valid pair is: ("set_terminate", "1")
+
+To disable terminate, the valid pair is: ("set_terminate", "0")
+
+Key: "set_terminate"
+
+Accepted values: ("0", "1")
diff --git a/examples/csharp/HelloPhi/HelloPhi.csproj b/examples/csharp/HelloPhi/HelloPhi.csproj
index 1ee51abefd..3c5855bda3 100644
--- a/examples/csharp/HelloPhi/HelloPhi.csproj
+++ b/examples/csharp/HelloPhi/HelloPhi.csproj
@@ -10,9 +10,9 @@
-
-
-
+
+
+
diff --git a/nuget/PACKAGE.md b/nuget/PACKAGE.md
index 7f7c324db6..7daeed7977 100644
--- a/nuget/PACKAGE.md
+++ b/nuget/PACKAGE.md
@@ -121,11 +121,11 @@ This implementation checks if a number is prime by iterating only up to the squa
## Source code repository
ONNX Runtime is an open source project. See:
-* (https://github.com/microsoft/onnxruntime)[https://github.com/microsoft/onnxruntime]
-* (https://github.com/microsoft/onnxruntime-genai)[https://github.com/microsoft/onnxruntime-genai]
+* (ONNX Runtime)[https://github.com/microsoft/onnxruntime]
+* (ONNX Runtime GenAI)[https://github.com/microsoft/onnxruntime-genai]
## Documentation
-See (https://onxxruntime.ai/docs/genai)[https://onxxruntime.ai/docs/genai]
+See (ONNX Runtime GenAI Documentation)[https://onxxruntime.ai/docs/genai]
diff --git a/src/generators.cpp b/src/generators.cpp
index 1d5e74ff43..ee51b73c66 100644
--- a/src/generators.cpp
+++ b/src/generators.cpp
@@ -30,6 +30,11 @@ std::string CurrentModulePath() {
}
#endif
+void ThrowErrorIfSessionTerminated(bool is_session_terminated) {
+ if (is_session_terminated)
+ throw std::runtime_error("Session in Terminated state, exiting!");
+}
+
namespace Generators {
#if USE_CUDA
@@ -284,6 +289,7 @@ Generator::Generator(const Model& model, const GeneratorParams& params) : model_
}
void Generator::ComputeLogits() {
+ ThrowErrorIfSessionTerminated(state_->session_terminated_);
if (computed_logits_)
throw std::runtime_error("ComputeLogits called again without calling GenerateNextToken first");
@@ -301,7 +307,25 @@ void Generator::ComputeLogits() {
search_->ApplyRepetitionPenalty(search.repetition_penalty);
}
+void Generator::SetRuntimeOption(const char* key, const char* value) {
+ // TODO: Need a better way to handle different keys
+ // We can create a config manager to host all configurations and do comparison at that point
+ if (strcmp(key, "terminate_session") == 0) {
+ if (strcmp(value, "0") == 0) {
+ state_->UnsetTerminate();
+ } else if (strcmp(value, "1") == 0) {
+ state_->SetTerminate();
+ } else {
+ // Value not expected
+ throw std::runtime_error(std::string("terminate_session key value unexpected: ") + value);
+ }
+ } else {
+ throw std::runtime_error(std::string("SetRuntimeOption key is not expected: ") + key);
+ }
+}
+
bool Generator::IsDone() const {
+ ThrowErrorIfSessionTerminated(state_->session_terminated_);
if (computed_logits_)
throw std::runtime_error("IsDone() can't be called in the middle of processing logits");
@@ -313,7 +337,12 @@ bool Generator::IsDone() const {
return is_done;
}
+bool Generator::IsSessionTerminated() const {
+ return state_->session_terminated_;
+}
+
void Generator::GenerateNextToken() {
+ ThrowErrorIfSessionTerminated(state_->session_terminated_);
if (!computed_logits_)
throw std::runtime_error("Must call ComputeLogits before GenerateNextToken");
computed_logits_ = false;
diff --git a/src/generators.h b/src/generators.h
index c36ed6b336..31ddbf84cc 100644
--- a/src/generators.h
+++ b/src/generators.h
@@ -40,6 +40,8 @@ using cudaStream_t = void*;
#include "runtime_settings.h"
#include "tensor.h"
+void ThrowErrorIfSessionTerminated(bool is_session_terminated);
+
namespace Generators {
struct Model;
struct State;
@@ -108,7 +110,9 @@ struct Generator : LeakChecked {
Generator(const Model& model, const GeneratorParams& params);
bool IsDone() const;
+ void SetRuntimeOption(const char* key, const char* value);
void ComputeLogits();
+ bool IsSessionTerminated() const;
void GenerateNextToken();
DeviceMemorySpan GetSequence(size_t index) const;
diff --git a/src/models/model.cpp b/src/models/model.cpp
index a5f549fb15..0c092cfb5b 100644
--- a/src/models/model.cpp
+++ b/src/models/model.cpp
@@ -65,7 +65,18 @@ void State::Run(OrtSession& session, int new_batch_size) {
}
}
+void State::SetTerminate() {
+ session_terminated_ = true;
+ run_options_->SetTerminate();
+}
+
+void State::UnsetTerminate() {
+ session_terminated_ = false;
+ run_options_->UnsetTerminate();
+}
+
OrtValue* State::GetInput(const char* name) {
+ ThrowErrorIfSessionTerminated(session_terminated_);
for (size_t i = 0; i < input_names_.size(); i++) {
if (std::strcmp(input_names_[i], name) == 0) {
return inputs_[i];
@@ -75,6 +86,7 @@ OrtValue* State::GetInput(const char* name) {
}
OrtValue* State::GetOutput(const char* name) {
+ ThrowErrorIfSessionTerminated(session_terminated_);
for (size_t i = 0; i < output_names_.size(); i++) {
if (std::strcmp(output_names_[i], name) == 0) {
return outputs_[i];
diff --git a/src/models/model.h b/src/models/model.h
index 9eed5df75d..a5e1325627 100644
--- a/src/models/model.h
+++ b/src/models/model.h
@@ -34,6 +34,9 @@ struct State {
virtual const CapturedGraphInfo* GetCapturedGraphInfo() const { return nullptr; }
virtual void Finalize() {}
+ void SetTerminate();
+ void UnsetTerminate();
+ mutable bool session_terminated_{};
OrtValue* GetInput(const char* name);
virtual OrtValue* GetOutput(const char* name);
diff --git a/src/ort_genai.h b/src/ort_genai.h
index 424e7aad77..23e9040edb 100644
--- a/src/ort_genai.h
+++ b/src/ort_genai.h
@@ -258,6 +258,10 @@ struct OgaGenerator : OgaAbstract {
return OgaGenerator_IsDone(this);
}
+ bool IsSessionTerminated() const {
+ return OgaGenerator_IsSessionTerminated(this);
+ }
+
void ComputeLogits() {
OgaCheckResult(OgaGenerator_ComputeLogits(this));
}
@@ -266,6 +270,10 @@ struct OgaGenerator : OgaAbstract {
OgaCheckResult(OgaGenerator_GenerateNextToken(this));
}
+ void SetRuntimeOption(const char* key, const char* value) {
+ OgaCheckResult(OgaGenerator_SetRuntimeOption(this, key, value));
+ }
+
size_t GetSequenceCount(size_t index) const {
return OgaGenerator_GetSequenceCount(this, index);
}
diff --git a/src/ort_genai_c.cpp b/src/ort_genai_c.cpp
index a8f259fdec..a7e3b41f59 100644
--- a/src/ort_genai_c.cpp
+++ b/src/ort_genai_c.cpp
@@ -271,6 +271,10 @@ bool OGA_API_CALL OgaGenerator_IsDone(const OgaGenerator* generator) {
return reinterpret_cast(generator)->IsDone();
}
+bool OGA_API_CALL OgaGenerator_IsSessionTerminated(const OgaGenerator* generator) {
+ return reinterpret_cast(generator)->IsSessionTerminated();
+}
+
OgaResult* OGA_API_CALL OgaGenerator_ComputeLogits(OgaGenerator* generator) {
OGA_TRY
reinterpret_cast(generator)->ComputeLogits();
@@ -285,6 +289,13 @@ OgaResult* OGA_API_CALL OgaGenerator_GenerateNextToken(OgaGenerator* generator)
OGA_CATCH
}
+OgaResult* OGA_API_CALL OgaGenerator_SetRuntimeOption(OgaGenerator* generator, const char* key, const char* value) {
+ OGA_TRY
+ reinterpret_cast(generator)->SetRuntimeOption(key, value);
+ return nullptr;
+ OGA_CATCH
+}
+
OgaResult* OGA_API_CALL OgaGenerator_GetOutput(const OgaGenerator* oga_generator, const char* name, OgaTensor** out) {
OGA_TRY
auto& generator = *reinterpret_cast(oga_generator);
diff --git a/src/ort_genai_c.h b/src/ort_genai_c.h
index b00494d7d6..4fb0aa5497 100644
--- a/src/ort_genai_c.h
+++ b/src/ort_genai_c.h
@@ -278,6 +278,7 @@ OGA_EXPORT void OGA_API_CALL OgaDestroyGenerator(OgaGenerator* generator);
* \return True if the generator has finished generating all the sequences, false otherwise.
*/
OGA_EXPORT bool OGA_API_CALL OgaGenerator_IsDone(const OgaGenerator* generator);
+OGA_EXPORT bool OGA_API_CALL OgaGenerator_IsSessionTerminated(const OgaGenerator* generator);
/*
* \brief Computes the logits from the model based on the input ids and the past state. The computed logits are stored in the generator.
@@ -287,6 +288,8 @@ OGA_EXPORT bool OGA_API_CALL OgaGenerator_IsDone(const OgaGenerator* generator);
OGA_EXPORT OgaResult* OGA_API_CALL OgaGenerator_ComputeLogits(OgaGenerator* generator);
OGA_EXPORT OgaResult* OGA_API_CALL OgaGenerator_GenerateNextToken(OgaGenerator* generator);
+OGA_EXPORT OgaResult* OGA_API_CALL OgaGenerator_SetRuntimeOption(OgaGenerator* generator, const char* key, const char* value);
+
/*
* \brief Returns a copy of the model output identified by the given name as an OgaTensor on CPU. The buffer is owned by returned OgaTensor
* and will be released when the OgaTensor is destroyed
diff --git a/src/python/py/models/builder.py b/src/python/py/models/builder.py
index 7951de8e27..76e73bf6bd 100644
--- a/src/python/py/models/builder.py
+++ b/src/python/py/models/builder.py
@@ -40,10 +40,10 @@ def __init__(self, config, io_dtype, onnx_dtype, ep, cache_dir, extra_options):
self.io_dtype = io_dtype # {'fp16', 'fp32'}
self.onnx_dtype = onnx_dtype # {"int4", "fp16", "fp32"}
self.quant_type = config.quantization_config["quant_method"] if hasattr(config, "quantization_config") else None
- self.adapter_path = extra_options["adapter_path"] if "adapter_path" in extra_options else None
+ self.adapter_path = extra_options.get("adapter_path", None)
self.cache_dir = cache_dir
- self.filename = extra_options["filename"] if "filename" in extra_options else "model.onnx"
+ self.filename = extra_options.get("filename", "model.onnx")
self.hf_token = parse_hf_token(extra_options.get("hf_token", "true"))
self.extra_options = extra_options
@@ -54,7 +54,7 @@ def __init__(self, config, io_dtype, onnx_dtype, ep, cache_dir, extra_options):
self.nodes = []
# EP-specific variables
- enable_cuda_graph = "1" if "enable_cuda_graph" in extra_options else "0"
+ enable_cuda_graph = extra_options.get("enable_cuda_graph", "0")
self.ep = ep
self.ep_attrs = {
"cpu": {},
@@ -147,6 +147,7 @@ def __init__(self, config, io_dtype, onnx_dtype, ep, cache_dir, extra_options):
}
# LayerNorm-specific variables
+ epsilon = config.rms_norm_eps if hasattr(config, "rms_norm_eps") else 1e-06
self.layernorm_attrs = {
"simple": True, # Use SimplifiedLayerNorm/SkipSimplifiedLayerNorm vs. LayerNorm/SkipLayerNorm
"first_layernorm": True, # 1st LayerNorm = LayerNorm, then SkipLayerNorm for all subsequent LayerNorms
@@ -156,6 +157,7 @@ def __init__(self, config, io_dtype, onnx_dtype, ep, cache_dir, extra_options):
"output_0": "", # Output 0 for LayerNorm and SkipLayerNorm
"output_3": "", # Output 3 for SkipLayerNorm
"add_offset": 0, # Offset value for LayerNorm weight
+ "epsilon": epsilon, # Epsilon value to avoid `sqrt(0)` in LayerNorm
}
# MatMul-specific variables
@@ -212,6 +214,8 @@ def __init__(self, config, io_dtype, onnx_dtype, ep, cache_dir, extra_options):
}
# Attention-specific variables (MHA, GQA, GQA + Rot.Emb., etc.)
+ softcap = config.attn_logit_softcapping if hasattr(config, "attn_logit_softcapping") else 0.0 # default is 0.0 in GroupQueryAttention kernel
+
# Block-sparse attention-specific variables
sparse_block_size = config.blocksparse_block_size if hasattr(config, "blocksparse_block_size") else 0
kernel_block_size = config.blocksparse_triton_kernel_block_size if hasattr(config, "blocksparse_triton_kernel_block_size") else 0
@@ -224,6 +228,7 @@ def __init__(self, config, io_dtype, onnx_dtype, ep, cache_dir, extra_options):
"v_path": "", # V path to attention
"op_type": "MultiHeadAttention", # Attention op to use
"scale": 1 / np.sqrt(self.head_size), # Scale value after calculating Q x K' in attention
+ "softcap": softcap, # Softcap value to prevent values from exploding in attention
"use_rotemb_in_attn": False, # Use rotary embeddings within attention (instead of a separate RotaryEmbedding op)
"use_packed_matmul": False, # Use packed MatMul (instead of 3 separate MatMuls for Q/K/V)
"block_sparse": { # Block-sparse attention-specific variables
@@ -288,8 +293,9 @@ def __init__(self, config, io_dtype, onnx_dtype, ep, cache_dir, extra_options):
# Quantization-specific variables (INT4, INT8, etc.)
self.quant_attrs = {
"int4": {
- "block_size": int(extra_options["int4_block_size"]) if "int4_block_size" in extra_options else 32,
- "accuracy_level": int(extra_options["int4_accuracy_level"]) if "int4_accuracy_level" in extra_options else 0, # Default is 0 for non-QDQ formats, default is 4 for QDQ formats
+ "accuracy_level": int(extra_options.get("int4_accuracy_level", 0)), # Default is 0 for non-QDQ formats, default is 4 for QDQ formats
+ "block_size": int(extra_options.get("int4_block_size", 32)),
+ "op_types_to_quantize": extra_options.get("int4_op_types_to_quantize", ("MatMul", )),
},
"use_qdq": False, # Use QDQ format
}
@@ -346,7 +352,7 @@ def make_genai_config(self, model_name_or_path, extra_kwargs, out_dir):
"no_repeat_ngram_size": config.no_repeat_ngram_size if hasattr(config, "no_repeat_ngram_size") else 0,
"num_beams": config.num_beams if hasattr(config, "num_beams") else 1,
"num_return_sequences": config.num_return_sequences if hasattr(config, "num_return_sequences") else 1,
- "past_present_share_buffer": self.past_present_share_buffer,
+ "past_present_share_buffer": False if "config_only" in self.extra_options else self.past_present_share_buffer,
"repetition_penalty": config.repetition_penalty if hasattr(config, "repetition_penalty") else 1.0,
"temperature": config.temperature if hasattr(config, "temperature") else 1.0,
"top_k": 1,
@@ -401,7 +407,8 @@ def save_model(self, out_dir):
# Quantize ONNX model to desired precision
# TODO: Replace by quantizing the MatMuls as they are created
- if self.onnx_dtype == "int4" and self.quant_type is None:
+ 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
+ if self.onnx_dtype == "int4" and not already_quantized_in_qdq_format:
model = self.to_int4(model)
# Save ONNX model with only one external data file and delete any existing duplicate copies
@@ -432,6 +439,7 @@ def to_int4(self, model):
accuracy_level=self.quant_attrs["int4"]["accuracy_level"],
nodes_to_exclude=[],
quant_format=QuantFormat.QDQ if self.quant_attrs["use_qdq"] else QuantFormat.QOperator,
+ op_types_to_quantize=self.quant_attrs["int4"]["op_types_to_quantize"],
)
quant.process()
return quant.model.model
@@ -969,7 +977,7 @@ def make_layernorm(self, layer_id, layernorm, skip, simple, location):
name = f"/model/layers.{layer_id}/{location}_layernorm/{'Skip' if skip else ''}LayerNorm"
op_type = f"{'Skip' if skip else ''}{'Simplified' if simple else ''}LayerNormalization"
- kwargs = {"epsilon": 9.999999747378752e-06}
+ kwargs = {"epsilon": self.layernorm_attrs["epsilon"]}
if not skip:
kwargs.update({"axis": -1, "stash_type": 1})
@@ -1381,7 +1389,7 @@ def make_group_query_attention(self, name, **kwargs):
self.make_node(
"GroupQueryAttention", inputs=inputs, outputs=outputs, name=name, domain="com.microsoft",
num_heads=self.num_attn_heads, kv_num_heads=self.num_kv_heads, scale=self.attention_attrs["scale"], # local_window_size=self.window_size, # Disable sliding window attribute temporarily
- do_rotary=self.attention_attrs["use_rotemb_in_attn"], rotary_interleaved=self.rotemb_attrs["interleaved"],
+ softcap=self.attention_attrs["softcap"], do_rotary=self.attention_attrs["use_rotemb_in_attn"], rotary_interleaved=self.rotemb_attrs["interleaved"],
)
self.make_value_info(output, self.io_dtype, shape=['batch_size', 'sequence_length', self.head_size * self.num_attn_heads])
@@ -3004,6 +3012,12 @@ def make_layer(self, layer_id, layer):
super().make_layer(layer_id, layer)
def check_extra_options(kv_pairs):
+ 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("/"):
+ op_types_to_quantize += (op_type, )
+ kv_pairs["int4_op_types_to_quantize"] = op_types_to_quantize
+
if "use_8bits_moe" in kv_pairs:
assert(kv_pairs["use_8bits_moe"] == "1" or kv_pairs["use_8bits_moe"] == "0"), "use_8bits_moe must be 0 or 1."
@@ -3181,12 +3195,15 @@ def get_args():
nargs='+',
help=textwrap.dedent("""\
Key value pairs for various options. Currently supports:
- int4_block_size = 16/32/64/128/256: Specify the block_size for int4 quantization.
int4_accuracy_level = 1/2/3/4: Specify the minimum accuracy level for activation of MatMul in int4 quantization.
4 is int8, which means input A of int4 quantized MatMul is quantized to int8 and input B is upcasted to int8 for computation.
3 is bf16.
2 is fp16.
1 is fp32.
+ int4_block_size = 16/32/64/128/256: Specify the block_size for int4 quantization.
+ int4_op_types_to_quantize = MatMul/Gather: Specify op types to target for int4 quantization.
+ Use this option when you want to quantize specific ops.
+ Separate the op types with a '/' when passing them here (e.g. int4_op_types_to_quantize=MatMul/Gather)
num_hidden_layers = Manually specify the number of layers in your ONNX model (for unit testing purposes).
filename = Filename for ONNX model (default is 'model.onnx').
For models with multiple components, each component is exported to its own ONNX model.
@@ -3199,13 +3216,13 @@ 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.
- enable_cuda_graph = 1 : The model can use CUDA graph capture for CUDA execution provider. If enabled, all nodes being placed on the CUDA EP
+ enable_cuda_graph = 1: The model can use CUDA graph capture for CUDA execution provider. 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.
- use_8bits_moe = 1 : Use 8-bit quantization for MoE layers. Default is using 4-bit quantization.
+ use_8bits_moe = 1: Use 8-bit quantization for MoE layers. Default is using 4-bit quantization.
hf_token = false/token: Use this to disable authentication with Hugging Face or provide a custom authentication token that differs from the one stored in your environment. Default behavior is to use the authentication token stored by `huggingface-cli login`.
If you have already authenticated via `huggingface-cli login`, you do not need to use this flag because Hugging Face has already stored your authentication token for you.
- use_qdq = 1 : Use the QDQ decomposition for quantized MatMul instead of the MatMulNBits operator.
+ use_qdq = 1: Use the QDQ decomposition for quantized MatMul instead of the MatMulNBits operator.
adapter_path = Path to folder on disk containing the adapter files (adapter_config.json and adapter model weights).
"""),
)
diff --git a/src/python/py/models/quantized_model.py b/src/python/py/models/quantized_model.py
index b41b16d90d..ed0a5beb67 100644
--- a/src/python/py/models/quantized_model.py
+++ b/src/python/py/models/quantized_model.py
@@ -117,6 +117,18 @@ def __init__(self, quant_type, input_path, bits, group_size, q_size, kv_size, in
# transformer.rotary_pos_emb.inv_freq in ChatGLM3.
# Skip rotary embedding weights since they can be re-calculated when looping through the model
continue
+ elif name == "lm_head.qweight" or name == "transformer.output_layer.qweight":
+ self._initialize_quantized_lm_head(bits, group_size)
+ self.lm_head.qweight = tensor
+ elif name == "lm_head.qzeros" or name == "transformer.output_layer.qzeros":
+ self._initialize_quantized_lm_head(bits, group_size)
+ self.lm_head.qzeros = tensor
+ elif name == "lm_head.scales" or name == "transformer.output_layer.scales":
+ self._initialize_quantized_lm_head(bits, group_size)
+ self.lm_head.scales = tensor
+ elif name == "lm_head.g_idx" or name == "transformer.output_layer.g_idx":
+ self._initialize_quantized_lm_head(bits, group_size)
+ self.lm_head.g_idx = tensor
else:
if name.startswith("transformer.encoder"):
# Chatglm3, e.g., transformer.encoder.layers.0.input_layernorm.weight
@@ -326,7 +338,7 @@ def __init__(self, quant_type, input_path, bits, group_size, q_size, kv_size, in
raise NotImplementedError(f"{name} in your quantized model is not recognized.")
# Set LM head weights + biases if not already set
- if self.lm_head.weight is None:
+ if isinstance(self.lm_head, TensorModule) and self.lm_head.weight is None:
# Embedding and LM head share same weights + biases (lm_head.weight == embedding.weight and lm_head.bias == embedding.bias)
self.lm_head.weight = self.embedding.weight
if self.lm_head.bias is not None:
@@ -339,10 +351,31 @@ def __init__(self, quant_type, input_path, bits, group_size, q_size, kv_size, in
# Set properties of each layer based on quantization type
self.set_properties()
+ def _initialize_quantized_lm_head(self, bits, group_size):
+ """
+ Initialize `QuantizedTensorModule` for LM head if not already set
+ """
+ if isinstance(self.lm_head, TensorModule):
+ assert self.lm_head.weight is None
+ assert self.lm_head.bias is None
+ if not isinstance(self.lm_head, QuantizedTensorModule):
+ self.lm_head = QuantizedTensorModule(bits, group_size)
+
def set_properties(self):
"""
Set in_features, out_features, and g_idx based on quantization type
"""
+ if isinstance(self.lm_head, QuantizedTensorModule):
+ if self.quant_type == "awq":
+ self.lm_head.out_features = self.lm_head.scales.shape[1]
+ self.lm_head.in_features = self.lm_head.qweight.shape[0]
+ # Set g_idx if not already set
+ self.lm_head.g_idx = self.lm_head.g_idx if self.lm_head.g_idx is not None else torch.tensor([i // self.lm_head.group_size for i in range(self.lm_head.in_features)], dtype=torch.int32)
+ elif self.quant_type == "gptq":
+ self.lm_head.out_features = self.lm_head.qweight.shape[1]
+ self.lm_head.in_features = self.lm_head.g_idx.shape[0]
+ else:
+ raise NotImplementedError(f"The {self.quant_type} quantization method is not recognized.")
for module in self.layers:
if self.quant_type == "awq":
# Set in_features and out_features
@@ -582,6 +615,13 @@ def __init__(self, quant_type, input_path, bits, group_size, q_size, kv_size, in
# Set `g_idx` to None since it's not used in `MatMulNBits`
q_tensors.g_idx = None
+ if isinstance(self.lm_head, QuantizedTensorModule) and self.lm_head.qweight is not None:
+ self.unpack(self.lm_head)
+ self.repack(self.lm_head)
+
+ # Set `g_idx` to None since it's not used in `MatMulNBits`
+ self.lm_head.g_idx = None
+
def unpack_qweight(self, module):
"""
Unpack `qweight` to standard format
@@ -604,12 +644,12 @@ def reverse_reorder_tensor(self, tensor, bits):
"""
compress_ratio = 32 // bits
assert tensor.shape[-1] % compress_ratio == 0
-
+
if bits == 4:
order_map = [0, 2, 4, 6, 1, 3, 5, 7]
else:
raise NotImplementedError(f"Unpacking for {bits}-bit quantization is not currently supported.")
-
+
order_tensor = torch.tensor(order_map, dtype=torch.int32).reshape(1, -1)
order_tensor = order_tensor.repeat(tensor.shape[1] // compress_ratio, 1)
order_tensor = order_tensor + torch.arange(0, tensor.shape[1], compress_ratio, dtype=torch.int32).reshape(-1, 1)
@@ -652,7 +692,16 @@ def __init__(self, quant_type, input_path, bits, group_size, use_g_idx, q_size,
if not use_g_idx:
# Set `g_idx` to None since it's not used in `MatMulNBits`
q_tensors.g_idx = None
-
+
+ if isinstance(self.lm_head, QuantizedTensorModule) and self.lm_head.qweight is not None:
+ self.handle_qzeros(self.lm_head)
+ self.unpack(self.lm_head)
+ self.repack(self.lm_head)
+
+ if not use_g_idx:
+ # Set `g_idx` to None since it's not used in `MatMulNBits`
+ self.lm_head.g_idx = None
+
def handle_qzeros(self, module):
"""
Re-pack `qzeros` to handle extra `-1`s
diff --git a/test/c_api_tests.cpp b/test/c_api_tests.cpp
index 019d037242..a93f7c645c 100644
--- a/test/c_api_tests.cpp
+++ b/test/c_api_tests.cpp
@@ -5,6 +5,8 @@
#include
#include
#include "../src/span.h"
+#include
+#include
#ifndef MODEL_PATH
#define MODEL_PATH "../../test/test_models/"
@@ -283,6 +285,55 @@ TEST(CAPITests, GetOutputCAPI) {
generator->GenerateNextToken();
}
+TEST(CAPITests, SetTerminate) {
+#if TEST_PHI2
+
+ auto GeneratorSetTerminateCall = [](OgaGenerator* generator) {
+ // Set Terminate
+ generator->SetRuntimeOption("terminate_session", "1");
+ };
+
+ auto GenerateOutput = [](OgaGenerator* generator, std::unique_ptr tokenizer_stream) {
+ try {
+ while (!generator->IsDone()) {
+ generator->ComputeLogits();
+ generator->GenerateNextToken();
+ }
+ }
+ catch (const std::exception& e) {
+ EXPECT_EQ(generator->IsSessionTerminated(), true);
+ std::cout << "Session Terminated: " << e.what() << std::endl;
+ }
+ };
+
+ auto model = OgaModel::Create(PHI2_PATH);
+ auto tokenizer = OgaTokenizer::Create(*model);
+ auto tokenizer_stream = OgaTokenizerStream::Create(*tokenizer);
+
+ const char* input_string = "She sells sea shells by the sea shore.";
+ auto input_sequences = OgaSequences::Create();
+ tokenizer->Encode(input_string, *input_sequences);
+ auto params = OgaGeneratorParams::Create(*model);
+ params->SetInputSequences(*input_sequences);
+ params->SetSearchOption("max_length", 40);
+
+ auto generator = OgaGenerator::Create(*model, *params);
+ EXPECT_EQ(generator->IsSessionTerminated(), false);
+ std::vector threads;
+ threads.push_back(std::thread(GenerateOutput, generator.get(), std::move(tokenizer_stream)));
+ threads.push_back(std::thread(GeneratorSetTerminateCall, generator.get()));
+
+ for (auto& th : threads) {
+ std::cout << "Waiting for threads completion" << std::endl;
+ th.join(); // Wait for each thread to finish
+ }
+ EXPECT_EQ(generator->IsSessionTerminated(), true);
+ // Unset terminate
+ generator->SetRuntimeOption("terminate_session", "0");
+ EXPECT_EQ(generator->IsSessionTerminated(), false);
+#endif
+}
+
#if TEST_PHI2
struct Phi2Test {
diff --git a/test/python/cpu/ort/requirements.txt b/test/python/cpu/ort/requirements.txt
index b465e7571a..80f6942c88 100644
--- a/test/python/cpu/ort/requirements.txt
+++ b/test/python/cpu/ort/requirements.txt
@@ -1,2 +1,2 @@
--i https://aiinfra.pkgs.visualstudio.com/PublicPackages/_packaging/ORT-Nightly/pypi/simple/
+# -i https://aiinfra.pkgs.visualstudio.com/PublicPackages/_packaging/ORT-Nightly/pypi/simple/
onnxruntime==1.20.0
diff --git a/test/python/cuda/ort/requirements.txt b/test/python/cuda/ort/requirements.txt
index a52982a2aa..8d73885d7a 100644
--- a/test/python/cuda/ort/requirements.txt
+++ b/test/python/cuda/ort/requirements.txt
@@ -1,2 +1,2 @@
--i https://aiinfra.pkgs.visualstudio.com/PublicPackages/_packaging/ORT-Nightly/pypi/simple/
+# -i https://aiinfra.pkgs.visualstudio.com/PublicPackages/_packaging/ORT-Nightly/pypi/simple/
onnxruntime-gpu==1.20.0
diff --git a/test/python/directml/ort/requirements.txt b/test/python/directml/ort/requirements.txt
index f2a3667dfb..73fdc91672 100644
--- a/test/python/directml/ort/requirements.txt
+++ b/test/python/directml/ort/requirements.txt
@@ -1,2 +1,2 @@
--i https://aiinfra.pkgs.visualstudio.com/PublicPackages/_packaging/ORT-Nightly/pypi/simple/
+# -i https://aiinfra.pkgs.visualstudio.com/PublicPackages/_packaging/ORT-Nightly/pypi/simple/
onnxruntime-directml==1.20.0
diff --git a/test/python/macos/ort/requirements.txt b/test/python/macos/ort/requirements.txt
index f75baa617d..5dfa4f3a8a 100644
--- a/test/python/macos/ort/requirements.txt
+++ b/test/python/macos/ort/requirements.txt
@@ -1,2 +1,2 @@
--i https://aiinfra.pkgs.visualstudio.com/PublicPackages/_packaging/ORT-Nightly/pypi/simple/
+# -i https://aiinfra.pkgs.visualstudio.com/PublicPackages/_packaging/ORT-Nightly/pypi/simple/
onnxruntime==1.20.0
\ No newline at end of file