Derive Qwen3-TTS codec conv encoder from config instead of hardcoding widths - #552
Merged
Merged
Conversation
The Mimi-style conv encoder in the Qwen3-TTS codec tokenizer hardcoded its channel widths (1->64->128->256->512->1024->512) and kernels, while the transformer that consumes its output was sized from `codec_encoder.hidden_size`. Nothing tied the two together, so any config whose `hidden_size` was not 512 built a structurally malformed graph: a LayerNormalization with `hidden_size`-wide weights applied to a 512-wide tensor, then `Incompatible MatMul contraction dimensions: 512 vs 32` at the first q_proj. This was invisible because `SymbolicShapeInferencePass` deliberately swallows inference errors, and the tiny test config in `TestBuildCodecGraph` used `hidden_size=32` — so the codec tests were building a broken graph and passing. HF's `MimiEncoder.__init__` derives the whole stack from config, so the hardcoded widths were an inlined specialization rather than a fixed property of the checkpoint. Parameterize it the same way: - Add the missing HF fields to `CodecEncoderConfig` (`audio_channels`, `num_filters`, `num_residual_layers`, `kernel_size`, `last_kernel_size`, `residual_kernel_size`, `compress`, `upsampling_ratios`) and read them in `ArchitectureConfig.from_transformers`. - Build `_MimiConvEncoder` by HF's derivation, keeping the parameterless ELU slots so `encoder.encoder.layers.*` numbering still matches the checkpoint. Also fixes the encoder RVQ widths, which had the same root cause: it projected from `codebook_dim` instead of `hidden_size` and halved the codebook dim. The real checkpoint has `input_proj [256, 512, 1]` and `embed_sum [2048, 256]` with `encoder_config.codebook_dim=256`, so the previous code built `Conv1d(256 -> 128)` on a 512-wide tensor and could not have loaded those weights. The tiny test config now shrinks only widths (`num_filters=4`), keeping the real depth so layer numbering is exercised. New tests assert the default config reproduces the checkpoint's conv weight *names* and shapes, that the tiny config keeps the same names, and that indices track `upsampling_ratios`/`num_residual_layers`. Symbolic-shape-inference failures across the codec tests: 6 before, 0 after (measured by instrumenting `infer_symbolic_shapes`, since the pass swallows them). Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com>
Performance Comparison
|
🏗️ Architecture Diff
No architecture changes detected. ✅ Legend: ⚪ No change · 🔵 Minor (attrs/inits) · 🟡 Moderate (nodes added/removed) · 🔴 Major (interface changed) |
The RVQ width fix is a weight-loading correctness fix: before it, the encoder built Conv1d(256->128) projections while the checkpoint has input_proj [256, 512, 1] and embed_sum [2048, 256], so the real weights could not have loaded at all. Assert the derived projection and codebook shapes for the default config, the same way the conv-stack test guards layer naming. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com>
Contributor
There was a problem hiding this comment.
Pull request overview
This PR makes the Qwen3-TTS codec encoder derive convolutional and RVQ dimensions from configuration, improving checkpoint compatibility and preventing malformed graphs.
Changes:
- Added Mimi encoder configuration fields and extraction.
- Parameterized convolutional layers, kernels, residual blocks, and output widths.
- Corrected RVQ projection dimensions.
- Added regression tests for shapes and layer naming.
Reviewed changes
Copilot reviewed 4 out of 4 changed files in this pull request and generated 3 comments.
| File | Summary and review status |
|---|---|
tests/build_graph_test.py |
Adds regression coverage for derived stack shapes, indices, and output widths. |
src/mobius/models/qwen3_tts_tokenizer.py |
Implements configurable convolution and RVQ dimensions. Moderate follow-ups: assert RVQ projection shapes and handle non-mono audio_channels in CodecTask. |
src/mobius/_configs/_sub_configs.py |
Adds Mimi encoder configuration fields and defaults. |
src/mobius/_configs/_base.py |
Extracts the new fields from Transformers configs. Moderate follow-up: add coverage using non-default nested encoder values, including upsampling_ratios. |
Suppressed comments (1)
src/mobius/_configs/_sub_configs.py:155
audio_channelsis now used to size the first convolution, butCodecTask._build_encoderstill declareswaveformas[batch, 1, audio_len](src/mobius/tasks/_codec.py:91-93). For any non-monoCodecEncoderConfig, the emitted Conv expectsaudio_channelsinputs while the graph supplies one, so this configuration path builds an invalid encoder. Either thread this field into the task input shape or reject values other than 1.
audio_channels: int = 1
💡 Configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
Two review findings, both valid: 1. `audio_channels` sized `encoder.layers.0.conv` but `CodecTask. _build_encoder` always declared `waveform` as `[batch, 1, audio_len]`, so any non-mono config fed a one-channel graph input into a conv expecting more. Thread the configured channel count into the task input shape (falling back to 1 when there is no codec_encoder, as for Mimi), and add a build test asserting the input channel dim follows the config. 2. The new nested-config extraction was untested — the existing tests construct `CodecEncoderConfig` directly, so a wrong getattr key or default would still pass. Add extraction tests covering the default values, non-default values including `upsampling_ratios`, and the fallback path where a config omits the conv fields entirely. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com>
Contributor
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 6 out of 6 changed files in this pull request and generated 1 comment.
Suppressed comments (1)
src/mobius/models/qwen3_tts_tokenizer.py:360
- The new config-driven stack is covered only through
named_parameters()shape/name assertions; none of the added tests executes the exported encoder. Because the original malformed graph was hidden by swallowed symbolic-inference errors, a graph can still build in these tests yet fail at ORT load/run. Add a tinyOnnxModelSessiontest using the generated package (including a non-default ratio/width configuration) to verify that the derived padding, strides, and full encoder path are executable.
for ratio in reversed(upsampling_ratios):
current_scale = scaling * num_filters
for _ in range(num_residual_layers):
layers.append(
_EncoderResBlock(
| codes: (B, 16, T) int64 audio codes. | ||
| """ | ||
| # 1. Conv encoder: (B, 1, samples) -> (B, 512, T') | ||
| # 1. Conv encoder: (B, 1, samples) -> (B, hidden_size, T') |
justinchuby
added a commit
that referenced
this pull request
Aug 23, 2026
…ive contract (#554) Fixes six long-standing schema-conformance failures caused by mobius's legacy onnx-genai emitters drifting from the runtime's published JSON schema. ## Why the drift was invisible The conformance tests searched a couple of hard-coded local onnx-genai checkout paths and `pytest.skip`ped when none was found. CI has no such checkout, so **six tests never ran there**; on the machines that did have one, the result depended on whatever revision that clone happened to sit on. Two upstream contract redesigns accumulated unnoticed. ## The two contract redesigns 1. **`pipeline` is now `PipelineSpec`** — a single property, `workflow` (a typed SSA graph), with `additionalProperties: false`. The legacy emitters produced `{models, dataflow, strategy, phases}`, all of which the schema now rejects. 2. **`speculative` is now `SpeculativeContract`** — requires `{proposer, target, vocabulary, max_proposal_width}` and forbids every field of the old flat MTP block. The flat `SpeculatorConfig` that block was modelled on still exists upstream, but it describes a HuggingFace `config.json` section, not `InferenceMetadata.speculative`. ## What was migrated ### Diffusion — `build_diffusion_pipeline_metadata` Now emits the denoise loop explicitly: the solver, sigma/alpha schedule, timestep table, guidance combine and output clamp are real ONNX components built from mobius's policy library and shipped with the document. It **cannot** delegate to `build_diffusion_workflow_metadata`: the ComfyUI conversion path has no component graphs at all (it deliberately does not build or export them), and that builder derives everything from live `ir.Model`s. So the workflow is built directly, reusing `workflow_metadata`'s `_invoke` / `_publish_workflow_v1` and `mobius.generation`'s solver builders. The schedule is derived from the scheduler's own betas through the same diffusers-compatible helpers the package exporter uses, so a ComfyUI conversion and a package export of the same checkpoint describe the same dynamics. img2img's `start_step` lowers to a sliced schedule; the VAE `scaling_factor` and the sigma-space initial-state scale are emitted as explicit components. Three cases now **fail closed** instead of being silently mis-described: - an ancestral sampler (mobius ships no stochastic solver), - Karras/exponential sigma spacing (the workflow ships the sigma table as a constant, so a `use_karras_sigmas` hint is no longer sufficient), - a latent-only graph with no VAE decode. ### MTP — `write_mtp_speculator_metadata` Emits a `SpeculativeContract` anchored to the backbone's workflow: it registers the head as a workflow component, names the target by its declared `logits` port role, states the hidden handoff as `port_bindings.target_hidden_context` plus a `hidden_states` role on the target output, and completes the rollback capacity its own claim requires. It is declared `block`, not `chained`: a chained proposer must expose a `logits_output` carrying the next-token distribution, and this sidecar emits only `mtp_hidden` — the runtime decodes it through the target's shared LM head, which is why that initializer is listed in `shared_weights`. ### VLM `build_native_vlm_package_metadata` produces mobius's **internal** structural descriptor, not a publishable document — `build_vlm_workflow_metadata` already consumes it and republishes only `preprocessing` under a real `pipeline.workflow`. The actual bug was that `write_native_vlm_package_metadata` wrote that descriptor to `inference_metadata.yaml`. It now writes the workflow document, and the tests validate the published document instead of the descriptor. ### CI visibility The upstream schema is vendored under `src/mobius/integrations/onnx_genai/_schema/` and is the default, so conformance never skips and drift becomes a test failure. A local checkout is no longer consulted implicitly — one that is ahead of or behind `main` reintroduces exactly the machine-dependent result that hid this bug. Set `ONNX_GENAI_SCHEMA` to validate against a specific revision. Three further conformance tests (codec, speech-to-text, duplex workflows) were skipping for the same reason and now run. ## Tests Migrated in the same change rather than loosened: - `TestMtpSpeculatorMetadata::test_exact_schema_keys_and_values` asserts the new contract; `test_no_legacy_field_names` is repointed at the now-legacy field names. - A new MTP test anchors against a workflow `write_decoder_workflow_metadata` actually emits (which ships ~11 generated policy components), so target selection can't regress to "the only ONNX component". - New tests pin that the solver schedule comes from the scheduler's betas rather than a placeholder ramp. - Diffusion/ComfyUI tests assert workflow structure instead of the removed `strategy` block; scheduler facts that were duplicated into the document are asserted on the parsed `ComfyUIWorkflow` where they actually live. ## Verification Rebased onto current `main` (includes #552 and #553). ``` python -m pytest tests/build_graph_test.py tests/cli_test.py src/ -q \ -k "not phi4mm and not apply_weights_unknown" -n auto ``` | | before (`main` at branch point) | after | |---|---|---| | full suite | 8 failed / 4613 passed | **0 failed / 4641 passed** | | `src/mobius/integrations/onnx_genai/` with `ONNX_GENAI_SCHEMA` at upstream `main` | 7 failed / 240 passed | **258 passed / 1 skipped** | Six of those eight failures are the schema-drift ones fixed here. The other two — `qwen_image_test.py::test_deterministic_l4_l5_image_edit_golden` and `TestNativeVlmPackageMetadata::test_cached_gemma_processor_matches_emitted_patch_budget` — were fixed by #553, which is now merged; this branch preserves both of its changes through the rebase. `lintrunner` is clean. ## Sequencing note `workflow_metadata.py` and `auto_export.py` are **not modified** (both are read from only), so this does not conflict with #551. --------- Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com> Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>
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.
The bug
_MimiConvEncoderinsrc/mobius/models/qwen3_tts_tokenizer.pytook no config at all and hardcoded its channel widths and kernels (1→64→128→256→512→1024→512). Its output fedCodecEncoderTransformerModel(hidden_size=config.codec_encoder.hidden_size), which has no input projection — so the conv output width andhidden_sizehad to agree, and nothing checked that they did.When they disagreed, the builder happily emitted a malformed graph:
Why it was invisible
SymbolicShapeInferencePasscatches inference errors and logs"Symbolic shape inference failed (upstream bug); skipping"— a deliberate generic guard. The tiny config inTestBuildCodecGraph._codec_config()setcodec_encoder.hidden_size=32, so the codec tests were building a broken graph and passing.The fix
The hardcoded widths were not dictated by the checkpoint — HF's
MimiEncoder.__init__derives the entire stack from config, and mobius's literals were just that derivation inlined withnum_filters=64,upsampling_ratios=[8,6,5,4],num_residual_layers=1,kernel_size=7,last_kernel_size=3. So rather than adding a validation guard, this parameterizes the stack and eliminates the mismatch class:CodecEncoderConfig: added the missing HF fields (audio_channels,num_filters,num_residual_layers,kernel_size,last_kernel_size,residual_kernel_size,compress,upsampling_ratios) with HF's own defaults, extracted inArchitectureConfig.from_transformers._MimiConvEncoder: builds by HF's derivation. The parameterless ELU slots are preserved soencoder.encoder.layers.*numbering still matches the checkpoint exactly.kernel = ratio * 2and the trailing conv's out-channels ishidden_size, both derived._EncoderResBlock: inner width is nowdim // compresswith a configurable kernel.Also fixed: the encoder RVQ could not load the real checkpoint
This one is a weight-loading correctness fix, not just a shape-consistency fix, which is why it belongs here despite being adjacent to the original scope.
_EncoderSplitRVQusedinput_dim=codebook_dimanddim=codebook_dim // 2. The real checkpoint hasencoder.quantizer.*.input_proj.weight [256, 512, 1]andcodebook.embed_sum [2048, 256], withencoder_config.codebook_dim=256. So the old code builtConv1d(256→128)projections and 128-wide codebooks: the encoder could never have loaded the real checkpoint's weights, independently of the malformed-graph issue. It only happened to be self-consistent on theenc is Nonefallback path, wherecodebook_dimdefaulted to 512.Now it projects
hidden_size → codebook_dim, matching HFMimiResidualVectorQuantizer.The encoder/decoder asymmetry is intentional — please don't unify these paths
The two sides interpret
codebook_dimdifferently, and both are correct:input_projcodebook_dim=256(Mimi semantics: the codebook dim itself)hidden_size=512[256, 512, 1]codebook_dim=512(the RVQ input dim)512 // 2 = 256[256, 512, 1]So the decoder's
// 2convention inSplitResidualVectorQuantizeris right and is left untouched; only the encoder path was misreading its field. A future "cleanup" that unifies the two would reintroduce this bug.Verification
Symbolic-shape-inference failures across
TestBuildCodecGraph(measured by temporarily instrumentinginfer_symbolic_shapes, since the pass swallows them — instrumentation not committed):Default config reproduces the real checkpoint exactly — verified against the safetensors header of
Qwen/Qwen3-TTS-Tokenizer-12Hz: all 14 conv-stack weights match by name and shape, including the ResBlocklayers.N.block.1/.3names and the0,1,3,4,6,7,9,10,12,14numbering, as do both RVQinput_proj/output_projpairs and the codebooks. This is asserted by tests, so indexing or width drift (which would otherwise surface far away at weight-load time) fails loudly.Tiny test config keeps the real depth and only shrinks widths (
num_filters=4,hidden_size=32→1→4→8→16→32→64→32), so it exercises the same weight-name structure the real model depends on. Codec test runtime: 2.05s (4 tests) → ~2.1s (9 tests).Suite:
pytest tests/build_graph_test.py tests/cli_test.py src/ -k "not phi4mm and not apply_weights_unknown" -n auto→ 8 failed, 4613 passed — the same 8 pre-existing failures onmain(onnx_genai jsonschema drift + one qwen_image golden).lintrunnerclean.New tests
test_default_config_conv_stack_matches_checkpoint— asserts derived parameter names and shapes against the real checkpoint.test_tiny_config_keeps_checkpoint_layer_names— tiny config has identical naming, narrower tensors.test_conv_stack_indices_follow_config— 2 ratios → final conv atlayers.8; 4 ratios × 2 residual layers →layers.18.test_hidden_size_drives_conv_output_width— direct regression guard for the original bug.test_default_config_rvq_projections_match_checkpoint— assertsinput_proj (256, 512, 1),output_proj (512, 256, 1)and(2048, 256)codebooks, guarding the weight-loading fix.Decoder conv stack
Checked: the codec decoder does not have the mirrored conv problem.
Qwen3TTSCodecDecoderModelalready derives everything from config (latent_dim,decoder_dim,upsample_rates,upsampling_ratios), andCodecDecoderTransformerModelhas aninput_proj/output_projpair that adaptslatent_dim ↔ hidden_size, so no width can be stranded. No change needed there.Scope
Limited to the codec tokenizer model, its config, and its tests. Does not touch
reuse.py,_ssm.py,_mamba_block.py, the onnx_genai integration (covered by #551), orSymbolicShapeInferencePass.