|
|
import time |
|
|
import unittest |
|
|
|
|
|
from parameterized import parameterized |
|
|
|
|
|
from transformers import AutoModelForCausalLM, AutoTokenizer, GenerationConfig |
|
|
from transformers.testing_utils import require_flash_attn, require_torch_gpu, slow |
|
|
|
|
|
|
|
|
_TEST_PROMPTS = [ |
|
|
"A man is a walking his dog down the street, and a the turn he sees", |
|
|
"Describe a fruit that is of orange color and round. It is a sweet fruit and a great source of Vitamine C. The fruit I'm thinking of is an", |
|
|
"A plane is flying high in the sky, out of the window are clouds and mountains. Where could the plane be located?", |
|
|
"Please fill in the form to", |
|
|
"For safety reasons, the train is stopped in the middle of the", |
|
|
] |
|
|
|
|
|
_EXPECTED_OUTPUTS = [ |
|
|
"a woman standing on the sidewalk, looking at him. He is immediately drawn to her and feels a strong attraction. He walks up to her and strikes up a conversation, and they quickly discover that they have a lot in common. They exchange numbers and", |
|
|
"orange.\n\n## Step 1: Identify the key characteristics of the fruit\nThe fruit is described as being orange in color and round in shape.\n\n## Step 2: Determine the taste and nutritional value of the fruit\nThe fruit is described as sweet", |
|
|
"This riddle is a classic example of a lateral thinking puzzle, which requires the test-taker to think creatively and consider multiple possibilities. The answer is not a straightforward one, and it requires some lateral thinking to arrive at the correct solution.", |
|
|
"get in touch with us. We will respond to your message as soon as possible.\n\n[Your Name]\n[Your Email]\n[Your Phone Number]\n[Your Message]\n\nWe are looking forward to hearing from you!\n\n[Insert Contact Information]\n\nNote:", |
|
|
"track. The train is stopped for 30 minutes. The train is moving at a speed of 60 km/h. How many kilometers does the train travel in 30 minutes?\n## Step 1: Convert the speed from km/h to km/min", |
|
|
] |
|
|
|
|
|
|
|
|
@slow |
|
|
@require_flash_attn |
|
|
@require_torch_gpu |
|
|
class TestBatchGeneration(unittest.TestCase): |
|
|
@classmethod |
|
|
def setUpClass(cls): |
|
|
cls.model = AutoModelForCausalLM.from_pretrained( |
|
|
"meta-llama/Llama-3.2-3b-Instruct", torch_dtype="bfloat16", device_map="auto" |
|
|
).eval() |
|
|
|
|
|
cls.tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-3.2-3b-Instruct", padding_side="left") |
|
|
|
|
|
if cls.tokenizer.pad_token is None: |
|
|
cls.tokenizer.pad_token = cls.tokenizer.eos_token |
|
|
cls.model.config.pad_token_id = cls.model.config.eos_token_id |
|
|
|
|
|
cls.model.use_cache = False |
|
|
|
|
|
@parameterized.expand( |
|
|
[ |
|
|
("eager_paged", 64, 128, 64), |
|
|
("sdpa_paged", 32, 256, 128), |
|
|
("paged_attention", 16, 512, 256), |
|
|
("flex_paged", 64, 128, 64), |
|
|
] |
|
|
) |
|
|
def test_generate_batch_consistency(self, attn_impl, num_blocks, block_size, max_batch_tokens): |
|
|
self.model.config.attn_implementation = attn_impl |
|
|
|
|
|
generation_config = GenerationConfig( |
|
|
max_new_tokens=50, |
|
|
top_k=0, |
|
|
eos_token_id=self.tokenizer.eos_token_id, |
|
|
pad_token_id=self.tokenizer.pad_token_id, |
|
|
use_cache=False, |
|
|
num_blocks=num_blocks, |
|
|
block_size=block_size, |
|
|
max_batch_tokens=max_batch_tokens, |
|
|
) |
|
|
|
|
|
tokenized = self.tokenizer(_TEST_PROMPTS, truncation=True, max_length=512) |
|
|
batch_inputs = list(tokenized["input_ids"]) |
|
|
|
|
|
start = time.time() |
|
|
batch_outputs = self.model.generate_batch( |
|
|
inputs=batch_inputs, |
|
|
generation_config=generation_config, |
|
|
) |
|
|
end = time.time() |
|
|
print( |
|
|
f"\n[{attn_impl}] Batch took {end - start:.2f}s with config: blocks={num_blocks}, block_size={block_size}, max_batch_tokens={max_batch_tokens}" |
|
|
) |
|
|
|
|
|
for i, req_id in enumerate(batch_outputs): |
|
|
generated = self.tokenizer.decode(batch_outputs[req_id].static_outputs, skip_special_tokens=False).strip() |
|
|
expected = _EXPECTED_OUTPUTS[i].strip() |
|
|
self.assertTrue( |
|
|
generated.startswith(expected), |
|
|
msg=f"[{attn_impl}] Mismatch in request {i}:\nExpected start: {expected}\nGot: {generated}", |
|
|
) |
|
|
|
