main
py 45 lines 1.31 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """Demonstrates how to achieve reproducibility in vLLM.
4
5 Main article: https://docs.vllm.ai/en/latest/usage/reproducibility.html
6 """
7
8 import os
9 import random
10
11 from vllm import LLM, SamplingParams
12
13 # Either:
14 ## Turn off multiprocessing to make the scheduling deterministic, or
15 os.environ["VLLM_ENABLE_V1_MULTIPROCESSING"] = "0"
16 ## Enable batch invariance to get consistent results regardless of scheduling.
17 os.environ["VLLM_BATCH_INVARIANT"] = "1"
18
19 prompts = [
20 "Hello, my name is",
21 "The president of the United States is",
22 "The capital of France is",
23 "The future of AI is",
24 ]
25 sampling_params = SamplingParams(temperature=0.8, top_p=0.95)
26
27
28 def main():
29 llm = LLM(model="facebook/opt-125m")
30 outputs = llm.generate(prompts, sampling_params)
31 print("-" * 50)
32 for output in outputs:
33 prompt = output.prompt
34 generated_text = output.outputs[0].text
35 print(f"Prompt: {prompt!r}\nGenerated text: {generated_text!r}")
36 print("-" * 50)
37
38 # Try generating random numbers outside vLLM
39 # The same number is output across runs, meaning that the random state
40 # in the user code has been updated by vLLM
41 print(random.randint(0, 100))
42
43
44 if __name__ == "__main__":
45 main()