test_generation_models.py 5.4 KB
Newer Older
1
2
3
4
5
6
7
8
9
10
11
12
13
# Copyright 2023-2024 SGLang Team
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
14
15
16
"""
Usage:

Kiv Chen's avatar
Kiv Chen committed
17
18
19
To test a specific model locally:
1. Add it to ALL_MODELS, for example, `ModelCase("Qwen/Qwen2-1.5B")`
2. Run `ONLY_RUN=Qwen/Qwen2-1.5B python3 -m unittest test_generation_models.TestGenerationModels`
20
21
22
"""

import dataclasses
23
import multiprocessing as mp
24
import os
Kiv Chen's avatar
Kiv Chen committed
25
import random
26
import unittest
27
from typing import List
28
29
30

import torch

31
32
33
34
35
36
from sglang.test.runners import (
    DEFAULT_PROMPTS,
    HFRunner,
    SRTRunner,
    check_close_model_outputs,
)
37
from sglang.test.test_utils import CustomTestCase, is_in_ci
38

39

40
41
42
43
44
45
46
@dataclasses.dataclass
class ModelCase:
    model_path: str
    tp_size: int = 1
    prefill_tolerance: float = 5e-2
    decode_tolerance: float = 5e-2
    rouge_l_tolerance: float = 1
47
    skip_long_prompt: bool = False
48
    trust_remote_code: bool = False
49
50


51
# Popular models that run on the CI
52
CI_MODELS = [
53
    ModelCase("meta-llama/Llama-3.1-8B-Instruct"),
54
    ModelCase("google/gemma-2-2b"),
55
]
56

Kiv Chen's avatar
Kiv Chen committed
57
58
59
# the complete set of models to test sglang's generation model
ALL_MODELS = [
    *CI_MODELS,
60
    ModelCase("Qwen/Qwen2-1.5B"),
61
    ModelCase("Qwen/Qwen2.5-14B-Instruct"),
62
63
    ModelCase("HuggingFaceTB/SmolLM-135M-Instruct", skip_long_prompt=True),
    ModelCase("allenai/OLMo-1B-0724-hf", decode_tolerance=8e-2, skip_long_prompt=True),
64
65
66
    ModelCase(
        "THUDM/glm-4-9b-chat", tp_size=2, trust_remote_code=True, skip_long_prompt=True
    ),
Chayenne's avatar
Chayenne committed
67
    ModelCase("openai-community/gpt2"),
Kiv Chen's avatar
Kiv Chen committed
68
    ModelCase("microsoft/Phi-3-small-8k-instruct", trust_remote_code=True),
Jani Monoses's avatar
Jani Monoses committed
69
    ModelCase("allenai/OLMo-2-1124-7B-Instruct", skip_long_prompt=True),
70
    ModelCase("ibm-granite/granite-3.0-2b-instruct", skip_long_prompt=True),
71
72
73
74
75
76
    ModelCase(
        "microsoft/Phi-3.5-MoE-instruct",
        tp_size=2,
        trust_remote_code=True,
        skip_long_prompt=True,
    ),
77
]
78

79
TORCH_DTYPES = [torch.float16]
80
81


82
class TestGenerationModels(CustomTestCase):
83

84
85
    @classmethod
    def setUpClass(cls):
86
        mp.set_start_method("spawn", force=True)
87

88
    def assert_close_logits_and_output_strs(
89
        self,
90
91
92
        prompts: List[str],
        model_case: ModelCase,
        torch_dtype: torch.dtype,
93
    ) -> None:
94
95
96
97
98
99
100
        model_path = model_case.model_path
        prefill_tolerance, decode_tolerance, rouge_l_tolerance = (
            model_case.prefill_tolerance,
            model_case.decode_tolerance,
            model_case.rouge_l_tolerance,
        )
        max_new_tokens = 32
101

102
        with HFRunner(
103
104
105
            model_path,
            torch_dtype=torch_dtype,
            model_type="generation",
106
            trust_remote_code=model_case.trust_remote_code,
107
        ) as hf_runner:
108
            hf_outputs = hf_runner.forward(prompts, max_new_tokens=max_new_tokens)
109
110
111

        with SRTRunner(
            model_path,
112
            tp_size=model_case.tp_size,
113
            torch_dtype=torch_dtype,
114
            model_type="generation",
115
            trust_remote_code=model_case.trust_remote_code,
116
        ) as srt_runner:
117
            srt_outputs = srt_runner.forward(prompts, max_new_tokens=max_new_tokens)
118

119
120
121
122
123
124
125
        check_close_model_outputs(
            hf_outputs=hf_outputs,
            srt_outputs=srt_outputs,
            prefill_tolerance=model_case.prefill_tolerance,
            decode_tolerance=model_case.decode_tolerance,
            rouge_l_tolerance=model_case.rouge_l_tolerance,
            debug_text=f"model_path={model_path} prompts={prompts}",
126
        )
127

Kiv Chen's avatar
Kiv Chen committed
128
    @unittest.skipIf(not is_in_ci(), "Local test should run all models")
129
130
    def test_ci_models(self):
        for model_case in CI_MODELS:
131
            for torch_dtype in TORCH_DTYPES:
Kiv Chen's avatar
Kiv Chen committed
132
                prompts = DEFAULT_PROMPTS
133
134
135
136
137
138

                # Skip long prompts for models that do not have a long context
                if model_case.skip_long_prompt:
                    prompts = [p for p in DEFAULT_PROMPTS if len(p) < 1000]

                # Assert the logits and output strs are close
139
                self.assert_close_logits_and_output_strs(
140
                    prompts, model_case, torch_dtype
141
142
                )

Kiv Chen's avatar
Kiv Chen committed
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
    @unittest.skipIf(is_in_ci(), "CI only runs selected models for simplicity")
    def test_all_models(self):
        for model_case in ALL_MODELS:
            for torch_dtype in TORCH_DTYPES:
                if (
                    "ONLY_RUN" in os.environ
                    and os.environ["ONLY_RUN"] != model_case.model_path
                ):
                    continue

                # Skip long prompts for models that do not have a long context
                prompts = DEFAULT_PROMPTS
                if model_case.skip_long_prompt:
                    prompts = [p for p in DEFAULT_PROMPTS if len(p) < 1000]

                # Assert the logits and output strs are close
                self.assert_close_logits_and_output_strs(
                    prompts, model_case, torch_dtype
                )
162

163

164
if __name__ == "__main__":
Mingyi's avatar
Mingyi committed
165
    unittest.main()