Fix on "cache position" for assisted generation #30068

zucchini-nlp · 2024-04-05T11:59:12Z

What does this PR do?

Fixes #30042 and fixes #30068. As it was reported in the issue, assisted decoding generation does not match greedy. The problem was in assuming there is only one new token generated when "update_kwargs_for_generation". That does not affect other model kwargs, because assisted generation handles them internally. But for "cache position" we need to update it to the new length, which is "accepted num candidates".

I added a small test for assisted generation + cache position, passes for all models where "cache position" is supported. Tested if anything breaks when 0 candidates accepted or when all of them accepted, seems to work for me.

Btw, I want to notice again that making a universal "update_kwargs_for_generation" would be very nice, so that it does not assume one token being added. I can work on it as a new PR, maybe this can help our future refactoring a bit :)

Update: After a little explorationn I found that using only "update_model_kwargs" to get rid of extra function is feasible, but that needs two separate methods: for pre-fill and for decoding

HuggingFaceDocBuilderDev · 2024-04-05T12:20:43Z

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

gante

LGTM, but a question: how does the new test differ from the existing one? (I may be missing something :) )

gante · 2024-04-18T14:08:42Z

Related change: 82be569

zucchini-nlp · 2024-04-19T08:00:33Z

src/transformers/generation/utils.py

+            cur_len = model_kwargs["inputs_embeds"].shape[1]
+        else:
+            cur_len = input_ids.shape[-1]
+        model_kwargs["cache_position"] = torch.arange(past_length, cur_len, device=input_ids.device)


It was not working in case when past key values are passed into the "generate" in Jamba, because Jamba apparently makes attention mask shape based on cache position values. Tested all generation tests + LLama/Jamba/Gemma tests, all are passing on my end

zucchini-nlp · 2024-04-19T08:01:41Z

src/transformers/models/jamba/modeling_jamba.py

@@ -276,6 +278,14 @@ def reorder_cache(self, beam_idx: torch.LongTensor):
            device = self.ssm_states[layer_idx].device
            self.ssm_states[layer_idx] = self.ssm_states[layer_idx].index_select(0, beam_idx.to(device))

+    def get_seq_length(self, layer_idx: Optional[int] = 0) -> int:


This was causing errors since cache at layer_id=0 is always length=0. Current code is a workaround, I could not find a better way for Jamba yet...

This looks good to me 👍

gante

In general looks great to me 🙌 I haven't approved because there is something wrong with CI.

I've also added a few improvement suggestions

src/transformers/generation/utils.py

gante · 2024-04-19T10:00:01Z

src/transformers/generation/utils.py

