[ROCm][CI] Set "highest" matmul precision for reference hf_runner in `test_bert_for_masked_lm` (#48784)
Signed-off-by: Micah Williamson <micah.williamson@amd.com>
Micah Williamson committed
Jul 16, 2026 at 00:49 UTC
7dc2698632867e2313625103674459c02f42f8a9
1 file changed
+6
tests/models/language/pooling/test_token_classification.py
+6
@@ -272,6 +272,11 @@ def test_bert_for_masked_lm(
272
if current_platform.is_rocm():
273
hf_model_kwargs["attn_implementation"] = "eager"
274
275
+ # Run hf_runner reference with "highest" fp32 precision to match
276
+ # default behvior of vLLM. This is needed on ROCm since the
277
+ # pooling tests set matmul precision to "high" in conftest.py
278
+ prev_matmul_precision = torch.get_float32_matmul_precision()
279
+ torch.set_float32_matmul_precision("highest")
280
with hf_runner(
281
model,
282
dtype=dtype,
@@ -285,6 +290,7 @@ def test_bert_for_masked_lm(
290
inputs = hf_model.wrap_device(inputs)
291
output = hf_model.model(**inputs)
292
hf_outputs.append(softmax(output.logits[0]))
293
+ torch.set_float32_matmul_precision(prev_matmul_precision)
294
295
# Compare the per-token vocabulary distributions position by position.
296
for hf_output, vllm_output in zip(hf_outputs, vllm_outputs):