fix(onnx): export vocabulary-quantized models with their token mapping and weights - #364
Merged
stephantul merged 2 commits intoSep 12, 2026
Conversation
Codecov Report✅ All modified and coverable lines are covered by tests.
... and 2 files with indirect coverage changes 🚀 New features to boost your workflow:
|
stephantul
approved these changes
Sep 11, 2026
Contributor
|
@serhiizghama thanks, very cool! |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Exporting a vocabulary-quantized model to ONNX gives a graph that can't actually run.
TorchStaticModelbuilds its embedding table frommodel.embedding, which for a quantized model has one row per cluster, and then gathers it with the raw token ids — so any id at or above the cluster count blows up at inference (Gather ... indices element out of data boundsfrom onnxruntime,IndexErrorfrom the torch module itself), andweightswere never applied at all. The export itself succeeds because the dummy inputs are all zeros.StaticModel._encode_helperlooks ids up throughtoken_mappingfirst and multiplies byweights; the exporter now does the same. Both are kept as buffers, so the exported graph stays cluster-sized instead of materializing a full vocab × dim table.Reproduced with
quantize_model(model, vocabulary_quantization=3)and with a classifier trained on such a model and exported viato_pipeline()— both go through the same module. Added a test for each; they fail on main with the out-of-bounds gather and matchencode/predict_probawith the fix.