A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 15, 14:45 UTC
[Ray] Improve documentation on batch inference (#16609)
Signed-off-by: Richard Liaw <rliaw@berkeley.edu>
Richard Liaw committed
Apr 16, 2025 at 22:19 UTC
8cac35ba435906fb7eb07e44fe1a8c26e8744f4e
2 files changed
+90
-109
examples/offline_inference/batch_llm_inference.py
new
+90
@@ -0,0 +1,90 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+"""
3
+This example shows how to use Ray Data for data parallel batch inference.
4
+
5
+Ray Data is a data processing framework that can handle large datasets
6
+and integrates tightly with vLLM for data-parallel inference.
7
+
8
+As of Ray 2.44, Ray Data has a native integration with
9
+vLLM (under ray.data.llm).
10
+
11
+Ray Data provides functionality for:
12
+* Reading and writing to cloud storage (S3, GCS, etc.)
13
+* Automatic sharding and load-balancing across a cluster
14
+* Optimized configuration of vLLM using continuous batching
15
+* Compatible with tensor/pipeline parallel inference as well.
16
+
17
+Learn more about Ray Data's LLM integration:
18
+https://docs.ray.io/en/latest/data/working-with-llms.html
19
+"""
20
+import ray
21
+from packaging.version import Version
22
+from ray.data.llm import build_llm_processor, vLLMEngineProcessorConfig
23
+
24
+assert Version(ray.__version__) >= Version(
25
+ "2.44.1"), "Ray version must be at least 2.44.1"
26
+
27
+# Uncomment to reduce clutter in stdout
28
+# ray.init(log_to_driver=False)
29
+# ray.data.DataContext.get_current().enable_progress_bars = False
30
+
31
+# Read one text file from S3. Ray Data supports reading multiple files
32
+# from cloud storage (such as JSONL, Parquet, CSV, binary format).
33
+ds = ray.data.read_text("s3://anonymous@air-example-data/prompts.txt")
34
+print(ds.schema())
35
+
36
+size = ds.count()
37
+print(f"Size of dataset: {size} prompts")
38
+
39
+# Configure vLLM engine.
40
+config = vLLMEngineProcessorConfig(
41
+ model_source="unsloth/Llama-3.1-8B-Instruct",
42
+ engine_kwargs={
43
+ "enable_chunked_prefill": True,
44
+ "max_num_batched_tokens": 4096,
45
+ "max_model_len": 16384,
46
+ },
47
+ concurrency=1, # set the number of parallel vLLM replicas
48
+ batch_size=64,
49
+)
50
+
51
+# Create a Processor object, which will be used to
52
+# do batch inference on the dataset
53
+vllm_processor = build_llm_processor(
54
+ config,
55
+ preprocess=lambda row: dict(
56
+ messages=[{
57
+ "role": "system",
58
+ "content": "You are a bot that responds with haikus."
59
+ }, {
60
+ "role": "user",
61
+ "content": row["text"]
62
+ }],
63
+ sampling_params=dict(
64
+ temperature=0.3,
65
+ max_tokens=250,
66
+ )),
67
+ postprocess=lambda row: dict(
68
+ answer=row["generated_text"],
69
+ **row # This will return all the original columns in the dataset.
70
+ ),
71
+)
72
+
73
+ds = vllm_processor(ds)
74
+
75
+# Peek first 10 results.
76
+# NOTE: This is for local testing and debugging. For production use case,
77
+# one should write full result out as shown below.
78
+outputs = ds.take(limit=10)
79
+
80
+for output in outputs:
81
+ prompt = output["prompt"]
82
+ generated_text = output["generated_text"]
83
+ print(f"Prompt: {prompt!r}")
84
+ print(f"Generated text: {generated_text!r}")
85
+
86
+# Write inference output data out as Parquet files to S3.
87
+# Multiple files would be written to the output destination,
88
+# and each task would write one or more files separately.
89
+#
90
+# ds.write_parquet("s3://<your-output-bucket>")
examples/offline_inference/distributed.py
deleted
-109
@@ -1,109 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-"""
3
-This example shows how to use Ray Data for running offline batch inference
4
-distributively on a multi-nodes cluster.
5
-
6
-Learn more about Ray Data in https://docs.ray.io/en/latest/data/data.html
7
-"""
8
-
9
-from typing import Any
10
-
11
-import numpy as np
12
-import ray
13
-from packaging.version import Version
14
-from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy
15
-
16
-from vllm import LLM, SamplingParams
17
-
18
-assert Version(ray.__version__) >= Version(
19
- "2.22.0"), "Ray version must be at least 2.22.0"
20
-
21
-# Create a sampling params object.
22
-sampling_params = SamplingParams(temperature=0.8, top_p=0.95)
23
-
24
-# Set tensor parallelism per instance.
25
-tensor_parallel_size = 1
26
-
27
-# Set number of instances. Each instance will use tensor_parallel_size GPUs.
28
-num_instances = 1
29
-
30
-
31
-# Create a class to do batch inference.
32
-class LLMPredictor:
33
-
34
- def __init__(self):
35
- # Create an LLM.
36
- self.llm = LLM(model="meta-llama/Llama-2-7b-chat-hf",
37
- tensor_parallel_size=tensor_parallel_size)
38
-
39
- def __call__(self, batch: dict[str, np.ndarray]) -> dict[str, list]:
40
- # Generate texts from the prompts.
41
- # The output is a list of RequestOutput objects that contain the prompt,
42
- # generated text, and other information.
43
- outputs = self.llm.generate(batch["text"], sampling_params)
44
- prompt: list[str] = []
45
- generated_text: list[str] = []
46
- for output in outputs:
47
- prompt.append(output.prompt)
48
- generated_text.append(' '.join([o.text for o in output.outputs]))
49
- return {
50
- "prompt": prompt,
51
- "generated_text": generated_text,
52
- }
53
-
54
-
55
-# Read one text file from S3. Ray Data supports reading multiple files
56
-# from cloud storage (such as JSONL, Parquet, CSV, binary format).
57
-ds = ray.data.read_text("s3://anonymous@air-example-data/prompts.txt")
58
-
59
-
60
-# For tensor_parallel_size > 1, we need to create placement groups for vLLM
61
-# to use. Every actor has to have its own placement group.
62
-def scheduling_strategy_fn():
63
- # One bundle per tensor parallel worker
64
- pg = ray.util.placement_group(
65
- [{
66
- "GPU": 1,
67
- "CPU": 1
68
- }] * tensor_parallel_size,
69
- strategy="STRICT_PACK",
70
- )
71
- return dict(scheduling_strategy=PlacementGroupSchedulingStrategy(
72
- pg, placement_group_capture_child_tasks=True))
73
-
74
-
75
-resources_kwarg: dict[str, Any] = {}
76
-if tensor_parallel_size == 1:
77
- # For tensor_parallel_size == 1, we simply set num_gpus=1.
78
- resources_kwarg["num_gpus"] = 1
79
-else:
80
- # Otherwise, we have to set num_gpus=0 and provide
81
- # a function that will create a placement group for
82
- # each instance.
83
- resources_kwarg["num_gpus"] = 0
84
- resources_kwarg["ray_remote_args_fn"] = scheduling_strategy_fn
85
-
86
-# Apply batch inference for all input data.
87
-ds = ds.map_batches(
88
- LLMPredictor,
89
- # Set the concurrency to the number of LLM instances.
90
- concurrency=num_instances,
91
- # Specify the batch size for inference.
92
- batch_size=32,
93
- **resources_kwarg,
94
-)
95
-
96
-# Peek first 10 results.
97
-# NOTE: This is for local testing and debugging. For production use case,
98
-# one should write full result out as shown below.
99
-outputs = ds.take(limit=10)
100
-for output in outputs:
101
- prompt = output["prompt"]
102
- generated_text = output["generated_text"]
103
- print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")
104
-
105
-# Write inference output data out as Parquet files to S3.
106
-# Multiple files would be written to the output destination,
107
-# and each task would write one or more files separately.
108
-#
109
-# ds.write_parquet("s3://<your-output-bucket>")