Skip to content

[GPT Big Code] Fix attention scaling #40041

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
vasqu merged 6 commits into from
Aug 11, 2025
Merged

[GPT Big Code] Fix attention scaling #40041

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
d345abb
fix
vasqu Aug 8, 2025
85103d8
update integration tests
vasqu Aug 11, 2025
e9a33c7
fmt
vasqu Aug 11, 2025
4b1ef2d
add regression test
vasqu Aug 11, 2025
1968f43
Merge branch 'main' into fix-gpt-bigcode
vasqu Aug 11, 2025
2ff1ad4
flaky ci
vasqu Aug 11, 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
2 changes: 1 addition & 1 deletion src/transformers/models/gpt_bigcode/modeling_gpt_bigcode.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -144,7 +144,7 @@ def __init__(self, config, is_cross_attention=False, layer_idx=None):
)

self.scale_attn_weights = config.scale_attn_weights
self.scaling = self.head_dim**0.5 if config.scale_attn_weights else 1.0
self.scaling = self.head_dim**-0.5 if config.scale_attn_weights else 1.0
self.is_cross_attention = is_cross_attention

self.layer_idx = layer_idx
Expand Down
23 changes: 20 additions & 3 deletions tests/models/gpt_bigcode/test_modeling_gpt_bigcode.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -29,6 +29,7 @@
import torch

from transformers import (
AutoTokenizer,
GPT2TokenizerFast,
GPTBigCodeForCausalLM,
GPTBigCodeForSequenceClassification,
Expand Down Expand Up @@ -510,7 +511,7 @@ def test_generate_simple(self):
output_sequence = model.generate(input_ids)
output_sentence = tokenizer.decode(output_sequence[0], skip_special_tokens=True)

expected_output = """def print_hello_world():\n print("Hello World!")\n\n\ndef print_hello_"""
expected_output = 'def print_hello_world():\n print("Hello World!")\n\n\ndef print_hello_world_with_args(name' # fmt: skip
self.assertEqual(output_sentence, expected_output)

def test_generate_batched(self):
Expand All @@ -527,11 +528,27 @@ def test_generate_batched(self):
outputs = tokenizer.batch_decode(outputs, skip_special_tokens=True)

expected_output = [
'def print_hello_world():\n print("Hello World!")\n\n\ndef print_hello_',
'def say_hello():\n print("Hello, World!")\n\n\nsay_hello()',
'def print_hello_world():\n print("Hello World!")\n\n\ndef print_hello_world_with_args(name',
'def say_hello():\n print("Hello, World!")\n\n\nsay_hello()\n',
]
self.assertListEqual(outputs, expected_output)

def test_newline_regression(self):
"""Added to prevent regressions regarding attention (scaling) indicated by excessive newlines"""
tokenizer = AutoTokenizer.from_pretrained("bigcode/tiny_starcoder_py")
model = GPTBigCodeForCausalLM.from_pretrained("bigcode/tiny_starcoder_py").to(torch_device)

input_ids = tokenizer(
"Analyze the impact of the COVID-19 pandemic on global economic structures and future business models.\n",
return_tensors="pt",
).input_ids.to(torch_device)

output_sequence = model.generate(input_ids, max_new_tokens=20, do_sample=False)
output_sentence = tokenizer.decode(output_sequence[0], skip_special_tokens=True)

expected_output = 'Analyze the impact of the COVID-19 pandemic on global economic structures and future business models.\n\nThe impact of the COVID-19 pandemic on global economic structures and future business' # fmt: skip
self.assertEqual(output_sentence, expected_output)


@require_torch
class GPTBigCodeMQATest(unittest.TestCase):
Expand Down