[core] overhaul memory profiling and fix backward compatibility #10511
Conversation
Signed-off-by: youkaichao <[email protected]>
youkaichao
requested review from
zhuohan123,
alexm-neuralmagic,
comaniac and
njhill
as code owners
|
👋 Hi! Thank you for contributing to the vLLM project. Once the PR is approved and ready to go, your PR reviewer(s) can run CI to test the changes comprehensively before merging. To run CI, PR reviewers can do one of these:
🚀 |
Signed-off-by: youkaichao <[email protected]>
youkaichao
requested review from
DarkLight1337,
robertgshaw2-neuralmagic and
simon-mo
as code owners
Signed-off-by: youkaichao <[email protected]>
Signed-off-by: youkaichao <[email protected]>
| | cuda memory | | ||
| | | torch memory | | | ||
| Before profiling: | --------- | +++++++++ | | | ||
| During profiling (peak): | --------- | +++++++++++++ | *** | | ||
| After profiling: | --------- | +++++++++++ | *** | |
There was a problem hiding this comment.
The alignment here is a bit off.
| (result.before_profile.cuda_memory_in_bytes) / (1024**3), | ||
| (result.before_profile.torch_memory_in_bytes) / (1024**3), | ||
| (result.after_profile.cuda_memory_in_bytes) / (1024**3), | ||
| (result.after_profile.torch_memory_in_bytes) / (1024**3), | ||
| result.torch_peak_memory_in_bytes / (1024**3), | ||
| result.non_torch_memory_in_bytes / (1024**3), | ||
| available_kv_cache_memory / (1024**3), |
There was a problem hiding this comment.
Btw we can import GiB_bytes from vllm.utils.
|
cc @joerunde |
| self.model_runner.profile_run() | ||
| torch.cuda.synchronize() | ||
| gc.collect() | ||
| torch.cuda.empty_cache() |
There was a problem hiding this comment.
Should we encapsulate more of the code into the profiler? The device sychronization, garbage collection / cache clean up, and the call to torch.cuda.reset_peak_memory_stats() could be moved into the context.
| result.after_profile.measure() | ||
|
|
||
| result.torch_peak_memory_in_bytes = torch.cuda.memory_stats( | ||
| )["allocated_bytes.all.peak"] |
There was a problem hiding this comment.
This torch peak memory will include bytes from other objects in the Python process other than the vLLM worker / LLM instance that we want to profile. Hence, the gpu_memory_utilization parameter would limit the total usage of the python process, not the usage of just the LLM instance. To fix, we need to measure the peak relative to a baseline torch memory recorded as part of the profiling.
There was a problem hiding this comment.
oh yeah, got your point. I think we need to record the baseline before we load the model.
|
This pull request has merge conflicts that must be resolved before it can be |
fixes #10451 , and clearly explain the memory classification and the procedure.
I also added the initial pytorch memory, to be aligned with the pytorch memory profiler.
the profiling procedure is extracted into
vllm/utils, so that we can use it later in v1 too.