test_dummy_grok_models.py 856 Bytes
Newer Older
saienduri's avatar
saienduri committed
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
import unittest

from sglang.test.test_utils import CustomTestCase, is_in_ci, run_bench_one_batch


class TestDummyGrok1(CustomTestCase):
    def test_dummy_grok_1(self):
        output_throughput = run_bench_one_batch(
            None,
            [
                "--model",
                "/dummy-grok",
                "--tokenizer-path",
                "Xenova/grok-1-tokenizer",
                "--batch-size",
                "2",
                "--tp",
                "2",
                "--quantization",
                "fp8",
                "--load-format",
                "dummy",
                "--json-model-override-args",
                '{"num_hidden_layers": 2}',
            ],
        )

        if is_in_ci():
            assert output_throughput > 0, f"{output_throughput=}"


if __name__ == "__main__":
    unittest.main()