diff --git a/src/python/py/models/builder.py b/src/python/py/models/builder.py index 1effaf4401..30ef8741ee 100644 --- a/src/python/py/models/builder.py +++ b/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) } @@ -799,24 +802,20 @@ def make_matmul_op(self, matmul, basename, root_input, **kwargs): def make_matmul_float(self, matmul, name, root_input, **kwargs): weight = name[1:].replace("/", ".") + ".weight" - self.make_external_tensor(matmul.weight.detach().T.to(torch.bfloat16).contiguous(), weight) + self.make_external_tensor(matmul.weight.detach().T.to( self.to_torch_dtype[self.io_dtype]).contiguous(), weight) last_dim = matmul.weight.shape[0] - first_dim = matmul.weight.shape[1] output = "logits" if kwargs.get("logits", False) else f"{name}/output_0" - cast_input = f"{root_input}/Cast" - cast_output = f"{output}/Cast" - - if kwargs.get("first_cast", True): - self.make_node("Cast", inputs=[root_input], outputs=[cast_input], name=f"{name}/Cast_Input", to=TensorProto.BFLOAT16) - self.make_value_info(cast_input, TensorProto.BFLOAT16, shape=["batch_size", "sequence_length", first_dim]) - - self.make_node("MatMul", inputs=[cast_input, weight], outputs=[cast_output], name=name) - self.make_value_info(output, self.io_dtype, shape=['batch_size', 'sequence_length', last_dim]) - - self.make_node("Cast", inputs=[cast_output], outputs=[output], name=f"{name}/Cast_Output", to=self.io_dtype) - self.make_value_info(cast_output, TensorProto.BFLOAT16, 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 @@ -1068,7 +1067,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"] @@ -1107,6 +1106,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 "gemma-3" in self.model_name_or_path and self.io_dtype == TensorProto.BFLOAT16: + 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,9 +1238,9 @@ def make_rotary_embedding_caches(self, **kwargs): # Slice cos/sin caches from (M, H) to (M, H/2) hidden_dim = cos_cache.shape[-1] cos_cache = cos_cache.squeeze()[:, : (hidden_dim // 2)].detach() - cos_cache = cos_cache.to(self.to_torch_dtype[TensorProto.BFLOAT16]) + cos_cache = cos_cache.to(self.to_torch_dtype[self.io_dtype]) sin_cache = sin_cache.squeeze()[:, : (hidden_dim // 2)].detach() - sin_cache = sin_cache.to(self.to_torch_dtype[TensorProto.BFLOAT16]) + sin_cache = sin_cache.to(self.to_torch_dtype[self.io_dtype]) # Slice cos/sin caches from (M, H/2) to (M, R/2) if partial rotary embeddings are used if self.rotemb_attrs["partial_rotary_factor"] != 1.0: @@ -1196,23 +1263,15 @@ def make_rotary_embedding(self, name, root_input, **kwargs): cos_cache_name, sin_cache_name = self.make_rotary_embedding_caches() num_heads = self.num_kv_heads if "k_rotary" in name else self.num_attn_heads - casted_root = f"{name}/root_input_bf16" - self.make_node( - "Cast", - inputs=[root_input], - outputs=[casted_root], - name=f"{name}/Cast_to_bf16", - to=TensorProto.BFLOAT16, - ) + inputs = [root_input, kwargs.pop("position_ids"), cos_cache_name, sin_cache_name] - inputs = [casted_root, 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", interleaved=self.rotemb_attrs["interleaved"], num_heads=(0 if self.rotemb_attrs["partial_rotary_factor"] == 1.0 else num_heads), # default is 0 in RotaryEmbedding kernel rotary_embedding_dim=self.rotemb_attrs["rotary_embedding_dim"], ) - self.make_value_info(output, TensorProto.BFLOAT16, shape=['batch_size', 'sequence_length', self.head_size * num_heads]) + self.make_value_info(output, self.io_dtype, shape=['batch_size', 'sequence_length', self.head_size * num_heads]) def make_rotary_embedding_multi_cache(self, **kwargs): cos_cache_name = kwargs.get("cos_cache_name", "cos_cache") @@ -1320,12 +1379,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().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().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().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 @@ -1339,12 +1409,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().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().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().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 @@ -1540,7 +1621,7 @@ def make_attention_op(self, name, **kwargs): if op_type == "MultiHeadAttention": self.make_multi_head_attention(name, add_qk=f"{self.mask_attrs['mask_name']}/output_0", **kwargs) elif op_type == "GroupQueryAttention": - self.make_group_query_attention_with_bf16(name, seqlens_k=f"{self.mask_attrs['seqlens_k']}/output_0", total_seq_len=f"{self.mask_attrs['total_seq_len']}/output_0", **kwargs) + self.make_group_query_attention(name, seqlens_k=f"{self.mask_attrs['seqlens_k']}/output_0", total_seq_len=f"{self.mask_attrs['total_seq_len']}/output_0", **kwargs) elif op_type == "SparseAttention": self.make_sparse_attention(name, block_row_indices=self.mask_attrs['block_row_indices'], block_col_indices=self.mask_attrs['block_col_indices'], key_total_seq_lens=f"{self.mask_attrs['key_total_seq_lens']}/output_0", total_seq_len=f"{self.mask_attrs['total_seq_len']}/output_0", **kwargs) else: @@ -1576,84 +1657,6 @@ def make_group_query_attention(self, name, **kwargs): ) self.make_value_info(output, self.io_dtype, shape=['batch_size', 'sequence_length', self.head_size * self.num_attn_heads]) - def make_group_query_attention_with_bf16(self, name, **kwargs): - raw_inputs = [ - kwargs["v_path"], - kwargs.get("past_k", ""), kwargs.get("past_v", ""), - kwargs.get("cos_cache", ""), kwargs.get("sin_cache", ""), - ] - - int64_inputs = [kwargs.get("seqlens_k", ""), kwargs.get("total_seq_len", "")] - - bf16_inputs = [kwargs["q_path"], kwargs["k_path"]] - for inp in raw_inputs: - if inp == "": - bf16_inputs.append("") - continue - bf16_inp = f"{inp}_bf16" - self.make_node( - "Cast", - inputs=[inp], - outputs=[bf16_inp], - name=f"{name}/cast_{inp}_to_bf16", - to=TensorProto.BFLOAT16 - ) - bf16_inputs.append(bf16_inp) - - bf16_inputs[5:5] = int64_inputs - - out_bf16 = f"{name}/output_0_bf16" - pk_bf16 = f"{kwargs['present_k']}_bf16" if kwargs.get("present_k") else "" - pv_bf16 = f"{kwargs['present_v']}_bf16" if kwargs.get("present_v") else "" - bf16_outputs = [out_bf16, pk_bf16, pv_bf16] - - self.make_node( - "GroupQueryAttention", - inputs=bf16_inputs, - outputs=bf16_outputs, - name=name, - domain="com.microsoft", - num_heads=self.num_attn_heads, - kv_num_heads=self.num_kv_heads, - scale=self.attention_attrs["scale"], - local_window_size=self.window_size, - softcap=self.attention_attrs["softcap"], - do_rotary=self.attention_attrs["use_rope_in_attn"], - rotary_interleaved=self.rotemb_attrs["interleaved"], - ) - - self.make_node( - "Cast", - inputs=[out_bf16], - outputs=[f"{name}/output_0"], - name=f"{name}/Cast_output_to_fp32", - to=TensorProto.FLOAT, - ) - - if pk_bf16: - self.make_node( - "Cast", - inputs=[pk_bf16], - outputs=[kwargs["present_k"]], - name=f"{name}/Cast_present_k_to_fp32", - to=TensorProto.FLOAT, - ) - - if pv_bf16: - self.make_node( - "Cast", - inputs=[pv_bf16], - outputs=[kwargs["present_v"]], - name=f"{name}/Cast_present_v_to_fp32", - to=TensorProto.FLOAT, - ) - - self.make_value_info( - f"{name}/output_0", - TensorProto.FLOAT, - shape=['batch_size', 'sequence_length', self.head_size * self.num_attn_heads] - ) - def make_sparse_attention(self, name, **kwargs): inputs = [ kwargs["q_path"], kwargs["k_path"], kwargs["v_path"], @@ -2207,24 +2210,6 @@ def make_gelu(self, layer_id, root_input, activation): return gelu_name - def make_gelu_bf16(self, layer_id, root_input, activation): - gelu_name = f"/model/layers.{layer_id}/mlp/act_fn/{activation}" - cast_to_bf16 = f"{gelu_name}/cast_to_bf16" - gelu_in = f"{cast_to_bf16}/output" - gelu_out_bf16 = f"{gelu_name}/gelu_bf16" - cast_to_fp32 = f"{gelu_name}/cast_to_fp32" - final_output = f"{gelu_name}/output_0" - - self.make_node("Cast", inputs=[root_input], outputs=[gelu_in], name=cast_to_bf16, to=TensorProto.BFLOAT16) - self.make_node(activation, inputs=[gelu_in], outputs=[gelu_out_bf16], name=gelu_name, domain="com.microsoft") - - self.make_value_info(gelu_out_bf16, TensorProto.BFLOAT16, shape=['batch_size', 'sequence_length', self.intermediate_size]) - - self.make_node("Cast",inputs=[gelu_out_bf16], outputs=[final_output], name=cast_to_fp32, to=TensorProto.FLOAT) - self.make_value_info(final_output, self.io_dtype, shape=['batch_size', 'sequence_length', self.intermediate_size]) - - return gelu_name - def make_relu(self, layer_id, root_input, activation): relu_name = f"/model/layers.{layer_id}/mlp/act_fn/{activation}" output = f"{relu_name}/output_0" @@ -2245,7 +2230,7 @@ def make_activation(self, layer_id, root_input): if self.activation in {"silu", "swish", "swiglu"}: output_name = self.make_activation_with_mul(layer_id, root_input, activation="Sigmoid", domain=None) elif self.activation in {"gelu_new", "gelu_fast", "gelu_pytorch_tanh"}: - output_name = self.make_gelu_bf16(layer_id, root_input, activation="FastGelu") + output_name = self.make_gelu(layer_id, root_input, activation="FastGelu") elif self.activation in {"gelu"}: output_name = self.make_gelu(layer_id, root_input, activation="Gelu") elif self.activation in {"gegelu", "geglu"}: