From 816b96a58a4596cb13fd120cd603d5238b4e623f Mon Sep 17 00:00:00 2001 From: ViperEkura <3081035982@qq.com> Date: Sat, 5 Sep 2026 01:48:06 +0800 Subject: [PATCH] test: compare logits_positions rows with assert_close - sliced M=2 and full M=5 lm_head projections may pick different GEMM kernels whose accumulation order differs in the last float32 bits, so torch.equal flakes by machine and thread count --- tests/module/test_model_forward.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/tests/module/test_model_forward.py b/tests/module/test_model_forward.py index 6e3efd3..7e2d9ed 100644 --- a/tests/module/test_model_forward.py +++ b/tests/module/test_model_forward.py @@ -157,7 +157,9 @@ def test_forward_logits_positions_projects_only_requested_rows(): assert full["logits"].shape == (5, config.vocab_size) assert sliced["logits"].shape == (2, config.vocab_size) - assert torch.equal(sliced["logits"], full["logits"][last_rows]) + # Projecting M=2 rows vs M=5 rows may pick different GEMM kernels and + # differ in the last float32 bits — assert_close, not bit equality. + torch.testing.assert_close(sliced["logits"], full["logits"][last_rows]) assert torch.equal(sliced["hidden_states"], full["hidden_states"][last_rows])