Skip to content

fix(vibevoice): preserve quantization metadata in sanitize() for quantized model loading#604

Merged
lucasnewman merged 1 commit into
Blaizzy:mainfrom
korale77:fix/vibevoice-quantized-loading
Mar 25, 2026
Merged

fix(vibevoice): preserve quantization metadata in sanitize() for quantized model loading#604
lucasnewman merged 1 commit into
Blaizzy:mainfrom
korale77:fix/vibevoice-quantized-loading

Conversation

@korale77

Copy link
Copy Markdown
Contributor

Context

Loading quantized VibeVoice TTS models (8/6/5/4-bit from mlx-community) fails with:

ValueError: Expected shape (151936, 896) but received shape (151936, 224)
for parameter language_model.embed_tokens.weight

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() checks f"{p}.scales" in weights to 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 .scales and .biases keys through sanitization, matching the existing pattern in chatterbox/s3gen.

Changes in the codebase

  • mlx_audio/tts/models/vibevoice/vibevoice.py: In sanitize(), preserve quantization metadata keys (.scales, .biases) that aren't in curr_shapes, instead of dropping them. This allows apply_quantization() to find the metadata and properly quantize the model before weight loading.
  • mlx_audio/tts/tests/test_models.py: Add test_sanitize_preserves_quantization_metadata to verify .scales and .biases keys survive sanitization.

Changes outside the codebase

None.

Additional information

Checklist

  • Tests added/updated
  • Documentation updated
  • Issue referenced (e.g., "Closes #...")

…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).

@lucasnewman lucasnewman left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks!

@lucasnewman
lucasnewman merged commit 1c936dd into Blaizzy:main Mar 25, 2026
10 checks passed
@korale77
korale77 deleted the fix/vibevoice-quantized-loading branch March 26, 2026 02:06
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants