Fix accuracy issues with Gemma models

src/python/py/models/builder.py

@@ -100,9 +100,12 @@ def __init__(self, config, io_dtype, onnx_dtype, ep, cache_dir, extra_options):
 
         # Map output names to their types and shapes
         self.output_names = ["logits"]
+
+        logits_type = TensorProto.FLOAT if "gemma-3" in self.model_name_or_path else self.io_dtype
+
         self.output_types = {
             "hidden_states": self.io_dtype,                                                                      # For standard models where you want to remove the language modeling head from the model (note that `hidden_states` is written this way to match Hugging Face format)
-            "logits": self.io_dtype,                                                                             # For standard models
+            "logits": logits_type,                                                                               # For standard models
             "present.key": self.io_dtype,                                                                        # For standard models (note that `present.key` is written this way to match Hugging Face format)
             "present.value": self.io_dtype,                                                                      # For standard models (note that `present.value` is written this way to match Hugging Face format)
         }
@@ -816,8 +819,16 @@ def make_matmul_float(self, matmul, name, root_input, **kwargs):
 
         last_dim = matmul.weight.shape[0]
         output = "logits" if kwargs.get("logits", False) else f"{name}/output_0"
-        self.make_node("MatMul", inputs=[root_input, weight], outputs=[output], name=name)
-        self.make_value_info(output, self.io_dtype, shape=['batch_size', 'sequence_length', last_dim])
+
+        if output != "logits" or "gemma-3" not in self.model_name_or_path:
+            self.make_node("MatMul", inputs=[root_input, weight], outputs=[output], name=name)
+            self.make_value_info(output, self.io_dtype, shape=['batch_size', 'sequence_length', last_dim])
+        else:
+            cast_name = name + "/Cast"
+            cast_name_input = name + "/Cast"
+            self.make_node("MatMul", inputs=[root_input, weight], outputs=[cast_name_input], name=name)
+            self.make_node("Cast", inputs=[cast_name_input], outputs=[output], name=cast_name, to=TensorProto.FLOAT)
+            self.make_value_info(output, TensorProto.FLOAT, shape=['batch_size', 'sequence_length', last_dim])
 
         return name
 
@@ -1042,7 +1053,7 @@ def make_embedding(self, embedding):
         self.layernorm_attrs["root_input"] = layernorm_attrs_value
         self.layernorm_attrs["skip_input"] = layernorm_attrs_value
 
-    def make_layernorm(self, layer_id, layernorm, skip, simple, location):
+    def make_layernorm_default(self, layer_id, layernorm, skip, simple, location):
         root_input = self.layernorm_attrs["root_input"]
         skip_input = self.layernorm_attrs["skip_input"]
 
@@ -1081,6 +1092,74 @@ def make_layernorm(self, layer_id, layernorm, skip, simple, location):
             # Assign output 3 of current SkipLayerNorm as root input to next SkipLayerNorm
             self.layernorm_attrs["root_input"] = output_3
 
