Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

[core] further polish memory profiling #12126

Merged
merged 6 commits into from
Jan 18, 2025
Merged

[core] further polish memory profiling #12126

merged 6 commits into from
Jan 18, 2025

Conversation

youkaichao
Copy link
Member

@youkaichao youkaichao commented Jan 16, 2025 •
edited by github-actions bot
Loading

Improve upon #11809

current main branch, vllm serve meta-llama/Llama-3.2-11B-Vision --load-format dummy --max-model-len 65536 --max-num-seqs 128 --enforce-eager will fail.

after this PR, it can work on H100-80G now:

INFO 01-16 09:40:00 worker.py:238] the current vLLM instance can use total_gpu_memory (79.22GiB) x gpu_memory_utilization (0.90) = 71.29GiB
INFO 01-16 09:40:00 worker.py:238] model weights take 19.86GiB; non_torch_memory takes 0.16GiB; PyTorch activation peak memory takes 36.63GiB; the rest of the memory reserved for KV Cache is 14.65GiB.
INFO 01-16 09:40:00 executor_base.py:95] # CUDA blocks: 6001, # CPU blocks: 1638

@youkaichao
fix
fa32104 
Signed-off-by: youkaichao <[email protected]>
@youkaichao
fix
41bf4f2 
Signed-off-by: youkaichao <[email protected]>
@youkaichao
fix
01b4364 
Signed-off-by: youkaichao <[email protected]>
@youkaichao
fix
42ee511 
Signed-off-by: youkaichao <[email protected]>
@youkaichao
fix
16779bd 
Signed-off-by: youkaichao <[email protected]>
@youkaichao
add tests
7e36d4a 
Signed-off-by: youkaichao <[email protected]>
@github-actions GitHub Actions
Copy link

👋 Hi! Thank you for contributing to the vLLM project.
Just a reminder: PRs would not trigger full CI run by default. Instead, it would only run fastcheck CI which starts running only a small and essential subset of CI tests to quickly catch errors. You can run other CI tests on top of those by going to your fastcheck build on Buildkite UI (linked in the PR checks section) and unblock them. If you do not have permission to unblock, ping simon-mo or khluu to add you in our Buildkite org.

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:

  • Add ready label to the PR
  • Enable auto-merge.

🚀

youkaichao
youkaichao commented Jan 16, 2025
View reviewed changes
tests/test_utils.py

# load weights

weights = torch.randn(128, 1024, 1024, device='cuda', dtype=torch.float32)

weights_memory_in_bytes = 128 * 1024 * 1024 * 4 # 512 MiB
Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

refactor in this PR: remote the _in_bytes in variable name to make the name shorter.

youkaichao
youkaichao commented Jan 16, 2025
View reviewed changes
tests/test_utils.py
assert abs(non_torch_ratio - 1) <= 0.05
assert abs(torch_peak_ratio - 1) <= 0.05
Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

now this becomes accurate.

youkaichao
youkaichao commented Jan 16, 2025
View reviewed changes
vllm/utils.py
Comment on lines +1937 to +1943
# we measure the torch peak memory usage via allocated_bytes,
# rather than `torch.cuda.memory_reserved()` .
# After `torch.cuda.reset_peak_memory_stats()`,
# `torch.cuda.memory_reserved()` will keep growing, and only shrink
# when we call `torch.cuda.empty_cache()` or OOM happens.
self.torch_peak = torch.cuda.memory_stats().get(
"allocated_bytes.all.peak", 0)
Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

this is the key change, reported by @gshtras

youkaichao
youkaichao commented Jan 16, 2025
View reviewed changes
vllm/worker/worker.py
Comment on lines +141 to +142
torch.cuda.reset_peak_memory_stats()
self.baseline_snapshot = MemorySnapshot()
Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

another key change: we also measure the non-torch memory before creating the vllm instance.

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Do we need to do this in v1/worker/gpu_worker.py as well?

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

v1 does not use this memory_profiling utility yet. welcome to port it to v1 code path!

@gshtras
Copy link
Contributor

gshtras commented Jan 16, 2025

Ran some tests and can confirm that this fixes the original issue we had with Llama3.2 90B model peak memory jumping from 98GB to 160GB

tlrmchlsmth
tlrmchlsmth approved these changes Jan 17, 2025
View reviewed changes
Copy link
Collaborator

@tlrmchlsmth tlrmchlsmth left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for wrestling this into a better state!

@youkaichao youkaichao enabled auto-merge (squash) January 17, 2025 15:48
@github-actions github-actions bot added the ready ONLY add when PR is ready to merge/full CI is needed label Jan 17, 2025
@youkaichao
Copy link
Member Author

failed tests look unrelated, merging.

@youkaichao youkaichao disabled auto-merge January 18, 2025 04:25
@youkaichao youkaichao merged commit da02cb4 into vllm-project:main Jan 18, 2025
67 of 70 checks passed
@youkaichao youkaichao deleted the fix_profiling branch January 18, 2025 04:25
joennlae pushed a commit to 44ai-labs/vllm that referenced this pull request Jan 19, 2025
@youkaichao
[core] further polish memory profiling (vllm-project#12126)
bb23d3c 
Signed-off-by: youkaichao <[email protected]>
joennlae pushed a commit to 44ai-labs/vllm that referenced this pull request Jan 19, 2025
@youkaichao
[core] further polish memory profiling (vllm-project#12126)
a6ac593 
Signed-off-by: youkaichao <[email protected]>
lckr pushed a commit to lckr/vllm that referenced this pull request Jan 19, 2025
@youkaichao
[core] further polish memory profiling (vllm-project#12126)
914bc14 
Signed-off-by: youkaichao <[email protected]>
@kzawora-intel kzawora-intel mentioned this pull request Jan 21, 2025
Open
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers

@tlrmchlsmth tlrmchlsmth tlrmchlsmth approved these changes

@zhuohan123 zhuohan123 Awaiting requested review from zhuohan123 zhuohan123 is a code owner

@comaniac comaniac Awaiting requested review from comaniac comaniac is a code owner

@njhill njhill Awaiting requested review from njhill njhill is a code owner

Assignees
No one assigned
Labels
ready ONLY add when PR is ready to merge/full CI is needed
Projects
None yet
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.
3 participants
@youkaichao @gshtras @tlrmchlsmth