mirror of
https://github.com/huggingface/transformers.git
synced 2025-07-03 21:00:08 +06:00
209 lines
8.9 KiB
Python
209 lines
8.9 KiB
Python
# coding=utf-8
|
|
# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
"""Testing suite for the PyTorch Gemma2 model."""
|
|
|
|
import unittest
|
|
|
|
from pytest import mark
|
|
|
|
from transformers import AutoModelForCausalLM, AutoTokenizer, Gemma2Config, is_torch_available, pipeline
|
|
from transformers.testing_utils import (
|
|
require_flash_attn,
|
|
require_read_token,
|
|
require_torch,
|
|
require_torch_gpu,
|
|
slow,
|
|
torch_device,
|
|
)
|
|
|
|
from ...models.gemma.test_modeling_gemma import GemmaModelTest, GemmaModelTester
|
|
from ...test_configuration_common import ConfigTester
|
|
|
|
|
|
if is_torch_available():
|
|
import torch
|
|
|
|
from transformers import (
|
|
Gemma2ForCausalLM,
|
|
Gemma2ForSequenceClassification,
|
|
Gemma2ForTokenClassification,
|
|
Gemma2Model,
|
|
)
|
|
|
|
|
|
class Gemma2ModelTester(GemmaModelTester):
|
|
if is_torch_available():
|
|
config_class = Gemma2Config
|
|
model_class = Gemma2Model
|
|
for_causal_lm_class = Gemma2ForCausalLM
|
|
for_sequence_class = Gemma2ForSequenceClassification
|
|
for_token_class = Gemma2ForTokenClassification
|
|
|
|
|
|
@require_torch
|
|
class Gemma2ModelTest(GemmaModelTest, unittest.TestCase):
|
|
all_model_classes = (
|
|
(Gemma2Model, Gemma2ForCausalLM, Gemma2ForSequenceClassification, Gemma2ForTokenClassification)
|
|
if is_torch_available()
|
|
else ()
|
|
)
|
|
all_generative_model_classes = ()
|
|
pipeline_model_mapping = (
|
|
{
|
|
"feature-extraction": Gemma2Model,
|
|
"text-classification": Gemma2ForSequenceClassification,
|
|
"token-classification": Gemma2ForTokenClassification,
|
|
"text-generation": Gemma2ForCausalLM,
|
|
"zero-shot": Gemma2ForSequenceClassification,
|
|
}
|
|
if is_torch_available()
|
|
else {}
|
|
)
|
|
test_headmasking = False
|
|
test_pruning = False
|
|
_is_stateful = True
|
|
model_split_percents = [0.5, 0.6]
|
|
_torch_compile_test_ckpt = "google/gemma-2-9b"
|
|
|
|
def setUp(self):
|
|
self.model_tester = Gemma2ModelTester(self)
|
|
self.config_tester = ConfigTester(self, config_class=Gemma2Config, hidden_size=37)
|
|
|
|
@unittest.skip("Failing because of unique cache (HybridCache)")
|
|
def test_model_outputs_equivalence(self, **kwargs):
|
|
pass
|
|
|
|
@unittest.skip("Gemma2's eager attn/sdpa attn outputs are expected to be different")
|
|
def test_eager_matches_sdpa_inference(self):
|
|
pass
|
|
|
|
@unittest.skip("Gemma2's eager attn/sdpa attn outputs are expected to be different")
|
|
def test_sdpa_equivalence(self):
|
|
pass
|
|
|
|
def test_eager_attention_loaded_by_default(self):
|
|
"""Gemma 2 + SDPA = inferior results, because of the logit softcapping. Eager is the default."""
|
|
config, _ = self.model_tester.prepare_config_and_inputs_for_common()
|
|
|
|
# Usually we enable SDPA by default, but not for Gemma2
|
|
model = Gemma2Model(config)
|
|
self.assertTrue(model.config._attn_implementation == "eager")
|
|
|
|
# We can still force SDPA
|
|
config._attn_implementation = "sdpa"
|
|
model = Gemma2Model(config)
|
|
self.assertTrue(model.config._attn_implementation == "sdpa")
|
|
|
|
|
|
@slow
|
|
@require_torch_gpu
|
|
class Gemma2IntegrationTest(unittest.TestCase):
|
|
input_text = ["Hello I am doing", "Hi today"]
|
|
# This variable is used to determine which CUDA device are we using for our runners (A10 or T4)
|
|
# Depending on the hardware we get different logits / generations
|
|
cuda_compute_capability_major_version = None
|
|
|
|
@classmethod
|
|
def setUpClass(cls):
|
|
if is_torch_available() and torch.cuda.is_available():
|
|
# 8 is for A100 / A10 and 7 for T4
|
|
cls.cuda_compute_capability_major_version = torch.cuda.get_device_capability()[0]
|
|
|
|
@require_read_token
|
|
def test_model_9b_bf16(self):
|
|
model_id = "google/gemma-2-9b"
|
|
EXPECTED_TEXTS = [
|
|
"<bos>Hello I am doing a project on the 1918 flu pandemic and I am trying to find out how many",
|
|
"<pad><pad><bos>Hi today I'm going to be talking about the history of the United States. The United States of America",
|
|
]
|
|
|
|
model = AutoModelForCausalLM.from_pretrained(
|
|
model_id, low_cpu_mem_usage=True, torch_dtype=torch.bfloat16, attn_implementation="eager"
|
|
).to(torch_device)
|
|
|
|
tokenizer = AutoTokenizer.from_pretrained(model_id)
|
|
inputs = tokenizer(self.input_text, return_tensors="pt", padding=True).to(torch_device)
|
|
|
|
output = model.generate(**inputs, max_new_tokens=20, do_sample=False)
|
|
output_text = tokenizer.batch_decode(output, skip_special_tokens=False)
|
|
|
|
self.assertEqual(output_text, EXPECTED_TEXTS)
|
|
|
|
@require_read_token
|
|
def test_model_9b_fp16(self):
|
|
model_id = "google/gemma-2-9b"
|
|
EXPECTED_TEXTS = [
|
|
"<bos>Hello I am doing a project on the 1918 flu pandemic and I am trying to find out how many",
|
|
"<pad><pad><bos>Hi today I'm going to be talking about the history of the United States. The United States of America",
|
|
]
|
|
|
|
model = AutoModelForCausalLM.from_pretrained(
|
|
model_id, low_cpu_mem_usage=True, torch_dtype=torch.float16, attn_implementation="eager"
|
|
).to(torch_device)
|
|
|
|
tokenizer = AutoTokenizer.from_pretrained(model_id)
|
|
inputs = tokenizer(self.input_text, return_tensors="pt", padding=True).to(torch_device)
|
|
|
|
output = model.generate(**inputs, max_new_tokens=20, do_sample=False)
|
|
output_text = tokenizer.batch_decode(output, skip_special_tokens=False)
|
|
|
|
self.assertEqual(output_text, EXPECTED_TEXTS)
|
|
|
|
@require_read_token
|
|
def test_model_9b_pipeline_bf16(self):
|
|
# See https://github.com/huggingface/transformers/pull/31747 -- pipeline was broken for Gemma2 before this PR
|
|
model_id = "google/gemma-2-9b"
|
|
# EXPECTED_TEXTS should match the same non-pipeline test, minus the special tokens
|
|
EXPECTED_TEXTS = [
|
|
"Hello I am doing a project on the 1918 flu pandemic and I am trying to find out how many",
|
|
"Hi today I'm going to be talking about the history of the United States. The United States of America",
|
|
]
|
|
|
|
model = AutoModelForCausalLM.from_pretrained(model_id, low_cpu_mem_usage=True, torch_dtype=torch.bfloat16).to(
|
|
torch_device
|
|
)
|
|
tokenizer = AutoTokenizer.from_pretrained(model_id)
|
|
pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)
|
|
|
|
output = pipe(self.input_text, max_new_tokens=20, do_sample=False, padding=True)
|
|
|
|
self.assertEqual(output[0][0]["generated_text"], EXPECTED_TEXTS[0])
|
|
self.assertEqual(output[1][0]["generated_text"], EXPECTED_TEXTS[1])
|
|
|
|
@require_read_token
|
|
@require_flash_attn
|
|
@require_torch_gpu
|
|
@mark.flash_attn_test
|
|
@slow
|
|
def test_model_9b_flash_attn(self):
|
|
# See https://github.com/huggingface/transformers/issues/31953 --- flash attn was generating garbage for gemma2, especially in long context
|
|
model_id = "google/gemma-2-9b"
|
|
EXPECTED_TEXTS = [
|
|
'<bos>Hello I am doing a project on the 1918 flu pandemic and I am trying to find out how many people died in the United States. I have found a few sites that say 500,000 but I am not sure if that is correct. I have also found a site that says 675,000 but I am not sure if that is correct either. I am trying to find out how many people died in the United States. I have found a few',
|
|
"<pad><pad><bos>Hi today I'm going to be talking about the history of the United States. The United States of America is a country in North America. It is the third largest country in the world by total area and the third most populous country with over 320 million people. The United States is a federal republic consisting of 50 states and a federal district. The 48 contiguous states and the district of Columbia are in central North America between Canada and Mexico. The state of Alaska is in the"
|
|
] # fmt: skip
|
|
|
|
model = AutoModelForCausalLM.from_pretrained(
|
|
model_id, attn_implementation="flash_attention_2", torch_dtype="float16"
|
|
).to(torch_device)
|
|
tokenizer = AutoTokenizer.from_pretrained(model_id)
|
|
inputs = tokenizer(self.input_text, return_tensors="pt", padding=True).to(torch_device)
|
|
|
|
output = model.generate(**inputs, max_new_tokens=100, do_sample=False)
|
|
output_text = tokenizer.batch_decode(output, skip_special_tokens=False)
|
|
print(output_text)
|
|
|
|
self.assertEqual(output_text, EXPECTED_TEXTS)
|