[TPU] Call torch._sync(param) during weight loading

[WoosukKwon]

During weight loading, we often do something like:

narrowed_tensor = param.data.narrow(0, offset, len)
narrowed_tensor.copy_(real_weight)

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,

• Before this PR: 3.4 GB allocated after weight loading
• After this PR: 2.0 GB allocated after weight loading

[github-actions]

👋 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.

🚀

[WoosukKwon]

Thanks @JackCaoG for finding out the bug and providing the solution.

[JackCaoG]

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.

[WoosukKwon]

Thanks for the elaboration. Fixed the comment!

[JackCaoG]

lgtm

[mgoin]

Thanks for referencing the CT issue, LGTM!