biondizzle/vllm
1
0
Fork 0
Code Issues Pull Requests Actions 2 Packages Projects Releases Wiki Activity

[Misc/Testing] Use torch.testing.assert_close (#7324)

Browse Source
This commit is contained in:
jon-chuang
2024-08-15 21:24:04 -07:00
committed by GitHub
parent e165528778
commit 50b8d08dbd
25 changed files with 197 additions and 188 deletions
Download Patch File Download Diff File Expand all files Collapse all files

10
tests/kernels/test_moe.py
View File

@@ -50,7 +50,7 @@ def test_fused_moe(
score = torch.randn((m, e), device='cuda', dtype=dtype)
triton_output = fused_moe(a, w1, w2, score, topk, renormalize=False)
torch_output = torch_moe(a, w1, w2, score, topk)
assert torch.allclose(triton_output, torch_output, atol=1e-2, rtol=0)
torch.testing.assert_close(triton_output, torch_output, atol=1e-2, rtol=0)
@pytest.mark.parametrize("dtype",
@@ -95,7 +95,7 @@ def test_mixtral_moe(dtype: torch.dtype):
torch.bfloat16: 1e-2,
}
assert torch.allclose(hf_states.flatten(0, 1),
vllm_states,
rtol=mixtral_moe_tol[dtype],
atol=mixtral_moe_tol[dtype])
torch.testing.assert_close(hf_states.flatten(0, 1),
vllm_states,
rtol=mixtral_moe_tol[dtype],
atol=mixtral_moe_tol[dtype])