Cast vision/audio inputs f32 to model dtype for GenAI compatibility - #265
Conversation
🏗️ Architecture Diff
No architecture changes detected. ✅ Legend: ⚪ No change · 🔵 Minor (attrs/inits) · 🟡 Moderate (nodes added/removed) · 🔴 Major (interface changed) |
Performance Comparison
|
Codecov Report❌ Patch coverage is 📢 Thoughts on this report? Let us know! |
titaiwangms
left a comment
titaiwangms
left a comment
There was a problem hiding this comment.
Review summary
Reviewed by Opus 4.7 + GPT-5.3-Codex (code) + GPT-5.5 (adversarial) + Claude Sonnet 4.6 (readability). Consensus: the direction is right, but the fix is incomplete and may not actually produce a loadable model in BF16/F16 multimodal builds.
The PR keeps selected encoder parameters at float32, but does not change the encoder graph's I/O dtype or the downstream embedding model's input dtype. That's three independent problems.
🔴 Major — Encoder graph I/O still uses config.dtype, mixing types with f32 weights
The tasks declare encoder graph inputs with dtype=config.dtype (e.g. _vision_language_3model.py:86-91, _phi4mm_multimodal.py:122-125, _gemma4.py:301-305). After this PR, in a bf16 build:
pixel_values/input_featuresgraph input → bf16- First Conv/MatMul/Gemm weight in the encoder → float32
ONNX numeric op type constraints generally require matching tensor types. The likely failure mode is an invalid graph or ORT load-time type-check failure — not a safe fallback to f32 compute.
Suggested direction: make precision a per-component contract, not a parameter-only exclusion. Encoder graph inputs, constants, and outputs should all be FLOAT, with an explicit Cast at the next-component boundary.
🔴 Major — Embedding model still expects bf16/f16 multimodal features
Even if the encoder is fully corrected to float32, ORT GenAI passes its outputs as image_features / audio_features into a separate embedding session whose input is declared at config.dtype (_phi4mm_multimodal.py:166-174, _gemma4.py:394-407). genai_config.json (integrations/ort_genai/genai_config.py:279-296) only wires names; it does not insert dtype conversions.
Suggested fix: either (a) declare embedding *_features inputs as FLOAT and Cast to decoder dtype inside the embedding graph before token-merge, or (b) keep embedding in float32 and Cast inputs_embeds at the decoder boundary.
🔴 Major — Hardcoded ("vision_encoder", "audio_encoder") misses real models
The exclude list only matches literal nn.Module attribute names. Several existing multimodal/speech models bypass the fix:
| Model | Module attr | Excluded? |
|---|---|---|
| Gemma4 / Qwen-VL / Gemma3 / MLLaMA / LLaVA / DeepSeek-OCR2 / InternVL / BLIP2 | vision_encoder, audio_encoder |
✅ |
| Phi4-MM | speech_encoder (models/phi.py:981, remapped via tasks/_phi4mm_multimodal.py:59) |
❌ |
| Qwen3-ASR | audio_tower (models/qwen3_asr.py:348, remapped via tasks/_speech_language.py:51) |
❌ |
| FunASR | audio_tower (models/fun_asr.py:511, remapped via tasks/_fun_asr_speech_language.py:54) |
❌ |
Suggested fix: derive exclusions from the resolved task's component-attribute mapping (or have each task declare a component_dtypes / encoder-attrs policy) instead of guessing module layout.
🟡 Minor — No tests covering the new behavior
Existing _cast_module_dtype tests (_exporter_test.py:181) only exercise full-module casting. Please add tests asserting (per affected model class) that encoder initializer dtypes remain FLOAT while decoder initializers are BF16/F16. Phi4-MM, Qwen3-ASR, and FunASR are the highest-value cases since they would have silently regressed.
🟡 Minor — Silent precision-policy override
Passing --dtype bf16 previously implied uniform bf16. After this PR it silently means "mixed precision". Worth at least a logger.info(...) listing which sub-modules were kept at float32, and ideally a named policy flag (keep_multimodal_encoders_fp32=True) exposed in the public API + docs.
🟢 Minor / Nit — Readability
- Hardcoded tuple
("vision_encoder", "audio_encoder")duplicates information already in_MODEL_ROLE_MAP(lines 111–115). Either co-locate as_ENCODER_SUBMODULE_ATTRS = frozenset(...)or derive from the role map. exclude=exclude_from_cast or Noneis a no-op:_cast_module_dtypealready guards withif exclude:(truthy check passes bothNoneandset()). Drop theor None, or change the signature toexclude: set[str] = frozenset().- Parameter name
excludeis weaker than the call-siteexclude_from_cast. Considerexcluded_attrs. - Docstring on
excludemixes contract + specific use case. Keep the rationale at the call site. excludesilently ignores typos in attr names — consider validating or warning.
✅ What's right
- Casting before graph construction (vs post-export mutation) is the right architectural layer.
id(param)set is the correct way to identify shared parameters under weight tying.keyword-only(*) marker on the newexcludearg is good API hygiene.- Inline comment at the call site cites the concrete ORT op gap (OneHot/Mul/ReduceMean) — exactly the "why" the project style guide asks for.
- Confirmed (
_optimizations.py:285-291): GQA / QKV-packing fusions are gated tomodel_role == "decoder", so passing bf16 tooptimize_model()for the f32 vision graph does not trigger decoder-only fusions on it.
Recommendation
Request changes. The two BF16 graph-boundary issues (encoder I/O dtype, embedding-input dtype) are likely to make exported BF16 multimodal models fail to load in ORT, which would silently invalidate the PR's "BF16 export succeeds" claim once the embedding/vision sessions are actually run end-to-end. Worth verifying with an ORT GenAI smoke test on Gemma4-bf16 (load all three sub-models in one session pipeline) before merging.
GenAI's image/audio processors output f32 data. Instead of keeping entire vision/audio encoders at f32 (wasteful), keep only the graph input at f32 and add a Cast to the model dtype at the graph entry. Vision/audio weights use the requested f16/bf16 for memory efficiency. Reverts the exclude-based approach in favor of Cast-at-input: - _cast_module_dtype: simplified back to original (no exclude param) - Gemma4 task: vision/audio inputs always f32 + Cast to config.dtype - Result: f32 input → Cast → f16/bf16 weights → f16/bf16 compute 13 gemma4 tests pass. Signed-off-by: Justin Chu <justinchu@microsoft.com>
daebaf3 to
10ccf59
Compare
Address review feedback on PR #265: extend the f32 encoder input pattern (FLOAT input + Cast at graph entry) to all multimodal task files, not just Gemma4. Affected tasks: - VisionLanguageTask (base _build_vision for LLaVA, Aya Vision, etc.) - QwenVLTask (Qwen2.5-VL, Qwen3-VL, Qwen3.5-VL) - PixtralVLTask (Pixtral) - Phi4MMMultiModalTask (vision + speech encoders) - SpeechLanguageTask (Qwen3-ASR) - FunASRSpeechLanguageTask (Fun-ASR) Encoder graph inputs (pixel_values, input_features, audio_embeds) are declared as FLOAT because ORT GenAI's image/audio processors output f32. A Cast at graph entry converts to the target dtype for the encoder's internal computation with cast weights. Add 4 new tests verifying encoder inputs stay FLOAT in f16/bf16 builds for both generic VL (LLaVA) and Gemma4 (vision + audio). Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <justinchu@microsoft.com>
|
Addressed the review feedback in commit 76e37fa: What was fixedExtended the f32 encoder input pattern (FLOAT input + Cast at graph entry) to all multimodal task files:
Embedding boundary dtypeEmbedding model
Tests added4 new tests verify encoder inputs stay FLOAT in f16/bf16 builds:
Test results
|
justinchuby
changed the title
|
@titaiwangms PTAL |
|
Merged LMK if follow ups needed |
…ill (#271) Update the multimodal-models skill to document why vision/audio encoder graphs accept f32 inputs and Cast to model dtype at graph entry. ## What this adds New section in `.agents/skills/multimodal-models/SKILL.md`: - **Why:** ORT GenAI image/audio processors always output f32, regardless of model dtype - **How:** Encoder graph adds `Cast(f32 → model_dtype)` at entry; weights remain in f16/bf16 - **When:** Handled automatically by mobius with `--runtime ort-genai` - **Error without it:** `Type Error: Type parameter (T) bound to different types` This documents the pattern introduced in PR #265. --------- Signed-off-by: Justin Chu <justinchu@microsoft.com> Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
2 participants
Summary
Keep vision/audio encoder graph inputs as float32 (matching GenAI image/audio processor output) and add a Cast(f32→model_dtype) at the graph entry. Encoder weights use the requested f16/bf16 dtype for memory efficiency.
Problem
GenAI multimodal processor outputs float32 pixel values and audio features. When building with
--dtype f16or--dtype bf16, the vision/audio encoder graph inputs were also cast to f16/bf16, creating a type mismatch: GenAI sends f32, encoder expects f16, ORT produces all-zero output (silent failure).Solution: Cast at graph entry
Instead of keeping entire vision/audio encoders at f32 (wasteful), add a
Castop at the start of each encoder graph:This gives both GenAI compatibility (f32 input) and memory efficiency (f16/bf16 weights).
Changes
_gemma4.py: Vision/audio encoder inputs alwaysir.DataType.FLOAT, withop.Casttoconfig.dtypewhen dtype ≠ f32_builder.py: No changes needed — all parameters cast to requested dtype as beforeTesting
13 gemma4 tests pass. Verified f16 build: vision input=f32, first node=Cast, weights=f16.