[Bugfix] use float32 precision in samplers/test_logprobs.py for comparing with HF (#6409)

Signed-off-by: Thomas Parnell <tpa@zurich.ibm.com>
This commit is contained in:
Thomas Parnell
2024-07-15 19:14:49 +02:00
committed by GitHub
parent eaec4b9153
commit 4ef95b0f06
2 changed files with 8 additions and 1 deletions

View File

@@ -11,7 +11,8 @@ MODELS = ["facebook/opt-125m"]
@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])