#4873 - Fix BF16 2:4 sparsity check in Polygraphy - #4920

Fixes #4873. Thanks @v0ropaev for the diagnosis and @gabortar for the int32_data correction.

Overview

process_bf16_tensor misreads BF16 initializers. It's shared by inspect sparsity, surgeon prune, surgeon weight-strip and surgeon weight-reconstruct, so all four are affected. There are two bugs.

1. raw_data: The check loop reuses the variable i. But short2long also uses i to compute positions, so the check ends up reading a diagonal of elements instead of each group of 4. That gives wrong answers, and it often crashes with an IndexError near the end of the tensor.

2. int32_data: Each entry holds one BF16 value in its low 16 bits. The check only reads one entry per group of 2, and then it counts that value's two bytes as if they were two separate values. So a dense tensor full of 2.0 (0x4000) gets reported as sparse, because its low byte happens to be zero.

Changes made

Fix: Both formats now use the same loop over groups of 4. A small is_zero helper reads two bytes for raw_data and the low 16 bits for int32_data. Pruning output doesn't change. The int32_data branch still zeroes elements 0 and 2 of each group, which is exactly what the old loop did.

Test: I added tests/tools/test_sparse.py with 18 tests. They run on CPU and don't need TensorRT. They cover sparse tensors, dense all-2.0 tensors, tensors with only one zero in a group, and pruning followed by a re-check. Each case runs for both raw_data and int32_data, and with both stride layouts (pstride == 1 and pstride > 1). The tests check the actual sparse or not-sparse answer, not just whether the command exited successfully.

Output

# before (main)
~/TensorRT/tools/Polygraphy$ PYTHONPATH=. pytest --noconftest tests/tools/test_sparse.py -q --tb=no -rf 2>&1 | tail -12
FFFF.F.FF...F.F...     
=========================== short test summary info ============================
FAILED tests/tools/test_sparse.py::TestBF16SparsityCheck::test_sparse_tensor_is_sparse[pstride3-raw_data]
FAILED tests/tools/test_sparse.py::TestBF16SparsityCheck::test_sparse_tensor_is_sparse[pstride3-int32_data]
FAILED tests/tools/test_sparse.py::TestBF16SparsityCheck::test_sparse_tensor_is_sparse[pstride1-raw_data]
FAILED tests/tools/test_sparse.py::TestBF16SparsityCheck::test_sparse_tensor_is_sparse[pstride1-int32_data]
FAILED tests/tools/test_sparse.py::TestBF16SparsityCheck::test_dense_tensor_is_not_sparse[pstride3-int32_data]
FAILED tests/tools/test_sparse.py::TestBF16SparsityCheck::test_dense_tensor_is_not_sparse[pstride1-int32_data]
FAILED tests/tools/test_sparse.py::TestBF16SparsityCheck::test_one_zero_per_group_is_not_sparse[pstride3-raw_data]
FAILED tests/tools/test_sparse.py::TestBF16SparsityPrune::test_pruned_tensor_is_2_4_sparse[pstride3-raw_data]
FAILED tests/tools/test_sparse.py::TestBF16SparsityPrune::test_pruned_tensor_is_2_4_sparse[pstride1-raw_data]
9 failed, 9 passed in 0.15s

# after
18 passed in 0.27s