ViperEkura/AstrAI
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

test: compare logits_positions rows with assert_close

Browse Source 
- 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
This commit is contained in:
ViperEkura
2026-09-05 01:48:06 +08:00
parent 6709534d64
commit 816b96a58a
1 changed files with 3 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/module/test_model_forward.py
+3 -1
View File
@@ -157,7 +157,9 @@ def test_forward_logits_positions_projects_only_requested_rows():
assert full["logits"].shape == (5, config.vocab_size) assert full["logits"].shape == (5, config.vocab_size)
assert sliced["logits"].shape == (2, 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]) assert torch.equal(sliced["hidden_states"], full["hidden_states"][last_rows])