Add GLM-ASR Nano speech-language support - #509
Conversation
Performance Comparison
|
🏗️ Architecture Diff
No architecture changes detected. ✅ Legend: ⚪ No change · 🔵 Minor (attrs/inits) · 🟡 Moderate (nodes added/removed) · 🔴 Major (interface changed) |
There was a problem hiding this comment.
Pull request overview
Adds first-class support for the GLM-ASR-Nano-2512 speech-language architecture in Mobius, exporting a standardized 3-model ONNX package (audio_encoder → embedding → cached decoder) and extending runtime metadata/export flows to support pinned Hugging Face revisions and audio processor artifacts.
Changes:
- Introduces the
glmasrmodel implementation (partial-RoPE audio tower + 4-frame projector, audio-token embedding replacement, and cached Llama-like decoder) plus task wiring and registry/config extraction. - Extends ORT GenAI + ONNX GenAI export/metadata to include an explicit embedding stage and to propagate pinned HF revisions to config/tokenizer/processor/runtime asset downloads.
- Adds graph/unit tests plus end-to-end golden cases (L4 prefill + L5 cached generation) and schema support for strict
exact_matchgeneration assertions.
Reviewed changes
Copilot reviewed 32 out of 33 changed files in this pull request and generated 2 comments.
Show a summary per file
| File | Description |
|---|---|
| tests/e2e_golden_test.py | Updates speech-language golden runner for GLM-ASR preprocessing, mask shapes, audio feature handling, and optional exact-match transcript assertions. |
| tests/cli_test.py | Adds coverage ensuring --revision is forwarded to detection and build. |
| tests/build_graph_test.py | Adds graph-contract + pipeline smoke tests for GLM-ASR three-stage split and CUDA attention op preservation. |
| tests/_test_configs.py | Adds tiny override config for glmasr test cases. |
| testdata/golden/speech/glm-asr-nano.json | Adds pinned golden logits + prompt input IDs for L4 prefill. |
| testdata/golden/speech/glm-asr-nano_generation.json | Adds pinned golden generated tokens + transcript for L5 generation. |
| testdata/cases/speech/glm-asr-nano.yaml | Adds L4+L5 speech-language case definition for GLM-ASR-Nano with exact_match. |
| testdata/cases/schema.json | Extends generation schema with exact_match boolean flag. |
| src/mobius/tasks/_glmasr_speech_language.py | Implements GLM-ASR three-model export task and audio encoder graph boundary casting/IO contract. |
| src/mobius/tasks/init.py | Exports GlmAsrSpeechLanguageTask and registers the task name. |
| src/mobius/models/glm_asr.py | Adds GLM-ASR model implementation (audio tower + projector, embedding mixer, cached decoder) and weight routing. |
| src/mobius/models/glm_asr_test.py | Adds synthetic stage-by-stage parity tests vs HF for audio → embedding → decoder. |
| src/mobius/models/init.py | Exports GlmAsrForConditionalGeneration from the models package. |
| src/mobius/integrations/ort_genai/genai_config.py | Adds an explicit with_embedding() generator path for multimodal embedding stages. |
| src/mobius/integrations/ort_genai/auto_export.py | Adds GLM-ASR audio processor emission, embedding stage metadata, and pinned-revision propagation for hub assets. |
| src/mobius/integrations/ort_genai/auto_export_test.py | Adds tests for GLM-ASR processor/genai config emission and pinned-revision forwarding. |
| src/mobius/integrations/onnx_genai/inference_metadata.py | Adds optional revision propagation for processor/runtime asset loading and diffusers scheduler config fetch. |
| src/mobius/integrations/onnx_genai/auto_export.py | Propagates revision to tokenizer/audio processor emission and diffusers scheduler loading. |
| src/mobius/integrations/onnx_genai/auto_export_test.py | Adds regression coverage for pinned-revision propagation to audio processor emission. |
| src/mobius/_registry.py | Registers glmasr model type, default task mapping, and default example model ID. |
| src/mobius/_diffusers_builder.py | Propagates pinned revision for optional diffusers metadata downloads. |
| src/mobius/_diffusers_builder_test.py | Adds unit test ensuring optional diffusers metadata downloads use pinned revision. |
| src/mobius/_configs/per_model/_glmasr_audio.py | Adds GLM-ASR audio_config extraction hook for nested composite configs. |
| src/mobius/_configs/per_model/init.py | Registers the GLM-ASR audio config hook. |
| src/mobius/_configs/_sub_configs.py | Extends AudioConfig with additional encoder-specific fields needed by GLM-ASR. |
| src/mobius/_configs/_base.py | Introduces GlmAsrConfig with nested HF config unwrapping and dtype/audio_token_id handling. |
| src/mobius/_configs/init.py | Exports GlmAsrConfig. |
| src/mobius/_configs_test.py | Adds config extraction unit test for GLM-ASR nested audio/text config behavior. |
| src/mobius/main.py | Adds CLI --revision and propagates it into diffusers detection/build and runtime metadata generation. |
| scripts/generate_golden.py | Ensures hub calls use pinned revisions and stores reference logits as float32 for bf16 compatibility; adds GLM-ASR loading/input prep path. |
| README.md | Lists GLM-ASR under supported Speech-to-Text architectures. |
| docs/cli_reference.md | Documents the new --revision CLI option. |
Suppressed comments (1)
src/mobius/models/glm_asr.py:251
- GlmAsrMultiModalProjector hard-codes GELU and never consults GlmAsrConfig.projector_hidden_act. Since the config extracts and stores this field from HF (projector_hidden_act), ignoring it can produce an incorrect projector for non-default configs. Consider using get_activation(projector_hidden_act) to keep the ONNX graph faithful to the checkpoint config.
merged = op.Reshape(
audio_hidden_states,
op.Constant(value_ints=[0, -1, self._projector_input_size]),
)
return self.linear_2(op, op.Gelu(self.linear_1(op, merged)))
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
Implement the GLM-ASR audio encoder, projector, embedding mixer, and cached Llama decoder with pinned real-audio L4/L5 coverage. Add revision-pinned CLI/runtime metadata export and exact CUDA transcript validation. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <justinchu@microsoft.com>
Honor configured audio and projector activations and emit single-output encoder Attention nodes. Release staged ORT sessions so large CUDA components do not retain device weights, and reject unsupported ORT GenAI metadata rather than producing an unloadable package. Clean up rebase regressions in Diffusers and CLI tests. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <justinchu@microsoft.com>
9b538cf to
64c6e2d
Compare
|
@copilot please update from main |
Resolve conflicts where main independently added pinned HuggingFace revision plumbing to the CLI, ORT-GenAI artifact export, and the staged ORT session teardown. Keep main's revision kwarg handling, drop the duplicated --revision argument and duplicated build() kwarg introduced by the merge, and retain main's upstream asset patching. Signed-off-by: GitHub <noreply@github.com> Co-authored-by: justinchuby <11205048+justinchuby@users.noreply.github.com>
|
|
Merged Conflicts were where
Post-merge validation: |
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 32 out of 33 changed files in this pull request and generated no new comments.
Suppressed comments (2)
tests/build_graph_test.py:3170
- The assertion intended to ensure the audio encoder’s opset-24 Attention nodes don’t request any optional inputs (mask / past KV) is incomplete:
len(node.inputs) == 3 or node.inputs[3] is Nonedoesn’t check the remaining optional inputs (e.g., past key/value), so a regression could slip through. Consider asserting that all inputs after the first 3 areNone(e.g.,all(inp is None for inp in node.inputs[3:])).
src/mobius/integrations/ort_genai/auto_export_test.py:858 - Test name is misleading:
test_audio_glmasr_writes_whisper_feature_extraction_jsonasserts that_write_audio_processor_config()writesaudio_processor.json(notfeature_extraction.json). Renaming the test (and/or its docstring) would better reflect the file it validates and avoid confusion with the Gemma4audio_feature_extraction.jsonpath.
def test_audio_glmasr_writes_whisper_feature_extraction_json(self, tmp_path):
config = mock.MagicMock()
config.audio = mock.MagicMock()
config.model_type = "glmasr"
path = _write_audio_processor_config(config, str(tmp_path))
assert path is not None
assert path.endswith("audio_processor.json")
Rebasing 141 commits onto `79b48bc0` conflicted in the five files where this branch's workflow migration overlaps the two new model integrations. The conflicts were resolved in favour of this branch, which is right for the metadata architecture and wrong for everything else main had just added, so the endpoint was then checked against main hunk by hunk rather than trusted. Every `def`/`class` main added was confirmed present, which surfaced what the resolution had silently dropped. `--revision` was threaded through the asset writers by #509 and did not survive: `_write_clip_tokenizer`, `_write_hf_tokenizer` and `_copy_runtime_assets` lost the keyword from their signatures while their bodies still referenced it, and `write_onnx_genai_config` lost the parameter entirely. The last one is the dangerous shape — it has `**kwargs`, so a pinned revision would have been absorbed and ignored rather than raising, and the package would still build and still validate while its tokenizer came from whatever the branch tip happened to be. That is precisely the failure a pin exists to prevent. Restored at all nine sites, plus `_write_text_runtime_assets`, which is this branch's own helper and needed the same parameter to pass it on. Also restored: audio-processor emission for speech-language packages, which #509 added to the multimodal dispatch and this branch's rewrite of that branch did not carry. Expressed as `_has_audio_encoder` rather than main's inline `"audio_encoder" in pkg` so a package object without `keys()` is a False rather than a TypeError. Two of main's tests could not be taken verbatim. `test_revision_is_forwarded_to_detection_and_build` was overwritten wholesale by a test of this branch's and is restored byte-identically alongside it. `test_dispatch_audio_only_multimodal_pipeline` asserts on `pipeline.models`, the legacy composite ABI this branch replaces, so its still-relevant half — that the revision reaches the feature extractor — is covered by a new test against the workflow instead. `test_runtime_onnx_genai_routes_vlm_through_workflow_emitter` pins the writer call exactly and now expects the threaded `revision`. Verified: 4580 passed (up from 4541; the increase is main's new tests now running), 11/11 `validate_metadata` and 11/11 runtime conformance against ONNX GenAI `6e2ddc78`, lintrunner clean. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: justinchuby <justinchuby@users.noreply.github.com>
Summary
glmasrsupport forzai-org/GLM-ASR-Nano-2512as standardizedaudio_encoder→embedding→ cacheddecoderONNX models.--revisionpropagation, and documentation metadata.Immutable inputs
zai-org/GLM-ASR-Nano-2512@61ba4e0b3309b6656edea3e93e419f7bd5c61957eustlb/audio-samples@1045238c3afdc0254472ecdde5dae95370d621e4Architecture checklist
audio_encoder,embedding, anddecoderpackage keys[B,128,3000], 16 kHz, FFT 400, hop 1605120 → 4096 → 2048projectorValidation evidence
L1–L3 and regression gates
L4 CUDA prefill (rebased head)
[1,153,59264], correlation0.9991001, mean absolute difference0.02795, maximum absolute difference7.456; greedy token matches.L5 CUDA cached generation (rebased head)
Yesterday it was thirty five degrees in Barcelona, but today the temperature will go down to minus twenty degrees.CLI export and CUDA execution
(138, 2048)(1, 153, 2048)(1, 153, 59264)89BFCArenainitialization OOM (242745344-byte decoder allocation) and returns CUDA memory to baseline between stages.Formatting and review
lintrunner -a: passedWaivers and runtime constraints
rtol=2e-4,atol=2e-4for decoder logits and tighter tolerances for earlier stages. A full real FP32 CUDA checkpoint run is not practical on the available 8 GB RTX A1000.No corresponding Numpy type for Tensor Type. bfloat16), so executable L4/L5 uses FP16; BF16 output validation requires I/O binding or a runtime with BF16 host materialization.GroupQueryAttentionchanged greedy termination from the exact 23-token HF sequence to an earlier valid EOS at 21 tokens. GLM-ASR intentionally preserves standard opset-24Attentionfor exact transcript parity; graph coverage locks this dispatch decision.gemma4andphi4mm; loadingtype: glmasrfails as unsupported. Mobius now rejects--runtime ort-genaibefore saving rather than emitting unloadable metadata. Direct ONNX and ONNX GenAI metadata remain supported; native ORT GenAI packaging is deferred until the runtime registers a compatible GLM-ASR pipeline.NonZero, integerCumSum, and placeholder indexing). Eliminating them would require changing the feature-stripping/embedding contract.