test_generation_models.py 5.84 KB
Newer Older
1
2
3
4
5
6
7
8
"""
Usage:

To test a specific model:
1. Add it to ALL_OTHER_MODELS
2. Run `ONLY_RUN=Qwen/Qwen2-1.5B python3 -m unittest test_generation_models.TestGenerationModels.test_others`
"""

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
14

15
    http://www.apache.org/licenses/LICENSE-2.0
16

17
18
19
20
21
22
23
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.
"""

24
import dataclasses
25
import multiprocessing as mp
26
import os
27
import unittest
28
from typing import List
29
30
31
32

import torch

from sglang.test.runners import DEFAULT_PROMPTS, HFRunner, SRTRunner
33
from sglang.test.test_utils import calculate_rouge_l, is_in_ci
34

35

36
37
38
39
40
41
42
@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
43
    skip_long_prompt: bool = False
44
45


46
# Popular models that run on the CI
47
CI_MODELS = [
48
    ModelCase("meta-llama/Llama-3.1-8B-Instruct"),
49
    ModelCase("google/gemma-2-2b"),
50
]
51

52
# All other models that do not run on the CI
53
54
ALL_OTHER_MODELS = [
    ModelCase("Qwen/Qwen2-1.5B"),
55
    ModelCase("Qwen/Qwen2.5-14B-Instruct"),
56
57
    ModelCase("HuggingFaceTB/SmolLM-135M-Instruct", skip_long_prompt=True),
    ModelCase("allenai/OLMo-1B-0724-hf", decode_tolerance=8e-2, skip_long_prompt=True),
58
    ModelCase("THUDM/glm-4-9b-chat"),
59
]
60

61
TORCH_DTYPES = [torch.float16]
62
63


64
class TestGenerationModels(unittest.TestCase):
65
66
67
68
    @classmethod
    def setUpClass(cls):
        mp.set_start_method("spawn")

69
    def assert_close_logits_and_output_strs(
70
        self,
71
72
73
        prompts: List[str],
        model_case: ModelCase,
        torch_dtype: torch.dtype,
74
    ) -> None:
75
76
77
78
79
80
81
        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
82

83
        with HFRunner(
84
85
86
            model_path,
            torch_dtype=torch_dtype,
            model_type="generation",
87
        ) as hf_runner:
88
            hf_outputs = hf_runner.forward(prompts, max_new_tokens=max_new_tokens)
89
90
91

        with SRTRunner(
            model_path,
92
            tp_size=model_case.tp_size,
93
            torch_dtype=torch_dtype,
94
            model_type="generation",
95
        ) as srt_runner:
96
            srt_outputs = srt_runner.forward(prompts, max_new_tokens=max_new_tokens)
97
98

        for i in range(len(prompts)):
99
            # Compare input logprobs
100
101
            hf_logprobs = torch.Tensor(hf_outputs.top_input_logprobs[i])
            srt_logprobs = torch.Tensor(srt_outputs.top_input_logprobs[i])
102
103
104
105
106
            input_len = hf_logprobs.shape[0]
            print(
                "prefill logprobs max_diff", torch.max(abs(hf_logprobs - srt_logprobs))
            )
            if input_len <= 100:
107
108
109
110
111
                assert torch.all(abs(hf_logprobs - srt_logprobs) < prefill_tolerance), (
                    f"prefill logprobs are not all close with model_path={model_path} prompts={prompts} "
                    f"prefill_tolerance={prefill_tolerance}."
                    f"{hf_logprobs=}, {srt_logprobs=}"
                )
112

113
            # Compare output logprobs
114
115
            hf_logprobs = torch.Tensor(hf_outputs.top_output_logprobs[i])
            srt_logprobs = torch.Tensor(srt_outputs.top_output_logprobs[i])
116

117
            print(
118
                "decode logprobs max_diff", torch.max(abs(hf_logprobs - srt_logprobs))
119
120
            )
            if input_len <= 100:
121
122
123
124
125
                assert torch.all(abs(hf_logprobs - srt_logprobs) < decode_tolerance), (
                    f"decode logprobs are not all close with model_path={model_path} prompts={prompts} "
                    f"decode_tolerance={decode_tolerance}."
                    f"{hf_logprobs=}, {srt_logprobs=}"
                )
126

127
128
129
        # Compare output strings
        print(f"{hf_outputs.output_strs=}")
        print(f"{srt_outputs.output_strs=}")
130
131
132
        rouge_l_scores = calculate_rouge_l(
            hf_outputs.output_strs, srt_outputs.output_strs
        )
133
        print(f"{rouge_l_scores=}")
134
        assert all(
135
136
137
138
139
            score >= rouge_l_tolerance for score in rouge_l_scores
        ), f"Not all ROUGE-L scores are greater than rouge_l_tolerance={rouge_l_tolerance}"

    def test_ci_models(self):
        for model_case in CI_MODELS:
140
            for torch_dtype in TORCH_DTYPES:
141
142
143
144
145
146
147

                # 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
148
                self.assert_close_logits_and_output_strs(
149
                    prompts, model_case, torch_dtype
150
151
                )

152
    def test_others(self):
153
154
155
        if is_in_ci():
            return

156
        for model_case in ALL_OTHER_MODELS:
157
            # Only run a specified model
158
159
160
161
162
            if (
                "ONLY_RUN" in os.environ
                and os.environ["ONLY_RUN"] != model_case.model_path
            ):
                continue
163

164
            # Skip long prompts for models that do not have a long context
165
            prompts = DEFAULT_PROMPTS
166
            if model_case.skip_long_prompt:
167
168
169
170
                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.float16)
171

172

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