+    def make_layernorm_with_casting_to_fp32(self, layer_id, layernorm, skip, simple, location):
+        root_input = self.layernorm_attrs["root_input"]
+        skip_input = self.layernorm_attrs["skip_input"]
+
+        skip_path = "skip" if skip else "noskip"
+
+        casted_root = f"{root_input}.{layer_id}.{skip_path}root/Cast"
+        self.make_node("Cast", inputs=[root_input], outputs=[casted_root], name=f"/model/layers.{layer_id}/{location}_layernorm/CastInput0", to=TensorProto.FLOAT)
+        root_input = casted_root
+
+        self.make_value_info(casted_root, TensorProto.FLOAT, shape=["batch_size", "sequence_length", self.hidden_size],)
+
+        if skip:
+            casted_skip = f"{skip_input}.{layer_id}.{skip_path}skip/Cast"
+            self.make_node("Cast", inputs=[skip_input], outputs=[casted_skip], name=f"/model/layers.{layer_id}/{location}_layernorm/CastInput3", to=TensorProto.FLOAT)
+            skip_input = casted_skip
+
+            self.make_value_info(casted_skip, TensorProto.FLOAT, shape=["batch_size", "sequence_length", self.hidden_size],)
+
+        weight = f"model.layers.{layer_id}.{location}_layernorm.weight"
+        self.make_external_tensor(layernorm.weight.detach().to(self.to_torch_dtype[TensorProto.FLOAT]).contiguous() + self.layernorm_attrs["add_offset"], weight)
+        bias = f"model.layers.{layer_id}.{location}_layernorm.bias"
+        if not simple:
+            self.make_external_tensor(layernorm.bias.detach().to(self.to_torch_dtype[TensorProto.FLOAT]).contiguous(), bias)
+
+        inputs = [root_input, skip_input, weight] if skip else [root_input, weight]
+        if not simple:
+            inputs.append(bias)
+
+        name = f"/model/layers.{layer_id}/{location}_layernorm/{'Skip' if skip else ''}LayerNorm"
+        op_type = f"{'Skip' if skip else ''}{'Simplified' if simple else ''}LayerNormalization"
+        kwargs = {"epsilon": self.layernorm_attrs["epsilon"]}
+        if not skip:
+            kwargs.update({"axis": -1, "stash_type": 1})
+
+        output_0 = f"/model/layers.{layer_id}/{location}_layernorm/output_0"
+        output_3 = f"/model/layers.{layer_id}/{location}_layernorm/output_3"
+        raw_out0 = f"{output_0}/Cast"
+        raw_out3 = f"{output_3}/Cast"
+        if self.layernorm_attrs["last_layernorm"] and (self.include_hidden_states or self.exclude_lm_head):
+            raw_out0 = "hidden_states"
+        outputs = [raw_out0, "", "", raw_out3] if skip and not self.layernorm_attrs["last_layernorm"] else [raw_out0]
+
+        self.make_node(op_type, inputs=inputs, outputs=outputs, name=name, domain=("com.microsoft" if skip else None), **kwargs)
+
+        self.make_value_info(raw_out0, TensorProto.FLOAT, shape=['batch_size', 'sequence_length', self.hidden_size])
+        if skip and not self.layernorm_attrs["last_layernorm"]:
+            self.make_value_info(raw_out3, TensorProto.FLOAT, shape=['batch_size', 'sequence_length', self.hidden_size])
+
+        if not(skip and not self.layernorm_attrs["last_layernorm"]):
+            raw_out3 = None
+
+        self.make_node("Cast", inputs=[raw_out0], outputs=[output_0], name=f"{name}/CastOutput0", to=self.io_dtype,)
+        self.make_value_info(output_0, self.io_dtype, shape=["batch_size", "sequence_length", self.hidden_size],)
+        self.layernorm_attrs["output_0"] = output_0
+
+        if raw_out3 is not None:
+            self.make_node("Cast", inputs=[raw_out3], outputs=[output_3], name=f"{name}/CastOutput3", to=self.io_dtype,)
+            self.make_value_info(output_3, self.io_dtype, shape=["batch_size", "sequence_length", self.hidden_size],)
+            self.layernorm_attrs["output_3"] = output_3
+            self.layernorm_attrs["root_input"] = output_3
+
+    def make_layernorm(self, layer_id, layernorm, skip, simple, location):
+        if self.layernorm_attrs["use_fp32_layernorms"]:
+            self.make_layernorm_with_casting_to_fp32(layer_id, layernorm, skip, simple, location)
+        else:
+            self.make_layernorm_default(layer_id, layernorm, skip, simple, location)
+
     def make_mscale_su(self, mscale):
         if mscale <= 1.0:
             return 1.0
@@ -1171,6 +1250,7 @@ def make_rotary_embedding(self, name, root_input, **kwargs):
         num_heads = self.num_kv_heads if "k_rotary" in name else self.num_attn_heads
 
         inputs = [root_input, kwargs.pop("position_ids"), cos_cache_name, sin_cache_name]
