[Ray Integration] Integrate vllm with experimental accelerated DAG API

[rkooo567]

Hi, we have an experimental accelerated DAG developed based on rep ray-project/enhancements#48.

TL;DR is that we are implementing a compile DAG api that can reduce the control plane overhead from ray (See _compiled_dag_init_dag for details). Our microbenchmark shows 35x control plane overhead reduction in scatter gather type of workload (which is equivalent to how vllm implements tensor parallel. I.e., send a single input to N actors and get the result from all actors);

PR summary

• Support a path to run compiled DAG.
• Use pickle for serialization instead of Ray's default serializer (cloudpickle). pickle is much cheaper to serialize the input required by vllm. Cloudpickle is useful when you have large data that can zero-copy.

Benchmark

The experimental feature can be used with the nightly Ray for the evaluation. I ran the llama 7B benchmark on g5.12xlarge (4 A10 GPUs), TP=4, 10 iteration.

python benchmark_latency.py --use-ray-compiled-dag --tensor-parallel-size 4 --num-iters 10 --model "meta-llama/Llama-2-7b-hf"

And got the result

Tp4, pathway: 1.7389498465017823
Tp4, default: 2.0027053648009314

which is about 13% improvement.

Limitation

Note that the current nightly implementation has a several limitation. We are working on follow-ups, but the feature itself could be enabled by default after these features are implemented.

• Currently, if the actor raises an exception or a worker dies, it hangs forever. It will be fixed by https://github.com/ray-project/ray/pull/41943/files
• Currently, once DAG is initialized, you cannot run other actor tasks. It will also be fixed by https://github.com/ray-project/ray/pull/41943/files
• Currently, you can have only 1 DAG for a set of actors (however, the 1 DAG can be reused without a new compilation). We are planning to fix this for other use case, but this is not required from VLLM now IIUC.

Followup

I am planning to make a follow up PR for these 2 features after merging https://github.com/ray-project/ray/pull/41943/files. The ETA is 12/26 (when I come back from OOO). Please let me know if async engine is a high priority item for evaluation.

• Better error handling
• Support async engine

[rkooo567]

Q: Please let me know Where I should add tests. I was thinking to add a flag to test_models.py, but it seems like it doesn't really test ray config (where tp > 1).

[njhill]

Presumably this is obsolete now that #2221 is merged?

[rkooo567]

@njhill we decided to contribute the feature off by default. We will internally productionize it within anyscale and consider to reenable in the future.

[rkooo567]

Decided to close over #2471