Add unittests for Universal Assisted generation #8

gauravjain14 · 2024-12-12T02:43:21Z

Introduce Unit tests for Universal Assisted Generation

tests/test_universal_assisted_generation.py is intended to test the functionality introduced by universal assisted generation.

Note: All but test_basic_generation have been disabled for now.

Who can review?

@keyboardAnt @jmamou

gauravjain14 · 2024-12-12T02:46:34Z

Proposing to include some unittests to ensure functionality.

I am, however, encountering some errors in these dummy examples. Any inputs into what this test might be missing?

The following is the error I am seeing -

======================================================================
ERROR: test_basic_generation (__main__.TestUniversalSpeculativeDecoding.test_basic_generation)
Test basic speculative decoding works
----------------------------------------------------------------------
Traceback (most recent call last):
  File "/disk1/universal_assisted_generation/transformers/tests/test_universal_assisted_generation.py", line 45, in test_basic_generation
    candidates, scores = self.generator.get_candidates(input_ids)
                         ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/disk1/universal_assisted_generation/transformers/src/transformers/generation/candidate_generator.py", line 744, in get_candidates
    target_logits = self._atm_translator.get_target_logits(candidate_logits)
                    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/disk1/universal_assisted_generation/transformers/src/transformers/generation/candidate_generator.py", line 628, in get_target_logits
    .apply_(lambda x: self._assistant_to_target_input_ids[x])
     ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/disk1/universal_assisted_generation/transformers/src/transformers/generation/candidate_generator.py", line 628, in <lambda>
    .apply_(lambda x: self._assistant_to_target_input_ids[x])
                      ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~^^^
KeyError: 151665

----------------------------------------------------------------------
Ran 8 tests in 20.816s

FAILED (errors=1)

jmamou · 2024-12-12T06:33:55Z

Thanks @gauravjain14
For now, I propose to run the tests on #7 branch that contains bug fixes and not on the main branch

gauravjain14 · 2024-12-12T07:28:20Z

@jmamou I'll rebase this on that.

However, this error occurs on #7 as well.

keyboardAnt · 2024-12-12T07:33:09Z

Proposing to include some unittests to ensure functionality.

I am, however, encountering some errors in these dummy examples. Any inputs into what this test might be missing?

The following is the error I am seeing -


======================================================================

ERROR: test_basic_generation (__main__.TestUniversalSpeculativeDecoding.test_basic_generation)

Test basic speculative decoding works

----------------------------------------------------------------------

Traceback (most recent call last):

  File "/disk1/universal_assisted_generation/transformers/tests/test_universal_assisted_generation.py", line 45, in test_basic_generation

    candidates, scores = self.generator.get_candidates(input_ids)

                         ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^

  File "/disk1/universal_assisted_generation/transformers/src/transformers/generation/candidate_generator.py", line 744, in get_candidates

    target_logits = self._atm_translator.get_target_logits(candidate_logits)

                    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^

  File "/disk1/universal_assisted_generation/transformers/src/transformers/generation/candidate_generator.py", line 628, in get_target_logits

    .apply_(lambda x: self._assistant_to_target_input_ids[x])

     ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^

  File "/disk1/universal_assisted_generation/transformers/src/transformers/generation/candidate_generator.py", line 628, in <lambda>

    .apply_(lambda x: self._assistant_to_target_input_ids[x])

                      ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~^^^

KeyError: 151665



----------------------------------------------------------------------

Ran 8 tests in 20.816s



FAILED (errors=1)

It seems the drafter sampled a token that the translator does not include. Perhaps that token is not in the target vocabulary?

jmamou · 2024-12-12T12:27:51Z

@gauravjain14
Fixed in last push
#7

gauravjain14 · 2024-12-13T06:17:42Z

Rebased on @jmamou's changes in #7.

Removed some tests that seemed unnecessary. All tests pass.

