Skip to content

Reduce peak memory for prompt_logprobs requests - #907

Open
50h100a wants to merge 3 commits into
mainfrom
prompt_logprobs_mem
Open

50h100a wants to merge 3 commits into
mainfrom
prompt_logprobs_mem

Conversation

@50h100a

@50h100a 50h100a commented Dec 16, 2024

Copy link
Copy Markdown
Collaborator

First order of business, make prompt_logprobs "compatible" with prefix caching.
It can't take advantage of the caching, but at least it will run.

Second order of business, reduce the peak memory usage of the samplers.
This PR slightly reduces the memory load, but not nearly enough:
On single-GPU, sampling can still take dozens of gigabytes at peak memory. (8b model at 16k was >10gb)
On multi-GPU, sampling is no cheaper, and there's also a colossal memory spike when gathering the logits.

Thoughts:

  • In this PR, some operations are split into smaller batches. Can we split the entire sampling process the same way? Leaving it mostly unchanged, but only handling a fixed k of rows at a time?
  • No idea what the fix is for the gather spikes, deferring to @AlpinDale on that. That might not even be the specific issue, just where it ran out of VRAM for me, but there's something about multi-GPU that's aggravating the memory peaks.

@AlpinDale
AlpinDale self-requested a review December 19, 2024 17:35
@AlpinDale

Copy link
Copy Markdown
Member

Will probably need some restructuring after #925

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants