[generate] skip compilation on cpu offload

[gante]

What does this PR do?

torch.compile + model CPU offload is resulting in crashes. It should work in theory, but it's not working atm.

from transformers import AutoModelForCausalLM, AutoTokenizer

device_map = {"model.embed_tokens": 0, "model.layers.0": 0, "model.layers.1": "cpu", "model.norm": "cpu", "lm_head": 0}
model = AutoModelForCausalLM.from_pretrained(
    "hf-internal-testing/tiny-random-MistralForCausalLM", device_map=device_map
)
tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-MistralForCausalLM")
tokenized_inputs = tokenizer(["Hello world"], return_tensors="pt")
input_ids = tokenized_inputs.input_ids.to(0)

# Uses a compilable cache -> compilation happens under the hood
output = model.generate(input_ids, max_new_tokens=20, cache_implementation="static")

This PR:

1. Moves the logic to trigger "auto compile" into its own function
2. Disables "auto compile" when there is CPU offload (and disk offload too, which is not expected to support torch.compile)
3. Adds a test to prevent regressions

[github-actions]

Hi 👋, thank you for opening this pull request! The pull request is converted to draft by default. The CI will be paused while the PR is in draft mode. When it is ready for review, please click the Ready for review button (at the bottom of the PR page). This will assign reviewers and trigger CI.

[HuggingFaceDocBuilderDev]

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.

[SunMarc]

Thanks !

[SunMarc]

thanks for adding there cases, maybe we should also check for disk offload ?

[gante]

Good point! Adding too

[zucchini-nlp]

Nice, thanks for fixing!

One question, do we expose "auto-compile" in generate to users through config or is this still under the hood? We might raise a small warning in case users forced "auto-compile" but the model doesn't meet all criteria

[gante]

One question, do we expose "auto-compile" in generate to users through config or is this still under the hood?

There is generation_config.compile_config. Yes, agreed, we should throw a warning if it is set but we don't meet the conditions for compilation to happen -- adding a commit with it 👍