I have disabled this test for now -

    def test_long_sequence(self):
        if False:
            """Test handling of very long input sequences"""
            long_input = torch.ones((1, 2048), dtype=torch.long)
            self.generator.input_ids = long_input
            candidates, scores = self.generator.get_candidates(long_input)
            self.assertLessEqual(
                candidates.shape[1],
                self.main_model.config.max_position_embeddings,
            )

Let me know what y'all think about it. If we should have it I can enable it. Disabled due to the context length.

keyboardAnt

@gauravjain14, thank you, it looks good! I added minor comments.

tests/test_universal_assisted_generation.py

keyboardAnt · 2024-12-13T07:48:34Z

tests/test_universal_assisted_generation.py

+    @classmethod
+    def setUpClass(cls):
+        # Setup main and assistant models
+        cls.main_model = AutoModelForCausalLM.from_pretrained(


Does it take <5s to load this 1B model? (Please see @gante's comment)

No, it takes about 33 seconds on a T4 machine. I think we can just add the tag @slow as mentioned in the comment. Wdyt?

What about using smaller models? There are a few examples of fast models used in existing Hugging Face tests.

@slow tests run less frequently, so I suggest striving for faster tests.

@gauravj14 @gauravjain14
Models for testing: https://huggingface.co/hf-internal-testing. For example, hf-internal-testing/tiny-random-gpt2 as used here.

tests/test_universal_assisted_generation.py

keyboardAnt · 2024-12-13T08:00:40Z

tests/generation/test_candidate_generator.py

Wdyt about moving the content of this file to tests/generation/test_candidate_generator.py?

gauravjain14 · 2024-12-18T10:02:31Z

Here's a quick update -

I am running the test 'test_mismatched_vocabulariesand for certain tokens the test is failing with the cache andpast_key_values` being empty (or of size 0). I am looking into but in case either of you want to give it a try, this is one of the tokens I have seen this issue on -

input_ids = torch.tensor([[128245]])

The test -

transformers/tests/test_universal_assisted_generation.py

Line 59 in 8e3d814

def test_mismatched_vocabularies(self):

jmamou · 2024-12-18T10:33:31Z

input_ids = torch.tensor([[128245]])

According to the target tokenizer, token_id 128245 corresponds to the special token '<|reserved_special_token_237|>'. We currently don't handle the case when the original prompt contains only special tokens.
I will try to handle that case.

gauravjain14 · 2024-12-18T10:35:32Z

input_ids = torch.tensor([[128245]])

According to the target tokenizer, token_id 128245 corresponds to the special token '<|reserved_special_token_237|>'. We currently don't handle the case when the original prompt contains only special tokens. I will try to handle that case.

Got it. Thanks for the quick response on that.

So, how do you propose we handle this for now? Should we skip special tokens in the target vocab or you think this will be a quick fix?

jmamou · 2024-12-18T11:07:16Z

input_ids = torch.tensor([[128245]])

According to the target tokenizer, token_id 128245 corresponds to the special token '<|reserved_special_token_237|>'. We currently don't handle the case when the original prompt contains only special tokens. I will try to handle that case.

Got it. Thanks for the quick response on that.

So, how do you propose we handle this for now? Should we skip special tokens in the target vocab or you think this will be a quick fix?

Let's skip special tokens for now. Note that UAG does not handle that case also.

keyboardAnt

LGTM. I only added two small comments.

tests/test_universal_assisted_generation.py

…ate_generator.py`

keyboardAnt

@gauravjain14, Please see the minor comment below.

keyboardAnt · 2024-12-19T19:03:22Z

tests/generation/test_candidate_generator.py

 import threading
 import unittest
 import weakref
 from unittest.mock import MagicMock

+from zmq import device


@gauravjain14 please remove this line

gauravjain14 force-pushed the unit_tests_usd branch from 3d7a709 to e18a060 Compare December 13, 2024 06:15

keyboardAnt requested review from keyboardAnt and jmamou and removed request for jmamou December 13, 2024 07:37

keyboardAnt requested changes Dec 13, 2024

View reviewed changes

Add unittests for Universal Assisted generation

4e92e9c

Remove unused import and fix test_speculation_depth test

011f595

gauravjain14 force-pushed the unit_tests_usd branch from 8e3d814 to 1a7f420 Compare December 18, 2024 21:06

exclude special and reserved tokens from tokenizer for UAG

2652490

gauravjain14 force-pushed the unit_tests_usd branch from 1a7f420 to 2652490 Compare December 18, 2024 21:08

gauravjain14 requested a review from keyboardAnt December 18, 2024 21:10

keyboardAnt reviewed Dec 18, 2024

View reviewed changes

tests/test_universal_assisted_generation.py Outdated Show resolved Hide resolved

tests/test_universal_assisted_generation.py Outdated Show resolved Hide resolved

mv test_universal_assisted_generation.py to `generation/test_candid…

701edbb

…ate_generator.py`

gauravjain14 merged commit 7088978 into usd Dec 19, 2024

keyboardAnt deleted the unit_tests_usd branch December 19, 2024 19:02

keyboardAnt restored the unit_tests_usd branch December 19, 2024 19:02

keyboardAnt reviewed Dec 19, 2024

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Add unittests for Universal Assisted generation #8

Add unittests for Universal Assisted generation #8

gauravjain14 commented Dec 12, 2024

gauravjain14 commented Dec 12, 2024

jmamou commented Dec 12, 2024

gauravjain14 commented Dec 12, 2024

keyboardAnt commented Dec 12, 2024 •

edited

Loading

jmamou commented Dec 12, 2024

gauravjain14 commented Dec 13, 2024

keyboardAnt left a comment

keyboardAnt Dec 13, 2024

gauravjain14 Dec 15, 2024

keyboardAnt Dec 15, 2024 •

edited

Loading

keyboardAnt Dec 16, 2024 •

edited

Loading

keyboardAnt Dec 13, 2024

gauravjain14 commented Dec 18, 2024

jmamou commented Dec 18, 2024

gauravjain14 commented Dec 18, 2024

jmamou commented Dec 18, 2024

keyboardAnt left a comment

keyboardAnt left a comment •

edited

Loading

keyboardAnt Dec 19, 2024

Add unittests for Universal Assisted generation #8

Add unittests for Universal Assisted generation #8

Conversation

gauravjain14 commented Dec 12, 2024

Who can review?

gauravjain14 commented Dec 12, 2024

jmamou commented Dec 12, 2024

gauravjain14 commented Dec 12, 2024

keyboardAnt commented Dec 12, 2024 • edited Loading

jmamou commented Dec 12, 2024

gauravjain14 commented Dec 13, 2024

keyboardAnt left a comment

Choose a reason for hiding this comment

keyboardAnt Dec 13, 2024

Choose a reason for hiding this comment

gauravjain14 Dec 15, 2024

Choose a reason for hiding this comment

keyboardAnt Dec 15, 2024 • edited Loading

Choose a reason for hiding this comment

keyboardAnt Dec 16, 2024 • edited Loading

Choose a reason for hiding this comment

keyboardAnt Dec 13, 2024

Choose a reason for hiding this comment

gauravjain14 commented Dec 18, 2024

jmamou commented Dec 18, 2024

gauravjain14 commented Dec 18, 2024

jmamou commented Dec 18, 2024

keyboardAnt left a comment

Choose a reason for hiding this comment

keyboardAnt left a comment • edited Loading

Choose a reason for hiding this comment

keyboardAnt Dec 19, 2024

Choose a reason for hiding this comment

keyboardAnt commented Dec 12, 2024 •

edited

Loading

keyboardAnt Dec 15, 2024 •

edited

Loading

keyboardAnt Dec 16, 2024 •

edited

Loading

keyboardAnt left a comment •

edited

Loading