+
         output = f"{name}/output_0"
         self.make_node(
             "RotaryEmbedding", inputs=inputs, outputs=[output], name=name, domain="com.microsoft",
@@ -1285,12 +1365,23 @@ def make_qk_norm(self, layer_id, attention):
         q_reshape_1_output = f"{q_reshape_1_name}/output_0"
         self.make_reshape(q_reshape_1_name, q_reshape_1_inputs, dtype=self.io_dtype, shape=['batch_size', 'sequence_length * num_attention_heads', self.head_size])
 
-        # Make Q LayerNorm
         q_layernorm_name = f"/model/layers.{layer_id}/attn/q_norm/SimplifiedLayerNormalization"
         q_weight_name = f"model.layers.{layer_id}.attn.q_norm.layernorm.weight"
         q_layernorm_output = f"{q_layernorm_name}/output_0"
-        self.make_external_tensor((attention.q_norm.weight.detach().cpu().to(self.to_torch_dtype[self.io_dtype]) + self.layernorm_attrs["add_offset"]).contiguous(), q_weight_name)
-        self.make_node("SimplifiedLayerNormalization", inputs=[q_reshape_1_output, q_weight_name], outputs=[q_layernorm_output], name=q_layernorm_name, **layernorm_kwargs)
+
+        if "gemma-3" in self.model_name_or_path and self.io_dtype == TensorProto.BFLOAT16:
+            q_layernorm_cast_in = f"{q_layernorm_name}/Cast"
+            self.make_node("Cast", inputs=[q_reshape_1_output], outputs=[q_layernorm_cast_in], name=f"{q_layernorm_name}/CastIn", to=TensorProto.FLOAT)
+
+            q_layernorm_weight_fp32 = (attention.q_norm.weight.detach().cpu().to(self.to_torch_dtype[TensorProto.FLOAT]) + self.layernorm_attrs["add_offset"]).contiguous()
+            self.make_external_tensor(q_layernorm_weight_fp32, q_weight_name)
+
+            self.make_node("SimplifiedLayerNormalization", inputs=[q_layernorm_cast_in, q_weight_name], outputs=[f"{q_layernorm_name}/output/Cast"], name=q_layernorm_name, **layernorm_kwargs)
+            self.make_node("Cast", inputs=[f"{q_layernorm_name}/output/Cast"], outputs=[q_layernorm_output], name=f"{q_layernorm_name}/CastOut", to=self.io_dtype)
+        else:
+            self.make_external_tensor((attention.q_norm.weight.detach().cpu().to(self.to_torch_dtype[self.io_dtype]) + self.layernorm_attrs["add_offset"]).contiguous(), q_weight_name)
+            self.make_node("SimplifiedLayerNormalization", inputs=[q_reshape_1_output, q_weight_name], outputs=[q_layernorm_output], name=q_layernorm_name, **layernorm_kwargs)
+
         self.make_value_info(q_layernorm_output, dtype=self.io_dtype, shape=['batch_size', 'sequence_length * num_attention_heads', self.head_size])
 
         # Reshape Q path after LayerNorm from Bx(SxN)xH to BxSxD
@@ -1304,12 +1395,23 @@ def make_qk_norm(self, layer_id, attention):
         k_reshape_1_output = f"{k_reshape_1_name}/output_0"
         self.make_reshape(k_reshape_1_name, k_reshape_1_inputs, dtype=self.io_dtype, shape=['batch_size', 'sequence_length * num_key_value_heads', self.head_size])
 
-        # Make K LayerNorm
         k_layernorm_name = f"/model/layers.{layer_id}/attn/k_norm/SimplifiedLayerNormalization"
         k_weight_name = f"model.layers.{layer_id}.attn.k_norm.layernorm.weight"
         k_layernorm_output = f"{k_layernorm_name}/output_0"
-        self.make_external_tensor((attention.k_norm.weight.detach().cpu().to(self.to_torch_dtype[self.io_dtype]) + self.layernorm_attrs["add_offset"]).contiguous(), k_weight_name)
-        self.make_node("SimplifiedLayerNormalization", inputs=[k_reshape_1_output, k_weight_name], outputs=[k_layernorm_output], name=k_layernorm_name, **layernorm_kwargs)
+
+        if "gemma-3" in self.model_name_or_path and self.io_dtype == TensorProto.BFLOAT16:
+            k_layernorm_cast_in = f"{k_layernorm_name}/Cast"
+            self.make_node("Cast", inputs=[k_reshape_1_output], outputs=[k_layernorm_cast_in], name=f"{k_layernorm_name}/CastIn", to=TensorProto.FLOAT)
+
+            k_layernorm_weight_fp32 = (attention.k_norm.weight.detach().cpu().to(self.to_torch_dtype[TensorProto.FLOAT]) + self.layernorm_attrs["add_offset"]).contiguous()
+            self.make_external_tensor(k_layernorm_weight_fp32, k_weight_name)
+
+            self.make_node("SimplifiedLayerNormalization", inputs=[k_layernorm_cast_in, k_weight_name], outputs=[f"{k_layernorm_name}/output/Cast"], name=k_layernorm_name, **layernorm_kwargs)
+            self.make_node("Cast", inputs=[f"{k_layernorm_name}/output/Cast"], outputs=[k_layernorm_output], name=f"{k_layernorm_name}/CastOut", to=self.io_dtype)
+        else:
+            self.make_external_tensor((attention.k_norm.weight.detach().cpu().to(self.to_torch_dtype[self.io_dtype]) + self.layernorm_attrs["add_offset"]).contiguous(), k_weight_name)
+            self.make_node("SimplifiedLayerNormalization", inputs=[k_reshape_1_output, k_weight_name], outputs=[k_layernorm_output], name=k_layernorm_name, **layernorm_kwargs)                
+
         self.make_value_info(k_layernorm_output, dtype=self.io_dtype, shape=['batch_size', 'sequence_length * num_key_value_heads', self.head_size])
 
         # Reshape K path after LayerNorm from Bx(SxN)xH to BxSxD
@@ -1606,10 +1708,10 @@ def make_attention(self, layer_id, attention, root_input, **kwargs):
             q_matmul_name = self.make_matmul(attention.q_proj, q_matmul_basename, root_input)
             self.attention_attrs["q_path"] = f"{q_matmul_name}/output_0"
             k_matmul_basename = f"/model/layers.{layer_id}/attn/k_proj/MatMul"
-            k_matmul_name = self.make_matmul(attention.k_proj, k_matmul_basename, root_input)
+            k_matmul_name = self.make_matmul(attention.k_proj, k_matmul_basename, root_input, first_cast=False)
             self.attention_attrs["k_path"] = f"{k_matmul_name}/output_0"
             v_matmul_basename = f"/model/layers.{layer_id}/attn/v_proj/MatMul"
-            v_matmul_name = self.make_matmul(attention.v_proj, v_matmul_basename, root_input)
+            v_matmul_name = self.make_matmul(attention.v_proj, v_matmul_basename, root_input, first_cast=False)
             self.attention_attrs["v_path"] = f"{v_matmul_name}/output_0"
 
         # Make Add nodes (if bias exists)
@@ -1875,7 +1977,7 @@ def make_mlp_proj(self, layer_id, mlp, root_input):
 
         # Make Up proj nodes
         up_matmul_basename = f"/model/layers.{layer_id}/mlp/up_proj/MatMul"
-        up_matmul_name = self.make_matmul(mlp.up_proj, up_matmul_basename, root_input)
+        up_matmul_name = self.make_matmul(mlp.up_proj, up_matmul_basename, root_input, first_cast=False)
         up_name = up_matmul_name
         if up_bias_exists:
             up_add_name = f"/model/layers.{layer_id}/mlp/up_proj/Add"
@@ -2209,6 +2311,7 @@ def make_model(self, input_path):
                 from onnxruntime_genai.models.gguf_model import GGUFModel
             model = GGUFModel.from_pretrained(self.model_type, input_path, self.head_size, self.hidden_size, self.intermediate_size, self.num_attn_heads, self.num_kv_heads, self.vocab_size)
             self.layernorm_attrs["add_offset"] = 0  # add offset already done for GGUF models
+            self.layernorm_attrs["use_fp32_layernorms"] = 0
         elif self.quant_type is not None:
             # Load quantized PyTorch model
             try:
@@ -2232,7 +2335,7 @@ def make_model(self, input_path):
         for module in model.modules():
             if isinstance(module, torch.nn.Embedding) or (hasattr(model, "embedding") and module == model.embedding):
                 # Checks (Hugging Face logic) or (GGUF logic)
-                if not self.exclude_embeds:
+                if not self.exclude_embeds and module == model.language_model.model.embed_tokens:
                     # Embedding layer
                     print("Reading embedding layer")
                     self.make_embedding(module.weight.detach().cpu())
@@ -2262,7 +2365,15 @@ def make_model(self, input_path):
         del model
 
     def has_final_norm(self, module, orig_model):
-        if hasattr(orig_model, "base_model") and hasattr(orig_model.base_model, "model"):
+        # Find where the language model is stored to check attributes. Some classes
+        # store the language model in a different attribute than `model.model`.
+        if hasattr(orig_model, "language_model"):
+            # Model is multimodal
+            # Note: This case is checked first because the `language_model` attribute and the `base_model` attribute
+            # exist for both multimodal models and PEFT models. However they represent different classes and their attributes
+            # differ.
+            model = orig_model.language_model
+        elif hasattr(orig_model, "base_model") and hasattr(orig_model.base_model, "model"):
             # Model is from PEFT
             model = orig_model.base_model.model
         else:
@@ -2272,12 +2383,11 @@ def has_final_norm(self, module, orig_model):
         hf_norm = hasattr(model, "model") and hasattr(model.model, "norm") and module == model.model.norm
         hf_final_layernorm = hasattr(model, "model") and hasattr(model.model, "final_layernorm") and module == model.model.final_layernorm
         hf_transformer_final_layernorm = hasattr(model, "transformer") and hasattr(model.transformer, "encoder") and hasattr(model.transformer.encoder, "final_layernorm") and module == model.transformer.encoder.final_layernorm
-        hf_multimodal_final_layernorm = hasattr(model, "language_model") and hasattr(model.language_model, "model") and hasattr(model.language_model.model, "norm") and module == model.language_model.model.norm
 
         # GGUF names
         gguf_final_norm = hasattr(model, "final_norm") and module == model.final_norm
 
-        hf_names = [hf_norm, hf_final_layernorm, hf_transformer_final_layernorm, hf_multimodal_final_layernorm]
+        hf_names = [hf_norm, hf_final_layernorm, hf_transformer_final_layernorm]
         gguf_names = [gguf_final_norm]
         return any(hf_names + gguf_names)
 
@@ -3230,6 +3340,7 @@ def __init__(self, config, io_dtype, onnx_dtype, ep, cache_dir, extra_options):
         self.is_local = lambda layer_id: bool((layer_id + 1) % config.sliding_window_pattern)
         self.rope_local_theta = config.rope_local_base_freq
         self.make_rotary_embedding_multi_cache()
+        self.layernorm_attrs["use_fp32_layernorms"] = 1 if self.io_dtype == TensorProto.BFLOAT16 else 0        
 
     def make_attention_init(self):
         self.attention_attrs["q_norm"] = True
@@ -3369,8 +3480,9 @@ def create_model(model_name, input_path, output_dir, precision, execution_provid
             for key in text_config:
                 if not hasattr(config, key):
                     setattr(config, key, getattr(text_config, key))
-            extra_options["exclude_embeds"] = True
+            # extra_options["exclude_embeds"] = True
             onnx_model = Gemma3Model(config, io_dtype, precision, execution_provider, cache_dir, extra_options)
+            onnx_model.model_type = "gemma3_text"
         elif config.architectures[0] == "GraniteForCausalLM":
             onnx_model = GraniteModel(config, io_dtype, precision, execution_provider, cache_dir, extra_options)
         elif config.architectures[0] == "LlamaForCausalLM":
@@ -3549,4 +3661,4 @@ def get_args():
 if __name__ == '__main__':
     args = get_args()
     extra_options = parse_extra_options(args.extra_options)
-    create_model(args.model_name, args.input, args.output, args.precision, args.execution_provider, args.cache_dir, **extra_options)
+    create_model(args.model_name, args.input, args.output, args.precision, args.execution_provider, args.cache_dir, **extra_options)