Skip to content

Fix LM head for Gemma-2 #1420

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
Baiju Meswani (baijumeswani) merged 5 commits into from
Apr 25, 2025
Merged

Fix LM head for Gemma-2 #1420

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
eb6f0fb
Fix node names for logit softcapping
kunal-vaishnavi Apr 24, 2025
d3b5ee0
Remove extra multiply
kunal-vaishnavi Apr 24, 2025
4880922
Add missing final norm for Gemma-3 multimodal
kunal-vaishnavi Apr 24, 2025
2e00e80
Fix naming of model type
kunal-vaishnavi Apr 24, 2025
196c1c2
Add Gemma-3 system prompt to chat example
kunal-vaishnavi Apr 24, 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
6 changes: 5 additions & 1 deletion examples/python/model-chat.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -55,6 +55,8 @@ def main(args):
print("Using Chat Template for LLAMA 3, if you are using LLAMA 2 please pass the argument --chat_template '{input} [/INST]')")
elif model_type.startswith("qwen2"):
args.chat_template = '<|im_start|>user\n{input}<|im_end|>\n<|im_start|>assistant\n'
elif model_type == "gemma3_text":
args.chat_template = '<start_of_turn>user\n{system_prompt}{input}<end_of_turn>\n<start_of_turn>model\n'
else:
raise ValueError(f"Chat Template for model type {model_type} is not known. Please provide chat template using --chat_template")

Expand All @@ -81,6 +83,8 @@ def main(args):
print("Using System Prompt for LLAMA 3, if you are using LLAMA 2 please pass the argument --system_prompt '<s>[INST] <<SYS>>\\n{args.system_prompt}\\n<</SYS>>')")
elif model_type.startswith("qwen2"):
system_prompt = f"<|im_start|>system\n{args.system_prompt}<|im_end|>\n"
elif model_type == "gemma3_text":
system_prompt = f"{args.system_prompt}"
else:
system_prompt = args.system_prompt

Expand All @@ -100,7 +104,7 @@ def main(args):

if args.timings: started_timestamp = time.time()

prompt = f'{args.chat_template.format(input=text)}'
prompt = f'{args.chat_template.format(system_prompt=system_prompt, input=text)}'

Check warning

Code scanning / CodeQL

Unused named argument in formatting call

Surplus named argument for string format. An argument named 'system_prompt' is provided, but it is not required by [format "<|im_start|>user<|im_sep|> {input}<|im_end|> <|im_start|>assistant<|im_sep|>"](1).

Copilot Autofix

AI over 1 year ago

To fix the issue, we need to remove the surplus system_prompt argument from the format call on line 107. This involves:

  1. Verifying that the format string in args.chat_template does not include a placeholder for system_prompt.
  2. Removing the system_prompt=system_prompt argument from the format call, leaving only the required arguments.

This change ensures that the code is cleaner and avoids unnecessary arguments in the formatting call.

Suggested changeset 1
examples/python/model-chat.py

Autofix patch

Autofix patch
Run the following command in your local git repository to apply this patch
cat << 'EOF' | git apply
diff --git a/examples/python/model-chat.py b/examples/python/model-chat.py
--- a/examples/python/model-chat.py
+++ b/examples/python/model-chat.py
@@ -106,3 +106,3 @@
 
-        prompt = f'{args.chat_template.format(system_prompt=system_prompt, input=text)}'
+        prompt = f'{args.chat_template.format(input=text)}'
         input_tokens = tokenizer.encode(prompt)
EOF
@@ -106,3 +106,3 @@

prompt = f'{args.chat_template.format(system_prompt=system_prompt, input=text)}'
prompt = f'{args.chat_template.format(input=text)}'
input_tokens = tokenizer.encode(prompt)
Copilot is powered by AI and may make mistakes. Always verify output.
input_tokens = tokenizer.encode(prompt)

generator.append_tokens(input_tokens)
Expand Down
12 changes: 6 additions & 6 deletions src/python/py/models/builder.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -367,7 +367,7 @@ def make_genai_config(self, model_name_or_path, extra_kwargs, out_dir):
},
"eos_token_id": config.eos_token_id,
"pad_token_id": config.pad_token_id if hasattr(config, "pad_token_id") and config.pad_token_id is not None else config.eos_token_id[0] if isinstance(config.eos_token_id, list) else config.eos_token_id,
"type": self.model_type[ : self.model_type.find("For")].lower(),
"type": self.model_type[ : self.model_type.find("For") if "For" in self.model_type else len(self.model_type)].lower(),
"vocab_size": self.vocab_size,
},
"search": {
Expand Down Expand Up @@ -2143,14 +2143,14 @@ def make_lm_head(self, lm_head):

if softcap_exists:
# Add final logit softcapping (Div --> Tanh --> Mul)
div_name = "/lm_head/Div"
div_name = "/lm_head/softcap/Div"
div_inputs = [f"{lm_name}/output_0", f"/model/constants/{self.to_str_dtype[self.io_dtype]}/0D/{self.lm_head_attrs['softcap']}"]
self.make_div(div_name, div_inputs, dtype=self.io_dtype, shape=['batch_size', 'sequence_length', self.vocab_size])

tanh_name = "/lm_head/Tanh"
tanh_name = "/lm_head/softcap/Tanh"
self.make_tanh(tanh_name, f"{div_name}/output_0", dtype=self.io_dtype, shape=['batch_size', 'sequence_length', self.vocab_size])

mul_name = "/lm_head/Mul"
mul_name = "/lm_head/softcap/Mul"
mul_inputs = [f"{tanh_name}/output_0", f"/model/constants/{self.to_str_dtype[self.io_dtype]}/0D/{self.lm_head_attrs['softcap']}"]
mul_output = "logits"
self.make_node('Mul', inputs=mul_inputs, outputs=[mul_output], name=mul_name)
Expand Down Expand Up @@ -2251,9 +2251,10 @@ def has_final_norm(self, module, 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
return hf_norm or hf_final_layernorm or hf_transformer_final_layernorm or gguf_final_norm
return hf_norm or hf_final_layernorm or hf_transformer_final_layernorm or hf_multimodal_final_layernorm or gguf_final_norm

def make_preprocessing_nodes(self):
self.make_attention_mask_reformatting()
Expand Down Expand Up @@ -2754,7 +2755,6 @@ class Gemma2Model(GemmaModel):
def __init__(self, config, io_dtype, onnx_dtype, ep, cache_dir, extra_options):
super().__init__(config, io_dtype, onnx_dtype, ep, cache_dir, extra_options)
self.attention_attrs["scale"] = config.query_pre_attn_scalar ** -0.5
self.lm_head_attrs["scale"] = config.final_logit_softcapping if config.final_logit_softcapping is not None else 1.0
self.is_local = lambda layer_id: layer_id % 2 == 1

def make_layer(self, layer_id, layer):
Expand Down