Feature request
Currently, beam search effectively multiplies the batch size memory-wise and compute-wise by the batch size. If you have a batch size of 1 and a beam search of 8, model.forward sees 8 samples at once. This becomes an unnecessary problem for situations where needs that beam size but is only able to fit a smaller quantity of samples into memory.
That's why I suggest that a "beam search sub-batch" size be added to limit the number of samples (including beams) that are seen at once by the model. Eg, if the beam search batch size is 4, even if the main batch size is 1 and the beam search is 8, model.foward would only be called on 4 samples at once. Same for if the main batch size is 2 and the beam size is 8.
Motivation
With the size of LLM & their popularity, people are often stretching the capacities of their hardware, & in some cases (like majority voting for chain-of-thoughts), need beam search with some beam size.
If that beam size can't fit all at once in the memory, they are currently doomed, which is unnecessary, as the beams could be computed in a smaller quantity at once.
Feature request
Currently, beam search effectively multiplies the batch size memory-wise and compute-wise by the batch size. If you have a batch size of 1 and a beam search of 8,
model.forwardsees 8 samples at once. This becomes an unnecessary problem for situations where needs that beam size but is only able to fit a smaller quantity of samples into memory.That's why I suggest that a "beam search sub-batch" size be added to limit the number of samples (including beams) that are seen at once by the model. Eg, if the beam search batch size is 4, even if the main batch size is 1 and the beam search is 8,
model.fowardwould only be called on 4 samples at once. Same for if the main batch size is 2 and the beam size is 8.Motivation
With the size of LLM & their popularity, people are often stretching the capacities of their hardware, & in some cases (like majority voting for chain-of-thoughts), need beam search with some beam size.
If that beam size can't fit all at once in the memory, they are currently doomed, which is unnecessary, as the beams could be computed in a smaller quantity at once.