forked from Karylab-cklius/vllm
Richard Zou
8f4f8425d0
Summary ======= This PR is on the way to overlapping torch.compile and weight loading. See [design doc]([https://docs.google.com/document/d/1hssZeQv_lJlKqOr0vfpoqEUNn4ASGHVRB9a49yhcY6E/edit?tab=t.0](https://docs.google.com/document/d/1hssZeQv_lJlKqOr0vfpoqEUNn4ASGHVRB9a49yhcY6E/edit?tab=t.0)) and [proof-of-concept PR](https://github.com/vllm-project/vllm/pull/36072) and [RFC](https://github.com/vllm-project/vllm/issues/34956) - Adds `vllm compile <model> [options]` CLI command and `vllm.compile_model()` Python API - These APIs populate vLLM's torch.compile cache and do nothing else. A subsequent `vllm serve` call can read from the cache and perform a warm start. - They also use minimal GPU memory. This is accomplished through a combination of using FakeTensors (tensors with no storage that report device correctly) and Meta tensors (tensors with no storage that report device="meta"). Note that there is still some minimal GPU memory allocation (< 10 Mb, from GPUModelRunner runtime buffers), I did not go track down all of it, but I'm also not sure it matters. - In the future we can extend "vllm compile" to more than just torch.compile; for example, if vLLM uses triton kernels, or JIT'ed flashinfer kernels, `vllm compile` may also just compile those and saved the compiled artifacts somewhere. How it works ============ - `FakeModelLoader` wraps the real model loader. It initializes weights on `meta` device and runs weight post-processing on meta tensors. - Before torch.compile tracing, `swap_meta_params_to_fake()` converts meta params to FakeTensors. This is required so that torch.compile sees Tensors with the correct devices. - In theory we should also get the FakeModelLoader to give us FakeTensors, but FakeTensors do not yet support the type of Tensor subclasses that vLLM uses. - We raise the `CompilationDone` exception after cache artifacts are saved to avoid executing with fake tensors. (calling torch.compile performs both the compilation and an initial run of invoking the compiled artifact with the inputs) - `EngineCore` early-returns after `compile_or_warm_up_model`, skipping KV cache allocation, scheduler, and sampler setup Test plan ========= - Added tests for compile_only cold start followed by a warm start. The tests verify that the compile_only cold start uses no GPU memory, and the warm start does end up reading from the cache. Future work =========== In the following order: - add an option to overlap torch.compile and weight loading. The main process will do weight loading while spawning a new process to do compile-only work (that does not use gpu memory) - Extend this design to more weight loading schemes. For example, we currently support no weight processing. This will involve getting the additional weight processing to support meta tensors. Signed-off-by: Richard Zou <zou3519@gmail.com>
compile test folder structure
compile/test_*.py: various unit tests meant for testing particular code path/features. Future tests are most likely added here. New test files added here will be included in CI automaticallycompile/fullgraph/: full model tests, including all tests previously in compile/piecewise. These tests do not target particular features. New test files added here will be included in CI automaticallycompile/distributed/: tests that require multiple GPUs. New test files added here will NOT be included in CI automatically as these tests generally need to be manually configured to run in runners with particular number/type of GPUs.