Feature Request: Option to transfer logits to CPU during generation

[YunruiZhang]

Feature request

Currently, in model.generate, the Transformers implementation stores all logits on GPU until the generation is finished, at which point they are returned as a PyTorch tensor.

This design causes significant GPU memory usage, especially for long generations, since the logits for every token must remain in GPU memory until the end. As a result, the usable GPU memory is reduced, limiting batch size and sequence length when users want to return logits.

Proposed feature:
It would be very useful to add an option that transfers the logits to CPU at each step of generation (e.g., per token), and stores them as NumPy arrays (or CPU tensors). This would free up GPU memory during generation while still allowing users to access the logits afterwards.

Motivation

When using model.generate with output_scores=True, the logits for all generated tokens are accumulated on GPU until the generation finishes. For long sequences or larger models, this quickly consumes a large portion of GPU memory, which limits batch size, sequence length, and overall usability.

Your contribution

I’m happy to help, but I’m not very familiar with the current codebase.

[Rocketknight1]

cc @gante

[SamuelBarryCS]

Hey @YunruiZhang!
I opened a PR to offer that option here (default is of course to keep the current behavior i.e. logits stay on GPU), test non regression and benchmark memory usage w/ vs. wo/ offloading.
Would greatly appreciate your review!

[gante]

Hi @YunruiZhang @SamuelBarryCS 👋

The feature is desirable, but it will crash with an ongoing refactor if we address it now. We're refactoring output accumulation such that we can accumulate and output from generate arbitrary model outputs 🤗

The order of execution to enable this feature is

1. Abstract output accumulation into a separate function (Refactor output handling in generate for cleaner decoding methods #40887)
2. Accept arbitrary model output keys. This will be written by me, since it implies deprecations.
3. [this feature] enable offloading of optional generate outputs

This means I won't accept #40870 as is, as it would result in nasty merge conflicts :)

cc @manueldeprada

[SamuelBarryCS]

Hey @gante - thanks for your reply & understood :c
Do you want me to close the PR in this case or to leave it open and merge once the refactor in finished (I can handle merge conflicts) ? If this is not useful we can close I'm fine with it.