samplers/test_logprobs.py: use float32 precision.

Signed-off-by: Thomas Parnell <[email protected]>
vllm-project · Jul 13, 2024 · e3b40a7 · e3b40a7
1 parent d80aef3
commit e3b40a7
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/tests/samplers/test_logprobs.py b/tests/samplers/test_logprobs.py
@@ -11,7 +11,8 @@
 
 
 @pytest.mark.parametrize("model", MODELS)
-@pytest.mark.parametrize("dtype", ["half"])
+@pytest.mark.parametrize("dtype",
+                         ["float"])  # needed for comparing logprobs with HF
 @pytest.mark.parametrize("chunked_prefill_token_size", [1, 4, 16, -1])
 @pytest.mark.parametrize("num_top_logprobs", [6])  # 32000 == vocab_size
 @pytest.mark.parametrize("detokenize", [True, False])