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
[Misc] format and refactor some examples (#16252)
Signed-off-by: reidliu41 <reid201711@gmail.com> Co-authored-by: reidliu41 <reid201711@gmail.com>
Reid committed
Apr 8, 2025 at 18:42 UTC
7f00899ff7e192706a591d77af1e5fa51d22a8e2
13 files changed
+190
-127
examples/offline_inference/mistral-small.py
+4
-1
@@ -90,8 +90,9 @@ def run_simple_demo(args: argparse.Namespace):
90
},
91
]
92
outputs = llm.chat(messages, sampling_params=sampling_params)
93
-
93
+ print("-" * 50)
94
print(outputs[0].outputs[0].text)
95
+ print("-" * 50)
96
97
98
def run_advanced_demo(args: argparse.Namespace):
@@ -162,7 +163,9 @@ def run_advanced_demo(args: argparse.Namespace):
163
]
164
165
outputs = llm.chat(messages=messages, sampling_params=sampling_params)
166
+ print("-" * 50)
167
print(outputs[0].outputs[0].text)
168
+ print("-" * 50)
169
170
171
def main():
examples/offline_inference/multilora_inference.py
+2
@@ -61,6 +61,7 @@ def process_requests(engine: LLMEngine,
61
"""Continuously process a list of prompts and handle the outputs."""
62
request_id = 0
63
64
+ print("-" * 50)
65
while test_prompts or engine.has_unfinished_requests():
66
if test_prompts:
67
prompt, sampling_params, lora_request = test_prompts.pop(0)
@@ -75,6 +76,7 @@ def process_requests(engine: LLMEngine,
76
for request_output in request_outputs:
77
if request_output.finished:
78
print(request_output)
79
+ print("-" * 50)
80
81
82
def initialize_engine() -> LLMEngine:
examples/offline_inference/neuron.py
+33
-24
@@ -12,27 +12,36 @@ prompts = [
12
# Create a sampling params object.
13
sampling_params = SamplingParams(temperature=0.8, top_p=0.95)
14
15
-# Create an LLM.
16
-llm = LLM(
17
- model="TinyLlama/TinyLlama-1.1B-Chat-v1.0",
18
- max_num_seqs=8,
19
- # The max_model_len and block_size arguments are required to be same as
20
- # max sequence length when targeting neuron device.
21
- # Currently, this is a known limitation in continuous batching support
22
- # in transformers-neuronx.
23
- # TODO(liangfu): Support paged-attention in transformers-neuronx.
24
- max_model_len=1024,
25
- block_size=1024,
26
- # The device can be automatically detected when AWS Neuron SDK is installed.
27
- # The device argument can be either unspecified for automated detection,
28
- # or explicitly assigned.
29
- device="neuron",
30
- tensor_parallel_size=2)
31
-# Generate texts from the prompts. The output is a list of RequestOutput objects
32
-# that contain the prompt, generated text, and other information.
33
-outputs = llm.generate(prompts, sampling_params)
34
-# Print the outputs.
35
-for output in outputs:
36
- prompt = output.prompt
37
- generated_text = output.outputs[0].text
38
- print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")
15
+
16
+def main():
17
+ # Create an LLM.
18
+ llm = LLM(
19
+ model="TinyLlama/TinyLlama-1.1B-Chat-v1.0",
20
+ max_num_seqs=8,
21
+ # The max_model_len and block_size arguments are required to be same as
22
+ # max sequence length when targeting neuron device.
23
+ # Currently, this is a known limitation in continuous batching support
24
+ # in transformers-neuronx.
25
+ # TODO(liangfu): Support paged-attention in transformers-neuronx.
26
+ max_model_len=1024,
27
+ block_size=1024,
28
+ # ruff: noqa: E501
29
+ # The device can be automatically detected when AWS Neuron SDK is installed.
30
+ # The device argument can be either unspecified for automated detection,
31
+ # or explicitly assigned.
32
+ device="neuron",
33
+ tensor_parallel_size=2)
34
+ # Generate texts from the prompts. The output is a list of RequestOutput objects
35
+ # that contain the prompt, generated text, and other information.
36
+ outputs = llm.generate(prompts, sampling_params)
37
+ # Print the outputs.
38
+ print("-" * 50)
39
+ for output in outputs:
40
+ prompt = output.prompt
41
+ generated_text = output.outputs[0].text
42
+ print(f"Prompt: {prompt!r}\nGenerated text: {generated_text!r}")
43
+ print("-" * 50)
44
+
45
+
46
+if __name__ == "__main__":
47
+ main()
examples/offline_inference/neuron_int8_quantization.py
+37
-28
@@ -22,31 +22,40 @@ prompts = [
22
# Create a sampling params object.
23
sampling_params = SamplingParams(temperature=0.8, top_p=0.95)
24
25
-# Create an LLM.
26
-llm = LLM(
27
- model="TinyLlama/TinyLlama-1.1B-Chat-v1.0",
28
- max_num_seqs=8,
29
- # The max_model_len and block_size arguments are required to be same as
30
- # max sequence length when targeting neuron device.
31
- # Currently, this is a known limitation in continuous batching support
32
- # in transformers-neuronx.
33
- # TODO(liangfu): Support paged-attention in transformers-neuronx.
34
- max_model_len=2048,
35
- block_size=2048,
36
- # The device can be automatically detected when AWS Neuron SDK is installed.
37
- # The device argument can be either unspecified for automated detection,
38
- # or explicitly assigned.
39
- device="neuron",
40
- quantization="neuron_quant",
41
- override_neuron_config={
42
- "cast_logits_dtype": "bfloat16",
43
- },
44
- tensor_parallel_size=2)
45
-# Generate texts from the prompts. The output is a list of RequestOutput objects
46
-# that contain the prompt, generated text, and other information.
47
-outputs = llm.generate(prompts, sampling_params)
48
-# Print the outputs.
49
-for output in outputs:
50
- prompt = output.prompt
51
- generated_text = output.outputs[0].text
52
- print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")
25
+
26
+def main():
27
+ # Create an LLM.
28
+ llm = LLM(
29
+ model="TinyLlama/TinyLlama-1.1B-Chat-v1.0",
30
+ max_num_seqs=8,
31
+ # The max_model_len and block_size arguments are required to be same as
32
+ # max sequence length when targeting neuron device.
33
+ # Currently, this is a known limitation in continuous batching support
34
+ # in transformers-neuronx.
35
+ # TODO(liangfu): Support paged-attention in transformers-neuronx.
36
+ max_model_len=2048,
37
+ block_size=2048,
38
+ # ruff: noqa: E501
39
+ # The device can be automatically detected when AWS Neuron SDK is installed.
40
+ # The device argument can be either unspecified for automated detection,
41
+ # or explicitly assigned.
42
+ device="neuron",
43
+ quantization="neuron_quant",
44
+ override_neuron_config={
45
+ "cast_logits_dtype": "bfloat16",
46
+ },
47
+ tensor_parallel_size=2)
48
+ # Generate texts from the prompts. The output is a list of RequestOutput objects
49
+ # that contain the prompt, generated text, and other information.
50
+ outputs = llm.generate(prompts, sampling_params)
51
+ # Print the outputs.
52
+ print("-" * 50)
53
+ for output in outputs:
54
+ prompt = output.prompt
55
+ generated_text = output.outputs[0].text
56
+ print(f"Prompt: {prompt!r}\nGenerated text: {generated_text!r}")
57
+ print("-" * 50)
58
+
59
+
60
+if __name__ == "__main__":
61
+ main()
examples/offline_inference/prefix_caching.py
+59
-52
@@ -31,55 +31,62 @@ generating_prompts = [prefix + prompt for prompt in prompts]
31
# Create a sampling params object.
32
sampling_params = SamplingParams(temperature=0.0)
33
34
-# Create an LLM without prefix caching as a baseline.
35
-regular_llm = LLM(model="facebook/opt-125m", gpu_memory_utilization=0.4)
36
-
37
-print("Results without `enable_prefix_caching`")
38
-
39
-# Generate texts from the prompts. The output is a list of RequestOutput objects
40
-# that contain the prompt, generated text, and other information.
41
-outputs = regular_llm.generate(generating_prompts, sampling_params)
42
-
43
-regular_generated_texts = []
44
-# Print the outputs.
45
-for output in outputs:
46
- prompt = output.prompt
47
- generated_text = output.outputs[0].text
48
- regular_generated_texts.append(generated_text)
49
- print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")
50
-
51
-print("-" * 80)
52
-
53
-# Destroy the LLM object and free up the GPU memory.
54
-del regular_llm
55
-cleanup_dist_env_and_memory()
56
-
57
-# Create an LLM with prefix caching enabled.
58
-prefix_cached_llm = LLM(model="facebook/opt-125m",
59
- enable_prefix_caching=True,
60
- gpu_memory_utilization=0.4)
61
-
62
-# Warmup so that the shared prompt's KV cache is computed.
63
-prefix_cached_llm.generate(generating_prompts[0], sampling_params)
64
-
65
-# Generate with prefix caching.
66
-outputs = prefix_cached_llm.generate(generating_prompts, sampling_params)
67
-
68
-print("Results with `enable_prefix_caching`")
69
-
70
-cached_generated_texts = []
71
-# Print the outputs. You should see the same outputs as before.
72
-for output in outputs:
73
- prompt = output.prompt
74
- generated_text = output.outputs[0].text
75
- cached_generated_texts.append(generated_text)
76
- print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")
77
-
78
-print("-" * 80)
79
-
80
-# Compare the results and display the speedup
81
-generated_same = all([
82
- regular_generated_texts[i] == cached_generated_texts[i]
83
- for i in range(len(prompts))
84
-])
85
-print(f"Generated answers are the same: {generated_same}")
34
+
35
+def main():
36
+ # Create an LLM without prefix caching as a baseline.
37
+ regular_llm = LLM(model="facebook/opt-125m", gpu_memory_utilization=0.4)
38
+
39
+ print("Results without `enable_prefix_caching`")
40
+
41
+ # ruff: noqa: E501
42
+ # Generate texts from the prompts. The output is a list of RequestOutput objects
43
+ # that contain the prompt, generated text, and other information.
44
+ outputs = regular_llm.generate(generating_prompts, sampling_params)
45
+
46
+ regular_generated_texts = []
47
+ # Print the outputs.
48
+ print("-" * 50)
49
+ for output in outputs:
50
+ prompt = output.prompt
51
+ generated_text = output.outputs[0].text
52
+ regular_generated_texts.append(generated_text)
53
+ print(f"Prompt: {prompt!r}\nGenerated text: {generated_text!r}")
54
+ print("-" * 50)
55
+
56
+ # Destroy the LLM object and free up the GPU memory.
57
+ del regular_llm
58
+ cleanup_dist_env_and_memory()
59
+
60
+ # Create an LLM with prefix caching enabled.
61
+ prefix_cached_llm = LLM(model="facebook/opt-125m",
62
+ enable_prefix_caching=True,
63
+ gpu_memory_utilization=0.4)
64
+
65
+ # Warmup so that the shared prompt's KV cache is computed.
66
+ prefix_cached_llm.generate(generating_prompts[0], sampling_params)
67
+
68
+ # Generate with prefix caching.
69
+ outputs = prefix_cached_llm.generate(generating_prompts, sampling_params)
70
+
71
+ print("Results with `enable_prefix_caching`")
72
+
73
+ cached_generated_texts = []
74
+ # Print the outputs. You should see the same outputs as before.
75
+ print("-" * 50)
76
+ for output in outputs:
77
+ prompt = output.prompt
78
+ generated_text = output.outputs[0].text
79
+ cached_generated_texts.append(generated_text)
80
+ print(f"Prompt: {prompt!r}\nGenerated text: {generated_text!r}")
81
+ print("-" * 50)
82
+
83
+ # Compare the results and display the speedup
84
+ generated_same = all([
85
+ regular_generated_texts[i] == cached_generated_texts[i]
86
+ for i in range(len(prompts))
87
+ ])
88
+ print(f"Generated answers are the same: {generated_same}")
89
+
90
+
91
+if __name__ == "__main__":
92
+ main()
examples/offline_inference/reproduciblity.py
+14
-7
@@ -19,8 +19,6 @@ SEED = 42
19
# because it is almost impossible to make the scheduling deterministic in the
20
# online serving setting.
21
22
-llm = LLM(model="facebook/opt-125m", seed=SEED)
23
-
22
prompts = [
23
"Hello, my name is",
24
"The president of the United States is",
@@ -29,8 +27,17 @@ prompts = [
27
]
28
sampling_params = SamplingParams(temperature=0.8, top_p=0.95)
29
32
-outputs = llm.generate(prompts, sampling_params)
33
-for output in outputs:
34
- prompt = output.prompt
35
- generated_text = output.outputs[0].text
36
- print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")
30
+
31
+def main():
32
+ llm = LLM(model="facebook/opt-125m", seed=SEED)
33
+ outputs = llm.generate(prompts, sampling_params)
34
+ print("-" * 50)
35
+ for output in outputs:
36
+ prompt = output.prompt
37
+ generated_text = output.outputs[0].text
38
+ print(f"Prompt: {prompt!r}\nGenerated text: {generated_text!r}")
39
+ print("-" * 50)
40
+
41
+
42
+if __name__ == "__main__":
43
+ main()
examples/offline_inference/rlhf.py
+6
-2
@@ -85,11 +85,13 @@ sampling_params = SamplingParams(temperature=0)
85
86
outputs = ray.get(llm.generate.remote(prompts, sampling_params))
87
88
+print("-" * 50)
89
for output in outputs:
90
prompt = output.prompt
91
generated_text = output.outputs[0].text
91
- print(f"Prompt: {prompt!r}, "
92
+ print(f"Prompt: {prompt!r}\n"
93
f"Generated text: {generated_text!r}")
94
+ print("-" * 50)
95
96
# set up the communication between the training process
97
# and the inference engine.
@@ -120,8 +122,10 @@ assert all(ray.get(llm.collective_rpc.remote("check_weights_changed")))
122
# use the updated model to generate texts, they will be nonsense
123
# because the weights are all zeros.
124
outputs_updated = ray.get(llm.generate.remote(prompts, sampling_params))
125
+print("-" * 50)
126
for output in outputs_updated:
127
prompt = output.prompt
128
generated_text = output.outputs[0].text
126
- print(f"Prompt: {prompt!r}, "
129
+ print(f"Prompt: {prompt!r}\n"
130
f"Generated text: {generated_text!r}")
131
+ print("-" * 50)
examples/offline_inference/simple_profiling.py
+3
-1
@@ -32,10 +32,12 @@ if __name__ == "__main__":
32
llm.stop_profile()
33
34
# Print the outputs.
35
+ print("-" * 50)
36
for output in outputs:
37
prompt = output.prompt
38
generated_text = output.outputs[0].text
38
- print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")
39
+ print(f"Prompt: {prompt!r}\nGenerated text: {generated_text!r}")
40
+ print("-" * 50)
41
42
# Add a buffer to wait for profiler in the background process
43
# (in case MP is on) to finish writing profiling output.
examples/offline_inference/torchrun_example.py
+3
-1
@@ -36,11 +36,13 @@ llm = LLM(
36
outputs = llm.generate(prompts, sampling_params)
37
38
# all ranks will have the same outputs
39
+print("-" * 50)
40
for output in outputs:
41
prompt = output.prompt
42
generated_text = output.outputs[0].text
42
- print(f"Prompt: {prompt!r}, "
43
+ print(f"Prompt: {prompt!r}\n"
44
f"Generated text: {generated_text!r}")
45
+ print("-" * 50)
46
"""
47
Further tips:
48
examples/offline_inference/tpu.py
+19
-11
@@ -16,14 +16,22 @@ N = 1
16
# Currently, top-p sampling is disabled. `top_p` should be 1.0.
17
sampling_params = SamplingParams(temperature=0, top_p=1.0, n=N, max_tokens=16)
18
19
-# Set `enforce_eager=True` to avoid ahead-of-time compilation.
20
-# In real workloads, `enforace_eager` should be `False`.
21
-llm = LLM(model="Qwen/Qwen2-1.5B-Instruct",
22
- max_num_batched_tokens=64,
23
- max_num_seqs=4)
24
-outputs = llm.generate(prompts, sampling_params)
25
-for output, answer in zip(outputs, answers):
26
- prompt = output.prompt
27
- generated_text = output.outputs[0].text
28
- print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")
29
- assert generated_text.startswith(answer)
19
+
20
+def main():
21
+ # Set `enforce_eager=True` to avoid ahead-of-time compilation.
22
+ # In real workloads, `enforace_eager` should be `False`.
23
+ llm = LLM(model="Qwen/Qwen2-1.5B-Instruct",
24
+ max_num_batched_tokens=64,
25
+ max_num_seqs=4)
26
+ outputs = llm.generate(prompts, sampling_params)
27
+ print("-" * 50)
28
+ for output, answer in zip(outputs, answers):
29
+ prompt = output.prompt
30
+ generated_text = output.outputs[0].text
31
+ print(f"Prompt: {prompt!r}\nGenerated text: {generated_text!r}")
32
+ assert generated_text.startswith(answer)
33
+ print("-" * 50)
34
+
35
+
36
+if __name__ == "__main__":
37
+ main()
examples/offline_inference/vision_language.py
+4
@@ -1089,14 +1089,18 @@ def main(args):
1089
start_time = time.time()
1090
outputs = llm.generate(inputs, sampling_params=sampling_params)
1091
elapsed_time = time.time() - start_time
1092
+ print("-" * 50)
1093
print("-- generate time = {}".format(elapsed_time))
1094
+ print("-" * 50)
1095
1096
else:
1097
outputs = llm.generate(inputs, sampling_params=sampling_params)
1098
1099
+ print("-" * 50)
1100
for o in outputs:
1101
generated_text = o.outputs[0].text
1102
print(generated_text)
1103
+ print("-" * 50)
1104
1105
1106
if __name__ == "__main__":
examples/offline_inference/vision_language_embedding.py
+2
@@ -143,8 +143,10 @@ def run_encode(model: str, modality: QueryModality, seed: Optional[int]):
143
"multi_modal_data": mm_data,
144
})
145
146
+ print("-" * 50)
147
for output in outputs:
148
print(output.outputs.embedding)
149
+ print("-" * 50)
150
151
152
def main(args: Namespace):
examples/offline_inference/vision_language_multi_image.py
+4
@@ -644,9 +644,11 @@ def run_generate(model, question: str, image_urls: list[str],
644
},
645
sampling_params=sampling_params)
646
647
+ print("-" * 50)
648
for o in outputs:
649
generated_text = o.outputs[0].text
650
print(generated_text)
651
+ print("-" * 50)
652
653
654
def run_chat(model: str, question: str, image_urls: list[str],
@@ -687,9 +689,11 @@ def run_chat(model: str, question: str, image_urls: list[str],
689
chat_template=req_data.chat_template,
690
)
691
692
+ print("-" * 50)
693
for o in outputs:
694
generated_text = o.outputs[0].text
695
print(generated_text)
696
+ print("-" * 50)
697
698
699
def main(args: Namespace):