[TPU] Call torch._sync(param) during weight loading #9437
Conversation
|
👋 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:
🚀 |
|
Thanks @JackCaoG for finding out the bug and providing the solution. |
vllm/model_executor/utils.py
Outdated
| @@ -28,4 +29,22 @@ def set_weight_attrs( | |||
| for key, value in weight_attrs.items(): | |||
| assert not hasattr( | |||
| weight, key), (f"Overwriting existing tensor attribute: {key}") | |||
|
|
|||
| # NOTE(woosuk): For TPU, param.data.copy_(weight) happens lazily, | |||
There was a problem hiding this comment.
to be more accurate this is because in VLLM we do
narrowed_tensor = param.data.narrow(0, offset, len)
narrowed_tensor.copy_(real_weight)
narrowed_tensor and param.data share the same storage. With functionization, the in place update on the narrowed_tensor will lazily propagate to the base tensor which is param.data.
There was a problem hiding this comment.
Thanks for the elaboration. Fixed the comment!
|
lgtm |
WoosukKwon
changed the title
During weight loading, we often do something like:
expecting narrowed_tensor and param.data to share the same storage. However, on TPUs, narrowed_tensor will lazily propagate to the base tensor, which is param.data, leading to the redundant memory usage. This sometimes causes OOM errors during model loading.
This PR address this problem by adding a post-hook to call
torch._sync(param)after the weight loader of each param is called.When loading Llama3-8B (bf16) on v5e-8,