[core] improve cpu offloading implementation

[youkaichao]

make it friendly with torch.compile

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

🚀

[youkaichao]

it seems the original cuda tensor is still held alive somewhere. the weights are not offloaded actually.

[bdhirsh]

One thing to call out is that the monkey patching here will allow you to override torch.empty, but not any at::empty() calls that come from C++ anywhere in the dispatcher. I'm not sure if the particular code you're running is actually running into this, but the way we normally handle "factory functions that you want to override to return tensor subclasses" is with a TorchDispatchMode:

from torch.utils._python_dispatch import TorchDispatchMode

class OffloadedTensorMode(TorchDispatchMode):
    def __torch_dispatch__(self, func, types, args=(), kwargs=None):
        if kwargs is None:
            kwargs = {}
        rs = func(*args, **kwargs)
        if func is torch.ops.aten.empty.memory_format and rs.device != "cpu" and ...:
            rs = OffloadedTensor(rs)
        return rs

[youkaichao]

this is very helpful!

[bdhirsh]

I think if your tensor subclass internally holds another tensor (elem here), you probably want to user the "wrapper" subclass API. Example here

out_tensor = torch.Tensor._make_wrapper_subclass(cls, shape, **kwargs)
out_tensor.elem = weak_ref_tensor(elem)

Side note: I would probably call that constructor unconditionally, any reason you aren't doing it when requires_grad is None?

[youkaichao]

I copy the code from https://github.com/albanD/subclass_zoo/blob/ec47458346c2a1cfcd5e676926a4bbc6709ff62e/base_tensor.py#L11

[youkaichao]

need to have __new__ , similar to https://github.com/albanD/subclass_zoo/blob/ec47458346c2a1cfcd5e676926a4bbc6709ff62e/base_tensor.py#L11

we might also need https://github.com/albanD/subclass_zoo/blob/ec47458346c2a1cfcd5e676926a4bbc6709ff62e/base_tensor.py#L25

[bdhirsh]

We do generally have support for subclasses that implement both torch_function and torch_dispatch, although if you only need torch_dispatch then I agree that you probably want to disable torch_function as linked above.

Let me know if you have any other questions / would like to chat more about the subclass work you're doing!

[youkaichao]

maybe use torch.device("cpu") instead of "cpu"?

[bdhirsh]

hmm tensor.device != "cpu should generally be ok

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @youkaichao.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork