diff --git a/src/mobius/_builder.py b/src/mobius/_builder.py index 8c3f0830..3db1dd34 100644 --- a/src/mobius/_builder.py +++ b/src/mobius/_builder.py @@ -191,6 +191,9 @@ def forward(self, op, input_ids, attention_mask, if hasattr(config, "validate"): config.validate() dtype = getattr(config, "dtype", ir.DataType.FLOAT) + # Cast all parameters to the target dtype. Vision/audio encoder weights + # are included — their graph inputs are kept at f32 (matching GenAI's + # image processor output) with a Cast at the graph entry. _cast_module_dtype(module, dtype) resolved_task = get_task(task) capabilities = ep_registry.require(execution_provider) diff --git a/src/mobius/tasks/_fun_asr_speech_language.py b/src/mobius/tasks/_fun_asr_speech_language.py index a28b25be..7b9d3b72 100644 --- a/src/mobius/tasks/_fun_asr_speech_language.py +++ b/src/mobius/tasks/_fun_asr_speech_language.py @@ -91,14 +91,19 @@ def _build_audio_encoder( input_dim = (config.audio.input_size if config.audio else None) or 560 graph, builder = _make_graph(name="audio_encoder") + op = builder.op + # Audio encoder input is always f32 (matching audio processor output). + # Cast at graph entry for f16/bf16 builds. input_features = builder.input( "input_features", - dtype=config.dtype, + dtype=ir.DataType.FLOAT, shape=[batch, seq_len, input_dim], ) + if config.dtype and config.dtype != ir.DataType.FLOAT: + input_features = op.Cast(input_features, to=config.dtype) - audio_features = audio_encoder(builder.op, input_features) + audio_features = audio_encoder(op, input_features) builder.add_output(audio_features, "audio_features") return _make_model(graph) diff --git a/src/mobius/tasks/_gemma4.py b/src/mobius/tasks/_gemma4.py index c21cb972..4e18546c 100644 --- a/src/mobius/tasks/_gemma4.py +++ b/src/mobius/tasks/_gemma4.py @@ -298,11 +298,16 @@ def _build_vision( graph, builder = _make_graph(name="vision_encoder") op = builder.op + # Vision encoder input is always f32 (matching GenAI's image processor + # output). When the model uses f16/bf16, add a Cast at the graph entry + # so weights can stay at the requested dtype for memory efficiency. pixel_values = builder.input( "pixel_values", - dtype=config.dtype, + dtype=ir.DataType.FLOAT, shape=[batch, num_patches, pixel_dim], ) + if config.dtype and config.dtype != ir.DataType.FLOAT: + pixel_values = op.Cast(pixel_values, to=config.dtype) pixel_position_ids = builder.input( "pixel_position_ids", dtype=ir.DataType.INT64, @@ -352,9 +357,11 @@ def _build_audio( input_features = builder.input( "input_features", - dtype=config.dtype, + dtype=ir.DataType.FLOAT, # Always f32 (matching audio processor output) shape=[batch, time, input_size], ) + if config.dtype and config.dtype != ir.DataType.FLOAT: + input_features = op.Cast(input_features, to=config.dtype) input_features_mask = builder.input( "input_features_mask", dtype=ir.DataType.BOOL, diff --git a/src/mobius/tasks/_phi4mm_multimodal.py b/src/mobius/tasks/_phi4mm_multimodal.py index 2e148a17..f64eeb03 100644 --- a/src/mobius/tasks/_phi4mm_multimodal.py +++ b/src/mobius/tasks/_phi4mm_multimodal.py @@ -85,19 +85,24 @@ def _build_vision( image_size = (config.vision.image_size if config.vision else None) or 448 graph, builder = _make_graph(name="vision_encoder") + op = builder.op + # Vision encoder input is always f32 (matching GenAI's image processor + # output). Cast at graph entry for f16/bf16 builds. pixel_values = builder.input( "pixel_values", - dtype=config.dtype, + dtype=ir.DataType.FLOAT, shape=[batch, 3, image_size, image_size], ) + if config.dtype and config.dtype != ir.DataType.FLOAT: + pixel_values = op.Cast(pixel_values, to=config.dtype) image_sizes = builder.input( "image_sizes", dtype=ir.DataType.INT64, shape=[num_images, 2], ) - image_features = vision(builder.op, pixel_values, image_sizes=image_sizes) + image_features = vision(op, pixel_values, image_sizes=image_sizes) builder.add_output(image_features, "image_features") return _make_model(graph) @@ -118,12 +123,17 @@ def _build_speech( input_size = (config.audio.input_size if config.audio else None) or 80 graph, builder = _make_graph(name="audio_encoder") + op = builder.op + # Audio encoder input is always f32 (matching audio processor output). + # Cast at graph entry for f16/bf16 builds. audio_embeds = builder.input( "audio_embeds", - dtype=config.dtype, + dtype=ir.DataType.FLOAT, shape=[batch, audio_seq_len, input_size], ) + if config.dtype and config.dtype != ir.DataType.FLOAT: + audio_embeds = op.Cast(audio_embeds, to=config.dtype) audio_sizes = builder.input( "audio_sizes", dtype=ir.DataType.INT64, @@ -136,7 +146,7 @@ def _build_speech( ) speech_out = speech( - builder.op, + op, audio_embeds, audio_sizes=audio_sizes, audio_projection_mode=audio_projection_mode, diff --git a/src/mobius/tasks/_speech_language.py b/src/mobius/tasks/_speech_language.py index 855af2b7..38a1dedc 100644 --- a/src/mobius/tasks/_speech_language.py +++ b/src/mobius/tasks/_speech_language.py @@ -83,14 +83,19 @@ def _build_audio_encoder( n_mels = (config.audio.num_mel_bins if config.audio else None) or 128 graph, builder = _make_graph(name="audio_encoder") + op = builder.op + # Audio encoder input is always f32 (matching audio processor output). + # Cast at graph entry for f16/bf16 builds. input_features = builder.input( "input_features", - dtype=config.dtype, + dtype=ir.DataType.FLOAT, shape=[batch, n_mels, mel_seq], ) + if config.dtype and config.dtype != ir.DataType.FLOAT: + input_features = op.Cast(input_features, to=config.dtype) - audio_features = audio_encoder(builder.op, input_features) + audio_features = audio_encoder(op, input_features) builder.add_output(audio_features, "audio_features") return _make_model(graph) diff --git a/src/mobius/tasks/_vision_language_3model.py b/src/mobius/tasks/_vision_language_3model.py index 00dae849..1c0c7db6 100644 --- a/src/mobius/tasks/_vision_language_3model.py +++ b/src/mobius/tasks/_vision_language_3model.py @@ -83,12 +83,17 @@ def _build_vision( image_size = (config.vision.image_size if config.vision else None) or 224 graph, builder = _make_graph(name="vision_encoder") + op = builder.op + # Vision encoder input is always f32 (matching GenAI's image processor + # output). Cast at graph entry for f16/bf16 builds. pixel_values = builder.input( "pixel_values", - dtype=config.dtype, + dtype=ir.DataType.FLOAT, shape=[batch, 3, image_size, image_size], ) - image_features = vision(builder.op, pixel_values=pixel_values) + if config.dtype and config.dtype != ir.DataType.FLOAT: + pixel_values = op.Cast(pixel_values, to=config.dtype) + image_features = vision(op, pixel_values=pixel_values) builder.add_output(image_features, "image_features") return _make_model(graph) @@ -134,11 +139,16 @@ def _build_vision( pixel_dim = in_channels * temporal_patch_size * patch_size * patch_size graph, builder = _make_graph(name="vision_encoder") + op = builder.op + # Vision encoder input is always f32 (matching GenAI's image processor + # output). Cast at graph entry for f16/bf16 builds. pixel_values = builder.input( "pixel_values", - dtype=config.dtype, + dtype=ir.DataType.FLOAT, shape=[total_patches, pixel_dim], ) + if config.dtype and config.dtype != ir.DataType.FLOAT: + pixel_values = op.Cast(pixel_values, to=config.dtype) image_grid_thw = builder.input( "image_grid_thw", dtype=ir.DataType.INT64, @@ -146,7 +156,7 @@ def _build_vision( ) image_features = vision( - builder.op, + op, pixel_values=pixel_values, image_grid_thw=image_grid_thw, ) @@ -209,12 +219,16 @@ def _build_vision( width = ir.SymbolicDim("width") graph, builder = _make_graph(name="vision_encoder") + op = builder.op + # Vision encoder input is always f32 (matching GenAI's image processor + # output). Cast at graph entry for f16/bf16 builds. pixel_values = builder.input( "pixel_values", - dtype=config.dtype, + dtype=ir.DataType.FLOAT, shape=[batch, 3, height, width], ) - op = builder.op + if config.dtype and config.dtype != ir.DataType.FLOAT: + pixel_values = op.Cast(pixel_values, to=config.dtype) image_features = vision( op, diff --git a/tests/build_graph_test.py b/tests/build_graph_test.py index 26c47140..78c87098 100644 --- a/tests/build_graph_test.py +++ b/tests/build_graph_test.py @@ -1857,6 +1857,117 @@ def test_dtype_casts_float_initializers(self, dtype_str, expected): f"Initializer '{name}' dtype is {init.dtype}, expected {expected_dtype}" ) + @pytest.mark.parametrize( + "dtype_str", + ["f16", "bf16"], + ) + def test_multimodal_encoder_inputs_are_float32(self, dtype_str): + """Vision/audio encoder graph inputs stay f32 with Cast at entry. + + When building multimodal models with f16/bf16, encoder graph inputs + (pixel_values, input_features) must remain FLOAT because ORT GenAI's + image/audio processors output f32. A Cast node at graph entry converts + to the target dtype for the encoder's internal computation. + """ + # Use a VL model with 3-model split (vision_encoder is separate) + config = _base_config( + vision=VisionConfig( + hidden_size=32, + intermediate_size=64, + num_hidden_layers=1, + num_attention_heads=2, + image_size=28, + patch_size=14, + norm_eps=1e-6, + ), + image_token_id=32000, + ) + config.dtype = DTYPE_MAP[dtype_str] + model_cls = registry.get("llava") + module = model_cls(config) + task = get_task("vision-language") + pkg = task.build(module, config) + + # Vision encoder pixel_values input must be FLOAT + vision_model = pkg["vision_encoder"] + pixel_values_input = vision_model.graph.inputs[0] + assert pixel_values_input.name == "pixel_values" + assert pixel_values_input.dtype == ir.DataType.FLOAT, ( + f"Vision encoder input dtype is {pixel_values_input.dtype}, " + f"expected FLOAT (Cast should handle conversion to {dtype_str})" + ) + + # First non-input node should be Cast to target dtype + first_node = next(iter(vision_model.graph)) + assert first_node.op_type == "Cast", ( + f"Expected Cast as first node, got {first_node.op_type}" + ) + + @pytest.mark.parametrize( + "dtype_str", + ["f16", "bf16"], + ) + def test_gemma4_encoder_inputs_are_float32(self, dtype_str): + """Gemma4 vision and audio encoder inputs stay f32 in bf16/f16 builds.""" + from mobius._configs import Gemma4AudioConfig, Gemma4Config + + config = Gemma4Config( + num_hidden_layers=2, + hidden_size=64, + intermediate_size=128, + num_attention_heads=4, + num_key_value_heads=1, + head_dim=16, + vocab_size=256, + rms_norm_eps=1e-6, + hidden_act="silu", + attn_qk_norm=True, + layer_types=["sliding_attention", "sliding_attention"], + sliding_window=8, + global_head_dim=16, + global_rope_theta=10_000.0, + global_partial_rotary_factor=0.25, + final_logit_softcapping=0.0, + hidden_size_per_layer_input=0, + image_token_id=255999, + pad_token_id=0, + tie_word_embeddings=True, + num_kv_shared_layers=1, + vision=VisionConfig( + hidden_size=32, + intermediate_size=64, + num_hidden_layers=1, + num_attention_heads=2, + patch_size=16, + norm_eps=1e-6, + ), + audio=Gemma4AudioConfig( + input_size=16, + hidden_size=32, + num_layers=1, + output_dim=64, + output_proj_dims=64, + audio_token_id=255998, + ), + dtype=DTYPE_MAP[dtype_str], + ) + model_cls = registry.get("gemma4") + module = model_cls(config) + task = get_task("gemma4") + pkg = task.build(module, config) + + # Vision encoder pixel_values must be FLOAT + vision_model = pkg["vision_encoder"] + pv_input = vision_model.graph.inputs[0] + assert pv_input.name == "pixel_values" + assert pv_input.dtype == ir.DataType.FLOAT + + # Audio encoder input_features must be FLOAT + audio_model = pkg["audio_encoder"] + af_input = audio_model.graph.inputs[0] + assert af_input.name == "input_features" + assert af_input.dtype == ir.DataType.FLOAT + class TestBuildGraphMultiModal: """Verify Phi4MM builds with Phi4MMMultiModalTask (4-model split)."""