add CFG for .generate()

[Vermeille]

This commit implements CFG

Fixes #24536 (I did not touch MusicGen)

Hope you enjoy it!

@sanchit-gandhi
@gante

[gante]

LGTM 👍

There are three missing bits before we can move forward:
1 - (I can take this one) adapt MusicGen, and merge its PR before this one
2 - Add a doc test using generate, in UnbatchedClassifierFreeGuidanceLogitsProcessor (example)
3 - Add a simple unit test (example; You may need to load a dummy model, example here)

Let me know if you'd want further pointers 🙌

[sanchit-gandhi]

Looks cool already @Vermeille! Left some questions about the implementation and default values below

[sanchit-gandhi]

Pinning use_cache=True could be unexpected for the user if they have set model.config.use_cache=False and then silently the model actually uses the k/v cache in the forward pass for the unconditional logits (requires higher memory usage, and so potentially silent OOMs)

[Vermeille]

Ditto. Any proposed way to fix it?

[sanchit-gandhi]

Possibly we could restrict the CFG logits processor to only run if use_cache is set to True? In a similar way to how we do for assisted_generation:

transformers/src/transformers/generation/utils.py

Lines 1494 to 1495 in 4957294

	if not model_kwargs["use_cache"]:
	raise ValueError("assisted generate requires `use_cache=True`")

This is not ideal though, since users should have the option of running the model with/without cache as they desire.

Perhaps instead we could add an attribute self.use_cache to the init of the logits processor, and set this based on whether the user has set use_cache to True / False? Along the lines of:

def CFGLogitsProcessor(LogitsProcessor):
    def __init__(..., use_cache):
        ...
        self.use_cache = use_cache

-> we'll set use_cache=model_kwargs["use_cache"] to get the correct argument here when we instantiate the logits processor.

WDYT @Vermeille @gante?

[gante]

I agree with requiring the use of cache to alleviate the implementation complexity -- we can revert it in the future if needed.

In practical terms, all modern models use a cache anyways ;)

[gante]

@Vermeille -- @sanchit-gandhi raises good points about the attention mask and taking the first item of the batch in the unconditional logits. As it stands, it will only work with batch size = 1, and our logits processors should be flexible wrt batch size :)

[Vermeille]

good catch with the batch size! As for the attention mask, could you guide me to a working solution with that? I'm quite unfamiliar with huggingface tbh.

[Vermeille]

Tests are on the way.

[Vermeille]

All right. we only need to address use_cache / attention_mask.

• use_cache: currently, the forward passes take care of automatically appending to the negative prompt. I don't think such a thing happens with use_cache=False so I gotta do the concat myself. probably meaning I have to make two branches based on the value of use_cache?
• attention_mask: Does it even make sense then to read out.logits[:, -1]? is -1 a valid index if that position has an attention_mask of 0 due to padding? If so, then I will concat a valid id to padding and the attention_mask will be something like [1, 1, 0, 0, 0, 0, 1, 1, 1, 1, 1], won't that screw up positional encoding with the "empty" slots?

Basically, I think .generate() had to answer the same questions so you guys will be able to answer them quite easily. Also I will need your guidance as for the API design to integrate this seamlessly.

[Vermeille]

@gante I think we're good. The failure looks totally unrelated.

[Vermeille]

Indeed. no more failed test.

[gante]

@Vermeille

• use_cache: currently, the forward passes take care of automatically appending to the negative prompt. I don't think such a thing happens with use_cache=False so I gotta do the concat myself. probably meaning I have to make two branches based on the value of use_cache?

As I've replied in the dedicated thread, don't worry about the uncached case :) Make sure an exception is thrown, though!
EDIT: I see that you've handled the uncached case. In that case, since you've already written the code, you can leave it be :)

• attention_mask: Does it even make sense then to read out.logits[:, -1]? is -1 a valid index if that position has an attention_mask of 0 due to padding? If so, then I will concat a valid id to padding and the attention_mask will be something like [1, 1, 0, 0, 0, 0, 1, 1, 1, 1, 1], won't that screw up positional encoding with the "empty" slots?

That's a non-issue: .generate() must always be used with left-padding, so you won't run into the case of picking a padded token with -1 indexing 🙌

[gante]

@Vermeille these two arguments are not part of the method configuration (similarly to the main prompt and its corresponding attention mask), so they have to be part of the signature of generate and passed into _get_logits_processor.

Should be an easy change 🤗

[HuggingFaceDocBuilderDev]

The documentation is not available anymore as the PR was closed or merged.

[Vermeille]

@gante looks like we're good now :)

[gante]

LGTM

(the music gen change has to be merged before this PR goes in, I'm working on it)

[gante]

(@sgugger this one possibly did not get through your notifications, gently pinging :) )

[sgugger]

Thanks for your PR! I agree with @sanchit-gandhi 's comment on model_kwargs but this can merged and iterated upon if you prefer it that way @gante .

[gante]

@Vermeille would you be able to retouch the tests? We can merge right after that change :)

[Vermeille]

I'm currently in vacations. What's the problem with the tests?

[gante]

@Vermeille there are a few patterns in the tests that we usually avoid in our codebase (like thin wrappers). However, it's a minor issue, and this feature is being requested by the community, so I'm favoring merging it now.

I understand the review process is long and somewhat tedious on your end. We err on the strict side, as we bear the cost of future maintenance. Thank you for collaborating with us, and looking forward to future contributions 🤗

Next steps: we will be communicating the feature on our end. Amplification and/or communication on your end will help bring awareness to the feature! 🔥

[sanchit-gandhi]

Thanks for your contribution @Vermeille and congrats on the PR!