Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
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
5 changes: 4 additions & 1 deletion src/transformers/generation/utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -2133,8 +2133,11 @@ def _valid_auto_compile_criteria(
generation_config.compile_config is not None and generation_config.compile_config._compile_all_devices
)
# Note: for some models that only use linear attention (e.g. Mamba), even a DynamicCache is compilable since all
# Encoder-decoder models hold that cache in a subcache, so we unwrap it to check the cache the decoder actually generates with
decoder_cache = cache.self_attention_cache if isinstance(cache, EncoderDecoderCache) else cache
# Note: for some models that only use linear attention (e.g. Mamba), even a DynamicCache is compilable since all
# layers are, but we don't want to ALWAYS compile when calling `generate`, so we check the type
using_compilable_cache = cache is not None and cache.is_compileable and type(cache) is not DynamicCache
using_compilable_cache = cache is not None and cache.is_compileable and type(decoder_cache) is not DynamicCache
can_compile = valid_hardware and using_compilable_cache

# Exception 1: Some quantization methods do not support compilation
Expand Down
17 changes: 17 additions & 0 deletions tests/generation/test_utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -5114,6 +5114,23 @@ def test_cpu_offload_doesnt_compile(self):
_ = model_cpu.generate(input_ids, **generate_kwargs)
self.assertFalse(hasattr(model_cpu, "_compiled_call"))

def test_compileable_default_cache_doesnt_compile_encoder_decoder(self):
"""Test that a compileable default cache doesn't trigger compilation on encoder-decoder models either"""
model = AutoModelForSeq2SeqLM.from_pretrained("hf-internal-testing/tiny-random-bart")
decoder_config = model.config.get_text_config(decoder=True)
# Linear attention layers are statically shaped, so the default `DynamicCache` is compileable (e.g. Mamba)
decoder_config.layer_types = ["linear_attention"] * decoder_config.num_hidden_layers
self_attention_cache = DynamicCache(config=decoder_config)
self.assertTrue(self_attention_cache.is_compileable)

cache = EncoderDecoderCache(
self_attention_cache, DynamicCache(config=model.config.get_text_config(decoder=True))
)
generation_config = GenerationConfig()
generation_config.compile_config = CompileConfig()
generation_config.compile_config._compile_all_devices = True # force compilation (e.g. fast CI, CPU)
self.assertFalse(model._valid_auto_compile_criteria({"past_key_values": cache}, generation_config))

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I guess it needs the static cache to be true?

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yup because otherwise the test passes on CPU regardless. Static cache alone doesn't repro the issue either (both main and the PR return True). The actual diff in behavior is with a compileable DynamicCache (main returns True but the fix returns False) so that's the behavior we're trying to catch.


def test_custom_generate_from_argument_in_generate(self):
"""Tests that the `custom_generate` argument is used when passed to `generate`"""
model = AutoModelForCausalLM.from_pretrained(
Expand Down
Loading