@@ -4663,20 +4676,21 @@ def _assisted_decoding(
            # we use this forward pass to also pick the subsequent logits in the original model.

            # 2.1. Prepare the model inputs
-            model_kwargs = _prepare_attention_mask(
-                model_kwargs, candidate_input_ids.shape[1], self.config.is_encoder_decoder
+            candidate_kwargs = copy.copy(model_kwargs)


Do we need the copy? It was originally added before the whole assistant run was abstracted into AssistedCandidateGenerator, which now holds its own set of kwargs.

If we don't need the copy, let's remove it and use model_kwargs

yes, if we remove copy if fails because the cache position is from prev model run, with all candidates. It works only if we discard all changes made by candidates, and extend cache position based on number of accepted tokens

Makes sense! Thank you for walking me through it :D

src/transformers/generation/utils.py

gante · 2024-04-19T10:17:30Z

src/transformers/models/jamba/modeling_jamba.py

@@ -276,6 +278,14 @@ def reorder_cache(self, beam_idx: torch.LongTensor):
            device = self.ssm_states[layer_idx].device
            self.ssm_states[layer_idx] = self.ssm_states[layer_idx].index_select(0, beam_idx.to(device))

+    def get_seq_length(self, layer_idx: Optional[int] = 0) -> int:


This looks good to me 👍

tests/generation/test_utils.py

gante

CI is happy now :D

Let's sort the comments, then we can tag a core maintainer.

gante

Thank you for iterating 💛

src/transformers/testing_utils.py

tests/generation/test_utils.py

gante · 2024-04-22T10:11:30Z

src/transformers/generation/utils.py

@@ -4663,20 +4676,21 @@ def _assisted_decoding(
            # we use this forward pass to also pick the subsequent logits in the original model.

            # 2.1. Prepare the model inputs
-            model_kwargs = _prepare_attention_mask(
-                model_kwargs, candidate_input_ids.shape[1], self.config.is_encoder_decoder
+            candidate_kwargs = copy.copy(model_kwargs)


Makes sense! Thank you for walking me through it :D

Co-authored-by: Joao Gante <joaofranciscocardosogante@gmail.com>

gante · 2024-04-22T10:45:13Z

Failing test fixed by #30389

amyeroberts

Thanks for fixing!

Just a small comment about model-specific login in the mixin

amyeroberts · 2024-04-22T12:16:53Z

tests/generation/test_utils.py

-            assistant_model = model
+            # test with the same assistant model or randomly init one
+            # in the first case all candidate tokens are accepted, in the second none is accepted
+            # case when some are accepted and some not is hard to reproduce, so let's hope this catches most errors :)


Does this mean we expect the second case to fail periodically?

No, we do not expect the test to fail but we expect that candidates will be discarded. We had some problems with getting correct cache positions when all tokens are accepted vs not all. So it's better to check that it works in both cases

amyeroberts · 2024-04-22T16:17:32Z

src/transformers/generation/utils.py

+            elif "gptbigcode" in self.__class__.__name__.lower() or (
+                self.config.architectures is not None and "gptbigcode" in self.config.architectures[0].lower()
+            ):
+                if self.config.multi_query:
+                    past_length = model_kwargs["past_key_values"][0].shape[1]
+                else:
+                    past_length = model_kwargs["past_key_values"][0].shape[2]


Rather than have this model specific code in the mixin side - it would be cleaner for inidividual models to implement their own _get_initial_cache_position wherever necessary

Hmm, This one is special only fot GPT bigcode, but I think we can make special method in bigcode's modeling and fallback to the general method in all other model types.

We can remove the gptbigcode logic here now

oppps, removed

zucchini-nlp · 2024-04-23T08:15:41Z

@amyeroberts can be re-reviewed, failing tests are due to connection errors and not related to the PR

amyeroberts

Thanks for fixing!

zucchini-nlp requested a review from gante April 5, 2024 11:59

gante reviewed Apr 18, 2024

View reviewed changes

zucchini-nlp commented Apr 19, 2024

View reviewed changes

clean commit history I hope

Verified

This commit was created on GitHub.com and signed with GitHub’s verified signature. The key has expired.

GPG key ID: 4AEE18F83AFDEB23
Expired

Verified
Learn about vigilant mode

c5d3c0f

zucchini-nlp force-pushed the assisted_gen branch from 4353bbc to c5d3c0f Compare April 19, 2024 08:06

zucchini-nlp mentioned this pull request Apr 19, 2024

assisted_decoding does not correctly update static KV cache #30312

Closed

get kv seq length correctly

640da9b

gante reviewed Apr 19, 2024

View reviewed changes

gante approved these changes Apr 19, 2024

View reviewed changes

PR suggestions

Loading
Loading status checks…

4db6130

zucchini-nlp requested a review from gante April 19, 2024 12:23

Merge branch 'huggingface:main' into assisted_gen

Loading
Loading status checks…

be6b477

gante approved these changes Apr 22, 2024

View reviewed changes

gante requested a review from amyeroberts April 22, 2024 10:12

Update src/transformers/testing_utils.py

54ea7ae

Co-authored-by: Joao Gante <joaofranciscocardosogante@gmail.com>

add comment

14f423d

amyeroberts reviewed Apr 22, 2024

View reviewed changes

zucchini-nlp and others added 2 commits April 22, 2024 20:01

give gpt bigcode it's own overriden method

298d1d7

Merge branch 'huggingface:main' into assisted_gen

edb8448

remove code

efd10cb

amyeroberts approved these changes Apr 23, 2024

View reviewed changes

zucchini-nlp merged commit 77b59dc into huggingface:main Apr 23, 2024
21 checks passed

zucchini-nlp mentioned this pull request May 2, 2024

Some functional problems in the implementation of Speculative Decoding #30608

Closed

4 tasks

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Fix on "cache position" for assisted generation #30068

Fix on "cache position" for assisted generation #30068

zucchini-nlp commented Apr 5, 2024 •

edited

Loading

HuggingFaceDocBuilderDev commented Apr 5, 2024

gante left a comment

gante commented Apr 18, 2024

zucchini-nlp Apr 19, 2024

zucchini-nlp Apr 19, 2024

gante Apr 19, 2024

gante left a comment

gante Apr 19, 2024

zucchini-nlp Apr 19, 2024

gante Apr 22, 2024

gante Apr 19, 2024

gante left a comment

gante left a comment

gante Apr 22, 2024

gante commented Apr 22, 2024

amyeroberts left a comment

amyeroberts Apr 22, 2024

zucchini-nlp Apr 22, 2024

amyeroberts Apr 22, 2024

zucchini-nlp Apr 22, 2024

amyeroberts Apr 23, 2024

zucchini-nlp Apr 23, 2024

zucchini-nlp commented Apr 23, 2024

amyeroberts left a comment

Fix on "cache position" for assisted generation #30068

Fix on "cache position" for assisted generation #30068

Conversation

zucchini-nlp commented Apr 5, 2024 • edited Loading

What does this PR do?

HuggingFaceDocBuilderDev commented Apr 5, 2024

gante left a comment

Choose a reason for hiding this comment

gante commented Apr 18, 2024

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

gante left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

gante left a comment

Choose a reason for hiding this comment

gante left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

gante commented Apr 22, 2024

amyeroberts left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

zucchini-nlp commented Apr 23, 2024

amyeroberts left a comment

Choose a reason for hiding this comment

zucchini-nlp commented Apr 5, 2024 •

edited

Loading