reproducibility.py 1.18 KB
Newer Older
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
5
6
7
8
"""
Demonstrates how to achieve reproducibility in vLLM.

Main article: https://docs.vllm.ai/en/latest/usage/reproducibility.html
"""

9
import os
10
import random
11
12
13

from vllm import LLM, SamplingParams

14
# Turn off multiprocessing to make the scheduling deterministic.
15
os.environ["VLLM_ENABLE_V1_MULTIPROCESSING"] = "0"
16

17
18
19
20
21
22
23
24
prompts = [
    "Hello, my name is",
    "The president of the United States is",
    "The capital of France is",
    "The future of AI is",
]
sampling_params = SamplingParams(temperature=0.8, top_p=0.95)

25
26

def main():
27
    llm = LLM(model="facebook/opt-125m")
28
29
30
31
32
33
34
35
    outputs = llm.generate(prompts, sampling_params)
    print("-" * 50)
    for output in outputs:
        prompt = output.prompt
        generated_text = output.outputs[0].text
        print(f"Prompt: {prompt!r}\nGenerated text: {generated_text!r}")
        print("-" * 50)

36
37
38
39
40
    # Try generating random numbers outside vLLM
    # The same number is output across runs, meaning that the random state
    # in the user code has been updated by vLLM
    print(random.randint(0, 100))

41
42
43

if __name__ == "__main__":
    main()