fix(vibevoice): preserve quantization metadata in sanitize() for quantized model loading#604
Merged
lucasnewman merged 1 commit intoMar 25, 2026
Conversation
…tized model loading sanitize() drops weight keys not found in the model's current parameter shapes. Since the model isn't quantized yet at sanitize time, quantization metadata keys (.scales, .biases) are silently removed. Later, apply_quantization() checks for these keys to decide which layers to quantize -- finds nothing -- skips quantization -- and loading fails with a shape mismatch. Preserve .scales and .biases keys through sanitization, matching the existing pattern in chatterbox/s3gen. Same class of bug as Blaizzy#584 (fish_qwen3_omni sanitize fix).
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.
Context
Loading quantized VibeVoice TTS models (8/6/5/4-bit from mlx-community) fails with:
The fp16 variant loads fine. This is the same class of bug fixed for fish_qwen3_omni in #584.
Description
sanitize()filters out weight keys not found in the model's current parameter shapes. At sanitize time the model hasn't been quantized yet, so quantization metadata keys (.scales,.biases) don't exist in the model and are silently dropped.Later,
apply_quantization()checksf"{p}.scales" in weightsto decide which layers to quantize. Since sanitize removed the scales, it skips quantization entirely. The model stays full-size but the weight file has quantized shapes, causing the mismatch.The fix preserves
.scalesand.biaseskeys through sanitization, matching the existing pattern inchatterbox/s3gen.Changes in the codebase
mlx_audio/tts/models/vibevoice/vibevoice.py: Insanitize(), preserve quantization metadata keys (.scales,.biases) that aren't incurr_shapes, instead of dropping them. This allowsapply_quantization()to find the metadata and properly quantize the model before weight loading.mlx_audio/tts/tests/test_models.py: Addtest_sanitize_preserves_quantization_metadatato verify.scalesand.biaseskeys survive sanitization.Changes outside the codebase
None.
Additional information
.endswith((".scales", ".biases"))pattern matches the existing approach inchatterbox/s3gen/s3gen.py.mlx-community/VibeVoice-Realtime-0.5B-8bit.Checklist