Skip to content

Add final norm for LoRA models #1446

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
kunal-vaishnavi merged 3 commits into from
May 7, 2025
Merged

Add final norm for LoRA models #1446

Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
62d0661
Fix missing final norm for LoRA models
kunal-vaishnavi May 4, 2025
d7aeb9a
Match against all name variants for LoRA models
kunal-vaishnavi May 4, 2025
19c35d5
Merge branch 'main' into kvaishnavi/fix-logits-with-adapter
kunal-vaishnavi May 5, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
89 changes: 36 additions & 53 deletions src/python/py/models/builder.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -917,7 +917,7 @@ def make_packed_matmul(self, q_matmul, k_matmul, v_matmul, basename, root_input,
else:
raise NotImplementedError(f"The {self.onnx_dtype} precision is not currently supported.")

def make_packed_matmul_fp16_or_fp32(self, q_matmul, k_matmul, v_matmul, name, root_input, **kwargs):
def make_packed_matmul_fp16_or_fp32(self, q_matmul, k_matmul, v_matmul, basename, root_input, **kwargs):
# N_q = num_attention_heads * head_size, N_kv = num_key_value_heads * head_size, H = hidden_size
# Combine 3 MatMuls of shape N_q x H, N_kv x H, N_kv x H into 1 packed MatMul of shape (N_q+N_kv+N_kv)xH
# Note: Packed MatMul is of shape (N_q+N_kv+N_kv)xH instead of Hx(N_q+N_kv+N_kv) because `make_matmul` will
Expand All @@ -930,7 +930,7 @@ class PackedMatMul:
def __init__(self):
self.weight = torch.concatenate([q_matmul.weight.detach().cpu(), k_matmul.weight.detach().cpu(), v_matmul.weight.detach().cpu()], dim=0).reshape(N_q + N_kv + N_kv, H)
matmul = PackedMatMul()
new_name = self.make_matmul(matmul, name, root_input, **kwargs)
new_name = self.make_matmul(matmul, basename, root_input, **kwargs)

return new_name

Expand All @@ -941,8 +941,6 @@ def make_packed_matmul_int4(self, q_matmul, k_matmul, v_matmul, basename, root_i
# print(f"Saving as {self.io_dtype} on-the-fly and quantizing to {self.onnx_dtype} at the end.")
return self.make_packed_matmul_fp16_or_fp32(q_matmul, k_matmul, v_matmul, basename, root_input, **kwargs)

name = f"{basename}NBits"

# Create dummy PackedMatMul class
class PackedMatMul:
def __init__(self):
Expand All @@ -956,34 +954,9 @@ def __init__(self):
self.bits = q_matmul.bits
self.group_size = q_matmul.group_size
matmul = PackedMatMul()
new_name = self.make_matmul_int4(matmul, basename, root_input, **kwargs)

# Input weights are quantized, save quantized MatMul numpy weights for onnx model
weight = name[1:].replace("/", ".") + ".qweight"
self.make_external_tensor(matmul.qweight.detach().numpy(), weight)
scales = name[1:].replace("/", ".") + ".scales"
self.make_external_tensor(matmul.scales.detach().numpy().astype(self.to_numpy_dtype[self.io_dtype]), scales)

inputs = [root_input, weight, scales]

if hasattr(matmul, "qzeros") and matmul.qzeros is not None:
zeros = name[1:].replace("/", ".") + ".qzeros"
self.make_external_tensor(matmul.qzeros.detach().numpy(), zeros)
inputs.append(zeros)

if hasattr(matmul, "g_idx") and matmul.g_idx is not None:
g_idx = name[1:].replace("/", ".") + ".g_idx"
self.make_external_tensor(matmul.g_idx.detach().numpy().astype(np.int32), g_idx)
inputs.append(g_idx)

output = "logits" if kwargs.get("logits", False) else f"{name}/output_0"
self.make_node(
"MatMulNBits", inputs=inputs, outputs=[output], name=name, domain="com.microsoft",
accuracy_level=self.quant_attrs["int4"]["accuracy_level"],
bits=matmul.bits, block_size=matmul.group_size, K=matmul.in_features, N=matmul.out_features,
)
self.make_value_info(output, self.io_dtype, shape=['batch_size', 'sequence_length', matmul.out_features])

return name
return new_name

def make_add_bias(self, add, name, root_input, **kwargs):
bias = name[1:].replace("/", ".") + ".bias"
Expand All @@ -992,7 +965,7 @@ def make_add_bias(self, add, name, root_input, **kwargs):
add_bias_inputs = [root_input, bias]
shape = ['batch_size', 'sequence_length', add.shape[0]]

if "logits" in kwargs:
if kwargs.get("logits", False):
output = "logits"
self.make_node("Add", inputs=add_bias_inputs, outputs=[output], name=name)
self.make_value_info(output, dtype=self.io_dtype, shape=shape)
Expand Down Expand Up @@ -2204,15 +2177,7 @@ def make_model(self, input_path):
from onnxruntime_genai.models.quantized_model import QuantModel
q_size = self.num_attn_heads * self.head_size
kv_size = self.num_kv_heads * self.head_size
model = QuantModel.from_pretrained(
self.quant_type,
input_path = input_path,
quant_attrs = self.quant_attrs,
q_size = q_size,
kv_size = kv_size,
intermediate_size = self.intermediate_size,
num_layers = self.num_layers,
)
model = QuantModel.from_pretrained(self.quant_type, input_path, self.quant_attrs, q_size, kv_size, self.intermediate_size, self.num_layers)
else:
# Load PyTorch model
extra_kwargs = {"num_hidden_layers": self.num_layers} if "num_hidden_layers" in self.extra_options else {}
Expand Down Expand Up @@ -2256,15 +2221,25 @@ def make_model(self, input_path):

del model

def has_final_norm(self, module, model):
# Hugging Face names
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
return hf_norm or hf_final_layernorm or hf_transformer_final_layernorm or hf_multimodal_final_layernorm or gguf_final_norm
def has_final_norm(self, module, orig_model):
if hasattr(orig_model, "base_model") and hasattr(orig_model.base_model, "model"):
# Model is from PEFT
model = orig_model.base_model.model
else:
model = orig_model

# Hugging Face names
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

Comment thread
kunal-vaishnavi marked this conversation as resolved.
hf_names = [hf_norm, hf_final_layernorm, hf_transformer_final_layernorm, hf_multimodal_final_layernorm]
gguf_names = [gguf_final_norm]
return any(hf_names + gguf_names)

def make_preprocessing_nodes(self):
self.make_attention_mask_reformatting()
Expand Down Expand Up @@ -3286,7 +3261,7 @@ def check_extra_options(kv_pairs):
"""
Check key-value pairs and set values correctly
"""
bools = ["int4_is_symmetric", "exclude_embeds", "exclude_lm_head", "include_hidden_states", "enable_cuda_graph", "use_8bits_moe", "use_qdq"]
bools = ["int4_is_symmetric", "exclude_embeds", "exclude_lm_head", "include_hidden_states", "enable_cuda_graph", "use_8bits_moe", "use_qdq", "use_webgpu_fp32"]
for key in bools:
if key in kv_pairs:
if kv_pairs[key] in {"false", "False", "0"}:
Expand Down Expand Up @@ -3347,6 +3322,15 @@ def parse_hf_token(hf_token):
return hf_token


def set_io_dtype(precision, execution_provider, extra_options):
if precision in {"int8", "fp32"} or (precision == "int4" and execution_provider == "cpu") or extra_options.get("use_webgpu_fp32", False):
# FP32 precision
return TensorProto.FLOAT

# FP16 precision
return TensorProto.FLOAT16


def create_model(model_name, input_path, output_dir, precision, execution_provider, cache_dir, **extra_options):
# Create cache and output directories
os.makedirs(output_dir, exist_ok=True)
Expand All @@ -3364,8 +3348,7 @@ def create_model(model_name, input_path, output_dir, precision, execution_provid
config.update(peft_config.__dict__)

# Set input/output precision of ONNX model
use_webgpu_fp32 = extra_options.get("use_webgpu_fp32", "0") == "1"
io_dtype = TensorProto.FLOAT if precision in {"int8", "fp32"} or (precision == "int4" and execution_provider == "cpu") or use_webgpu_fp32 else TensorProto.FLOAT16
io_dtype = set_io_dtype(precision, execution_provider, extra_options)

if "config_only" not in extra_options:
# List architecture options in alphabetical order
Expand Down